[Bugfix] MoRIIO toy P/D proxy: add /health (#45222)

Signed-off-by: Chaemin Lim <chaemin.lim@mangoboost.io>
Signed-off-by: Edwin Lim <edwinlim0919@gmail.com>
Co-authored-by: Edwin Lim <edwin.lim@mangoboost.io>
Co-authored-by: Jaeyoun Kim <jaeyoun.kim@mangoboost.io>
Co-authored-by: Edwin Lim <edwinlim0919@gmail.com>

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,

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"