Skip to content

Commit 4beebce

Browse files
authored
Merge branch 'main' into bugfix/MPT-24897/stream-raises-list-omits-truncation-and-retry-errors
2 parents 7cae6f4 + 7789cb2 commit 4beebce

3 files changed

Lines changed: 133 additions & 26 deletions

File tree

‎docs/streaming.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,10 @@ record rather than by the size of the export. A ten-million-record stream costs
108108
ten-record one. This holds in both wire formats — the envelope is tokenized incrementally
109109
rather than buffered.
110110

111+
Time scales the same way. A record that spans many body chunks is assembled once, when its
112+
line ends, so the cost of a record is linear in its own length: doubling the record size
113+
doubles the work rather than quadrupling it.
114+
111115
The buffering paths, for contrast:
112116

113117
- `iterate()` holds one page. It buffers each page response in full, deserializes it into a

‎mpt_api_client/http/jsonl_lines.py‎

Lines changed: 49 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,40 @@ def decode_record_line(line: str) -> dict[str, Any]:
2929
return record
3030

3131

32+
type _SplitResult = tuple[list[str], list[str]]
33+
34+
35+
def _split_chunk(pending: list[str], chunk: str, *, at_body_start: bool) -> _SplitResult:
36+
"""Split one body chunk into the lines it completes.
37+
38+
The fragments of the line still being assembled are kept as a list and joined once,
39+
when a line feed finally ends that line. Accumulating them into one string instead
40+
would copy and re-scan the whole partial record on every chunk, making a record that
41+
spans many chunks quadratic in its own length.
42+
43+
Args:
44+
pending: Fragments of the line still being assembled, in arrival order; the list
45+
is consumed here, so the caller must go on with the one returned.
46+
chunk: Next non-empty decoded chunk of the body.
47+
at_body_start: Whether this chunk opens the body, and so may carry the one
48+
tolerated byte order mark — `pending` is necessarily empty in that case.
49+
50+
Returns:
51+
The lines this chunk completes, without their line endings, and the fragments
52+
buffered for the next chunk.
53+
"""
54+
fragment = chunk.removeprefix(UTF8_BOM) if at_body_start else chunk
55+
if "\n" not in fragment:
56+
pending.append(fragment)
57+
return [], pending
58+
lines = fragment.split("\n")
59+
pending.append(lines[0])
60+
lines[0] = "".join(pending)
61+
trailing = lines.pop()
62+
completed = [line.removesuffix("\r") for line in lines]
63+
return completed, ([trailing] if trailing else [])
64+
65+
3266
def iter_jsonl_lines(text_chunks: Iterable[str]) -> Iterator[str]:
3367
"""Iterate the lines of a JSONL body, splitting on newlines alone.
3468
@@ -49,22 +83,16 @@ def iter_jsonl_lines(text_chunks: Iterable[str]) -> Iterator[str]:
4983
of a CRLF ending, so a final unterminated line is yielded as-is; a blank line
5084
is yielded as an empty string, for the caller to skip.
5185
"""
52-
pending = ""
86+
pending: list[str] = []
5387
at_body_start = True
5488
for chunk in text_chunks:
55-
if at_body_start and chunk:
56-
# `pending` is necessarily empty before the first non-empty chunk, so the
57-
# body's first character — the only place a BOM is tolerated — is its start.
58-
pending = chunk.removeprefix(UTF8_BOM)
59-
at_body_start = False
60-
else:
61-
pending += chunk
62-
lines = pending.split("\n")
63-
pending = lines.pop()
64-
for line in lines:
65-
yield line.removesuffix("\r")
89+
if not chunk:
90+
continue
91+
completed, pending = _split_chunk(pending, chunk, at_body_start=at_body_start)
92+
at_body_start = False
93+
yield from completed
6694
if pending:
67-
yield pending
95+
yield "".join(pending)
6896

6997

