[PD] Introduce runtime role switching between prefill and decode (#28403)

Signed-off-by: huanglong <huanglong@linux.alibaba.com> Signed-off-by: inkcherry <mingzhi.liu@amd.com> Co-authored-by: huanglong <huanglong@linux.alibaba.com> Co-authored-by: Shangming Cai <csmthu@gmail.com> Co-authored-by: Huang Long <121648372+LLLL114@users.noreply.github.com>
2026-09-18 01:45:12 +08:00
parent a98d921658
commit 1f60ddef5d
27 changed files with 1737 additions and 55 deletions
@@ -0,0 +1,84 @@
+import asyncio
+
+import pytest
+from sglang_router import mini_lb
+
+
+@pytest.mark.parametrize(
+    ("result", "restored"),
+    [
+        (
+            {
+                "success": False,
+                "message": "instance is not idle",
+                "safe_to_restore": True,
+            },
+            True,
+        ),
+        (
+            {
+                "success": False,
+                "message": "instance unhealthy, restart required",
+                "safe_to_restore": False,
+            },
+            False,
+        ),
+        ({"success": False, "message": "connection lost"}, False),
+    ],
+)
+def test_failed_role_switch_restores_only_healthy_worker(monkeypatch, result, restored):
+    worker_url = "http://prefill:8000"
+    load_balancer = mini_lb.MiniLoadBalancer.__new__(mini_lb.MiniLoadBalancer)
+    load_balancer.timeout = 1
+    load_balancer.prefill_urls = [worker_url]
+    load_balancer.prefill_bootstrap_ports = [8998]
+    load_balancer.decode_urls = ["http://decode:8000"]
+    monkeypatch.setattr(mini_lb, "lb", load_balancer)
+
+    async def post_role_switch(*_args, **_kwargs):
+        return 400, result
+
+    monkeypatch.setattr(mini_lb, "_post_role_switch", post_role_switch)
+    response = asyncio.run(
+        mini_lb.pd_role_switch(
+            {
+                "worker_url": worker_url,
+                "new_role": "decode",
+                "drain_timeout_secs": 0,
+            }
+        )
+    )
+
+    assert response.status_code == 400
+    assert (worker_url in load_balancer.prefill_urls) is restored
+
+
+def test_role_switch_forwards_decode_graph_requirements(monkeypatch):
+    worker_url = "http://prefill:8000"
+    load_balancer = mini_lb.MiniLoadBalancer.__new__(mini_lb.MiniLoadBalancer)
+    load_balancer.timeout = 1
+    load_balancer.prefill_urls = [worker_url]
+    load_balancer.prefill_bootstrap_ports = [8998]
+    load_balancer.decode_urls = ["http://decode:8000"]
+    monkeypatch.setattr(mini_lb, "lb", load_balancer)
+    sent_body = {}
+
+    async def post_role_switch(_worker_url, body):
+        sent_body.update(body)
+        return 200, {"success": True, "message": "ok"}
+
+    monkeypatch.setattr(mini_lb, "_post_role_switch", post_role_switch)
+    response = asyncio.run(
+        mini_lb.pd_role_switch(
+            {
+                "worker_url": worker_url,
+                "new_role": "decode",
+                "decode_cuda_graph_bs": [1, 2, 4],
+                "decode_cuda_graph_memory_gb": 1.25,
+            }
+        )
+    )
+
+    assert response.status_code == 200
+    assert sent_body["decode_cuda_graph_bs"] == [1, 2, 4]
+    assert sent_body["decode_cuda_graph_memory_gb"] == 1.25