autoinference-utils 0.2.8__tar.gz → 0.2.10__tar.gz

{autoinference_utils-0.2.8 → autoinference_utils-0.2.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: autoinference-utils
-Version: 0.2.8
+Version: 0.2.10
 Summary: Shared endpoint abstractions for autoinference deployments
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown

{autoinference_utils-0.2.8 → autoinference_utils-0.2.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "autoinference-utils"
-version = "0.2.8"
+version = "0.2.10"
 description = "Shared endpoint abstractions for autoinference deployments"
 readme = "README.md"
 requires-python = ">=3.10"

{autoinference_utils-0.2.8 → autoinference_utils-0.2.10}/src/autoinference_utils/endpoint.py

@@ -27,7 +27,7 @@ BENCH_MODE_ENV = "ENDPOINT_BENCH_MODE"
 DUMMY_WEIGHTS_ENV = "ENDPOINT_DUMMY"
 BENCH_MODE_PORT_OFFSET = 10000
 
-MODAL_FLASH_REQUEST_UUID_HEADER= "modal-flash-request-uuid"
+MODAL_FLASH_REQUEST_UUID_HEADER = "modal-flash-request-uuid"
 MODAL_SESSION_ID_HEADER = "modal-session-id"
 
 ENDPOINTS_REQUIRING_OAI_STREAM_COMPAT: tuple[str, ...] = ("TRTLLMEndpoint",)
@@ -63,9 +63,7 @@ def _flag_name(key: str) -> str:
     return key.partition("=")[0]
 
 
-def get_server_arg(
-    server_args: Mapping[str, str], flag: str
-) -> tuple[str, str] | None:
+def get_server_arg(server_args: Mapping[str, str], flag: str) -> tuple[str, str] | None:
     """Return the key and value setting a flag under either spelling."""
     for key, value in server_args.items():
         if _flag_name(key) == flag:
@@ -176,9 +174,7 @@ class SGLangEndpoint(Endpoint):
         self.disaggregation_mode = disaggregation_mode
         self.prefill_bootstrap_port = prefill_bootstrap_port
         self.launcher_module = launcher_module
-        self.extra_server_args = (
-            dict(extra_server_args) if extra_server_args else {}
-        )
+        self.extra_server_args = dict(extra_server_args) if extra_server_args else {}
         self.health_timeout = health_timeout
         self.health_poll_interval = health_poll_interval
         self.health_request_timeout = health_request_timeout
@@ -193,17 +189,16 @@ class SGLangEndpoint(Endpoint):
         self._bench_server: Optional[http.server.ThreadingHTTPServer] = None
 
         if self.disaggregation_mode not in (None, "prefill", "decode"):
-            raise ValueError(
-                "disaggregation_mode must be None, 'prefill', or 'decode'"
-            )
+            raise ValueError("disaggregation_mode must be None, 'prefill', or 'decode'")
 
         if os.environ.get(DUMMY_WEIGHTS_ENV) == "1":
-            if (
-                self.load_format is None
-                and not has_server_arg(self.extra_server_args, "--load-format")
+            if self.load_format is None and not has_server_arg(
+                self.extra_server_args, "--load-format"
             ):
                 self.load_format = "dummy"
-            if not has_server_arg(self.extra_server_args, "--model-loader-extra-config"):
+            if not has_server_arg(
+                self.extra_server_args, "--model-loader-extra-config"
+            ):
                 self.extra_server_args["--model-loader-extra-config"] = "{}"
 
         # Request logging enabled, log only metadata by default
@@ -214,7 +209,9 @@ class SGLangEndpoint(Endpoint):
             # Set request logging level
             level = get_server_arg(self.extra_server_args, "--log-requests-level")
             if level is None:
-                self.extra_server_args["--log-requests-level"] = str(self.log_requests_level)
+                self.extra_server_args["--log-requests-level"] = str(
+                    self.log_requests_level
+                )
             else:
                 key, value = level
                 print(
@@ -236,19 +233,21 @@ class SGLangEndpoint(Endpoint):
             self.log_request_headers = ",".join(sglang_log_request_headers)
             os.environ["SGLANG_LOG_REQUEST_HEADERS"] = self.log_request_headers
 
-
     def _build_cmd(self) -> list[str]:
         cmd = [
-            "python", "-m", self.launcher_module,
-            "--host", "0.0.0.0",
-            "--port", str(self.worker_port),
-            "--model-path", self.model_path,
+            "python",
+            "-m",
+            self.launcher_module,
+            "--host",
+            "0.0.0.0",
+            "--port",
+            str(self.worker_port),
+            "--model-path",
+            self.model_path,
         ]
 
         if self.speculative_model_path is not None:
-            cmd.extend(
-                ["--speculative-draft-model-path", self.speculative_model_path]
-            )
+            cmd.extend(["--speculative-draft-model-path", self.speculative_model_path])
         if self.load_format is not None:
             cmd.extend(["--load-format", self.load_format])
         if self.tp is not None:
@@ -273,14 +272,18 @@ class SGLangEndpoint(Endpoint):
                 raise ValueError("dist_init_host is required when nnodes > 1")
             cmd.extend(
                 [
-                    "--nnodes", str(self.nnodes),
-                    "--node-rank", str(self.node_rank),
+                    "--nnodes",
+                    str(self.nnodes),
+                    "--node-rank",
+                    str(self.node_rank),
                     "--dist-init-addr",
                     f"{self.dist_init_host}:{self.dist_init_port}",
                 ]
             )
 
-        merged = _merge_server_args(self.DEFAULT_OPERATIONAL_ARGS, self.extra_server_args)
+        merged = _merge_server_args(
+            self.DEFAULT_OPERATIONAL_ARGS, self.extra_server_args
+        )
         for key, value in merged.items():
             cmd.extend(server_arg_tokens(key, value))
 
@@ -345,9 +348,7 @@ class VLLMEndpoint(Endpoint):
         super().__init__(base_url=f"http://localhost:{vllm_port}")
         self.model = model
         self.worker_port = vllm_port
-        self.extra_server_args = (
-            dict(extra_server_args) if extra_server_args else {}
-        )
+        self.extra_server_args = dict(extra_server_args) if extra_server_args else {}
         self.health_timeout = health_timeout
         self.health_poll_interval = health_poll_interval
         self.health_request_timeout = health_request_timeout
@@ -360,10 +361,15 @@ class VLLMEndpoint(Endpoint):
 
     def _build_cmd(self) -> list[str]:
         cmd = [
-            "python", "-m", "vllm.entrypoints.openai.api_server",
-            "--host", "0.0.0.0",
-            "--port", str(self.worker_port),
-            "--model", self.model,
+            "python",
+            "-m",
+            "vllm.entrypoints.openai.api_server",
+            "--host",
+            "0.0.0.0",
+            "--port",
+            str(self.worker_port),
+            "--model",
+            self.model,
         ]
         for key, value in self.extra_server_args.items():
             cmd.extend(server_arg_tokens(key, value))
@@ -417,9 +423,7 @@ class TRTLLMEndpoint(Endpoint):
         super().__init__(base_url=f"http://localhost:{worker_port}")
         self.model = model
         self.worker_port = worker_port
-        self.extra_server_args = (
-            dict(extra_server_args) if extra_server_args else {}
-        )
+        self.extra_server_args = dict(extra_server_args) if extra_server_args else {}
         self.health_timeout = health_timeout
         self.health_poll_interval = health_poll_interval
         self.health_request_timeout = health_request_timeout
@@ -429,8 +433,10 @@ class TRTLLMEndpoint(Endpoint):
         cmd = [
             "trtllm-serve",
             self.model,
-            "--host", "0.0.0.0",
-            "--port", str(self.worker_port),
+            "--host",
+            "0.0.0.0",
+            "--port",
+            str(self.worker_port),
         ]
         for key, value in self.extra_server_args.items():
             cmd.extend(server_arg_tokens(key, value))
@@ -499,18 +505,30 @@ class RouterEndpoint(Endpoint):
 
     def _build_cmd(self) -> list[str]:
         cmd = [
-            "python", "-m", "sglang_router.launch_router",
-            "--host", f"[{self.host}]" if ":" in self.host else self.host,
-            "--port", str(self.router_port),
-            "--prefill-policy", "cache_aware",
-            "--decode-policy", "round_robin",
-            "--max-concurrent-requests", "128",
-            "--rate-limit-tokens-per-second", "0",
-            "--queue-size", "0",
-            "--health-check-timeout-secs", "600",
-            "--log-level", "info",
+            "python",
+            "-m",
+            "sglang_router.launch_router",
+            "--host",
+            f"[{self.host}]" if ":" in self.host else self.host,
+            "--port",
+            str(self.router_port),
+            "--prefill-policy",
+            "cache_aware",
+            "--decode-policy",
+            "round_robin",
+            "--max-concurrent-requests",
+            "128",
+            "--rate-limit-tokens-per-second",
+            "0",
+            "--queue-size",
+            "0",
+            "--health-check-timeout-secs",
+            "600",
+            "--log-level",
+            "info",
             "--disable-circuit-breaker",
-            "--request-timeout-secs", "3600",
+            "--request-timeout-secs",
+            "3600",
         ]
 
         if self.api_key is not None:
@@ -522,9 +540,7 @@ class RouterEndpoint(Endpoint):
         for role, node_ip in self.pd_config:
             node_url = _url(node_ip, self.worker_port)
             if role == "prefill":
-                cmd.extend(
-                    ["--prefill", node_url, str(self.prefill_bootstrap_port)]
-                )
+                cmd.extend(["--prefill", node_url, str(self.prefill_bootstrap_port)])
             elif role == "decode":
                 cmd.extend(["--decode", node_url])
             elif role == "worker":
@@ -576,6 +592,7 @@ class RouterEndpoint(Endpoint):
 # and this proxy listens on worker_port. POST /bench shells out to a benchmark
 # task via run_bench(); every other path is forwarded to the upstream server.
 
+
 def start_bench_proxy(
     *,
     listen_port: int,
@@ -587,17 +604,12 @@ def start_bench_proxy(
         (_BenchProxyHandler,),
         {"upstream_port": upstream_port, "upstream_base_url": upstream_base_url},
     )
-    server = http.server.ThreadingHTTPServer(
-        ("0.0.0.0", listen_port), handler_cls
-    )
+    server = http.server.ThreadingHTTPServer(("0.0.0.0", listen_port), handler_cls)
     thread = threading.Thread(
         target=server.serve_forever, daemon=True, name="bench-proxy"
     )
     thread.start()
-    print(
-        f"[bench-proxy] listening on :{listen_port}, "
-        f"forwarding to :{upstream_port}"
-    )
+    print(f"[bench-proxy] listening on :{listen_port}, forwarding to :{upstream_port}")
     return server
 
 
@@ -627,9 +639,7 @@ class _BenchProxyHandler(http.server.BaseHTTPRequestHandler):
         try:
             payload = self._read_json_body()
         except (ValueError, json.JSONDecodeError) as exc:
-            return self._send_json(
-                400, {"ok": False, "error": f"bad body: {exc}"}
-            )
+            return self._send_json(400, {"ok": False, "error": f"bad body: {exc}"})
 
         benchmark = payload.get("benchmark") or ""
         args = payload.get("args") or []
@@ -641,13 +651,9 @@ class _BenchProxyHandler(http.server.BaseHTTPRequestHandler):
         output_dir = payload.get("output_dir") or "/tmp/bench-output"
 
         if not benchmark:
-            return self._send_json(
-                400, {"ok": False, "error": "benchmark is required"}
-            )
+            return self._send_json(400, {"ok": False, "error": "benchmark is required"})
         if not isinstance(args, list):
-            return self._send_json(
-                400, {"ok": False, "error": "args must be a list"}
-            )
+            return self._send_json(400, {"ok": False, "error": "args must be a list"})
 
         rc, body = run_bench(benchmark, list(args), target, output_dir)
         status = 200 if rc == 0 else 500
@@ -699,9 +705,7 @@ class _BenchProxyHandler(http.server.BaseHTTPRequestHandler):
         return json.loads(raw)
 
     def _send_json(self, status: int, obj: Mapping[str, Any]):
-        self._send_raw(
-            status, "application/json", json.dumps(obj).encode()
-        )
+        self._send_raw(status, "application/json", json.dumps(obj).encode())
 
     def _send_raw(self, status: int, content_type: str, body: bytes):
         self.send_response(status)
@@ -726,18 +730,21 @@ def run_bench(
         return 500, json.dumps(
             {
                 "ok": False,
-                "error": (
-                    "autoinference.benchmarks not importable in container"
-                ),
+                "error": ("autoinference.benchmarks not importable in container"),
                 "volume_path": output_dir,
             }
         ).encode()
 
     cmd = [
-        sys.executable, "-m", "invoke",
-        "--search-root", search_root,
-        "-c", "tasks",
-        benchmark, *[str(a) for a in args],
+        sys.executable,
+        "-m",
+        "invoke",
+        "--search-root",
+        search_root,
+        "-c",
+        "tasks",
+        benchmark,
+        *[str(a) for a in args],
     ]
     env = {
         **os.environ,
@@ -788,9 +795,7 @@ def wait_ready(
 
     while time.time() < deadline:
         try:
-            error = _health_check(
-                url, request_timeout=request_timeout, process=process
-            )
+            error = _health_check(url, request_timeout=request_timeout, process=process)
         except subprocess.CalledProcessError as exc:
             print(
                 f"[endpoint] !!! server process exited with code "
@@ -804,13 +809,11 @@ def wait_ready(
         time.sleep(poll_interval)
 
     print(
-        f"[endpoint] !!! health check timed out after {timeout}s; "
-        f"last={last_error}",
+        f"[endpoint] !!! health check timed out after {timeout}s; last={last_error}",
         flush=True,
     )
     raise TimeoutError(
-        f"Health check timed out after {timeout}s for {url}. "
-        f"Last error: {last_error}"
+        f"Health check timed out after {timeout}s for {url}. Last error: {last_error}"
     )
 
 
@@ -839,9 +842,7 @@ def warmup_chat_completions(
             request_timeout=request_timeout,
             max_attempts=max_attempts_per_request,
             retry_delay=retry_delay,
-            description=(
-                f"warmup request {request_idx + 1}/{successful_requests}"
-            ),
+            description=(f"warmup request {request_idx + 1}/{successful_requests}"),
         )
 
 
@@ -863,9 +864,7 @@ def validate_embeddings_endpoint(
             1 if all(isinstance(item, int) for item in inputs) else len(inputs)
         )
     else:
-        raise ValueError(
-            "embedding validation payload must contain non-empty input"
-        )
+        raise ValueError("embedding validation payload must contain non-empty input")
 
     request_headers = {"Content-Type": "application/json"}
     if headers:
@@ -928,10 +927,7 @@ def start_heartbeat_thread(
             try:
                 error = health_check_fn()
             except subprocess.CalledProcessError as exc:
-                print(
-                    f"[heartbeat] server process exited with code "
-                    f"{exc.returncode}"
-                )
+                print(f"[heartbeat] server process exited with code {exc.returncode}")
                 on_failure()
                 return
             if error is None:
@@ -943,10 +939,7 @@ def start_heartbeat_thread(
                 f"({consecutive_failures}/{max_consecutive_failures})"
             )
             if consecutive_failures >= max_consecutive_failures:
-                print(
-                    "[heartbeat] sustained health-check failure, "
-                    "invoking on_failure"
-                )
+                print("[heartbeat] sustained health-check failure, invoking on_failure")
                 on_failure()
                 return
 
@@ -1031,8 +1024,7 @@ def _wait_ready_url(
         last_error = error
         time.sleep(poll_interval)
     raise TimeoutError(
-        f"Timed out after {timeout}s waiting for {url}. "
-        f"Last error: {last_error}"
+        f"Timed out after {timeout}s waiting for {url}. Last error: {last_error}"
     )
 
 
@@ -1050,9 +1042,7 @@ def _post_json(
     timeout: float | None = None,
 ) -> bytes:
     body = json.dumps(payload).encode("utf-8")
-    req = urllib.request.Request(
-        url, data=body, headers=dict(headers), method="POST"
-    )
+    req = urllib.request.Request(url, data=body, headers=dict(headers), method="POST")
     with urllib.request.urlopen(req, timeout=timeout) as resp:
         return resp.read()
 

{autoinference_utils-0.2.8 → autoinference_utils-0.2.10}/src/autoinference_utils/pd.py

@@ -5,6 +5,8 @@ from __future__ import annotations
 import asyncio
 import math
 import os
+import socket
+import struct
 import subprocess
 import threading
 import time
@@ -19,7 +21,52 @@ from .endpoint import (
     _local_base_url,
     _merge_server_args,
     _url,
+    server_arg_tokens,
 )
+from .router import verify_router
+
+# PDEndpoint owns the topology and every flag it mirrors in its own supervision
+# (drain deadline, health probes); recipes tune everything else via router_args.
+_MANAGED_ROUTER_FLAGS = frozenset(
+    {
+        "--host",
+        "--port",
+        "--pd-disaggregation",
+        "--prefill",
+        "--decode",
+        "--api-key",
+        "--shutdown-grace-period-secs",
+        "--health-check-interval-secs",
+        "--health-check-timeout-secs",
+        "--health-failure-threshold",
+    }
+)
+
+# Leave headroom inside Modal's 30-second exit-handler deadline.
+_SHUTDOWN_TIMEOUT = 25
+
+
+def _managed_router_tokens(router_args: Mapping[str, str]) -> list[str]:
+    """Rendered override tokens that would reach a managed flag.
+
+    Checks argv as the router parses it: values split into extra tokens, and
+    argparse accepts any unambiguous prefix of a long flag.
+    """
+    tokens = [
+        token
+        for key, value in router_args.items()
+        for token in server_arg_tokens(key, value)
+    ]
+    return sorted(
+        {
+            token
+            for token in tokens
+            if token.startswith("-")
+            and any(
+                flag.startswith(_flag_name(token)) for flag in _MANAGED_ROUTER_FLAGS
+            )
+        }
+    )
 
 
 def _configure_fabric() -> None:
@@ -61,6 +108,7 @@ class _PDRouter(RouterEndpoint):
         health_interval,
         health_request_timeout,
         health_failure_threshold,
+        router_args,
         **kwargs,
     ):
         super().__init__(health_path="/readiness", host="::", **kwargs)
@@ -69,6 +117,7 @@ class _PDRouter(RouterEndpoint):
         self.health_interval = health_interval
         self.health_request_timeout = health_request_timeout
         self.health_failure_threshold = health_failure_threshold
+        self.router_args = dict(router_args)
 
     def _build_cmd(self) -> list[str]:
         cmd = [
@@ -85,20 +134,21 @@ class _PDRouter(RouterEndpoint):
             cmd.extend([f"--{role}", _url(host, self.worker_port)])
             if role == "prefill":
                 cmd.append(str(self.prefill_bootstrap_port))
-        for flag, value in {
+        defaults = {
             "--prefill-policy": "cache_aware",
             "--decode-policy": "round_robin",
-            "--max-concurrent-requests": self.max_concurrent_requests,
-            "--queue-size": 0,
-            "--rate-limit-tokens-per-second": 0,
-            "--health-check-interval-secs": self.health_interval,
-            "--health-check-timeout-secs": self.health_request_timeout,
-            "--health-failure-threshold": self.health_failure_threshold,
-            "--shutdown-grace-period-secs": self.drain_timeout,
-            "--request-timeout-secs": 1800,
+            "--max-concurrent-requests": str(self.max_concurrent_requests),
+            "--queue-size": "0",
+            "--rate-limit-tokens-per-second": "0",
+            "--health-check-interval-secs": str(self.health_interval),
+            "--health-check-timeout-secs": str(self.health_request_timeout),
+            "--health-failure-threshold": str(self.health_failure_threshold),
+            "--shutdown-grace-period-secs": str(self.drain_timeout),
+            "--request-timeout-secs": "1800",
             "--log-level": "info",
-        }.items():
-            cmd.extend([flag, str(value)])
+        }
+        for flag, value in _merge_server_args(defaults, self.router_args).items():
+            cmd.extend(server_arg_tokens(flag, value))
         if self.api_key is not None:
             cmd.extend(["--api-key", self.api_key])
         return cmd
@@ -110,6 +160,8 @@ class PDEndpoint(Endpoint):
     Call start/stop from Modal enter/exit hooks. Model-specific arguments and
     warmup stay in the recipe. The image must support NIXL and graceful shutdown.
     The ratio is (prefill, decode) counts; prefill ranks precede decode ranks.
+    The router is the build pinned in autoinference_utils.router; router_args
+    override its default flags, and allow_custom_router skips the pin check.
     """
 
     def __init__(
@@ -128,6 +180,9 @@ class PDEndpoint(Endpoint):
         router_health_request_timeout: int = 5,
         health_failure_threshold: int = 3,
         health_failure_timeout: float = 30,
+        router_args: Mapping[str, str] | None = None,
+        allow_custom_router: bool = False,
+        shutdown_port: int | None = None,
     ):
         from modal.experimental import get_cluster_info
 
@@ -156,6 +211,10 @@ class PDEndpoint(Endpoint):
             raise ValueError(
                 "ratio must contain positive integer prefill and decode counts"
             )
+        if managed := _managed_router_tokens(router_args or {}):
+            raise ValueError(
+                f"router_args cannot set PDEndpoint-managed flags: {managed}"
+            )
         roles = ("prefill",) * ratio[0] + ("decode",) * ratio[1]
         cluster = get_cluster_info()
         hosts = cluster.container_ips
@@ -164,6 +223,13 @@ class PDEndpoint(Endpoint):
                 f"P/D ratio {ratio} requires a {len(roles)}-container Modal cluster"
             )
         self._host_ip = hosts[cluster.rank]
+        self._hosts = hosts
+        self._rank = cluster.rank
+        self._shutdown_port = (
+            router_port + 1 if shutdown_port is None else shutdown_port
+        )
+        if not 1 <= self._shutdown_port <= 65535:
+            raise ValueError("shutdown_port must be between 1 and 65535")
         self.role = roles[cluster.rank]
         super().__init__(_url(hosts[0], router_port))
         self.engine = engine
@@ -199,16 +265,22 @@ class PDEndpoint(Endpoint):
                 health_interval=health_interval,
                 health_request_timeout=router_health_request_timeout,
                 health_failure_threshold=health_failure_threshold,
+                router_args=router_args or {},
             )
             if cluster.rank == 0
             else None
         )
+        # Fail before any weights load, not after a 20-minute engine boot.
+        if self.router is not None and not allow_custom_router:
+            verify_router()
         self.drain_timeout = drain_timeout
         self.health_interval = health_interval
         self.health_failure_timeout = health_failure_timeout
         self._stopped = threading.Event()
         self._lock = threading.Lock()
         self._started = False
+        self._serving = False
+        self._shutdown_listener: socket.socket | None = None
         self._threads: list[threading.Thread] = []
 
     def start(self, *, warmup: Callable[[], None] | None = None) -> None:
@@ -223,6 +295,13 @@ class PDEndpoint(Endpoint):
             SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION="1",
         )
         try:
+            if self._rank == 0:
+                self._shutdown_listener = socket.socket(socket.AF_INET6)
+                self._shutdown_listener.setsockopt(
+                    socket.SOL_SOCKET, socket.SO_REUSEADDR, 1
+                )
+                self._shutdown_listener.bind(("::", self._shutdown_port))
+                self._shutdown_listener.listen(len(self._hosts) - 1)
             for endpoint in (self.engine, self.router):
                 if endpoint is not None:
                     endpoint.start()
@@ -234,6 +313,7 @@ class PDEndpoint(Endpoint):
                 for _, host in self.router.pd_config:
                     self._watch_health(_url(host, self.engine.worker_port) + "/health")
                 self._watch_health(self.router.base_url + "/health")
+            self._serving = True
         except BaseException:
             self.stop()
             raise
@@ -295,7 +375,7 @@ class PDEndpoint(Endpoint):
             if self._stopped.is_set():
                 return
             self._stopped.set()
-        deadline = time.monotonic() + self.drain_timeout + 30
+        deadline = time.monotonic() + _SHUTDOWN_TIMEOUT
         try:
             for endpoint in (self.router, self.engine):
                 if endpoint is None:
@@ -307,9 +387,58 @@ class PDEndpoint(Endpoint):
                         process.wait(timeout=max(0, deadline - time.monotonic()))
                     except subprocess.TimeoutExpired:
                         process.kill()
+                        process.wait()
                 endpoint.stop()
         finally:
             for thread in self._threads:
-                thread.join(
-                    timeout=self.engine.health_request_timeout + self.health_interval
+                thread.join(timeout=max(0, deadline - time.monotonic()))
+            if self._serving:
+                if self._rank == 0:
+                    self._wait_for_followers(deadline)
+                else:
+                    self._notify_leader(deadline)
+            if self._shutdown_listener is not None:
+                self._shutdown_listener.close()
+                self._shutdown_listener = None
+
+    def _notify_leader(self, deadline: float) -> None:
+        while True:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                print("[pd] Failed to notify rank zero of shutdown", flush=True)
+                return
+            try:
+                with socket.create_connection(
+                    (self._hosts[0], self._shutdown_port), timeout=min(remaining, 1)
+                ) as connection:
+                    connection.sendall(struct.pack("!I", self._rank))
+                return
+            except OSError:
+                time.sleep(min(0.1, max(0, deadline - time.monotonic())))
+
+    def _wait_for_followers(self, deadline: float) -> None:
+        listener = self._shutdown_listener
+        if listener is None:
+            return
+        pending = set(range(1, len(self._hosts)))
+        while pending:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                print(
+                    f"[pd] Timed out waiting for ranks {sorted(pending)} to stop",
+                    flush=True,
                 )
+                return
+            listener.settimeout(min(remaining, 1))
+            try:
+                connection, _ = listener.accept()
+            except TimeoutError:
+                continue
+            with connection:
+                connection.settimeout(min(remaining, 1))
+                try:
+                    payload = connection.recv(4, socket.MSG_WAITALL)
+                except TimeoutError:
+                    continue
+            if len(payload) == 4:
+                pending.discard(struct.unpack("!I", payload)[0])

autoinference_utils-0.2.10/src/autoinference_utils/router.py

@@ -0,0 +1,103 @@
+"""The pinned SGLang model gateway (SMG) wheel that PDEndpoint routes through.
+
+This module is the source of truth for the P/D router. The wheel is built once
+from the fork commit below and published as a GitHub release asset; bumping
+the pin here and releasing autoinference-utils moves every recipe that takes
+the new release. Install it inside an image after autoinference-utils:
+
+    image.run_commands("python -m autoinference_utils.router install")
+"""
+
+from __future__ import annotations
+
+import json
+import subprocess
+import sys
+from importlib.metadata import Distribution, PackageNotFoundError, distribution
+
+ROUTER_REPOSITORY = "https://github.com/modal-labs/sglang"
+# Kimi K3 production router: HRRN decode preallocation (SMG_PD_DECODE_HRRN)
+# and brace-free metrics HELP text (kimi-k3-sglang#158). The wheel is the one
+# extracted from the Kimi K3 production image built at this commit.
+ROUTER_COMMIT = "e3e602e3ae46fa4da7ec105d5dc3efcfc85c226b"
+# The wheel is x86-64 Linux and links against glibc 2.39 (Ubuntu 24.04, which
+# the lmsysorg/sglang and vllm base images use); uv refuses it on older images.
+ROUTER_WHEEL_URL = (
+    f"{ROUTER_REPOSITORY}/releases/download/smg-router-{ROUTER_COMMIT[:8]}/"
+    "sglang_router-0.3.2-cp38-abi3-manylinux_2_39_x86_64.whl"
+)
+ROUTER_WHEEL_SHA256 = "5b867c6c308a609395a6f5c146023274f570b7187b3a7b92ebb11658da52a6c7"
+
+
+def router_requirement() -> str:
+    """The hash-pinned direct requirement for the router wheel."""
+    return f"sglang-router @ {ROUTER_WHEEL_URL}#sha256={ROUTER_WHEEL_SHA256}"
+
+
+def verify_router(dist: Distribution | None = None) -> None:
+    """Raise unless the installed sglang-router came from the pinned wheel.
+
+    Checks the installer's PEP 610 record: the URL must match, and when the
+    installer recorded a digest it must match too. uv leaves archive_info empty
+    and enforces the digest at install time through the URL's sha256 fragment.
+    """
+    try:
+        dist = dist or distribution("sglang-router")
+        origin = json.loads(dist.read_text("direct_url.json") or "{}")
+    except PackageNotFoundError:
+        origin = None
+    except ValueError:
+        origin = {}
+    if origin is None:
+        found = "no sglang-router is installed"
+    else:
+        recorded = ((origin.get("archive_info") or {}).get("hashes") or {}).get(
+            "sha256"
+        )
+        if origin.get("url") == ROUTER_WHEEL_URL and recorded in (
+            None,
+            ROUTER_WHEEL_SHA256,
+        ):
+            return
+        found = f"installed sglang-router is {origin.get('url') or 'not a direct URL install'}"
+        if recorded is not None:
+            found += f" (sha256 {recorded})"
+    raise RuntimeError(
+        f"{found}; expected {ROUTER_WHEEL_URL} (sha256 {ROUTER_WHEEL_SHA256}). "
+        "Run `python -m autoinference_utils.router install` in the image after "
+        "autoinference-utils, or pass allow_custom_router=True."
+    )
+
+
+def install() -> None:
+    """Install the pinned wheel over whatever router the image ships."""
+    # --no-deps keeps the engine image's own web stack untouched.
+    subprocess.run(
+        [
+            "uv",
+            "pip",
+            "install",
+            "--python",
+            sys.executable,
+            "--no-cache",
+            "--no-deps",
+            "--reinstall-package",
+            "sglang-router",
+            router_requirement(),
+        ],
+        check=True,
+    )
+    verify_router()
+    # --no-deps means a pin whose dependencies the image lacks must fail here,
+    # at image build, not at container start.
+    subprocess.run(
+        [sys.executable, "-m", "sglang_router.launch_router", "--help"],
+        check=True,
+        stdout=subprocess.DEVNULL,
+    )
+
+
+if __name__ == "__main__":
+    if sys.argv[1:] != ["install"]:
+        raise SystemExit("usage: python -m autoinference_utils.router install")
+    install()