From 520a20ba4e2d2a15e58f34751f668bd8d046bc45 Mon Sep 17 00:00:00 2001 From: chaeminlim-mb Date: Wed, 15 Jul 2026 07:56:46 +0900 Subject: [PATCH] [Bugfix] MoRIIO toy P/D proxy: add /health (#45222) Signed-off-by: Chaemin Lim Signed-off-by: Edwin Lim Co-authored-by: Edwin Lim Co-authored-by: Jaeyoun Kim Co-authored-by: Edwin Lim --- .../moriio_toy_proxy_server.py | 16 +++++++- .../unit/test_moriio_toy_proxy_server.py | 40 +++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) create mode 100644 tests/v1/kv_connector/unit/test_moriio_toy_proxy_server.py diff --git a/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py b/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py index 07a462711d2..6f7d210ff3c 100644 --- a/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py +++ b/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py @@ -214,6 +214,14 @@ def example_round_robin_dp_loader(request_number, dp_size): return request_nums % dp_size +@app.route("/health", methods=["GET"]) +async def health(): + # Benchmark harnesses and load balancers often probe the proxy URL and + # treat a non-200 response as a dead service. Backends expose /health, but + # this proxy used to lack one; report the proxy process itself as healthy. + return ("ok", 200) + + @app.route("/v1/completions", methods=["POST"]) async def handle_completions_request(): return await handle_request("/completions", request) @@ -288,7 +296,13 @@ async def handle_request(api: str, request: Request): ) ) - req_data["max_tokens"] -= 1 + # max_completion_tokens takes precedence when present (that's the limit + # the backend enforces); fall back to max_tokens. If neither is set + # (e.g. benchmark clients that omit it), leave the request unchanged. + if "max_completion_tokens" in req_data: + req_data["max_completion_tokens"] -= 1 + elif "max_tokens" in req_data: + req_data["max_tokens"] -= 1 req_data["kv_transfer_params"] = { "do_remote_decode": False, diff --git a/tests/v1/kv_connector/unit/test_moriio_toy_proxy_server.py b/tests/v1/kv_connector/unit/test_moriio_toy_proxy_server.py new file mode 100644 index 00000000000..d1a270cac14 --- /dev/null +++ b/tests/v1/kv_connector/unit/test_moriio_toy_proxy_server.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import runpy +from pathlib import Path + +import pytest + +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip( + "MoRIIO toy proxy tests run only in the ROCm image.", + allow_module_level=True, + ) + +pytest.importorskip("quart") + +_PROXY_SERVER = ( + Path(__file__).parents[4] + / "examples" + / "disaggregated" + / "disaggregated_serving" + / "moriio_toy_proxy_server.py" +) + + +def _load_proxy(): + namespace = runpy.run_path(str(_PROXY_SERVER)) + return namespace["handle_request"].__globals__ + + +@pytest.mark.asyncio +async def test_health_route_returns_ok(): + proxy = _load_proxy() + + response = await proxy["app"].test_client().get("/health") + + assert response.status_code == 200 + assert await response.get_data(as_text=True) == "ok"