7098
async def aiter_jsonl_lines(text_chunks: AsyncIterable[str]) -> AsyncIterator[str]:
@@ -87,19 +115,14 @@ async def aiter_jsonl_lines(text_chunks: AsyncIterable[str]) -> AsyncIterator[st
87115
of a CRLF ending, so a final unterminated line is yielded as-is; a blank line
88116
is yielded as an empty string, for the caller to skip.
89117
"""
90-
pending = ""
118+
pending: list[str] = []
91119
at_body_start = True
92120
async for chunk in text_chunks:
93-
if at_body_start and chunk:
94-
# `pending` is necessarily empty before the first non-empty chunk, so the
95-
# body's first character — the only place a BOM is tolerated — is its start.
96-
pending = chunk.removeprefix(UTF8_BOM)
97-
at_body_start = False
98-
else:
99-
pending += chunk
100-
lines = pending.split("\n")
101-
pending = lines.pop()
102-
for line in lines:
103-
yield line.removesuffix("\r")
121+
if not chunk:
122+
continue
123+
completed, pending = _split_chunk(pending, chunk, at_body_start=at_body_start)
124+
at_body_start = False
125+
for line in completed:
126+
yield line
104127
if pending:
105-
yield pending
128+
yield "".join(pending)

‎tests/unit/http/test_jsonl_lines.py‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
import asyncio
22
import json
33
import re
4+
import time
45
from pathlib import Path
6+
from typing import NamedTuple
57

68
import pytest
79

@@ -103,6 +105,67 @@ def test_keeps_a_bom_after_the_body_start():
103105
assert result == ['{"id": 1}', '\ufeff{"id": 2}']
104106

105107

108+
LONG_RECORD_CHUNK_SIZE = 8 * 1024
109+
LONG_RECORD_SIZE = 16 * 1024 * 1024
110+
# The linear splitter needs a few milliseconds for a 16MB record, while re-splitting the
111+
# whole buffer on every chunk needed seconds, so this budget is generous on a slow
112+
# machine and still fails on a return to quadratic behaviour.
113+
LONG_RECORD_BUDGET_SECONDS = 2
114+
115+
116+
@pytest.fixture(scope="module")
117+
def long_record():
118+
# Built on demand rather than at import, so runs that do not select the tests below
119+
# do not pay for the 16MB body during collection.
120+
return json.dumps({"note": "a" * LONG_RECORD_SIZE})
121+
122+
123+
def chunked(body, size=LONG_RECORD_CHUNK_SIZE):
124+
offsets = range(0, len(body), size)
125+
return [body[offset : offset + size] for offset in offsets]
126+
127+
128+
class TimedLines(NamedTuple):
129+
lines: list[str]
130+
elapsed: float
131+
132+
133+
def timed_lines(chunks):
134+
started_at = time.perf_counter()
135+
lines = list(iter_jsonl_lines(chunks))
136+
return TimedLines(lines, time.perf_counter() - started_at)
137+
138+
139+
async def collected_lines(chunks):
140+
return [line async for line in aiter_jsonl_lines(chunks)]
141+
142+
143+
async def timed_async_lines(chunks):
144+
started_at = time.perf_counter()
145+
lines = await collected_lines(chunks)
146+
return TimedLines(lines, time.perf_counter() - started_at)
147+
148+
149+
def test_long_record_across_chunks_is_linear(long_record):
150+
chunks = chunked(f"{long_record}\n")
151+
152+
result = timed_lines(chunks)
153+
154+
assert result.lines == [long_record]
155+
assert result.elapsed < LONG_RECORD_BUDGET_SECONDS
156+
157+
158+
def test_unterminated_long_record_is_linear(long_record):
159+
# The body confirm_stream_format tolerates without a content type: a single long
160+
# line, unterminated, so every chunk stays buffered and nothing is flushed early.
161+
chunks = chunked(long_record)
162+
163+
result = timed_lines(chunks)
164+
165+
assert result.lines == [long_record]
166+
assert result.elapsed < LONG_RECORD_BUDGET_SECONDS
167+
168+
106169
def test_decode_record_line_returns_the_object():
107170
result = decode_record_line('{"id": "ID-1"}')
108171

@@ -156,6 +219,15 @@ async def test_async_preserves_standalone_trailing_cr():
156219
assert result == ['{"id": 1}\r']
157220

158221

222+
async def test_async_long_record_is_linear(long_record):
223+
chunks = async_chunks(chunked(f"{long_record}\n"))
224+
225+
result = await timed_async_lines(chunks)
226+
227+
assert result.lines == [long_record]
228+
assert result.elapsed < LONG_RECORD_BUDGET_SECONDS
229+
230+
159231
async def test_async_drops_a_bom_opening_the_body():
160232
chunks = async_chunks(['\ufeff{"id": 1}\n{"id": 2}\n'])
161233

@@ -164,6 +236,14 @@ async def test_async_drops_a_bom_opening_the_body():
164236
assert result == ['{"id": 1}', '{"id": 2}']
165237

166238

239+
async def test_async_drops_bom_after_empty_chunk():
240+
chunks = async_chunks(["", '\ufeff{"id": 1}\n'])
241+
242+
result = await collected_lines(chunks)
243+
244+
assert result == ['{"id": 1}']
245+
246+
167247
async def test_async_drops_the_bom_in_its_own_chunk():
168248
chunks = async_chunks(["\ufeff", '{"id": 1}\n'])
169249

0 commit comments

Comments
 (0)