[Router] Add load-aware prefill admission and bounded policy proposals (#37843)

Signed-off-by: Vincent Gao <vincentbo@linux.alibaba.com>
Co-authored-by: Kangyan Zhou <zky314343421@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Shangming Cai <csmthu@gmail.com>

experimental/sgl-router/tests/e2e/infra/model_pool.py

@@ -34,6 +34,18 @@ from .model_specs import get_model_spec
 logger = logging.getLogger(__name__)
 
 
+def _wait_for_process_group_exit(pgid: int, timeout: float) -> bool:
+    deadline = time.monotonic() + timeout
+    while True:
+        try:
+            os.killpg(pgid, 0)
+        except ProcessLookupError:
+            return True
+        if time.monotonic() >= deadline:
+            return False
+        time.sleep(0.1)
+
+
 def _get_open_port() -> int:
     """Allocate an ephemeral TCP port in the range [20000, 55535].
 
@@ -72,6 +84,7 @@ class ModelInstance:
     model_id: str
     gpu_ids: list[int] = field(default_factory=list)
     kv_events_endpoint: str | None = None
+    _shutdown_started: bool = field(default=False, init=False, repr=False)
 
     def __enter__(self) -> "ModelInstance":
         return self
@@ -80,16 +93,31 @@ class ModelInstance:
         self.shutdown()
 
     def shutdown(self) -> None:
-        if self.process is not None and self.process.poll() is None:
-            try:
-                self.process.send_signal(signal.SIGTERM)
-                try:
-                    self.process.wait(timeout=60)
-                except subprocess.TimeoutExpired:
-                    self.process.kill()
-                    self.process.wait()
-            except ProcessLookupError:
-                pass
+        if self.process is None or self._shutdown_started:
+            return
+        self._shutdown_started = True
+        pgid = self.process.pid
+
+        try:
+            os.killpg(pgid, signal.SIGTERM)
+        except ProcessLookupError:
+            return
+
+        try:
+            self.process.wait(timeout=60)
+        except subprocess.TimeoutExpired:
+            pass
+
+        if _wait_for_process_group_exit(pgid, timeout=30):
+            return
+
+        try:
+            os.killpg(pgid, signal.SIGKILL)
+        except ProcessLookupError:
+            return
+        self.process.wait()
+        if not _wait_for_process_group_exit(pgid, timeout=5):
+            raise RuntimeError(f"worker process group {pgid} did not exit")
 
 
 def spawn_worker(

experimental/sgl-router/tests/e2e/infra/test_model_pool.py

@@ -0,0 +1,50 @@
+import signal
+
+from infra import model_pool
+
+
+class _Process:
+    pid = 1234
+
+    def __init__(self):
+        self.wait_timeouts = []
+
+    def poll(self):
+        return None
+
+    def send_signal(self, sig):
+        raise AssertionError(f"signaled only the parent process: {sig}")
+
+    def wait(self, timeout=None):
+        self.wait_timeouts.append(timeout)
+        return 0
+
+
+def test_shutdown_waits_for_the_worker_process_group(monkeypatch):
+    process = _Process()
+    signals = []
+    probes = iter([True, True, False])
+
+    def killpg(pgid, sig):
+        signals.append((pgid, sig))
+        if sig == 0 and not next(probes):
+            raise ProcessLookupError
+
+    monkeypatch.setattr(model_pool.os, "killpg", killpg)
+    monkeypatch.setattr(model_pool.time, "sleep", lambda _: None)
+
+    instance = model_pool.ModelInstance(
+        url="http://127.0.0.1:30000",
+        port=30000,
+        process=process,
+        model_id="qwen3-0.6b",
+    )
+    instance.shutdown()
+
+    assert signals == [
+        (process.pid, signal.SIGTERM),
+        (process.pid, 0),
+        (process.pid, 0),
+        (process.pid, 0),
+    ]
+    assert process.wait_timeouts == [60]

experimental/sgl-router/tests/e2e/k8s_integration/conftest.py

@@ -11,6 +11,7 @@ Teardown: ./tests/e2e/k8s_integration/setup.sh teardown
 
 from __future__ import annotations
 
+import json
 import logging
 import socket
 import subprocess
@@ -97,6 +98,60 @@ def _wait_for_pod_ready(
     )
 
 
+def _wait_for_replacement_pod_ready(
+    old_pod: str,
+    selector: str,
+    namespace: str = NAMESPACE,
+    timeout: int = 120,
+    interval: float = 0.5,
+) -> str:
+    deadline = time.time() + timeout
+    last_observed = "no pods"
+
+    while time.time() < deadline:
+        result = _kubectl(
+            "get",
+            "pods",
+            "-n",
+            namespace,
+            "-l",
+            selector,
+            "-o",
+            "json",
+            check=False,
+        )
+        if getattr(result, "returncode", 0) == 0:
+            pods = json.loads(result.stdout or "{}").get("items", [])
+            names = [pod.get("metadata", {}).get("name", "") for pod in pods]
+            last_observed = ", ".join(filter(None, names)) or "no pods"
+
+            if old_pod not in names:
+                for pod in sorted(
+                    pods, key=lambda item: item.get("metadata", {}).get("name", "")
+                ):
+                    metadata = pod.get("metadata", {})
+                    status = pod.get("status", {})
+                    ready = any(
+                        condition.get("type") == "Ready"
+                        and condition.get("status") == "True"
+                        for condition in status.get("conditions", [])
+                    )
+                    if (
+                        metadata.get("name") != old_pod
+                        and not metadata.get("deletionTimestamp")
+                        and status.get("phase") == "Running"
+                        and ready
+                    ):
+                        return metadata["name"]
+
+        time.sleep(interval)
+
+    raise TimeoutError(
+        f"No ready replacement for pod {old_pod!r} with selector {selector!r} "
+        f"after {timeout}s; last observed: {last_observed}"
+    )
+
+
 def _wait_for_port(port: int, proc: subprocess.Popen, timeout: int = 15) -> None:
     """Poll until a TCP connection to localhost:port succeeds."""
     deadline = time.time() + timeout

experimental/sgl-router/tests/e2e/k8s_integration/test_conftest_helpers.py

@@ -0,0 +1,54 @@
+import json
+from types import SimpleNamespace
+
+import conftest as k8s_conftest
+
+
+def _pod(name: str, phase: str, ready: bool) -> dict:
+    return {
+        "metadata": {"name": name},
+        "status": {
+            "phase": phase,
+            "conditions": [
+                {
+                    "type": "Ready",
+                    "status": "True" if ready else "False",
+                }
+            ],
+        },
+    }
+
+
+def test_wait_for_replacement_pod_ignores_old_and_pending_pods(monkeypatch):
+    old_pod = "sgl-router-old"
+    new_pod = "sgl-router-new"
+    responses = iter(
+        [
+            [_pod(old_pod, "Running", True)],
+            [
+                _pod(old_pod, "Running", True),
+                _pod(new_pod, "Running", True),
+            ],
+            [_pod(new_pod, "Pending", False)],
+            [_pod(new_pod, "Running", True)],
+        ]
+    )
+    calls = []
+
+    def fake_kubectl(*args, **kwargs):
+        calls.append((args, kwargs))
+        return SimpleNamespace(stdout=json.dumps({"items": next(responses)}))
+
+    monkeypatch.setattr(k8s_conftest, "_kubectl", fake_kubectl)
+    monkeypatch.setattr(k8s_conftest.time, "sleep", lambda _: None)
+
+    replacement = k8s_conftest._wait_for_replacement_pod_ready(
+        old_pod,
+        "app=sgl-router",
+        timeout=5,
+        interval=0,
+    )
+
+    assert replacement == new_pod
+    assert len(calls) == 4
+    assert all("-o" in args and "json" in args for args, _ in calls)

experimental/sgl-router/tests/e2e/k8s_integration/test_lifecycle.py

@@ -13,10 +13,7 @@ by driving the deployment scale.
 
 from __future__ import annotations
 
-import logging
-
 import httpx
-import pytest
 from conftest import (
     NAMESPACE,
     _cleanup_port_forward,
@@ -24,7 +21,7 @@ from conftest import (
     _poll_until,
     _port_forward_start,
     _wait_for_deployment_ready,
-    logger,
+    _wait_for_replacement_pod_ready,
 )
 
 ROUTER_RESTART_PORT = 8092
@@ -132,7 +129,10 @@ class TestRouterRestart:
                 _cleanup_port_forward("router-restart-pre-kill", pf_holder[0])
                 pf_holder[0] = None
 
-            _wait_for_deployment_ready("sgl-router")
+            if old_pod:
+                _wait_for_replacement_pod_ready(old_pod, "app=sgl-router")
+            else:
+                _wait_for_deployment_ready("sgl-router")
 
             pf_holder[0] = _port_forward_start(
                 NAMESPACE, "sgl-router", ROUTER_RESTART_PORT, 8090