autoinference-utils 0.2.7__tar.gz → 0.2.9__tar.gz

{autoinference_utils-0.2.7 → autoinference_utils-0.2.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: autoinference-utils
-Version: 0.2.7
+Version: 0.2.9
 Summary: Shared endpoint abstractions for autoinference deployments
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
@@ -9,12 +9,6 @@ Description-Content-Type: text/markdown
 
 Shared endpoint abstractions (`SGLangEndpoint`, `VLLMEndpoint`) for autoinference deployments.
 
-For SGLang models whose custom Transformers code cannot load from symlinked
-Hugging Face snapshots, set `AUTOINFERENCE_MATERIALIZE_MODEL_PATH=1`. The endpoint
-copies local base/draft model code and metadata into unique temporary directories
-and symlinks top-level safetensors weights. Repository IDs pass through unchanged.
-Leave the setting unset to preserve the original paths.
-
 ## Publishing
 
 Bump `version` in `pyproject.toml` and merge to `main` — CI publishes automatically via trusted publishing.

autoinference_utils-0.2.9/README.md

@@ -0,0 +1,9 @@
+# autoinference-utils
+
+Shared endpoint abstractions (`SGLangEndpoint`, `VLLMEndpoint`) for autoinference deployments.
+
+## Publishing
+
+Bump `version` in `pyproject.toml` and merge to `main` — CI publishes automatically via trusted publishing.
+
+To publish manually: `cd autoinference_utils && uv build && uv publish`

{autoinference_utils-0.2.7 → autoinference_utils-0.2.9}/pyproject.toml

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

{autoinference_utils-0.2.7 → autoinference_utils-0.2.9}/src/autoinference_utils/endpoint.py

@@ -23,8 +23,6 @@ import urllib.request
 from abc import ABC
 from typing import Any, Callable, Literal, Mapping, Optional, Sequence
 
-from autoinference_utils.model_path import materialize_model_path
-
 BENCH_MODE_ENV = "ENDPOINT_BENCH_MODE"
 DUMMY_WEIGHTS_ENV = "ENDPOINT_DUMMY"
 BENCH_MODE_PORT_OFFSET = 10000
@@ -35,6 +33,14 @@ MODAL_SESSION_ID_HEADER = "modal-session-id"
 ENDPOINTS_REQUIRING_OAI_STREAM_COMPAT: tuple[str, ...] = ("TRTLLMEndpoint",)
 
 
+def _url(host: str, port: int) -> str:
+    return f"http://{'[' + host + ']' if ':' in host else host}:{port}"
+
+
+def _local_base_url(host: str, port: int) -> str:
+    return _url({"0.0.0.0": "127.0.0.1", "::": "::1"}.get(host, host), port)
+
+
 def server_arg_tokens(flag: str, value: str) -> list[str]:
     """Render one server arg as argv tokens.
 
@@ -150,18 +156,16 @@ class SGLangEndpoint(Endpoint):
         sglang_port = (
             worker_port + BENCH_MODE_PORT_OFFSET if self.bench_mode else worker_port
         )
-        super().__init__(base_url=f"http://localhost:{sglang_port}")
+        host_arg = get_server_arg(extra_server_args or {}, "--host")
+        host = (host_arg[1] or host_arg[0].partition("=")[2]) if host_arg else "0.0.0.0"
+        super().__init__(base_url=_local_base_url(host, sglang_port))
 
         self.worker_port = sglang_port
-        self.model_path = materialize_model_path(model_path)
+        self.model_path = model_path
         self.tp = tp
         self.ep = ep
         self.dp = dp
-        self.speculative_model_path = (
-            materialize_model_path(speculative_model_path)
-            if speculative_model_path is not None
-            else None
-        )
+        self.speculative_model_path = speculative_model_path
         self.load_format = load_format
         self.nnodes = nnodes
         self.node_rank = node_rank
@@ -286,7 +290,7 @@ class SGLangEndpoint(Endpoint):
         return cmd
 
     def health_check(self) -> str | None:
-        url = f"http://127.0.0.1:{self.worker_port}/health"
+        url = f"{self.base_url}/health"
         return _health_check(
             url,
             request_timeout=self.health_request_timeout,
@@ -300,6 +304,7 @@ class SGLangEndpoint(Endpoint):
         wait_ready(
             self._proc,
             port=self.worker_port,
+            base_url=self.base_url,
             timeout=self.health_timeout,
             poll_interval=self.health_poll_interval,
             request_timeout=self.health_request_timeout,
@@ -308,6 +313,7 @@ class SGLangEndpoint(Endpoint):
             self._bench_server = start_bench_proxy(
                 listen_port=self.listen_port,
                 upstream_port=self.worker_port,
+                upstream_base_url=self.base_url,
             )
 
     def stop(self):
@@ -474,13 +480,16 @@ class RouterEndpoint(Endpoint):
         api_key: Optional[str] = None,
         health_timeout: float = 10 * 60,
         health_poll_interval: float = 5.0,
+        health_path: str = "/health",
+        host: str = "0.0.0.0",
     ):
         self.bench_mode = os.environ.get(BENCH_MODE_ENV) == "1"
         self.listen_port = router_port
         actual_router_port = (
             router_port + BENCH_MODE_PORT_OFFSET if self.bench_mode else router_port
         )
-        super().__init__(base_url=f"http://localhost:{actual_router_port}")
+        self.host = host
+        super().__init__(base_url=_local_base_url(host, actual_router_port))
         self.pd_config = list(pd_config)
         self.worker_port = (
             worker_port + BENCH_MODE_PORT_OFFSET if self.bench_mode else worker_port
@@ -490,6 +499,7 @@ class RouterEndpoint(Endpoint):
         self.api_key = api_key
         self.health_timeout = health_timeout
         self.health_poll_interval = health_poll_interval
+        self.health_path = health_path
         self._proc: Optional[subprocess.Popen] = None
         self._bench_server: Optional[http.server.ThreadingHTTPServer] = None
 
@@ -499,7 +509,7 @@ class RouterEndpoint(Endpoint):
             "-m",
             "sglang_router.launch_router",
             "--host",
-            "0.0.0.0",
+            f"[{self.host}]" if ":" in self.host else self.host,
             "--port",
             str(self.router_port),
             "--prefill-policy",
@@ -528,7 +538,7 @@ class RouterEndpoint(Endpoint):
             cmd.append("--pd-disaggregation")
 
         for role, node_ip in self.pd_config:
-            node_url = f"http://{node_ip}:{self.worker_port}"
+            node_url = _url(node_ip, self.worker_port)
             if role == "prefill":
                 cmd.extend(["--prefill", node_url, str(self.prefill_bootstrap_port)])
             elif role == "decode":
@@ -543,7 +553,7 @@ class RouterEndpoint(Endpoint):
     def start(self):
         for _, node_ip in self.pd_config:
             _wait_ready_url(
-                f"http://{node_ip}:{self.worker_port}/health",
+                f"{_url(node_ip, self.worker_port)}/health",
                 timeout=self.health_timeout,
                 poll_interval=self.health_poll_interval,
             )
@@ -551,17 +561,25 @@ class RouterEndpoint(Endpoint):
         cmd = self._build_cmd()
         print(f"[router] starting: {shlex.join(cmd)}")
         self._proc = subprocess.Popen(cmd)
-        _wait_ready_url(
-            f"http://localhost:{self.router_port}/health",
-            timeout=self.health_timeout,
-            poll_interval=self.health_poll_interval,
-        )
+        self.wait_ready()
         if self.bench_mode:
             self._bench_server = start_bench_proxy(
                 listen_port=self.listen_port,
                 upstream_port=self.router_port,
+                upstream_base_url=self.base_url,
             )
 
+    def wait_ready(self) -> None:
+        assert self._proc is not None
+        wait_ready(
+            self._proc,
+            port=self.router_port,
+            base_url=self.base_url,
+            health_path=self.health_path,
+            timeout=self.health_timeout,
+            poll_interval=self.health_poll_interval,
+        )
+
     def stop(self):
         if self._bench_server is not None:
             self._bench_server.shutdown()
@@ -579,11 +597,12 @@ def start_bench_proxy(
     *,
     listen_port: int,
     upstream_port: int,
+    upstream_base_url: str | None = None,
 ) -> http.server.ThreadingHTTPServer:
     handler_cls = type(
         "BenchProxyHandler",
         (_BenchProxyHandler,),
-        {"upstream_port": upstream_port},
+        {"upstream_port": upstream_port, "upstream_base_url": upstream_base_url},
     )
     server = http.server.ThreadingHTTPServer(("0.0.0.0", listen_port), handler_cls)
     thread = threading.Thread(
@@ -596,6 +615,7 @@ def start_bench_proxy(
 
 class _BenchProxyHandler(http.server.BaseHTTPRequestHandler):
     upstream_port: int = 8000
+    upstream_base_url: str | None = None
 
     def log_message(self, format, *args):
         return
@@ -623,7 +643,11 @@ class _BenchProxyHandler(http.server.BaseHTTPRequestHandler):
 
         benchmark = payload.get("benchmark") or ""
         args = payload.get("args") or []
-        target = payload.get("target") or f"http://localhost:{self.upstream_port}"
+        target = (
+            payload.get("target")
+            or self.upstream_base_url
+            or f"http://localhost:{self.upstream_port}"
+        )
         output_dir = payload.get("output_dir") or "/tmp/bench-output"
 
         if not benchmark:
@@ -636,7 +660,8 @@ class _BenchProxyHandler(http.server.BaseHTTPRequestHandler):
         self._send_raw(status, "application/json", body)
 
     def _proxy(self, method: str):
-        url = f"http://localhost:{self.upstream_port}{self.path}"
+        base_url = self.upstream_base_url or f"http://localhost:{self.upstream_port}"
+        url = f"{base_url}{self.path}"
         length = int(self.headers.get("Content-Length", 0))
         body = self.rfile.read(length) if length else None
         forward_headers = {
@@ -759,12 +784,13 @@ def wait_ready(
     port: int,
     timeout: float,
     health_path: str = "/health",
+    base_url: str | None = None,
     poll_interval: float = 5.0,
     request_timeout: float = 5.0,
 ) -> None:
     """Poll an HTTP health endpoint until ready, raising if the process dies."""
     deadline = time.time() + timeout
-    url = f"http://127.0.0.1:{port}{health_path}"
+    url = f"{base_url or _local_base_url('127.0.0.1', port)}{health_path}"
     last_error = "no response yet"
 
     while time.time() < deadline:
@@ -794,6 +820,7 @@ def wait_ready(
 def warmup_chat_completions(
     *,
     port: int,
+    base_url: str | None = None,
     payload: Mapping[str, Any],
     headers: Mapping[str, str] | None = None,
     successful_requests: int = 3,
@@ -802,7 +829,7 @@ def warmup_chat_completions(
     retry_delay: float = 1.0,
 ) -> None:
     """Warm the OpenAI chat completions endpoint with strict retries."""
-    url = f"http://127.0.0.1:{port}/v1/chat/completions"
+    url = f"{base_url or _local_base_url('127.0.0.1', port)}/v1/chat/completions"
     request_headers = {"Content-Type": "application/json"}
     if headers:
         request_headers.update(headers)

autoinference_utils-0.2.9/src/autoinference_utils/pd.py

@@ -0,0 +1,372 @@
+"""SGLang p/d disaggregated serving in a multi-container static cluster."""
+
+from __future__ import annotations
+
+import asyncio
+import math
+import os
+import subprocess
+import threading
+import time
+from collections.abc import Callable, Mapping
+
+from .endpoint import (
+    Endpoint,
+    RouterEndpoint,
+    SGLangEndpoint,
+    _flag_name,
+    _health_check,
+    _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",
+    }
+)
+
+
+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:
+    if os.environ.get("FI_PROVIDER") == "efa":
+        os.environ["SGLANG_DISAGGREGATION_NIXL_BACKEND"] = "LIBFABRIC"
+    elif hca := os.environ.get("NCCL_IB_HCA", ""):
+        rails = [name for name in hca.lstrip("=").split(",") if name]
+        os.environ["UCX_NET_DEVICES"] = ",".join(f"{name}:1" for name in rails)
+
+
+async def _stop_container() -> None:
+    from modal.client import _Client
+    from modal_proto import api_pb2
+
+    async def stop():
+        client = await _Client.from_env()
+        await client.stub.ContainerStop(
+            api_pb2.ContainerStopRequest(task_id=os.environ["MODAL_TASK_ID"])
+        )
+
+    await asyncio.wait_for(stop(), timeout=3)
+
+
+def _terminate_container() -> None:
+    from modal._utils.async_utils import synchronize_api
+
+    try:
+        synchronize_api(_stop_container)()
+    finally:
+        os._exit(1)
+
+
+class _PDRouter(RouterEndpoint):
+    def __init__(
+        self,
+        *,
+        max_concurrent_requests,
+        drain_timeout,
+        health_interval,
+        health_request_timeout,
+        health_failure_threshold,
+        router_args,
+        **kwargs,
+    ):
+        super().__init__(health_path="/readiness", host="::", **kwargs)
+        self.max_concurrent_requests = max_concurrent_requests
+        self.drain_timeout = drain_timeout
+        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 = [
+            "python",
+            "-m",
+            "sglang_router.launch_router",
+            "--host",
+            "[::]",
+            "--port",
+            str(self.router_port),
+            "--pd-disaggregation",
+        ]
+        for role, host in self.pd_config:
+            cmd.extend([f"--{role}", _url(host, self.worker_port)])
+            if role == "prefill":
+                cmd.append(str(self.prefill_bootstrap_port))
+        defaults = {
+            "--prefill-policy": "cache_aware",
+            "--decode-policy": "round_robin",
+            "--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",
+        }
+        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
+
+
+class PDEndpoint(Endpoint):
+    """Own a local SGLang engine and the rank-zero router for P/D serving.
+
+    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__(
+        self,
+        engine: SGLangEndpoint,
+        *,
+        ratio: tuple[int, int],
+        prefill_args: Mapping[str, str] | None = None,
+        decode_args: Mapping[str, str] | None = None,
+        router_port: int = 9000,
+        max_concurrent_requests: int = 96,
+        api_key: str | None = None,
+        drain_timeout: int = 120,
+        health_interval: int = 10,
+        health_request_timeout: int = 10,
+        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,
+    ):
+        from modal.experimental import get_cluster_info
+
+        if not isinstance(engine, SGLangEndpoint):
+            raise TypeError("PDEndpoint requires a SGLangEndpoint")
+        if engine._proc is not None or engine.bench_mode:
+            raise ValueError("PDEndpoint requires an unstarted, non-benchmark engine")
+        if (
+            min(
+                health_interval,
+                health_request_timeout,
+                router_health_request_timeout,
+                health_failure_threshold,
+            )
+            <= 0
+            or drain_timeout < 0
+            or not math.isfinite(health_failure_timeout)
+            or health_failure_timeout <= 0
+        ):
+            raise ValueError(
+                "health settings must be positive and drain_timeout nonnegative"
+            )
+        if len(ratio) != 2 or any(
+            type(count) is not int or count < 1 for count in ratio
+        ):
+            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
+        if len(hosts) != len(roles) or not 0 <= cluster.rank < len(roles):
+            raise ValueError(
+                f"P/D ratio {ratio} requires a {len(roles)}-container Modal cluster"
+            )
+        self._host_ip = hosts[cluster.rank]
+        self.role = roles[cluster.rank]
+        super().__init__(_url(hosts[0], router_port))
+        self.engine = engine
+        engine.disaggregation_mode = self.role
+        engine.health_request_timeout = health_request_timeout
+        engine.extra_server_args = _merge_server_args(
+            engine.extra_server_args,
+            (prefill_args if self.role == "prefill" else decode_args) or {},
+        )
+        engine.extra_server_args = _merge_server_args(
+            engine.extra_server_args,
+            {
+                "--disaggregation-transfer-backend": "nixl",
+                "--host": "::",
+            },
+        )
+        engine.extra_server_args = {
+            key: value
+            for key, value in engine.extra_server_args.items()
+            if _flag_name(key) != "--disaggregation-mode"
+        }
+        engine.base_url = _local_base_url("::", engine.worker_port)
+        self.router = (
+            _PDRouter(
+                pd_config=list(zip(roles, hosts)),
+                worker_port=engine.worker_port,
+                router_port=router_port,
+                prefill_bootstrap_port=engine.prefill_bootstrap_port,
+                health_timeout=engine.health_timeout,
+                api_key=api_key,
+                max_concurrent_requests=max_concurrent_requests,
+                drain_timeout=drain_timeout,
+                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._threads: list[threading.Thread] = []
+
+    def start(self, *, warmup: Callable[[], None] | None = None) -> None:
+        """Run warmup on rank zero before supervising every worker's health."""
+        if self._started or self._stopped.is_set():
+            raise RuntimeError("PDEndpoint can only be started once")
+        self._started = True
+        os.environ.setdefault("SGLANG_HOST_IP", self._host_ip)
+        _configure_fabric()
+        os.environ.update(
+            SGLANG_GRACEFUL_SHUTDOWN_TIMEOUT=str(self.drain_timeout),
+            SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION="1",
+        )
+        try:
+            for endpoint in (self.engine, self.router):
+                if endpoint is not None:
+                    endpoint.start()
+                    self._watch_process(endpoint)
+            if self.router is not None:
+                if warmup is not None:
+                    warmup()
+                    self.router.wait_ready()
+                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")
+        except BaseException:
+            self.stop()
+            raise
+
+    def _watch_process(self, endpoint: SGLangEndpoint | RouterEndpoint) -> None:
+        process = endpoint._proc
+        assert process is not None
+
+        def wait():
+            code = process.wait()
+            if not self._stopped.is_set():
+                print(f"[pd] Serving process exited with code {code}", flush=True)
+                self._fail()
+
+        thread = threading.Thread(target=wait, daemon=True)
+        self._threads.append(thread)
+        thread.start()
+
+    def _watch_health(self, url: str) -> None:
+        def poll():
+            failed_since = None
+            while not self._stopped.wait(self.health_interval):
+                started = time.monotonic()
+                error = _health_check(
+                    url, request_timeout=self.engine.health_request_timeout
+                )
+                if error is None:
+                    failed_since = None
+                    continue
+                if failed_since is None:
+                    failed_since = started
+                elapsed = time.monotonic() - failed_since
+                print(f"[pd] {url}: {error}; unhealthy for {elapsed:.1f}s", flush=True)
+                if elapsed >= self.health_failure_timeout:
+                    self._fail()
+                    return
+
+        thread = threading.Thread(target=poll, daemon=True)
+        self._threads.append(thread)
+        thread.start()
+
+    def _fail(self) -> None:
+        with self._lock:
+            if self._stopped.is_set():
+                return
+            self._stopped.set()
+        try:
+            for endpoint in (self.router, self.engine):
+                if endpoint is not None and endpoint._proc is not None:
+                    try:
+                        endpoint._proc.kill()
+                    except ProcessLookupError:
+                        pass
+        finally:
+            _terminate_container()
+
+    def stop(self) -> None:
+        with self._lock:
+            if self._stopped.is_set():
+                return
+            self._stopped.set()
+        deadline = time.monotonic() + self.drain_timeout + 30
+        try:
+            for endpoint in (self.router, self.engine):
+                if endpoint is None:
+                    continue
+                process = endpoint._proc
+                if process is not None and process.poll() is None:
+                    process.terminate()
+                    try:
+                        process.wait(timeout=max(0, deadline - time.monotonic()))
+                    except subprocess.TimeoutExpired:
+                        process.kill()
+                endpoint.stop()
+        finally:
+            for thread in self._threads:
+                thread.join(
+                    timeout=self.engine.health_request_timeout + self.health_interval
+                )

autoinference_utils-0.2.9/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()

autoinference_utils-0.2.7/README.md

@@ -1,15 +0,0 @@
-# autoinference-utils
-
-Shared endpoint abstractions (`SGLangEndpoint`, `VLLMEndpoint`) for autoinference deployments.
-
-For SGLang models whose custom Transformers code cannot load from symlinked
-Hugging Face snapshots, set `AUTOINFERENCE_MATERIALIZE_MODEL_PATH=1`. The endpoint
-copies local base/draft model code and metadata into unique temporary directories
-and symlinks top-level safetensors weights. Repository IDs pass through unchanged.
-Leave the setting unset to preserve the original paths.
-
-## Publishing
-
-Bump `version` in `pyproject.toml` and merge to `main` — CI publishes automatically via trusted publishing.
-
-To publish manually: `cd autoinference_utils && uv build && uv publish`

autoinference_utils-0.2.7/src/autoinference_utils/model_path.py

@@ -1,38 +0,0 @@
-"""Prepare local model snapshots for Transformers' remote-code loader."""
-
-import os
-import shutil
-import tempfile
-from pathlib import Path
-
-MATERIALIZE_MODEL_PATH_ENV = "AUTOINFERENCE_MATERIALIZE_MODEL_PATH"
-
-
-def materialize_model_path(path: str) -> str:
-    """Copy snapshot metadata/code into a real directory; retain weight links.
-
-    Transformers can resolve Python-module symlinks into the HF blob directory,
-    losing the relative imports next to the original snapshot entry. This opt-in
-    keeps small files together without copying large safetensors shards. The
-    unique directory lasts for the container's lifetime.
-    """
-    source = Path(path)
-    if os.environ.get(MATERIALIZE_MODEL_PATH_ENV) != "1" or not source.is_dir():
-        return path
-
-    source = source.resolve()
-    destination = Path(tempfile.mkdtemp(prefix="autoinference-model-"))
-    try:
-        for entry in source.iterdir():
-            target = destination / entry.name
-            if entry.is_file() and entry.suffix == ".safetensors":
-                target.symlink_to(entry.resolve(strict=True))
-            elif entry.is_dir():
-                shutil.copytree(entry, target, symlinks=False)
-            else:
-                shutil.copy2(entry, target, follow_symlinks=True)
-    except BaseException:
-        shutil.rmtree(destination)
-        raise
-    print(f"Materialized model snapshot {source} -> {destination}")
-    return str(destination)

autoinference_utils-0.2.7/src/autoinference_utils/tests/test_model_path.py

@@ -1,84 +0,0 @@
-import shutil
-from pathlib import Path
-
-import pytest
-
-from autoinference_utils.endpoint import SGLangEndpoint
-from autoinference_utils.model_path import (
-    MATERIALIZE_MODEL_PATH_ENV,
-    materialize_model_path,
-)
-
-
-def snapshot(tmp_path):
-    blobs = tmp_path / "blobs"
-    source = tmp_path / "snapshots" / "revision"
-    blobs.mkdir()
-    source.mkdir(parents=True)
-    for name, content in {
-        "config.json": '{"model_type": "kimi"}',
-        "configuration_kimi.py": "from .helper import VALUE\n",
-        "helper.py": "VALUE = 42\n",
-        "model.safetensors": "weight placeholder",
-    }.items():
-        blob = blobs / ("hash-" + name)
-        blob.write_text(content)
-        (source / name).symlink_to(blob)
-    return source
-
-
-def test_opt_in_and_repository_ids(tmp_path, monkeypatch):
-    source = snapshot(tmp_path)
-    monkeypatch.delenv(MATERIALIZE_MODEL_PATH_ENV, raising=False)
-    assert materialize_model_path(str(source)) == str(source)
-    monkeypatch.setenv(MATERIALIZE_MODEL_PATH_ENV, "1")
-    assert materialize_model_path("moonshotai/Kimi-K3") == "moonshotai/Kimi-K3"
-
-
-def test_local_code_is_materialized_and_weights_stay_linked(tmp_path, monkeypatch):
-    source = snapshot(tmp_path)
-    monkeypatch.setenv(MATERIALIZE_MODEL_PATH_ENV, "1")
-    monkeypatch.setattr("tempfile.tempdir", str(tmp_path))
-    result = Path(materialize_model_path(str(source)))
-    assert result != source
-    for name in ("config.json", "configuration_kimi.py", "helper.py"):
-        assert not (result / name).is_symlink()
-        assert (result / name).resolve().parent == result
-        assert (result / name).read_bytes() == (source / name).read_bytes()
-    assert (result / "model.safetensors").is_symlink()
-    assert (result / "model.safetensors").resolve() == (
-        source / "model.safetensors"
-    ).resolve()
-    assert (source / "configuration_kimi.py").is_symlink()
-
-
-def test_endpoint_uses_materialized_base_and_draft(tmp_path, monkeypatch):
-    source = snapshot(tmp_path)
-    monkeypatch.setenv(MATERIALIZE_MODEL_PATH_ENV, "1")
-    monkeypatch.setattr("tempfile.tempdir", str(tmp_path))
-    endpoint = SGLangEndpoint(
-        model_path=str(source), speculative_model_path=str(source)
-    )
-    assert endpoint.model_path != str(source)
-    assert endpoint.speculative_model_path != str(source)
-    assert endpoint.model_path != endpoint.speculative_model_path
-    cmd = endpoint._build_cmd()
-    assert cmd[cmd.index("--model-path") + 1] == endpoint.model_path
-    assert (
-        cmd[cmd.index("--speculative-draft-model-path") + 1]
-        == endpoint.speculative_model_path
-    )
-
-
-def test_failed_copy_removes_partial_directory(tmp_path, monkeypatch):
-    source = snapshot(tmp_path)
-    monkeypatch.setenv(MATERIALIZE_MODEL_PATH_ENV, "1")
-    monkeypatch.setattr("tempfile.tempdir", str(tmp_path))
-
-    def fail(*args, **kwargs):
-        raise OSError("copy failed")
-
-    monkeypatch.setattr(shutil, "copy2", fail)
-    with pytest.raises(OSError, match="copy failed"):
-        materialize_model_path(str(source))
-    assert not list(tmp_path.glob("autoinference-model-*"))