fix(sglang): treat streamed usage metadata as optional (#11062)

Signed-off-by: Kirill Shkolenko <shkolenko.kirill@gmail.com>
Co-authored-by: Kirill Shkolenko <shkolenko.kirill@gmail.com>
Co-authored-by: Ryan McCormick <rmccormick@nvidia.com>

Author: 3 people

components/src/dynamo/sglang/request_handlers/llm/decode_handler.py

@@ -580,18 +580,23 @@ async def _process_token_stream(
                     # as nvext.routed_experts); disaggregated_params stays KV-transfer only.
                     out["engine_data"] = {"routed_experts": routed_experts}
                 if finish_reason:
-                    input_tokens = meta_info["prompt_tokens"]
-                    completion_tokens = meta_info["completion_tokens"]
-                    cached_tokens = meta_info["cached_tokens"]
+                    input_tokens = meta_info.get("prompt_tokens")
+                    completion_tokens = meta_info.get("completion_tokens")
+                    cached_tokens = meta_info.get("cached_tokens")
                     prefill_prompt_tokens_details = None
                     if cached_tokens is not None and cached_tokens > 0:
                         prefill_prompt_tokens_details = {"cached_tokens": cached_tokens}
-                    out["completion_usage"] = {
-                        "prompt_tokens": input_tokens,
-                        "completion_tokens": completion_tokens,
-                        "total_tokens": input_tokens + completion_tokens,
-                        "prompt_tokens_details": prefill_prompt_tokens_details,
-                    }
+                    if input_tokens is not None and completion_tokens is not None:
+                        completion_usage = {
+                            "prompt_tokens": input_tokens,
+                            "completion_tokens": completion_tokens,
+                            "total_tokens": input_tokens + completion_tokens,
+                        }
+                        if prefill_prompt_tokens_details is not None:
+                            completion_usage[
+                                "prompt_tokens_details"
+                            ] = prefill_prompt_tokens_details
+                        out["completion_usage"] = completion_usage
                     if metadata_uploader is not None:
                         try:
                             await metadata_uploader.upload_choice(output_idx, meta_info)

components/src/dynamo/sglang/tests/test_sglang_decode_handler.py

@@ -478,6 +478,53 @@ async def test_metadata_upload_normalizes_numpy_values(tmp_path):
     assert uploaded_array["data"] == expected.tobytes()
 
 
+@pytest.mark.asyncio
+async def test_process_token_stream_treats_completion_usage_as_optional():
+    handler = _new_decode_handler()
+
+    chunks = await _collect(
+        handler._process_token_stream(
+            _stream(
+                [
+                    {
+                        "index": 0,
+                        "output_ids": [],
+                        "meta_info": {
+                            "id": "request-1",
+                            "finish_reason": {"type": "stop"},
+                        },
+                    },
+                    {
+                        "index": 1,
+                        "output_ids": [],
+                        "meta_info": {
+                            "id": "request-1",
+                            "finish_reason": {"type": "stop"},
+                            "prompt_tokens": 2,
+                            "completion_tokens": 3,
+                        },
+                    },
+                ]
+            ),
+            _Context(),
+        )
+    )
+
+    assert chunks == [
+        {"index": 0, "finish_reason": "stop", "token_ids": []},
+        {
+            "index": 1,
+            "finish_reason": "stop",
+            "token_ids": [],
+            "completion_usage": {
+                "prompt_tokens": 2,
+                "completion_tokens": 3,
+                "total_tokens": 5,
+            },
+        },
+    ]
+
+
 @pytest.mark.asyncio
 async def test_process_token_stream_tracks_logprobs_per_choice_index():
     handler = _new_decode_handler()