Skip to content

fix(sagemaker): parse every streamed record however SageMaker cuts the payload parts - #47488

Closed
mateo-berri wants to merge 1 commit into
mainfrom
litellm_sagemaker_stream_framing
Closed

fix(sagemaker): parse every streamed record however SageMaker cuts the payload parts#47488
mateo-berri wants to merge 1 commit into
mainfrom
litellm_sagemaker_stream_framing

Conversation

@mateo-berri

@mateo-berri mateo-berri commented Oct 20, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

TLDR

Problem this solves:

  • SageMaker cuts a container's stream into parts at any byte
  • The decoder read each part as exactly one JSON record
  • vLLM streams log Unparseable JSON data remained: [DONE] at ERROR every time
  • Some streams silently lose the rest of the answer and the real usage

How it solves it:

  • Buffer the text across parts and peel off every complete JSON record
  • Skip SSE framing (data:, [DONE], whitespace) between records
  • Report leftover text only when a record is truly incomplete at stream end

User Flow

Before: an app streaming chat completions from a vLLM model on a SageMaker endpoint gets some answers cut short with no error, and the proxy log fills with ERROR lines

  1. The proxy admin adds a sagemaker_chat/<endpoint> deployment for an endpoint running a vLLM container
  2. The app sends POST https://litellm-domain/v1/chat/completions with "stream": true and "stream_options": {"include_usage": true}
  3. Every stream adds LiteLLM:ERROR ... Unparseable JSON data remained: [DONE] to the proxy log
  4. Now and then, an answer stops mid-number (..., 23, 24, 25, 2) yet still ends with finish_reason: "stop", a 200, and data: [DONE]
  5. The usage chunk on those streams holds LiteLLM's estimate of the cut text, not the container's token counts

After: the same streams arrive whole with the container's token counts and nothing in the error log

  1. The proxy admin adds a sagemaker_chat/<endpoint> deployment for an endpoint running a vLLM container
  2. The app sends POST https://litellm-domain/v1/chat/completions with "stream": true and "stream_options": {"include_usage": true}
  3. The proxy log shows no ERROR or WARNING line for the stream
  4. No answer is cut off mid-stream
  5. The usage chunk holds the token counts the container reported

Relevant issues

Replaces #31598 by @twerkmeister, which found the [DONE] log spam. That PR strips a trailing [DONE] at stream end, but the lost chunks come from SageMaker splitting or merging events across parts, so this PR fixes the parsing itself. Credit for the report and the first fix goes to that PR

Design decisions

The framing skipped between records is whitespace, data:, and [DONE] only, since the containers on this route (vLLM, TGI, LMI) send data-only SSE or JSON lines. A full SSE parser would also handle event:, id:, and : comment lines, which stall the stream today as they did before, but no container here sends them, so it stays out of scope

A record that fails validation on the async path is logged and drops the rest of its own part, as before. Decoding each part as UTF-8 on its own is unchanged, since #46884 and #46887 already cover a multi-byte character split across parts

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

  • I have added meaningful tests
  • The handful of test files covering my change pass locally, e.g. uv run pytest tests/unit/<your_test_file>.py -v. Leave the suites (make test-unit-*, make test-unit) to CI: it finishes in ~15 minutes where a laptop takes an hour or more
  • My PR passes all required CI/CD checks (e.g., lint, schema.d.ts sync check, etc.)
  • My PR's scope is as isolated as possible; it only solves 1 specific problem
  • Every automated review on my PR's latest commit has finished, and I have addressed or replied to each of its findings before requesting a maintainer review

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

Screenshots / Proof of Fix

Before is the merge base fb555b0 and After is this PR's head eb9f9e9. Each leg ran one proxy with 2 uvicorn workers and no database, from that commit's tree, with the same config and requests against a real SageMaker endpoint (ml.g5.xlarge, AWS vLLM 0.32.0 container, Qwen/Qwen3.5-0.8B) and a JumpStart TGI endpoint for the sagemaker/ route. Raw client output and both proxy consoles are in this gist

config.yaml and proxy command 
model_list:
  - model_name: qwen-sagemaker
    litellm_params:
      model: sagemaker_chat/litellm-qa-sm-vllm
      hf_model_name: Qwen/Qwen3.5-0.8B
      aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
      aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
      aws_region_name: us-west-2
  - model_name: mpt-sagemaker
    litellm_params:
      model: sagemaker/<JumpStart TGI endpoint>
      aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
      aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
      aws_region_name: us-west-2
