fix(sagemaker): forward stream events as they arrive to cut TTFT

[devin-ai-integration]

TLDR

Problem this solves:

• SageMaker streaming buffered tokens and inflated client TTFT
• tokens arrived in gap-then-burst waves instead of steadily

How it solves it:

• stop forcing httpx to buffer to a fixed 1024-byte threshold
• forward each decoded AWS event as soon as its bytes arrive

Relevant issues

Linear ticket

Resolves LIT-4313

Pre-Submission checklist

• I have added meaningful tests
• My PR passes all CI/CD checks (e.g., lint, format, unit tests)
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review

Screenshots / Proof of Fix

Live e2e against real SageMaker TGI endpoints (us-west-2), comparing a litellm proxy built from the merge-base (d25bac5a41, still chunk_size=1024) against this branch (5bc1df5e47), same endpoint and prompt, 3 runs each. A streaming client timestamps every SSE content delta as it arrives: "burst gaps" counts deltas that landed <1ms after the previous one (flushed together from the buffer), "arrival cadence" is the median gap between distinct arrivals

sagemaker_chat through proxy /v1/chat/completions (TinyLlama-1.1B-Chat on TGI 1.4.0 with MESSAGES_API_ENABLED=true, 80 tokens):

proxy code	median TTFT	burst gaps	arrival cadence
before	332 ms	52-60 / 79	21-26 ms
after	171 ms	7 / 79	8-9 ms

Native sagemaker/ through proxy (MPT-7B-Instruct on TGI 1.4.0, 25 tokens):

proxy code	median TTFT	burst gaps	arrival cadence
before	617 ms	19 / 24	123-136 ms
after	572 ms	0-6 / 24	31-33 ms

The sync SDK paths (litellm.completion(stream=True)) show the same collapse: native sagemaker/ goes from 19/24 burst gaps at ~126 ms cadence to 0/24 at ~33 ms, and sagemaker_chat from 13-16/24 to 2-3/24 at the model's ~8 ms token cadence. Together this exercises all four changed reads (sagemaker_chat sync/async, completion handler sync/async): with the fix every delta arrives at the provider's real token cadence instead of in 1024-byte flush waves. Native-path TTFT moves less because MPT-7B's own first token dominates it; the chat endpoint's TTFT halves. The temporary Messages API endpoint was deleted after the run

Earlier adapter-path proof with simulated arrival timing (pre-provisioning)

This drives real AWS event-stream frames through exactly the adapter path (httpx.Response.iter_bytes(...) -> AWSEventStreamDecoder.iter_bytes(...)) with a stream that emits one frame every 40ms, so only the network arrival timing is simulated; the buffering code under test is the real thing. Comparison captured at commit 0ff3ade by toggling the chunk_size the code used before vs after:

Before (chunk_size=1024), provider emits a frame every 40ms:

deltas decoded: 24
TTFT (first delta): 146 ms
gaps>1.5x provider cadence: 5      # deltas arrive in bursts of ~4 with ~160ms gaps

After (no fixed threshold):

deltas decoded: 24
TTFT (first delta): 0 ms
gaps>1.5x provider cadence: 0      # each delta flows through at the 40ms provider cadence

Type

🐛 Bug Fix

Changes

Both SageMaker streaming surfaces read the /invocations-response-stream body with response.iter_bytes(chunk_size=1024) / response.aiter_bytes(chunk_size=1024). httpx's ByteChunker withholds bytes until chunk_size accumulates before yielding, so when each AWS event frame is small (a single-token SSE delta is ~150-270 bytes) the client cannot see delta N until enough later frames arrive to cross 1024 bytes; AWSEventStreamDecoder then drains every complete frame from the released block in one pass, producing the ticket's wait-then-drain-a-burst pattern and inflating client-observed time-to-first-token

The fix drops the chunk_size argument on both reads so httpx yields transport bytes as they arrive and each decoded event is forwarded immediately. This matches the bedrock invoke path, whose stream_chunk_size defaults to None

sagemaker_chat (litellm/llms/sagemaker/chat/transformation.py):

-completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=1024))
+completion_stream = decoder.iter_bytes(response.iter_bytes())
-completion_stream = decoder.aiter_bytes(response.aiter_bytes(chunk_size=1024))
+completion_stream = decoder.aiter_bytes(response.aiter_bytes())

The native sagemaker completion path had the identical buffering on its real event-stream reads, so it is fixed the same way for consistency (litellm/llms/sagemaker/completion/handler.py, sync completion + async make_async_call):

-completion_stream = decoder.iter_bytes(sync_response.iter_bytes(chunk_size=1024))
+completion_stream = decoder.iter_bytes(sync_response.iter_bytes())
-completion_stream = decoder.aiter_bytes(response.aiter_bytes(chunk_size=1024))
+completion_stream = decoder.aiter_bytes(response.aiter_bytes())

