fix(batches): return original JSONL on unparseable row to avoid silent batch truncation

The streaming rewrite of replace_model_in_jsonl accumulated physical lines and
skipped a row on JSONDecodeError to support multi-line objects, but a genuinely
malformed or truncated row never completes: it poisons the buffer, swallows every
following row, and the function still returned the partial rewrite (the rows before
the bad one, already model-rewritten) as if the batch were complete. That turned the
pre-rewrite behavior of returning the original file unchanged (so the provider rejects
the bad batch loudly) into a silent partial submission.

Restore the original-content fallback: when an unparseable remainder is left after the
loop, return the original file_content (rewinding a consumed seekable source) instead of
the truncated output. The multi-line happy path is unchanged.
This commit is contained in:
mubashir1osmani 2026-06-22 17:15:52 -07:00
parent 935ee873ae
commit 809eda2d99
2 changed files with 67 additions and 0 deletions

View file

@ -139,9 +139,19 @@ def replace_model_in_jsonl(file_content: FileTypes, new_model_name: str) -> File
buffer = ""
if buffer.strip():
# A row never parsed (truncated/malformed, or it swallowed the rows
# that followed it). Returning the partial `output` would silently
# drop those rows; return the unchanged original so the provider
# rejects the batch loudly instead of accepting a truncated one.
verbose_logger.error(
f"error parsing trailing batch content: {buffer[:100]}..."
)
if hasattr(source, "seek"):
try:
source.seek(0) # type: ignore[attr-defined]
except (OSError, ValueError):
pass
return file_content
# If no valid JSON objects were found, return the original content
if not wrote_any:

View file

@ -233,6 +233,63 @@ def test_parse_jsonl_with_embedded_newlines_whitespace_only():
assert len(result) == 0
def test_replace_model_in_jsonl_malformed_middle_row_returns_original():
"""Regression: a malformed/truncated middle row must not silently drop the
rows that follow it. The streaming rewrite accumulates physical lines into a
buffer; a row that never parses poisons the buffer so every later valid row
is concatenated into it and dropped. Returning that partial rewrite would
ship a truncated batch with no error to the caller. Instead the original
content is returned unchanged so the provider rejects the bad batch loudly."""
content = (
b'{"custom_id":"a","body":{"model":"x"}}\n'
b'{"custom_id":"b","body":{"model":\n' # truncated, never completes
b'{"custom_id":"c","body":{"model":"x"}}\n'
)
result = replace_model_in_jsonl(content, "new-model")
assert (
result == content
), "must return the original unchanged, not a partial rewrite"
def test_replace_model_in_jsonl_malformed_row_seekable_handle_rewound():
"""When the source is a seekable handle that gets consumed during the failed
rewrite, it must be rewound to 0 so the caller can re-read the full original."""
content = (
b'{"custom_id":"a","body":{"model":"x"}}\n'
b'{"custom_id":"b","body":{"model":\n'
b'{"custom_id":"c","body":{"model":"x"}}\n'
)
handle = BytesIO(content)
result = replace_model_in_jsonl(handle, "new-model")
assert result is handle
assert handle.read() == content, "handle must be rewound for the caller to re-read"
def test_replace_model_in_jsonl_multi_row_rewrites_every_model():
"""Happy path: a well-formed multi-row file gets every row's model rewritten
and no row is dropped."""
content = (
b'{"custom_id":"a","body":{"model":"old1"}}\n'
b'{"custom_id":"b","body":{"model":"old2"}}\n'
b'{"custom_id":"c","body":{"model":"old3"}}\n'
)
result = replace_model_in_jsonl(content, "new-model")
assert isinstance(result, InMemoryFile)
rows = [
json.loads(line)
for line in result.getvalue().decode("utf-8").splitlines()
if line.strip()
]
assert [row["custom_id"] for row in rows] == ["a", "b", "c"]
assert all(row["body"]["model"] == "new-model" for row in rows)
def test_replace_model_in_jsonl_with_embedded_newlines():
"""Test that replace_model_in_jsonl works correctly with embedded newlines in content"""
# Create a JSONL with embedded newlines in the message content