From bbf989e89ba70334aec148668a547df9fd518c14 Mon Sep 17 00:00:00 2001 From: Calin Teodor <45018530+TeoSlayer@users.noreply.github.com> Date: Wed, 28 Oct 2026 09:05:17 -0700 Subject: [PATCH 1/2] fix: decode SGLang streaming responses --- .../endpoint_client/adapter_protocol.py | 9 ++-- src/inference_endpoint/sglang/adapter.py | 2 +- src/inference_endpoint/sglang/types.py | 6 +-- .../performance/sglang/test_sglang_adapter.py | 2 +- .../endpoint_client/test_adapter_protocol.py | 47 ++++++++++++++++--- 5 files changed, 49 insertions(+), 17 deletions(-) diff --git a/src/inference_endpoint/endpoint_client/adapter_protocol.py b/src/inference_endpoint/endpoint_client/adapter_protocol.py index a96faaa0a..1f58c59a3 100644 --- a/src/inference_endpoint/endpoint_client/adapter_protocol.py +++ b/src/inference_endpoint/endpoint_client/adapter_protocol.py @@ -127,12 +127,11 @@ def parse_sse_chunk(cls, buffer: bytes, end_pos: int) -> list[Any]: """ json_docs = cls.SSE_DATA_PATTERN.findall(buffer[:end_pos]) parsed: list[Any] = [] - # Note: if one frame is malformed, remaining frames are skipped - try: - for json_doc in json_docs: + for json_doc in json_docs: + try: content = cls.decode_sse_message(json_doc) if content is not None: parsed.append(content) - except (msgspec.DecodeError, msgspec.ValidationError) as exc: - logger.warning("skipping malformed SSE batch (%s)", type(exc).__name__) + except (msgspec.DecodeError, msgspec.ValidationError) as exc: + logger.warning("skipping malformed SSE frame (%s)", type(exc).__name__) return parsed diff --git a/src/inference_endpoint/sglang/adapter.py b/src/inference_endpoint/sglang/adapter.py index a6614fa35..c19e3d2a4 100644 --- a/src/inference_endpoint/sglang/adapter.py +++ b/src/inference_endpoint/sglang/adapter.py @@ -122,7 +122,7 @@ def decode_sse_message(cls, json_bytes: bytes) -> SGLangSSEDelta: # the output_tokens is the delta from the previous chunk token_delta = resp.output_ids - has_retractions = resp.meta_info.total_retractions > 0 + has_retractions = resp.meta_info.num_retractions > 0 return SGLangSSEDelta( text=total_text, token_delta=token_delta, diff --git a/src/inference_endpoint/sglang/types.py b/src/inference_endpoint/sglang/types.py index 9e75b3a4d..89f71eed8 100644 --- a/src/inference_endpoint/sglang/types.py +++ b/src/inference_endpoint/sglang/types.py @@ -52,13 +52,13 @@ class SGLangGenerateRequest( class MetaInfo(msgspec.Struct, frozen=True, kw_only=True, omit_defaults=True, gc=False): # type: ignore[call-arg] id: str - finish_reason: dict[str, Any] + finish_reason: dict[str, Any] | None prompt_tokens: int weight_version: str - total_retractions: int + num_retractions: int completion_tokens: int cached_tokens: int - e2e_latency: float + e2e_latency: float | None = None class SGLangGenerateResponse( diff --git a/tests/performance/sglang/test_sglang_adapter.py b/tests/performance/sglang/test_sglang_adapter.py index f614cc0d7..3e0fa90e8 100644 --- a/tests/performance/sglang/test_sglang_adapter.py +++ b/tests/performance/sglang/test_sglang_adapter.py @@ -64,7 +64,7 @@ def make_response_bytes(text: str, n_tokens: int) -> bytes: "finish_reason": {"type": "stop"}, "prompt_tokens": 10, "weight_version": "v1", - "total_retractions": 0, + "num_retractions": 0, "completion_tokens": n_tokens, "cached_tokens": 0, "e2e_latency": 0.1, diff --git a/tests/unit/endpoint_client/test_adapter_protocol.py b/tests/unit/endpoint_client/test_adapter_protocol.py index 591dc4236..b348b93ff 100644 --- a/tests/unit/endpoint_client/test_adapter_protocol.py +++ b/tests/unit/endpoint_client/test_adapter_protocol.py @@ -22,10 +22,11 @@ from inference_endpoint.config.schema import ModelParams from inference_endpoint.dataset_manager.transforms import Transform from inference_endpoint.endpoint_client.adapter_protocol import HttpRequestAdapter +from inference_endpoint.sglang.adapter import SGLangGenerateAdapter class _SimpleAdapter(HttpRequestAdapter): - """Minimal adapter: raises DecodeError for b'BAD', returns {"x": 1} otherwise.""" + """Decode JSON objects with integer values.""" @classmethod def dataset_transforms(cls, model_params: ModelParams) -> list[Transform]: @@ -41,15 +42,47 @@ def decode_response(cls, response_bytes: bytes, query_id: str) -> Any: @classmethod def decode_sse_message(cls, json_bytes: bytes) -> Any: - if json_bytes.strip() == b"BAD": - raise msgspec.DecodeError("bad frame") - return {"x": 1} + return msgspec.json.decode(json_bytes, type=dict[str, int]) @pytest.mark.unit -def test_parse_sse_chunk_skips_bad_frame_and_keeps_valid(): +@pytest.mark.parametrize("bad_frame", [b'{"ok":}', b'{"ok":"bad"}']) +def test_parse_sse_chunk_skips_bad_frame_and_keeps_valid(bad_frame): """A malformed SSE frame is skipped; surrounding valid frames are preserved.""" - buffer = b'data: {"ok":1}\n\ndata: BAD\n\ndata: {"ok":1}\n\n' + buffer = b'data: {"ok":1}\n\ndata: ' + bad_frame + b'\n\ndata: {"ok":2}\n\n' end_pos = len(buffer) result = _SimpleAdapter.parse_sse_chunk(buffer, end_pos) - assert result == [{"x": 1}, {"x": 1}] + assert result == [{"ok": 1}, {"ok": 2}] + + +@pytest.mark.unit +@pytest.mark.parametrize("finish_reason", [None, {"type": "stop"}]) +@pytest.mark.parametrize("latency", [{}, {"e2e_latency": None}, {"e2e_latency": 0.1}]) +@pytest.mark.parametrize("retractions", [0, 2]) +def test_sglang_response_metadata(finish_reason, latency, retractions): + response = msgspec.json.encode( + { + "text": "Hello", + "output_ids": [10], + "meta_info": { + "id": "sample", + "finish_reason": finish_reason, + "prompt_tokens": 3, + "weight_version": "v1", + "num_retractions": retractions, + "completion_tokens": 1, + "cached_tokens": 0, + **latency, + }, + } + ) + result = SGLangGenerateAdapter.decode_response(response, "sample") + assert result.get_response_output_string() == "Hello" + assert result.metadata == {"token_ids": [10], "n_tokens": 1} + frame = b"data: " + response + b"\n\n" + chunks = SGLangGenerateAdapter.parse_sse_chunk(frame, len(frame)) + assert len(chunks) == 1 + assert chunks[0].text == "Hello" + assert chunks[0].token_delta == [10] + assert chunks[0].total_completion_tokens == 1 + assert chunks[0].has_retractions is (retractions > 0) From a8a8d0aa2c9760282bef893f5399bdc4fd1cac2e Mon Sep 17 00:00:00 2001 From: Calin Teodor <45018530+TeoSlayer@users.noreply.github.com> Date: Thu, 5 Nov 2026 09:38:25 -0800 Subject: [PATCH 2/2] test: wait for benchmark startup before wedging metrics --- tests/integration/commands/test_sigint.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/integration/commands/test_sigint.py b/tests/integration/commands/test_sigint.py index 1197744dc..d6986981e 100644 --- a/tests/integration/commands/test_sigint.py +++ b/tests/integration/commands/test_sigint.py @@ -276,6 +276,7 @@ def test_sigint_grace_expiry_abandons_wedged_drain(mock_http_echo_server, tmp_pa """ report_dir = tmp_path / "report" config_path = tmp_path / "bench.yaml" + log_file = tmp_path / "run.log" _write_config( report_dir, mock_http_echo_server.url, @@ -286,10 +287,18 @@ def test_sigint_grace_expiry_abandons_wedged_drain(mock_http_echo_server, tmp_pa agg_pid: int | None = None try: with _benchmark_proc( - [_cli(), "benchmark", "from-config", "-c", str(config_path)] + [_cli(), "benchmark", "from-config", "-c", str(config_path)], + log_file, ) as proc: _wait_services_ready(proc, report_dir) - time.sleep(3.0) # comfortably inside the ~120 s performance phase + # Service readiness precedes client startup. Wait for the session + # before allowing live metrics ticks to reach the parent. + deadline = time.monotonic() + 60.0 + while "Starting phase:" not in log_file.read_text(errors="replace"): + assert proc.poll() is None, log_file.read_text(errors="replace") + assert time.monotonic() < deadline, "performance phase never started" + time.sleep(0.05) + time.sleep(3.0) agg_pid = _pid_of_child("metrics_aggregator", str(report_dir)) assert agg_pid is not None, "aggregator child not found"