From c0250ba6c4e89746a2eb62395857d1bb0ac5b612 Mon Sep 17 00:00:00 2001 From: AyushKashyapII Date: Fri, 25 Sep 2026 11:40:05 +0530 Subject: [PATCH 1/2] [Router][Feat] Add /reset_prefix_cache router passthrough Signed-off-by: AyushKashyapII --- src/tests/test_reset_prefix_cache.py | 138 ++++++++++++++++++ src/vllm_router/routers/main_router.py | 7 + .../services/request_service/request.py | 14 +- 3 files changed, 158 insertions(+), 1 deletion(-) create mode 100644 src/tests/test_reset_prefix_cache.py diff --git a/src/tests/test_reset_prefix_cache.py b/src/tests/test_reset_prefix_cache.py new file mode 100644 index 000000000..0ec07510b --- /dev/null +++ b/src/tests/test_reset_prefix_cache.py @@ -0,0 +1,138 @@ +"""Unit tests for route_sleep_wakeup_request's /reset_prefix_cache support. + +route_sleep_wakeup_request already backs /sleep, /wake_up, /is_sleeping. +These tests cover two things: the new /reset_prefix_cache endpoint relays +vLLM's real response body (unlike /sleep and /wake_up, whose upstream body +is always empty and safe to ignore), and that existing /sleep behavior is +provably unchanged by that change. +""" + +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from vllm_router.services.request_service.request import route_sleep_wakeup_request +from vllm_router.utils import SingletonABCMeta + + +class EndpointInfo: + def __init__(self, id_, url, pod_name): + self.Id = id_ + self.url = url + self.pod_name = pod_name + + +@pytest.fixture(autouse=True) +def cleanup_singletons(): + yield + for cls in list(SingletonABCMeta._instances.keys()): + del SingletonABCMeta._instances[cls] + + +def _build_request(query_params, body=b""): + request = MagicMock() + request.headers = {} + request.query_params = query_params + request.body = AsyncMock(return_value=body) + return request + + +def _mock_aiohttp_session(response_status, response_json=None): + """Build a mock aiohttp.ClientSession usable both as + `async with aiohttp.ClientSession() as client:` and, on the returned + client, `async with client.post(...) as response:`. + """ + mock_response = MagicMock() + mock_response.status = response_status + mock_response.raise_for_status = MagicMock() + mock_response.json = AsyncMock(return_value=response_json) + + mock_post_cm = MagicMock() + mock_post_cm.__aenter__ = AsyncMock(return_value=mock_response) + mock_post_cm.__aexit__ = AsyncMock(return_value=False) + + mock_client = MagicMock() + mock_client.post = MagicMock(return_value=mock_post_cm) + + mock_session_cm = MagicMock() + mock_session_cm.__aenter__ = AsyncMock(return_value=mock_client) + mock_session_cm.__aexit__ = AsyncMock(return_value=False) + + return mock_session_cm, mock_client, mock_response + + +@pytest.mark.asyncio +async def test_reset_prefix_cache_relays_real_response_body(): + """/reset_prefix_cache's real vLLM response body (`{"success": bool}`) + must be relayed to the caller, since `success: false` is a legitimate, + meaningful outcome (blocks still held) -- not swallowed into the old + canned `{"status": "success"}`. + """ + endpoint_info = EndpointInfo("X", "http://engine-x:8000", "pod-x") + service_discovery = MagicMock() + service_discovery.get_endpoint_info.return_value = [endpoint_info] + + mock_session_cm, mock_client, mock_response = _mock_aiohttp_session( + response_status=200, response_json={"success": False} + ) + + request = _build_request({"id": "X", "reset_external": "true"}) + + with ( + patch( + "vllm_router.services.request_service.request.get_service_discovery", + return_value=service_discovery, + ), + patch( + "vllm_router.services.request_service.request.aiohttp.ClientSession", + return_value=mock_session_cm, + ), + ): + response = await route_sleep_wakeup_request( + request, "/reset_prefix_cache", MagicMock() + ) + + # The real vLLM body was relayed, not the canned status. + assert json.loads(response.body) == {"success": False} + + # Correct upstream URL and forwarded params (id excluded). + call_args = mock_client.post.call_args + assert call_args.args[0] == "http://engine-x:8000/reset_prefix_cache" + assert call_args.kwargs["params"] == {"reset_external": "true"} + mock_response.json.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_sleep_still_returns_canned_status_unaffected(): + """Regression: /sleep's existing behavior (canned {"status": "success"}, + ignoring the empty upstream body) must be unchanged by the + /reset_prefix_cache relay logic added alongside it. + """ + endpoint_info = EndpointInfo("X", "http://engine-x:8000", "pod-x") + service_discovery = MagicMock() + service_discovery.get_endpoint_info.return_value = [endpoint_info] + service_discovery.add_sleep_label = MagicMock() + + mock_session_cm, mock_client, mock_response = _mock_aiohttp_session( + response_status=200, response_json=None + ) + + request = _build_request({"id": "X", "level": "2"}) + + with ( + patch( + "vllm_router.services.request_service.request.get_service_discovery", + return_value=service_discovery, + ), + patch( + "vllm_router.services.request_service.request.aiohttp.ClientSession", + return_value=mock_session_cm, + ), + ): + response = await route_sleep_wakeup_request(request, "/sleep", MagicMock()) + + assert json.loads(response.body) == {"status": "success"} + # The relay logic must never even attempt to parse /sleep's body. + mock_response.json.assert_not_called() + service_discovery.add_sleep_label.assert_called_once_with("pod-x") diff --git a/src/vllm_router/routers/main_router.py b/src/vllm_router/routers/main_router.py index 56bf83790..a8ad40ac1 100644 --- a/src/vllm_router/routers/main_router.py +++ b/src/vllm_router/routers/main_router.py @@ -120,6 +120,13 @@ async def route_is_sleeping(request: Request, background_tasks: BackgroundTasks) return await route_sleep_wakeup_request(request, "/is_sleeping", background_tasks) +@main_router.post("/reset_prefix_cache") +async def route_reset_prefix_cache(request: Request, background_tasks: BackgroundTasks): + return await route_sleep_wakeup_request( + request, "/reset_prefix_cache", background_tasks + ) + + @main_router.get("/version") async def show_version(): ver = {"version": __version__} diff --git a/src/vllm_router/services/request_service/request.py b/src/vllm_router/services/request_service/request.py index 44b8cb498..25c0a5923 100644 --- a/src/vllm_router/services/request_service/request.py +++ b/src/vllm_router/services/request_service/request.py @@ -1114,6 +1114,7 @@ async def route_sleep_wakeup_request( else: request_body = await request.body() response_status = None + response_content = None if request_body: req_data = json.loads(request_body) async with client.post( @@ -1127,6 +1128,8 @@ async def route_sleep_wakeup_request( ) as response: response.raise_for_status() response_status = response.status + if endpoint == "/reset_prefix_cache": + response_content = await response.json() pod_name = endpoints[0].pod_name if endpoint == "/sleep": @@ -1134,9 +1137,18 @@ async def route_sleep_wakeup_request( elif endpoint == "/wake_up": service_discovery.remove_sleep_label(pod_name) + # /reset_prefix_cache's response body is meaningful (success can + # legitimately be false while blocks are still held, e.g. by + # in-flight requests) and must be relayed to the caller, unlike + # /sleep and /wake_up whose upstream response body is always + # empty. return JSONResponse( status_code=response_status, - content={"status": "success"}, + content=( + response_content + if response_content is not None + else {"status": "success"} + ), headers={"X-Request-Id": request_id}, ) From 5a90126bc9e98d8874126cb32da18f8d0ca1affc Mon Sep 17 00:00:00 2001 From: AyushKashyapII Date: Fri, 25 Sep 2026 22:22:15 +0530 Subject: [PATCH 2/2] adding tests and considering changes between documented and what we can observe at runtime Signed-off-by: AyushKashyapII --- src/tests/test_reset_prefix_cache.py | 12 ++++++++-- .../services/request_service/request.py | 23 ++++++++----------- 2 files changed, 19 insertions(+), 16 deletions(-) diff --git a/src/tests/test_reset_prefix_cache.py b/src/tests/test_reset_prefix_cache.py index 0ec07510b..8cc2e7781 100644 --- a/src/tests/test_reset_prefix_cache.py +++ b/src/tests/test_reset_prefix_cache.py @@ -63,11 +63,17 @@ def _mock_aiohttp_session(response_status, response_json=None): @pytest.mark.asyncio -async def test_reset_prefix_cache_relays_real_response_body(): +@pytest.mark.parametrize("request_body", [b"", b"{}"]) +async def test_reset_prefix_cache_relays_real_response_body(request_body): """/reset_prefix_cache's real vLLM response body (`{"success": bool}`) must be relayed to the caller, since `success: false` is a legitimate, meaningful outcome (blocks still held) -- not swallowed into the old canned `{"status": "success"}`. + + Parameterized over an empty and a non-empty request body: a real + /reset_prefix_cache caller never sends a body, but nothing stops one + from doing so (e.g. a client that always attaches `{}`), and this + endpoint's relay logic must not depend on which branch that puts it in. """ endpoint_info = EndpointInfo("X", "http://engine-x:8000", "pod-x") service_discovery = MagicMock() @@ -77,7 +83,9 @@ async def test_reset_prefix_cache_relays_real_response_body(): response_status=200, response_json={"success": False} ) - request = _build_request({"id": "X", "reset_external": "true"}) + request = _build_request( + {"id": "X", "reset_external": "true"}, body=request_body + ) with ( patch( diff --git a/src/vllm_router/services/request_service/request.py b/src/vllm_router/services/request_service/request.py index 2f1749f26..aaa01540d 100644 --- a/src/vllm_router/services/request_service/request.py +++ b/src/vllm_router/services/request_service/request.py @@ -1121,21 +1121,16 @@ async def route_sleep_wakeup_request( request_body = await request.body() response_status = None response_content = None + + post_kwargs = {"headers": headers, "params": upstream_params} if request_body: - req_data = json.loads(request_body) - async with client.post( - url, json=req_data, headers=headers, params=upstream_params - ) as response: - response.raise_for_status() - response_status = response.status - else: - async with client.post( - url, headers=headers, params=upstream_params - ) as response: - response.raise_for_status() - response_status = response.status - if endpoint == "/reset_prefix_cache": - response_content = await response.json() + post_kwargs["json"] = json.loads(request_body) + + async with client.post(url, **post_kwargs) as response: + response.raise_for_status() + response_status = response.status + if endpoint == "/reset_prefix_cache": + response_content = await response.json() pod_name = endpoints[0].pod_name if endpoint == "/sleep":