From 27c91e6574c2efded57dced72f6a653d46bd712e Mon Sep 17 00:00:00 2001 From: mateo Date: Thu, 23 Jul 2026 04:03:25 +0000 Subject: [PATCH] fix(sagemaker): forward native streaming events as they arrive to cut TTFT Mirror the sagemaker_chat fix on the native sagemaker/ streaming path: the sync and async completion handlers read the invocations-response-stream body with iter_bytes(chunk_size=1024) / aiter_bytes(chunk_size=1024), so httpx withholds bytes until 1024 accumulate and tokens arrive in gap-then-burst waves. Drop the fixed chunk size so each decoded event is forwarded as its bytes arrive. Also add a boundary-agnostic decoder test proving frames reassemble correctly regardless of where transport reads split the stream. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/sagemaker/completion/handler.py | 4 +- .../test_sagemaker_chat_transformation.py | 26 +++++ .../test_sagemaker_completion_handler.py | 107 ++++++++++++++++++ 3 files changed, 135 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index 4b87271fd44..5012ed01c1e 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -216,7 +216,7 @@ class SagemakerLLM(BaseAWSLLM): decoder = AWSEventStreamDecoder(model="") - completion_stream = decoder.iter_bytes(sync_response.iter_bytes(chunk_size=1024)) + completion_stream = decoder.iter_bytes(sync_response.iter_bytes()) streaming_response = CustomStreamWrapper( completion_stream=completion_stream, model=model, @@ -358,7 +358,7 @@ class SagemakerLLM(BaseAWSLLM): raise SagemakerError(status_code=response.status_code, message=response.text) decoder = AWSEventStreamDecoder(model="") - completion_stream = decoder.aiter_bytes(response.aiter_bytes(chunk_size=1024)) + completion_stream = decoder.aiter_bytes(response.aiter_bytes()) return completion_stream diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py index 806d040caf6..54e6f95c795 100644 --- a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py @@ -207,3 +207,29 @@ def test_signed_body_includes_stream_flag(): ) assert signed_body is not None assert json.loads(signed_body)["stream"] is True + + +@pytest.mark.parametrize("split_size", [1, 3, 7, 64, 4096]) +def test_decoder_reassembles_frames_across_arbitrary_byte_boundaries(split_size): + """Correctness must not depend on chunk boundaries falling on frame edges. + + Removing `chunk_size=1024` lets httpx yield raw transport reads, so in + production a single read can straddle several frames or split one frame in + half. This re-chunks the concatenated stream at boundaries that deliberately + ignore frame edges and asserts every delta still decodes, in order, exactly + once - the guarantee botocore's EventStreamBuffer provides. + """ + from litellm.llms.sagemaker.chat.transformation import AWSEventStreamDecoder + + frames = _make_frames(24) + blob = b"".join(frames) + chunks = [blob[i : i + split_size] for i in range(0, len(blob), split_size)] + + decoder = AWSEventStreamDecoder(model="phi-4", is_messages_api=True) + texts = [ + _content_of(chunk) + for chunk in decoder.iter_bytes(iter(chunks)) + if chunk is not None and _content_of(chunk) is not None + ] + + assert texts == [f"token{i} " for i in range(len(frames))] diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py new file mode 100644 index 00000000000..d43c07e09bb --- /dev/null +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py @@ -0,0 +1,107 @@ +""" +Regression tests for LIT-4313: the native `sagemaker/` streaming path must +forward each AWS event-stream frame as it arrives instead of buffering to a +fixed 1024-byte threshold and then draining a burst of tokens. + +The buffering came from `response.aiter_bytes(chunk_size=1024)`: httpx's +ByteChunker withholds bytes until `chunk_size` accumulates, so the first token +could not be produced until enough later frames had arrived to cross 1024 bytes, +inflating TTFT and turning a steady provider stream into gap-then-burst delivery. +""" + +import binascii +import json +import struct +from typing import AsyncIterator +from unittest.mock import MagicMock + +import httpx +import pytest + +from litellm.llms.sagemaker.completion.handler import SagemakerLLM + + +def _encode_header(name: str, value: str) -> bytes: + name_b = name.encode("utf-8") + value_b = value.encode("utf-8") + return struct.pack("B", len(name_b)) + name_b + struct.pack("B", 7) + struct.pack(">H", len(value_b)) + value_b + + +def _encode_event_frame(payload: bytes) -> bytes: + """Encode one AWS event-stream message that botocore's EventStreamBuffer decodes.""" + headers = { + ":event-type": "PayloadPart", + ":content-type": "application/json", + ":message-type": "event", + } + headers_b = b"".join(_encode_header(k, v) for k, v in headers.items()) + total_len = 16 + len(headers_b) + len(payload) + prelude = struct.pack(">I", total_len) + struct.pack(">I", len(headers_b)) + prelude_crc = struct.pack(">I", binascii.crc32(prelude) & 0xFFFFFFFF) + message = prelude + prelude_crc + headers_b + payload + message_crc = struct.pack(">I", binascii.crc32(message) & 0xFFFFFFFF) + return message + message_crc + + +def _token_frame(text: str) -> bytes: + # SageMaker HF TGI streaming payloads are `{"token": {"text": ...}}` blobs. + sse = "data: " + json.dumps({"token": {"text": text}}) + "\n\n" + return _encode_event_frame(sse.encode("utf-8")) + + +def _make_frames(n: int) -> list[bytes]: + frames = [_token_frame(f"token{i} ") for i in range(n)] + assert all(len(f) < 1024 for f in frames) + return frames + + +class _CountingAsyncStream(httpx.AsyncByteStream): + """Yields provider frames one at a time and records how many have been pulled.""" + + def __init__(self, frames: list[bytes]) -> None: + self._frames = frames + self.consumed = 0 + + async def __aiter__(self) -> AsyncIterator[bytes]: + for frame in self._frames: + self.consumed += 1 + yield frame + + +class _FakeAsyncClient: + def __init__(self, response: httpx.Response) -> None: + self._response = response + + async def post(self, *args, **kwargs) -> httpx.Response: + return self._response + + +@pytest.mark.asyncio +async def test_async_native_streaming_forwards_each_frame_incrementally(): + """Each token must be emitted after exactly one newly-pulled source frame. + + With the old `chunk_size=1024` the httpx chunker would swallow several small + frames before yielding, so the first token would arrive only after `consumed` + had already crossed multiple frames, and tokens would then replay in a burst. + """ + frames = _make_frames(24) + stream = _CountingAsyncStream(frames) + response = httpx.Response(200, stream=stream) + + completion_stream = await SagemakerLLM().make_async_call( + api_base="https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/phi-4/invocations-response-stream", + headers={}, + data="", + logging_obj=MagicMock(), + client=_FakeAsyncClient(response), + ) + + consumed_at_token = [] + texts = [] + async for chunk in completion_stream: + if chunk is not None and chunk["text"]: + consumed_at_token.append(stream.consumed) + texts.append(chunk["text"]) + + assert texts == [f"token{i} " for i in range(len(frames))] + assert consumed_at_token == list(range(1, len(frames) + 1))