diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 29e0582ec5..c79b451749 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -214,6 +214,7 @@ steps: - vllm/v1/ - tests/v1 - tests/watermarking + - examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py device: cpu-small commands: # split the test to avoid interference diff --git a/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py b/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py index 6815f42061..bd63f04627 100644 --- a/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py +++ b/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py @@ -340,7 +340,10 @@ async def _handle_completions(api: str, request: Request): ): yield chunk - return StreamingResponse(generate_stream(), media_type="application/json") + media_type = ( + "text/event-stream" if req_data.get("stream") else "application/json" + ) + return StreamingResponse(generate_stream(), media_type=media_type) except Exception as e: import sys diff --git a/tests/v1/kv_connector/unit/test_mooncake_proxy.py b/tests/v1/kv_connector/unit/test_mooncake_proxy.py new file mode 100644 index 0000000000..2d4d93a887 --- /dev/null +++ b/tests/v1/kv_connector/unit/test_mooncake_proxy.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import asyncio +import itertools +import runpy +from pathlib import Path +from unittest.mock import AsyncMock + +import httpx +import pytest + +pytestmark = pytest.mark.cpu_test + + +@pytest.mark.asyncio +@pytest.mark.parametrize("endpoint", ["/v1/completions", "/v1/chat/completions"]) +@pytest.mark.parametrize("stream", [False, True]) +async def test_response_media_type_matches_payload(monkeypatch, endpoint, stream): + """Clients must select SSE parsing for streaming completion responses.""" + path = ( + Path(__file__).parents[4] + / "examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py" + ) + proxy = runpy.run_path(str(path))["_handle_completions"].__globals__ + app = proxy["app"] + app.state.ready = asyncio.Event() + app.state.ready.set() + app.state.prefill_iterator = itertools.cycle([({}, 0)]) + app.state.decode_iterator = itertools.cycle([0]) + app.state.decode_clients = [{}] + payload = b'{"choices":[{"text":"hello"}]}' + body = b"data: " + payload + b"\n\ndata: [DONE]\n\n" if stream else payload + + async def decode_response(*args, **kwargs): + yield body + + monkeypatch.setitem(proxy, "send_request_to_service", AsyncMock()) + monkeypatch.setitem(proxy, "stream_service_response", decode_response) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://proxy" + ) as client: + response = await client.post(endpoint, json={"stream": stream}) + + expected_type = "text/event-stream" if stream else "application/json" + assert response.headers["content-type"].split(";")[0] == expected_type + assert response.content == body diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index 4ac3eaf7c5..af111483ef 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -805,6 +805,23 @@ class TestMultiConnectorStats: child_metrics.observe.assert_called_once_with({"transfer_duration": [1.5]}, 2) + @pytest.mark.parametrize("with_exporter", [False, True]) + def test_prom_metrics_skips_stats_only_connectors(self, with_exporter): + """Stats-only children must not prevent supported metrics from exporting.""" + child_metrics = MagicMock() + metrics = object.__new__(MultiKVConnectorPromMetrics) + metrics._prom_metrics = ( + {"MooncakeStoreConnector": child_metrics} if with_exporter else {} + ) + payload = {"MooncakeConnector": {"bytes_transferred": [1024]}} + if with_exporter: + payload["MooncakeStoreConnector"] = {"load_get_bytes": [2048]} + + metrics.observe(payload, engine_idx=2) + + if with_exporter: + child_metrics.observe.assert_called_once_with({"load_get_bytes": [2048]}, 2) + def test_aggregate_same_connector(self): """Test aggregating stats from the same connector type.""" stats1 = MultiKVConnectorStats( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index bc59f3cb27..70f9454a62 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py @@ -125,11 +125,9 @@ class MultiKVConnectorPromMetrics(KVConnectorPromMetrics): def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): for connector_id, stats_data in transfer_stats_data.items(): - assert connector_id in self._prom_metrics, ( - f"{connector_id} is not contained in the list of registered connectors " - f"with Prometheus metrics support: {self._prom_metrics.keys()}" - ) - self._prom_metrics[connector_id].observe(stats_data, engine_idx) + prom_metrics = self._prom_metrics.get(connector_id) + if prom_metrics is not None: + prom_metrics.observe(stats_data, engine_idx) class MultiConnector(KVConnectorBase_V1, SupportsHMA):