On the ticket's second proposed root cause (supports_stream_param_in_request_body returning False so "stream": true is missing from the signed body): that does not reproduce. stream already reaches the signed body today because it flows through optional_params into the transformed request ({**optional_params} in OpenAIGPTConfig.transform_request) and is serialized by _sign_request before SigV4 signing. Flipping the property is a no-op for the request that actually goes out (the signed path sends signed_json_body, and _add_stream_param_to_request_body only mutates the post-sign data dict that sagemaker_chat never sends), so I left the property alone rather than make a change with no runtime effect that also breaks an existing backwards-compat assertion. A regression test pins that stream survives signing

Tests in tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py feed real AWS event-stream frames (one small SSE delta each) through the actual get_sync_custom_stream_wrapper / get_async_custom_stream_wrapper wiring using a stream that counts how many source frames have been pulled. They prove the first delta is emitted after exactly one frame is pulled (not held until a 1024-byte buffer fills), that delta i is emitted after exactly i+1 frames (steady, not held-and-replayed as a burst), and that stream: true survives SigV4 signing. A parametrized decoder test also re-chunks the concatenated stream at boundaries that deliberately ignore frame edges (1, 3, 7, 64, 4096 bytes) and asserts every delta still decodes in order exactly once, since removing the fixed chunk size lets httpx hand the decoder arbitrary transport-sized reads that can straddle or split frames

tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py covers the native sagemaker async path through make_async_call with an injected client, asserting each token is emitted after exactly one newly-pulled frame. The buffering tests fail on the pre-fix code and pass after, so they lock in the regression

Docs are updated separately in litellm-docs to distinguish native sagemaker_chat streaming from the legacy sagemaker/ fake-streaming path

Final Attestation

• The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

Link to Devin session: https://app.devin.ai/sessions/aef0886b17894acd972bab04f108708e
Requested by: @mateo-berri

[devin-ai-integration]

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

• Address comments on this PR. Add '(aside)' to your comment to have me ignore it.
• Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

• Disable automatic comment, CI, and merge conflict monitoring

[greptile-apps]

Greptile Summary

This PR fixes SageMaker streaming latency by dropping the chunk_size=1024 argument from all four iter_bytes/aiter_bytes calls across both the sagemaker_chat and native sagemaker/ completion paths. It also extracts a make_sync_call() helper on SagemakerLLM that mirrors the existing make_async_call() pattern.

• Core fix (transformation.py, handler.py): Removing chunk_size=1024 lets httpx yield transport bytes as they arrive instead of accumulating 1 KB before releasing, which was causing gap-then-burst delivery and inflating client-observed TTFT when individual AWS event-stream frames were smaller than the threshold (~150–270 bytes for a single SSE delta).
• Refactor (handler.py): make_sync_call() is extracted to keep the sync path testable via client injection; logging_obj is accepted for interface parity with make_async_call but is not wired to the raw httpx call (consistent with the original inline code).
• Tests (test_sagemaker_chat_transformation.py, test_sagemaker_completion_handler.py): New mock-only regression tests verify per-frame incremental delivery for all four code paths and parametrized chunk-boundary reassembly for the decoder; no real network calls.

Confidence Score: 5/5

Safe to merge — the production changes are four one-line argument deletions, each with dedicated regression tests that fail on the old code and pass after the fix.

The change is minimal: dropping a chunk_size argument from httpx streaming calls to restore the library's default (yield bytes as they arrive). The refactored make_sync_call mirrors the existing make_async_call pattern exactly. New mock-only tests cover all four code paths and lock in the streaming cadence guarantee. No auth, database, or routing logic is touched.

No files require special attention.

Important Files Changed

Filename	Overview
litellm/llms/sagemaker/chat/transformation.py	Drops chunk_size=1024 from both sync iter_bytes() and async aiter_bytes() calls — minimal, correct change that lets httpx yield transport bytes as they arrive
litellm/llms/sagemaker/completion/handler.py	Extracts a new make_sync_call() method mirroring make_async_call(), removes chunk_size=1024 from async path; logging_obj is accepted but unused in make_sync_call, consistent with the original inline code which also never passed it to the raw httpx client
tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py	New regression tests: mock-only (no network calls), verify per-frame incremental delivery for both sync and async paths; test helper code is duplicated verbatim with the completion handler test file
tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py	New regression tests for native sagemaker/ completion path; covers sync, async, and 500-error paths; helper classes are identical copies of those in the chat transformation test file

_{Reviews (3): Last reviewed commit: "test(sagemaker): assert make_sync_call m..." | Re-trigger Greptile}

[codecov]

Codecov Report

❌ Patch coverage is 83.33333% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
litellm/llms/sagemaker/completion/handler.py	80.00%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

[shivamrawat1]

@greptile review again

[codspeed-hq]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing litellm_lit_4313_sagemaker_chat_streaming_ttft (5bc1df5) with litellm_internal_staging (ba86889)}

[mateo-berri]

@greptileai

[mateo-berri]

LGTM; thanks!