[Bench] Fix bench_serving missing reasoning_content stream chunks (#23954)
This commit is contained in:
@@ -462,10 +462,22 @@ async def async_request_openai_chat_completions(
|
|||||||
pass
|
pass
|
||||||
else:
|
else:
|
||||||
data = json.loads(chunk)
|
data = json.loads(chunk)
|
||||||
|
# Check for usage info in final chunks. OpenAI-compatible
|
||||||
|
# servers may emit usage-only chunks with choices=[].
|
||||||
|
output_len = (data.get("usage") or {}).get(
|
||||||
|
"completion_tokens", output_len
|
||||||
|
)
|
||||||
|
|
||||||
# Check if this chunk contains content
|
choices = data.get("choices") or []
|
||||||
delta = data.get("choices", [{}])[0].get("delta", {})
|
if not choices:
|
||||||
content = delta.get("content", "")
|
continue
|
||||||
|
|
||||||
|
# Reasoning models stream thoughts via
|
||||||
|
# `reasoning_content`; count them like content.
|
||||||
|
delta = choices[0].get("delta") or {}
|
||||||
|
content = (delta.get("reasoning_content") or "") + (
|
||||||
|
delta.get("content") or ""
|
||||||
|
)
|
||||||
|
|
||||||
if content:
|
if content:
|
||||||
timestamp = time.perf_counter()
|
timestamp = time.perf_counter()
|
||||||
@@ -484,11 +496,6 @@ async def async_request_openai_chat_completions(
|
|||||||
most_recent_timestamp = timestamp
|
most_recent_timestamp = timestamp
|
||||||
generated_text += content
|
generated_text += content
|
||||||
|
|
||||||
# Check for usage info in final chunk
|
|
||||||
output_len = (data.get("usage") or {}).get(
|
|
||||||
"completion_tokens", output_len
|
|
||||||
)
|
|
||||||
|
|
||||||
output.generated_text = generated_text
|
output.generated_text = generated_text
|
||||||
output.success = True
|
output.success = True
|
||||||
output.latency = latency
|
output.latency = latency
|
||||||
|
|||||||
@@ -0,0 +1,210 @@
|
|||||||
|
"""Unit tests for bench_serving streaming with reasoning_content chunks.
|
||||||
|
|
||||||
|
Reasoning models (DeepSeek-R1, MiMo, Qwen3 reasoning, Kimi-K2, ...) stream their
|
||||||
|
chain-of-thought via OpenAI's `delta.reasoning_content` field. Without explicit
|
||||||
|
support, bench_serving only inspects `delta.content` and silently reports zero
|
||||||
|
TTFT / ITL and an empty `generated_text`, which then retokenizes to 0 tokens
|
||||||
|
even though the backend completed real work.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
import socket
|
||||||
|
import threading
|
||||||
|
import time
|
||||||
|
import unittest
|
||||||
|
from argparse import Namespace
|
||||||
|
from http.server import BaseHTTPRequestHandler, HTTPServer
|
||||||
|
|
||||||
|
from sglang.bench_serving import (
|
||||||
|
RequestFuncInput,
|
||||||
|
async_request_openai_chat_completions,
|
||||||
|
set_global_args,
|
||||||
|
)
|
||||||
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
|
register_cpu_ci(est_time=10, suite="stage-a-test-cpu")
|
||||||
|
|
||||||
|
|
||||||
|
def _free_port() -> int:
|
||||||
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||||
|
s.bind(("127.0.0.1", 0))
|
||||||
|
return s.getsockname()[1]
|
||||||
|
|
||||||
|
|
||||||
|
class _SSEHandler(BaseHTTPRequestHandler):
|
||||||
|
"""Streams a fixed sequence of OpenAI-compatible SSE chunks per test."""
|
||||||
|
|
||||||
|
chunks: list = []
|
||||||
|
chunk_delay_s: float = 0.02
|
||||||
|
|
||||||
|
def do_POST(self): # noqa: N802 (BaseHTTPRequestHandler interface)
|
||||||
|
length = int(self.headers.get("Content-Length", "0"))
|
||||||
|
if length:
|
||||||
|
self.rfile.read(length)
|
||||||
|
self.send_response(200)
|
||||||
|
self.send_header("Content-Type", "text/event-stream")
|
||||||
|
self.send_header("Cache-Control", "no-cache")
|
||||||
|
self.end_headers()
|
||||||
|
for chunk in self.chunks:
|
||||||
|
self.wfile.write(b"data: " + json.dumps(chunk).encode() + b"\n\n")
|
||||||
|
self.wfile.flush()
|
||||||
|
time.sleep(self.chunk_delay_s)
|
||||||
|
self.wfile.write(b"data: [DONE]\n\n")
|
||||||
|
self.wfile.flush()
|
||||||
|
|
||||||
|
def log_message(self, fmt, *args): # silence access logs
|
||||||
|
return
|
||||||
|
|
||||||
|
|
||||||
|
def _make_chunk(content=None, reasoning_content=None, completion_tokens=None):
|
||||||
|
delta = {}
|
||||||
|
if content is not None:
|
||||||
|
delta["content"] = content
|
||||||
|
if reasoning_content is not None:
|
||||||
|
delta["reasoning_content"] = reasoning_content
|
||||||
|
chunk = {"choices": [{"index": 0, "delta": delta}]}
|
||||||
|
if completion_tokens is not None:
|
||||||
|
chunk["usage"] = {"completion_tokens": completion_tokens}
|
||||||
|
return chunk
|
||||||
|
|
||||||
|
|
||||||
|
class TestBenchServingReasoningStream(CustomTestCase):
|
||||||
|
@classmethod
|
||||||
|
def setUpClass(cls):
|
||||||
|
set_global_args(
|
||||||
|
Namespace(
|
||||||
|
disable_stream=False,
|
||||||
|
disable_ignore_eos=True,
|
||||||
|
print_requests=False,
|
||||||
|
tokenizer="",
|
||||||
|
header=None,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
def _run(self, chunks):
|
||||||
|
port = _free_port()
|
||||||
|
|
||||||
|
class Handler(_SSEHandler):
|
||||||
|
pass
|
||||||
|
|
||||||
|
Handler.chunks = list(chunks)
|
||||||
|
server = HTTPServer(("127.0.0.1", port), Handler)
|
||||||
|
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
||||||
|
thread.start()
|
||||||
|
try:
|
||||||
|
req = RequestFuncInput(
|
||||||
|
prompt="hello",
|
||||||
|
api_url=f"http://127.0.0.1:{port}/v1/chat/completions",
|
||||||
|
prompt_len=1,
|
||||||
|
output_len=64,
|
||||||
|
model="dummy-model",
|
||||||
|
lora_name="",
|
||||||
|
image_data=None,
|
||||||
|
extra_request_body={},
|
||||||
|
)
|
||||||
|
return asyncio.run(async_request_openai_chat_completions(req))
|
||||||
|
finally:
|
||||||
|
server.shutdown()
|
||||||
|
server.server_close()
|
||||||
|
|
||||||
|
def test_reasoning_only_stream_populates_metrics(self):
|
||||||
|
chunks = [
|
||||||
|
_make_chunk(reasoning_content="Let "),
|
||||||
|
_make_chunk(reasoning_content="me "),
|
||||||
|
_make_chunk(reasoning_content="think."),
|
||||||
|
_make_chunk(completion_tokens=3),
|
||||||
|
]
|
||||||
|
out = self._run(chunks)
|
||||||
|
|
||||||
|
self.assertTrue(out.success, msg=f"request failed: {out.error}")
|
||||||
|
self.assertEqual(out.generated_text, "Let me think.")
|
||||||
|
self.assertGreater(out.ttft, 0.0)
|
||||||
|
self.assertEqual(len(out.itl), 2, msg="should record ITL for chunks 2..N")
|
||||||
|
for v in out.itl:
|
||||||
|
self.assertGreater(v, 0.0)
|
||||||
|
self.assertEqual(out.text_chunks, ["me ", "think."])
|
||||||
|
self.assertEqual(out.output_len, 3)
|
||||||
|
|
||||||
|
def test_reasoning_then_content_accounts_both(self):
|
||||||
|
chunks = [
|
||||||
|
_make_chunk(reasoning_content="step1 "),
|
||||||
|
_make_chunk(reasoning_content="step2 "),
|
||||||
|
_make_chunk(content="answer "),
|
||||||
|
_make_chunk(content="here"),
|
||||||
|
_make_chunk(completion_tokens=4),
|
||||||
|
]
|
||||||
|
out = self._run(chunks)
|
||||||
|
|
||||||
|
self.assertTrue(out.success, msg=f"request failed: {out.error}")
|
||||||
|
self.assertEqual(out.generated_text, "step1 step2 answer here")
|
||||||
|
self.assertGreater(out.ttft, 0.0)
|
||||||
|
self.assertEqual(len(out.itl), 3)
|
||||||
|
self.assertEqual(out.text_chunks, ["step2 ", "answer ", "here"])
|
||||||
|
self.assertEqual(out.output_len, 4)
|
||||||
|
|
||||||
|
def test_single_delta_preserves_reasoning_before_content(self):
|
||||||
|
chunks = [
|
||||||
|
_make_chunk(content="answer", reasoning_content="thought "),
|
||||||
|
_make_chunk(completion_tokens=2),
|
||||||
|
]
|
||||||
|
out = self._run(chunks)
|
||||||
|
|
||||||
|
self.assertTrue(out.success, msg=f"request failed: {out.error}")
|
||||||
|
self.assertEqual(out.generated_text, "thought answer")
|
||||||
|
self.assertGreater(out.ttft, 0.0)
|
||||||
|
self.assertEqual(out.output_len, 2)
|
||||||
|
|
||||||
|
def test_usage_only_stream_chunk_does_not_break(self):
|
||||||
|
chunks = [
|
||||||
|
_make_chunk(reasoning_content="thinking"),
|
||||||
|
{"choices": [], "usage": {"completion_tokens": 1}},
|
||||||
|
]
|
||||||
|
out = self._run(chunks)
|
||||||
|
|
||||||
|
self.assertTrue(out.success, msg=f"request failed: {out.error}")
|
||||||
|
self.assertEqual(out.generated_text, "thinking")
|
||||||
|
self.assertGreater(out.ttft, 0.0)
|
||||||
|
self.assertEqual(out.output_len, 1)
|
||||||
|
|
||||||
|
def test_content_only_stream_unchanged(self):
|
||||||
|
chunks = [
|
||||||
|
_make_chunk(content="hi "),
|
||||||
|
_make_chunk(content="there"),
|
||||||
|
_make_chunk(completion_tokens=2),
|
||||||
|
]
|
||||||
|
out = self._run(chunks)
|
||||||
|
|
||||||
|
self.assertTrue(out.success, msg=f"request failed: {out.error}")
|
||||||
|
self.assertEqual(out.generated_text, "hi there")
|
||||||
|
self.assertGreater(out.ttft, 0.0)
|
||||||
|
self.assertEqual(len(out.itl), 1)
|
||||||
|
self.assertEqual(out.text_chunks, ["there"])
|
||||||
|
self.assertEqual(out.output_len, 2)
|
||||||
|
|
||||||
|
def test_null_reasoning_field_does_not_break(self):
|
||||||
|
# Mirrors sglang's _StreamDelta: reasoning_content is always emitted,
|
||||||
|
# serialized as null when only content is present.
|
||||||
|
chunks = [
|
||||||
|
{
|
||||||
|
"choices": [
|
||||||
|
{"index": 0, "delta": {"content": "ok", "reasoning_content": None}}
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"choices": [
|
||||||
|
{"index": 0, "delta": {"content": None, "reasoning_content": None}}
|
||||||
|
],
|
||||||
|
"usage": {"completion_tokens": 1},
|
||||||
|
},
|
||||||
|
]
|
||||||
|
out = self._run(chunks)
|
||||||
|
|
||||||
|
self.assertTrue(out.success, msg=f"request failed: {out.error}")
|
||||||
|
self.assertEqual(out.generated_text, "ok")
|
||||||
|
self.assertGreater(out.ttft, 0.0)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
Reference in New Issue
Block a user