[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>
chaeminlim-mb committed
Jul 15, 2026 at 07:56 UTC
520a20ba4e2d2a15e58f34751f668bd8d046bc45
2 files changed
+55
-1
examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
+15
-1
@@ -214,6 +214,14 @@ def example_round_robin_dp_loader(request_number, dp_size):
214
return request_nums % dp_size
215
216
217
+@app.route("/health", methods=["GET"])
218
+async def health():
219
+ # Benchmark harnesses and load balancers often probe the proxy URL and
220
+ # treat a non-200 response as a dead service. Backends expose /health, but
221
+ # this proxy used to lack one; report the proxy process itself as healthy.
222
+ return ("ok", 200)
223
+
224
+
225
@app.route("/v1/completions", methods=["POST"])
226
async def handle_completions_request():
227
return await handle_request("/completions", request)
@@ -288,7 +296,13 @@ async def handle_request(api: str, request: Request):
296
)
297
)
298
291
- req_data["max_tokens"] -= 1
299
+ # max_completion_tokens takes precedence when present (that's the limit
300
+ # the backend enforces); fall back to max_tokens. If neither is set
301
+ # (e.g. benchmark clients that omit it), leave the request unchanged.
302
+ if "max_completion_tokens" in req_data:
303
+ req_data["max_completion_tokens"] -= 1
304
+ elif "max_tokens" in req_data:
305
+ req_data["max_tokens"] -= 1
306
307
req_data["kv_transfer_params"] = {
308
"do_remote_decode": False,
tests/v1/kv_connector/unit/test_moriio_toy_proxy_server.py
new
+40
@@ -0,0 +1,40 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import runpy
5
+from pathlib import Path
6
+
7
+import pytest
8
+
9
+from vllm.platforms import current_platform
10
+
11
+if not current_platform.is_rocm():
12
+ pytest.skip(
13
+ "MoRIIO toy proxy tests run only in the ROCm image.",
14
+ allow_module_level=True,
15
+ )
16
+
17
+pytest.importorskip("quart")
18
+
19
+_PROXY_SERVER = (
20
+ Path(__file__).parents[4]
21
+ / "examples"
22
+ / "disaggregated"
23
+ / "disaggregated_serving"
24
+ / "moriio_toy_proxy_server.py"
25
+)
26
+
27
+
28
+def _load_proxy():
29
+ namespace = runpy.run_path(str(_PROXY_SERVER))
30
+ return namespace["handle_request"].__globals__
31
+
32
+
33
+@pytest.mark.asyncio
34
+async def test_health_route_returns_ok():
35
+ proxy = _load_proxy()
36
+
37
+ response = await proxy["app"].test_client().get("/health")
38
+
39
+ assert response.status_code == 200
40
+ assert await response.get_data(as_text=True) == "ok"