[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

@@ -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"
+            status_codes = asyncio.run(
-            json_data = {
+                _send_disaggregation_warmup_requests(
-                "sampling_params": {
+                    server_args=server_args,
-                    "temperature": 0.0,
+                    url=url,
-                    "max_new_tokens": 8,
+                    headers=headers,
-                    "ignore_eos": True,
+                    ssl_verify=ssl_verify,
-                },
+                    timeout=warmup_timeout if warmup_timeout > 0 else 1800,
-                "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,
             )
-            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
 
 

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()