Fix LineIterator hanging forever when the stream ends without a trailing newline

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.

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)