From f2c875d1c8d1a35392992a9611d4bf5dd84d8e5c Mon Sep 17 00:00:00 2001 From: weireweire Date: Wed, 15 Jul 2026 15:59:41 +0800 Subject: [PATCH] [PD] Route PD server warmup to every DP rank (#30748) Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com> --- python/sglang/srt/entrypoints/http_server.py | 86 ++++++++++++------- .../entrypoints/test_http_server_warmup.py | 83 ++++++++++++++++++ 2 files changed, 139 insertions(+), 30 deletions(-) create mode 100644 test/registered/unit/entrypoints/test_http_server_warmup.py diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index a0d48b8f7..dbc1748be 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -21,6 +21,7 @@ import asyncio import dataclasses import logging import os +import ssl import tempfile import threading import time @@ -39,6 +40,7 @@ from typing import ( Union, ) +import aiohttp import numpy as np import requests import uvicorn @@ -2013,6 +2015,46 @@ def _admin_api_key_missing_response( MINIMUM_PNG_PICTURE_BASE64 = "iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAYAAABzenr0AAAACXBIWXMAAA7EAAAOxAGVKw4bAAAAbUlEQVRYhe3VsQ2AMAxE0Y/lIgNQULD/OqyCMgCihCKSG4yRuKuiNH6JLsoEbMACOGBcua9HOR7Y6w6swBwMy0qLTpkeI77qdEBpBFAHBBDAGH8WrwJKI4AAegUCfAKgEgpQDvh3CR3oQCuav58qlAw73kKCSgAAAABJRU5ErkJggg==" +async def _send_disaggregation_warmup_requests( + server_args: ServerArgs, + url: str, + headers: Dict[str, str], + ssl_verify: Union[bool, str], + timeout: int, +) -> List[int]: + ssl_context = ( + ssl_verify + if isinstance(ssl_verify, bool) + else ssl.create_default_context(cafile=ssl_verify) + ) + + async def send_request(session: aiohttp.ClientSession, dp_rank: int) -> int: + json_data = { + "sampling_params": { + "temperature": 0.0, + "max_new_tokens": 8, + "ignore_eos": True, + }, + "bootstrap_host": FAKE_BOOTSTRAP_HOST, + "bootstrap_room": dp_rank, + "input_ids": [10, 11, 12, 13], + "routed_dp_rank": dp_rank, + } + async with session.post( + url + "/generate", json=json_data, ssl=ssl_context + ) as response: + await response.read() + return response.status + + async with aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout(total=timeout), + headers=headers, + ) as session: + return await asyncio.gather( + *(send_request(session, dp_rank) for dp_rank in range(server_args.dp_size)) + ) + + def _execute_server_warmup(server_args: ServerArgs): headers = {} url = server_args.url() @@ -2124,42 +2166,28 @@ def _execute_server_warmup(server_args: ServerArgs): else: logger.info(f"Start of pd disaggregation warmup ...") - request_name = "/generate" - json_data = { - "sampling_params": { - "temperature": 0.0, - "max_new_tokens": 8, - "ignore_eos": True, - }, - "bootstrap_host": [FAKE_BOOTSTRAP_HOST] * server_args.dp_size, - # This is a hack to ensure fake transfer is enabled during prefill warmup - # ensure each dp rank has a unique bootstrap_room during prefill warmup - "bootstrap_room": [ - i * (2**63 // server_args.dp_size) + (i % server_args.tp_size) - for i in range(server_args.dp_size) - ], - "input_ids": [[10, 11, 12, 13]] * server_args.dp_size, - } - res = requests.post( - url + request_name, - json=json_data, - headers=headers, - timeout=( - warmup_timeout if warmup_timeout > 0 else 1800 - ), # because of deep gemm precache is very long if not precache. - verify=ssl_verify, + status_codes = asyncio.run( + _send_disaggregation_warmup_requests( + server_args=server_args, + url=url, + headers=headers, + ssl_verify=ssl_verify, + timeout=warmup_timeout if warmup_timeout > 0 else 1800, + ) ) - if res.status_code == 200: + failed_status_codes = [code for code in status_codes if code != 200] + if not failed_status_codes: logger.info( - f"Disaggregation warmup request completed with status {res.status_code}, resp: {res.json()}" + "Disaggregation warmup requests completed for all %s DP ranks", + server_args.dp_size, ) logger.info("End of disaggregation warmup") _global_state.tokenizer_manager.server_status = ServerStatus.Up else: logger.info( - "Disaggregation warmup failed (mode=%s), status code: %s", + "Disaggregation warmup failed (mode=%s), status codes: %s", server_args.disaggregation_mode, - res.status_code, + failed_status_codes, ) _global_state.tokenizer_manager.server_status = ServerStatus.UnHealthy @@ -2169,8 +2197,6 @@ def _execute_server_warmup(server_args: ServerArgs): kill_process_tree(os.getpid()) return False - # Debug print - # logger.info(f"warmup request returns: {res.json()=}") return success diff --git a/test/registered/unit/entrypoints/test_http_server_warmup.py b/test/registered/unit/entrypoints/test_http_server_warmup.py new file mode 100644 index 000000000..cf26faeb6 --- /dev/null +++ b/test/registered/unit/entrypoints/test_http_server_warmup.py @@ -0,0 +1,83 @@ +import asyncio +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST +from sglang.srt.entrypoints.http_server import ( + _send_disaggregation_warmup_requests, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class TestDisaggregationServerWarmup(unittest.IsolatedAsyncioTestCase): + async def test_sends_concurrent_scalar_request_to_each_dp_rank(self): + server_args = SimpleNamespace(dp_size=4) + all_started = asyncio.Event() + calls = [] + sessions = [] + + class Response: + status = 200 + + async def __aenter__(self): + if len(calls) == server_args.dp_size: + all_started.set() + await asyncio.wait_for(all_started.wait(), timeout=5) + return self + + async def __aexit__(self, *args): + pass + + async def read(self): + return b"" + + class Session: + def __init__(self, **kwargs): + self.kwargs = kwargs + sessions.append(self) + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + def post(self, *args, **kwargs): + calls.append((args, kwargs)) + return Response() + + with patch("sglang.srt.entrypoints.http_server.aiohttp.ClientSession", Session): + status_codes = await _send_disaggregation_warmup_requests( + server_args=server_args, + url="http://localhost:30000", + headers={"Authorization": "Bearer token"}, + ssl_verify=False, + timeout=123, + ) + + self.assertEqual(status_codes, [200] * server_args.dp_size) + self.assertEqual(len(calls), server_args.dp_size) + self.assertEqual(len(sessions), 1) + self.assertEqual( + sessions[0].kwargs["headers"], {"Authorization": "Bearer token"} + ) + self.assertEqual(sessions[0].kwargs["timeout"].total, 123) + + calls_by_rank = { + kwargs["json"]["routed_dp_rank"]: (args, kwargs) for args, kwargs in calls + } + self.assertEqual(set(calls_by_rank), set(range(server_args.dp_size))) + + for dp_rank, (args, kwargs) in calls_by_rank.items(): + self.assertEqual(args, ("http://localhost:30000/generate",)) + self.assertEqual(kwargs["json"]["input_ids"], [10, 11, 12, 13]) + self.assertEqual(kwargs["json"]["bootstrap_host"], FAKE_BOOTSTRAP_HOST) + self.assertEqual(kwargs["json"]["bootstrap_room"], dp_rank) + self.assertFalse(kwargs["ssl"]) + + +if __name__ == "__main__": + unittest.main()