[PD] Route PD server warmup to every DP rank (#30748)

Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>

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