diff --git a/sagemaker-core/src/sagemaker/core/iterators.py b/sagemaker-core/src/sagemaker/core/iterators.py index 60914cbdd0..37f39fce6f 100644 --- a/sagemaker-core/src/sagemaker/core/iterators.py +++ b/sagemaker-core/src/sagemaker/core/iterators.py @@ -175,7 +175,15 @@ def __next__(self): chunk = next(self.byte_iterator) except StopIteration: if self.read_pos < self.buffer.getbuffer().nbytes: - continue + # Stream ended with a trailing partial line (no "\n"). + # `continue` alone would spin forever here: byte_iterator + # is already exhausted, so it keeps raising StopIteration + # and read_pos/buffer never change. Return the remainder + # once, so the next call correctly raises StopIteration. + self.buffer.seek(self.read_pos) + remainder = self.buffer.read() + self.read_pos += len(remainder) + return remainder raise if "PayloadPart" not in chunk: # handle API response errors and force terminate. diff --git a/sagemaker-core/tests/unit/test_iterators.py b/sagemaker-core/tests/unit/test_iterators.py index 02ed29e2be..3bbefda59e 100644 --- a/sagemaker-core/tests/unit/test_iterators.py +++ b/sagemaker-core/tests/unit/test_iterators.py @@ -232,3 +232,26 @@ def test_line_iterator_incomplete_line_at_end(): # After consuming all complete lines, should raise StopIteration with pytest.raises(StopIteration): next(iterator) + + +def test_line_iterator_no_trailing_newline_at_end(): + """Test LineIterator returns the final chunk even if it has no trailing "\n". + + Regression test: previously this hung forever instead of returning or + raising, because the StopIteration handler for a leftover partial line + just did `continue`, re-reading the same unterminated line and calling + next() on an already-exhausted iterator every time. + """ + mock_stream = [ + {"PayloadPart": {"Bytes": b'{"outputs": [" first"]}\n'}}, + {"PayloadPart": {"Bytes": b'{"outputs": [" second"]}'}}, # no trailing \n + ] + iterator = LineIterator(mock_stream) + + line1 = next(iterator) + line2 = next(iterator) + assert line1 == b'{"outputs": [" first"]}' + assert line2 == b'{"outputs": [" second"]}' + + with pytest.raises(StopIteration): + next(iterator)