A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 05:00 UTC
[ROCm] Profiler api support for ROCm MORI toy proxy server in PD Disaggregation (#40264)
Signed-off-by: Tej Kiran <kiran.tej@amd.com>
tej committed
May 7, 2026 at 03:58 UTC
8a4888be21219bf081b5bc7e01f8ecf2af9057a8
1 file changed
+73
-1
examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
+73
-1
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import argparse
4
import asyncio
5
import copy
6
import logging
@@ -7,6 +8,7 @@ import os
8
import socket
9
import threading
10
import uuid
11
+from urllib.parse import urlparse
12
13
import aiohttp
14
import msgpack
@@ -336,11 +338,81 @@ async def handle_request(api: str, request: Request):
338
)
339
340
341
+async def send_profile_cmd(req_data: dict, profiler_cmd: str):
342
+ assert profiler_cmd in {"start", "stop"}
343
+
344
+ with _list_lock:
345
+ p_instances = list(prefill_instances)
346
+ d_instances = list(decode_instances)
347
+
348
+ if not p_instances and not d_instances:
349
+ raise RuntimeError(
350
+ "Service Unavailable: No prefill or decode instances are registered."
351
+ )
352
+
353
+ headers = {
354
+ "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
355
+ }
356
+
357
+ tasks = []
358
+
359
+ async with aiohttp.ClientSession(
360
+ timeout=aiohttp.ClientTimeout(total=60)
361
+ ) as session:
362
+ for instances in (p_instances, d_instances):
363
+ for inst in instances:
364
+ _p = urlparse(inst["request_address"])
365
+ url = f"http://{_p.hostname}:{_p.port}/{profiler_cmd}_profile"
366
+
367
+ tasks.append(
368
+ session.post(
369
+ url,
370
+ json=req_data,
371
+ headers=headers,
372
+ )
373
+ )
374
+
375
+ responses = await asyncio.gather(*tasks, return_exceptions=True)
376
+
377
+ for r in responses:
378
+ if isinstance(r, Exception):
379
+ raise r
380
+ if r.status >= 400:
381
+ msg = await r.text()
382
+ raise RuntimeError(f"{profiler_cmd}_profile failed: {r.status}, {msg}")
383
+
384
+ return await responses[0].json()
385
+
386
+
387
+@app.post("/start_profile")
388
+async def start_profile():
389
+ try:
390
+ req_data = await request.get_json()
391
+ return await send_profile_cmd(req_data, "start")
392
+ except Exception as e:
393
+ logger.exception("start_profile failed: %s", e)
394
+ return await make_response((str(e), 500))
395
+
396
+
397
+@app.post("/stop_profile")
398
+async def stop_profile():
399
+ try:
400
+ req_data = await request.get_json()
401
+ return await send_profile_cmd(req_data, "stop")
402
+ except Exception as e:
403
+ logger.exception("stop_profile failed: %s", e)
404
+ return await make_response((str(e), 500))
405
+
406
+
407
if __name__ == "__main__":
408
+ parser = argparse.ArgumentParser()
409
+ parser.add_argument("--port", type=int, default=10001)
410
+ args = parser.parse_args()
411
+
412
t = start_service_discovery("0.0.0.0", 36367)
413
app.debug = True
414
app.config["BODY_TIMEOUT"] = 360000
415
app.config["RESPONSE_TIMEOUT"] = 360000
416
345
- app.run(host="0.0.0.0", port=10001)
417
+ app.run(host="0.0.0.0", port=args.port)
418
t.join()