python -m litellm.proxy.proxy_cli --config config.yaml --port $P --num_workers 2
Case Before After
/v1/chat/completions stream 200, ERROR logged 200, clean log
/v1/responses stream 200, ERROR logged 200, clean log
/v1/messages stream 200, ERROR logged 200, clean log
sagemaker/ TGI stream 200, clean log 200, clean log
30 concurrent streams 30 ERROR lines 0 ERROR lines
SDK sync stream, 5 runs 5 ERROR lines 0 ERROR lines
Whole proxy run 183 ERROR lines 0 ERROR lines

Every stream on both legs ended with the container's own usage chunk, so no answer was cut off. Before logged one ERROR per vLLM stream, 183 over 184 requests, and After logged none over 94

Single chat stream

curl -sN http://127.0.0.1:$P/v1/chat/completions -H "Authorization: Bearer $KEY" -H "Content-Type: application/json" \
  -d '{"model":"qwen-sagemaker","messages":[{"role":"user","content":"Count from 1 to 5, digits only."}],"max_tokens":24,"stream":true,"stream_options":{"include_usage":true}}'

Both legs stream the same answer 1, 2, 3, 4, 5 with finish_reason: stop, the container's usage (23 in, 14 out), and data: [DONE]. Before, the proxy console prints

INFO:     127.0.0.1:54460 - "POST /v1/chat/completions HTTP/1.1" 200 OK
21:35:26 - LiteLLM:ERROR: common_utils.py:262 - Warning: Unparseable JSON data remained: [DONE]

After, it prints

