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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: autoinference-utils
-Version: 0.2.7
+Version: 0.2.8
 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.8/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.8}/pyproject.toml

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

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

@@ -23,18 +23,24 @@ 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
 
-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",)
 
 
+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.
 
@@ -57,7 +63,9 @@ 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:
@@ -150,18 +158,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
@@ -170,7 +176,9 @@ 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
@@ -185,16 +193,17 @@ 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
@@ -205,9 +214,7 @@ 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(
@@ -229,21 +236,19 @@ 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:
@@ -268,25 +273,21 @@ 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))
 
         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 +301,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 +310,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):
@@ -342,7 +345,9 @@ 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
@@ -355,15 +360,10 @@ 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,7 +417,9 @@ 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
@@ -427,10 +429,8 @@ 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))
@@ -474,13 +474,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,35 +493,24 @@ 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
 
     def _build_cmd(self) -> list[str]:
         cmd = [
-            "python",
-            "-m",
-            "sglang_router.launch_router",
-            "--host",
-            "0.0.0.0",
-            "--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:
@@ -528,9 +520,11 @@ 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)])
+                cmd.extend(
+                    ["--prefill", node_url, str(self.prefill_bootstrap_port)]
+                )
             elif role == "decode":
                 cmd.extend(["--decode", node_url])
             elif role == "worker":
@@ -543,7 +537,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 +545,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()
@@ -574,28 +576,34 @@ 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,
     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
     )
-    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}, forwarding to :{upstream_port}")
+    print(
+        f"[bench-proxy] listening on :{listen_port}, "
+        f"forwarding to :{upstream_port}"
+    )
     return server
 
 
 class _BenchProxyHandler(http.server.BaseHTTPRequestHandler):
     upstream_port: int = 8000
+    upstream_base_url: str | None = None
 
     def log_message(self, format, *args):
         return
@@ -619,24 +627,35 @@ 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 []
-        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:
-            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
         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 = {
@@ -680,7 +699,9 @@ 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)
@@ -705,21 +726,18 @@ 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,
@@ -759,17 +777,20 @@ 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:
         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 "
@@ -783,17 +804,20 @@ def wait_ready(
         time.sleep(poll_interval)
 
     print(
-        f"[endpoint] !!! health check timed out after {timeout}s; last={last_error}",
+        f"[endpoint] !!! health check timed out after {timeout}s; "
+        f"last={last_error}",
         flush=True,
     )
     raise TimeoutError(
-        f"Health check timed out after {timeout}s for {url}. Last error: {last_error}"
+        f"Health check timed out after {timeout}s for {url}. "
+        f"Last error: {last_error}"
     )
 
 
 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 +826,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)
@@ -815,7 +839,9 @@ 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}"
+            ),
         )
 
 
@@ -837,7 +863,9 @@ 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:
@@ -900,7 +928,10 @@ def start_heartbeat_thread(
             try:
                 error = health_check_fn()
             except subprocess.CalledProcessError as exc:
-                print(f"[heartbeat] server process exited with code {exc.returncode}")
+                print(
+                    f"[heartbeat] server process exited with code "
+                    f"{exc.returncode}"
+                )
                 on_failure()
                 return
             if error is None:
@@ -912,7 +943,10 @@ 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
 
@@ -997,7 +1031,8 @@ def _wait_ready_url(
         last_error = error
         time.sleep(poll_interval)
     raise TimeoutError(
-        f"Timed out after {timeout}s waiting for {url}. Last error: {last_error}"
+        f"Timed out after {timeout}s waiting for {url}. "
+        f"Last error: {last_error}"
     )
 
 
@@ -1015,7 +1050,9 @@ 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/src/autoinference_utils/pd.py

@@ -0,0 +1,315 @@
+"""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,
+)
+
+
+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,
+        **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
+
+    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))
+        for flag, value in {
+            "--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,
+            "--log-level": "info",
+        }.items():
+            cmd.extend([flag, str(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.
+    """
+
+    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,
+    ):
+        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"
+            )
+        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,
+            )
+            if cluster.rank == 0
+            else None
+        )
+        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.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-*"))