From bacf59bd497faae7b486bbd8ebe027651354bc93 Mon Sep 17 00:00:00 2001 From: Yinon Kahta Date: Thu, 27 Aug 2026 13:05:57 +0300 Subject: [PATCH] feat(pointfive): add gzipped ndjson batch encoding --- litellm/integrations/pointfive/payload.py | 53 ++++++++++++ .../integrations/pointfive/test_payload.py | 80 +++++++++++++++++++ 2 files changed, 133 insertions(+) create mode 100644 litellm/integrations/pointfive/payload.py create mode 100644 tests/test_litellm/integrations/pointfive/test_payload.py diff --git a/litellm/integrations/pointfive/payload.py b/litellm/integrations/pointfive/payload.py new file mode 100644 index 00000000000..e3362eda4ee --- /dev/null +++ b/litellm/integrations/pointfive/payload.py @@ -0,0 +1,53 @@ +"""Turns buffered log records into the gzipped NDJSON objects that get uploaded.""" + +import gzip +from collections.abc import Iterator, Mapping, Sequence +from itertools import accumulate, groupby, islice +from typing import Final + +from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + +_NEWLINE_BYTES: Final = 1 + + +def serialize_records(records: Sequence[Mapping[str, object]]) -> tuple[str, ...]: + """Serialize each record to one JSON line.""" + return tuple(safe_dumps(record) for record in records) + + +def _encoded_size(line: str) -> int: + return len(line.encode("utf-8")) + _NEWLINE_BYTES + + +def _object_indices(sizes: Sequence[int], max_bytes: int) -> Iterator[int]: + """Number each line with the object it belongs to, opening a new one on overflow.""" + + def advance(state: tuple[int, int], size: int) -> tuple[int, int]: + index, used = state + return (index + 1, size) if used and used + size > max_bytes else (index, used + size) + + return (index for index, _ in islice(accumulate(sizes, advance, initial=(0, 0)), 1, None)) + + +def chunk_lines(lines: Sequence[str], max_bytes: int) -> tuple[tuple[str, ...], ...]: + """ + Group serialized lines into objects of at most ``max_bytes`` uncompressed. + + A line above the bound on its own still becomes its own object. A record cannot be + split, and holding it back would stall every record queued behind it. + """ + sizes: Final = tuple(_encoded_size(line) for line in lines) + numbered: Final = zip(_object_indices(sizes, max_bytes), lines, strict=True) + return tuple(tuple(line for _, line in group) for _, group in groupby(numbered, lambda pair: pair[0])) + + +async def encode_lines(lines: Sequence[str]) -> bytes: + """ + Join lines as NDJSON and gzip them off the event loop. + + An object can be several megabytes, and compressing that inline would block the + proxy for as long as it takes. + """ + compress: Final = asyncify(gzip.compress) + return await compress("\n".join(lines).encode("utf-8")) diff --git a/tests/test_litellm/integrations/pointfive/test_payload.py b/tests/test_litellm/integrations/pointfive/test_payload.py new file mode 100644 index 00000000000..d310941e396 --- /dev/null +++ b/tests/test_litellm/integrations/pointfive/test_payload.py @@ -0,0 +1,80 @@ +import gzip +import json + +import pytest + +from litellm.integrations.pointfive.payload import chunk_lines, encode_lines, serialize_records + +UNBOUNDED = 10_000_000 + + +def test_each_record_becomes_one_json_line(): + lines = serialize_records([{"id": "a"}, {"id": "b"}, {"id": "c"}]) + + assert len(lines) == 3 + assert [json.loads(line)["id"] for line in lines] == ["a", "b", "c"] + + +def test_non_serializable_values_do_not_raise(): + """An odd payload must not kill the flush.""" + lines = serialize_records([{"id": "a", "when": object()}]) + + assert json.loads(lines[0])["id"] == "a" + + +def test_records_that_fit_stay_in_one_object(): + lines = serialize_records([{"id": f"r{i}"} for i in range(50)]) + + assert chunk_lines(lines, UNBOUNDED) == (lines,) + + +def test_objects_are_capped_by_uncompressed_size(): + lines = serialize_records([{"id": f"r{i}", "blob": "x" * 100} for i in range(10)]) + line_bytes = len(lines[0].encode("utf-8")) + 1 + + chunks = chunk_lines(lines, line_bytes * 3) + + assert [len(chunk) for chunk in chunks] == [3, 3, 3, 1] + + +def test_oversized_single_record_is_sent_alone_not_stalled(): + """A record too big for the cap must still go out, or it blocks everything behind it.""" + lines = serialize_records([{"id": "small"}, {"id": "huge", "blob": "x" * 5000}, {"id": "small2"}]) + + chunks = chunk_lines(lines, 200) + + assert sum(len(chunk) for chunk in chunks) == 3 + huge = [chunk for chunk in chunks if any("huge" in line for line in chunk)] + assert len(huge) == 1 + assert len(huge[0]) == 1 + + +def test_no_records_produces_no_objects(): + assert chunk_lines((), UNBOUNDED) == () + + +def test_every_record_appears_exactly_once(): + lines = serialize_records([{"id": f"r{i}"} for i in range(37)]) + + chunks = chunk_lines(lines, len(lines[0]) * 4) + + assert [line for chunk in chunks for line in chunk] == list(lines) + + +@pytest.mark.asyncio +async def test_encode_lines_round_trips_through_gzip(): + lines = serialize_records([{"id": f"r{i}"} for i in range(5)]) + + encoded = await encode_lines(lines) + + assert encoded[:2] == b"\x1f\x8b" + assert gzip.decompress(encoded).decode("utf-8") == "\n".join(lines) + + +@pytest.mark.asyncio +async def test_encode_lines_compresses_repetitive_records(): + lines = serialize_records([{"id": f"r{i}", "model": "gpt-4o", "cost": 0.01} for i in range(200)]) + + encoded = await encode_lines(lines) + + assert len(encoded) < len(gzip.decompress(encoded)) / 2