INFO:     127.0.0.1:42220 - "POST /v1/chat/completions HTTP/1.1" 200 OK
Full client output, Before (fb555b0) 
data: {"id":"chatcmpl-93b9711235c3848c","object":"chat.completion.chunk","created":1792470925,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"role":"assistant","content":""}}]}
data: {"id":"chatcmpl-93b9711235c3848c","object":"chat.completion.chunk","created":1792470925,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":"1"}}]}
data: {"id":"chatcmpl-93b9711235c3848c","object":"chat.completion.chunk","created":1792470925,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":","}}]}
data: {"id":"chatcmpl-93b9711235c3848c","object":"chat.completion.chunk","created":1792470925,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":" "}}]}
data: {"id":"chatcmpl-93b9711235c3848c","object":"chat.completion.chunk","created":1792470925,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":"2"}}]}
data: {"id":"chatcmpl-93b9711235c3848c","object":"chat.completion.chunk","created":1792470925,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":","}}]}
data: {"id":"chatcmpl-93b9711235c3848c","object":"chat.completion.chunk","created":1792470925,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":" "}}]}
data: {"id":"chatcmpl-93b9711235c3848c","object":"chat.completion.chunk","created":1792470925,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":"3"}}]}
data: {"id":"chatcmpl-93b9711235c3848c","object":"chat.completion.chunk","created":1792470925,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":","}}]}
data: {"id":"chatcmpl-93b9711235c3848c","object":"chat.completion.chunk","created":1792470925,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":" "}}]}
data: {"id":"chatcmpl-93b9711235c3848c","object":"chat.completion.chunk","created":1792470925,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":"4"}}]}
data: {"id":"chatcmpl-93b9711235c3848c","object":"chat.completion.chunk","created":1792470925,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":","}}]}
data: {"id":"chatcmpl-93b9711235c3848c","object":"chat.completion.chunk","created":1792470925,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":" "}}]}
data: {"id":"chatcmpl-93b9711235c3848c","object":"chat.completion.chunk","created":1792470925,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":"5"}}]}
data: {"id":"chatcmpl-93b9711235c3848c","object":"chat.completion.chunk","created":1792470925,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}
data: {"id":"chatcmpl-93b9711235c3848c","created":1792470925,"model":"qwen-sagemaker","object":"chat.completion.chunk","system_fingerprint":"vllm-0.32.0-8fa2052a","choices":[{"index":0,"delta":{}}],"usage":{"completion_tokens":14,"prompt_tokens":23,"total_tokens":37,"completion_tokens_details":{"reasoning_tokens":0,"text_tokens":14}}}
data: [DONE]
HTTP 200
Full client output, After (eb9f9e9) 
data: {"id":"chatcmpl-9900bca6e1b21735","object":"chat.completion.chunk","created":1792470926,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"role":"assistant","content":""}}]}
data: {"id":"chatcmpl-9900bca6e1b21735","object":"chat.completion.chunk","created":1792470926,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":"1"}}]}
data: {"id":"chatcmpl-9900bca6e1b21735","object":"chat.completion.chunk","created":1792470926,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":","}}]}
data: {"id":"chatcmpl-9900bca6e1b21735","object":"chat.completion.chunk","created":1792470926,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":" "}}]}
data: {"id":"chatcmpl-9900bca6e1b21735","object":"chat.completion.chunk","created":1792470926,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":"2"}}]}
data: {"id":"chatcmpl-9900bca6e1b21735","object":"chat.completion.chunk","created":1792470926,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":","}}]}
data: {"id":"chatcmpl-9900bca6e1b21735","object":"chat.completion.chunk","created":1792470926,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":" "}}]}
data: {"id":"chatcmpl-9900bca6e1b21735","object":"chat.completion.chunk","created":1792470926,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":"3"}}]}
data: {"id":"chatcmpl-9900bca6e1b21735","object":"chat.completion.chunk","created":1792470926,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":","}}]}
data: {"id":"chatcmpl-9900bca6e1b21735","object":"chat.completion.chunk","created":1792470926,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":" "}}]}
data: {"id":"chatcmpl-9900bca6e1b21735","object":"chat.completion.chunk","created":1792470926,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":"4"}}]}
data: {"id":"chatcmpl-9900bca6e1b21735","object":"chat.completion.chunk","created":1792470926,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":","}}]}
data: {"id":"chatcmpl-9900bca6e1b21735","object":"chat.completion.chunk","created":1792470926,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":" "}}]}
data: {"id":"chatcmpl-9900bca6e1b21735","object":"chat.completion.chunk","created":1792470926,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{"content":"5"}}]}
data: {"id":"chatcmpl-9900bca6e1b21735","object":"chat.completion.chunk","created":1792470926,"model":"qwen-sagemaker","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}
data: {"id":"chatcmpl-9900bca6e1b21735","created":1792470926,"model":"qwen-sagemaker","object":"chat.completion.chunk","system_fingerprint":"vllm-0.32.0-8fa2052a","choices":[{"index":0,"delta":{}}],"usage":{"completion_tokens":14,"prompt_tokens":23,"total_tokens":37,"completion_tokens_details":{"reasoning_tokens":0,"text_tokens":14}}}
data: [DONE]
HTTP 200

/v1/responses and /v1/messages streams

Commands and output 
curl -sN http://127.0.0.1:$P/v1/responses -H "Authorization: Bearer $KEY" -H "Content-Type: application/json" \
  -d '{"model":"qwen-sagemaker","input":"Write the numbers from 1 to 30 separated by commas.","max_output_tokens":120,"stream":true}'
curl -sN http://127.0.0.1:$P/v1/messages -H "Authorization: Bearer $KEY" -H "Content-Type: application/json" -H "anthropic-version: 2023-06-01" \
  -d '{"model":"qwen-sagemaker","max_tokens":120,"messages":[{"role":"user","content":"Write the numbers from 1 to 30 separated by commas."}],"stream":true}'

Both legs return the same text ending ..., 28, 29, 30. /v1/responses ends with response.completed carrying usage 26 in and 110 out, and /v1/messages ends with message_delta carrying stop_reason: end_turn and the same usage. Before, the proxy console prints LiteLLM:ERROR: common_utils.py:262 - Warning: Unparseable JSON data remained: [DONE] after each request. After, it prints only the access lines

Full output is in the gist as before-responses, after-responses, before-messages, and after-messages

sagemaker/ TGI stream

Command and output 
curl -sN http://127.0.0.1:$P/v1/chat/completions -H "Authorization: Bearer $KEY" -H "Content-Type: application/json" \
  -d '{"model":"mpt-sagemaker","messages":[{"role":"user","content":"Write the numbers from 1 to 30 separated by commas."}],"max_tokens":120,"stream":true,"stream_options":{"include_usage":true}}'

Both legs return the same answer with finish_reason: stop, usage 20 in and 26 out, and data: [DONE], with no ERROR line on either proxy (before, after)

30 concurrent streams

Command and output 
for i in $(seq 1 30); do
  curl -sN http://127.0.0.1:$P/v1/chat/completions -H "Authorization: Bearer $KEY" -H "Content-Type: application/json" \
    -d '{"model":"qwen-sagemaker","messages":[{"role":"user","content":"Write the numbers from 1 to 30 separated by commas."}],"max_tokens":120,"stream":true,"stream_options":{"include_usage":true}}' \
    -w '\nHTTP %{http_code}\n' > burst-$i.sse &
done; wait
for i in $(seq 1 30); do echo "$(grep -o 'HTTP [0-9]*' burst-$i.sse) $(grep -o '"prompt_tokens":[0-9]*' burst-$i.sse | tail -1)"; done | sort | uniq -c

Output on both legs

     30 HTTP 200 "prompt_tokens":26

Every stream on both legs ends with the container's own usage chunk, so none was cut off. The Before proxy console prints the [DONE] ERROR line 30 times, and the After console prints 30 access lines and nothing else (before, after)

An earlier Before run on a single-worker proxy at d56ff28, whose SageMaker stream decoder is byte-identical to fb555b0, saw 3 of 30 answers stop mid-stream (..., 23, 24, 25, 2) while still ending in finish_reason: "stop", a 200, and data: [DONE], with usage LiteLLM estimated from the cut text (prompt 20 instead of the container's 26). Its console showed the lost rest of each answer as Unparseable JSON data remained: data: {"id":"chatcmpl-...","choices":[{"index":0,"delta":{"content":"6"}}]}{"id":...}{...}. The 2-worker runs above did not hit a cut-off on either leg. The unit and integration tests in this PR replay those part boundaries and fail on main

SDK sync stream

Code and output 
import litellm

for model in ("sagemaker_chat/litellm-qa-sm-vllm", "sagemaker/<JumpStart TGI endpoint>"):
    for run in range(5):
        stream = litellm.completion(
            model=model,
            messages=[{"role": "user", "content": "Write the numbers from 1 to 30 separated by commas."}],
            max_tokens=120,
            stream=True,
            stream_options={"include_usage": True},
            aws_region_name="us-west-2",
        )
        chunks = list(stream)
        text = "".join(c.choices[0].delta.content or "" for c in chunks if c.choices)
        finish = next((c.choices[0].finish_reason for c in reversed(chunks) if c.choices and c.choices[0].finish_reason), None)
        usage = next((c.usage for c in reversed(chunks) if getattr(c, "usage", None)), None)
        print(f"{model.split('/')[0]} run {run}: finish={finish} usage={usage and (usage.prompt_tokens, usage.completion_tokens)} tail={text.strip()[-24:]!r}")

Before (fb555b0)

21:35:37 - LiteLLM:ERROR: common_utils.py:206 - Warning: Unparseable JSON data remained: [DONE]
21:35:38 - LiteLLM:ERROR: common_utils.py:206 - Warning: Unparseable JSON data remained: [DONE]
21:35:38 - LiteLLM:ERROR: common_utils.py:206 - Warning: Unparseable JSON data remained: [DONE]
21:35:38 - LiteLLM:ERROR: common_utils.py:206 - Warning: Unparseable JSON data remained: [DONE]
21:35:39 - LiteLLM:ERROR: common_utils.py:206 - Warning: Unparseable JSON data remained: [DONE]
sagemaker_chat run 0: finish=stop usage=(26, 110) tail=', 25, 26, 27, 28, 29, 30'
sagemaker_chat run 1: finish=stop usage=(26, 110) tail=', 25, 26, 27, 28, 29, 30'
sagemaker_chat run 2: finish=stop usage=(26, 3) tail='45'
sagemaker_chat run 3: finish=stop usage=(26, 110) tail=', 25, 26, 27, 28, 29, 30'
sagemaker_chat run 4: finish=stop usage=(26, 110) tail=', 25, 26, 27, 28, 29, 30'
sagemaker run 0: finish=stop usage=(20, 26) tail='rom smallest to largest.'
sagemaker run 1: finish=stop usage=(20, 26) tail='rom smallest to largest.'
sagemaker run 2: finish=stop usage=(20, 26) tail='rom smallest to largest.'
sagemaker run 3: finish=stop usage=(20, 26) tail='rom smallest to largest.'
sagemaker run 4: finish=stop usage=(20, 26) tail='rom smallest to largest.'

After (eb9f9e9)

sagemaker_chat run 0: finish=stop usage=(26, 110) tail=', 25, 26, 27, 28, 29, 30'
sagemaker_chat run 1: finish=stop usage=(26, 104) tail=', 25, 26, 27, 28, 29, 30'
sagemaker_chat run 2: finish=stop usage=(26, 110) tail=', 25, 26, 27, 28, 29, 30'
sagemaker_chat run 3: finish=stop usage=(26, 110) tail=', 25, 26, 27, 28, 29, 30'
sagemaker_chat run 4: finish=stop usage=(26, 110) tail=', 25, 26, 27, 28, 29, 30'
sagemaker run 0: finish=stop usage=(20, 26) tail='rom smallest to largest.'
sagemaker run 1: finish=stop usage=(20, 26) tail='rom smallest to largest.'
sagemaker run 2: finish=stop usage=(20, 26) tail='rom smallest to largest.'
sagemaker run 3: finish=stop usage=(20, 26) tail='rom smallest to largest.'
sagemaker run 4: finish=stop usage=(20, 26) tail='rom smallest to largest.'

The short vLLM answer in Before run 2 carries the container's own usage (3 tokens out), so it is the model's answer, not a cut-off

What surprised the run

  • Live cut-offs are rare and did not recur here
  • TGI route was already clean, and stays clean
  • /v1/messages start event has input_tokens 0, both legs, unchanged

Type

🐛 Bug Fix

Caveats (if any)

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

@mateo-berri
fix(sagemaker): parse every streamed record however SageMaker cuts th…
eb9f9e9 
…e payload parts

SageMaker cuts a container's response stream into PayloadParts at arbitrary
byte offsets, so one part can hold several SSE events or JSON lines, or only
part of one. The decoder treated each part as one record, so vLLM's [DONE]
was logged at ERROR on every stream and a split at an unlucky byte dropped
the rest of the answer.

Peel records off a running buffer with JSONDecoder.raw_decode, skipping SSE
framing and [DONE] between them.
@mateo-berri

mateo-berri commented Oct 20, 2026

Copy link
Copy Markdown
Contributor Author

@litellm-stampbot

@litellm-stampbot

litellm-stampbot commented Oct 20, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

litellm-stampbot: risk high

Not approved at eb9f9e9. risk is high, a human reviewer is needed.

Rules (stampbot/risk_rules.yml):

  • low: no rule flagged it

Cursor Bugbot: Medium Risk at eb9f9e9. Changes shared SageMaker event-stream decoding used by chat streaming; behavior is better defined and heavily tested, but incorrect framing assumptions could still affect live vLLM/TGI/LMI endpoints. Bugbot run requested, finished in 7 min.
Bugbot Low is low, Medium and High are high. The judge model was not called.

Pushes do not re-run this check. Comment @litellm-stampbot again to re-check the new head.

@litellm-stampbot

litellm-stampbot commented Oct 20, 2026

Copy link
Copy Markdown
Contributor

bugbot run

@codspeed

codspeed Bot commented Oct 20, 2026

Copy link
Copy Markdown
Contributor

Merging this PR will not alter performance

⚠️ 1 benchmark spent significant time in system calls

System calls cannot be consistently instrumented, so they are not included in the measure, which understates the real cost. Please switch to the Walltime instrument to accurately measure system calls.

Measurement and system calls

✅ 32 untouched benchmarks

Comparing litellm_sagemaker_stream_framing (eb9f9e9) with main (2d181e6)1

Open in CodSpeed

Footnotes

  1. No successful run was found on main (8884f46) during the generation of this report, so 2d181e6 was used instead as the comparison base. There might be some changes unrelated to this pull request in this report. ↩

cursor[bot]
cursor Bot reviewed Oct 20, 2026
View reviewed changes

@cursor cursor Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

✅ Bugbot reviewed your changes and found no new issues!

Comment @cursor review or bugbot run to trigger another review on this PR

Reviewed by Cursor Bugbot for commit eb9f9e9. Configure here.

@litellm-stampbot litellm-stampbot added the risk:high Sensitive, broad, or high-risk change label Oct 20, 2026
@codecov

codecov Bot commented Oct 20, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@mateo-berri

mateo-berri commented Oct 20, 2026

Copy link
Copy Markdown
Contributor Author

Superseded by #47495, which carries the same fix and tests after a smaller opening commit that stampbot could approve

@mateo-berri
mateo-berri deleted the litellm_sagemaker_stream_framing branch October 20, 2026 09:40
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@cursor cursor[bot] cursor[bot] left review comments

Assignees

No one assigned

Labels

risk:high Sensitive, broad, or high-risk change

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@mateo-berri @litellm-stampbot