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

{autoinference_utils-0.2.6 → autoinference_utils-0.2.8}/.gitignore

@@ -150,6 +150,7 @@ activemq-data/
 
 # Environments
 .env
+.env.dev
 .envrc
 .venv
 .autoinference/
@@ -241,6 +242,22 @@ Network Trash Folder
 Temporary Items
 .apdisk
 
-# Ephemeral files from Experiments
+# Ephemeral files from benchmark runs
 benchmark_results/
 src/autoinference/deployments/
+
+# Modal Skills
+.agents/skills/modal/
+.agents/skills/.modal-skill-*
+
+# cook/etc is an extension point: local tools live there untracked.
+packages/cook/src/cook/etc/*
+# Local tools that existed before /etc are explicitly added back
+!packages/cook/src/cook/etc/__init__.py
+!packages/cook/src/cook/etc/tests/
+!packages/cook/src/cook/etc/inspect_hf_model.py
+!packages/cook/src/cook/etc/download.py
+!packages/cook/src/cook/etc/quantize.py
+!packages/cook/src/cook/etc/graft_mtp.py
+!packages/cook/src/cook/etc/tune_moe.py
+!packages/cook/src/cook/etc/tool_calling.py

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

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

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

@@ -1,6 +1,6 @@
 [project]
 name = "autoinference-utils"
-version = "0.2.6"
+version = "0.2.8"
 description = "Shared endpoint abstractions for autoinference deployments"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -10,7 +10,7 @@ requires = ["hatchling"]
 build-backend = "hatchling.build"
 
 [tool.hatch.build.targets.wheel]
-packages = ["autoinference_utils"]
+packages = ["src/autoinference_utils"]
 
 [tool.hatch.build.targets.sdist]
-include = ["autoinference_utils", "README.md", "pyproject.toml"]
+include = ["src", "README.md", "pyproject.toml"]

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

@@ -33,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,7 +158,9 @@ 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 = model_path
@@ -277,7 +287,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,
@@ -291,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,
@@ -299,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):
@@ -462,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
@@ -478,13 +493,14 @@ 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",
+            "--host", f"[{self.host}]" if ":" in self.host else self.host,
             "--port", str(self.router_port),
             "--prefill-policy", "cache_aware",
             "--decode-policy", "round_robin",
@@ -504,7 +520,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)]
@@ -521,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,
             )
@@ -529,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()
@@ -556,11 +580,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
@@ -578,6 +603,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
@@ -609,6 +635,7 @@ class _BenchProxyHandler(http.server.BaseHTTPRequestHandler):
         args = payload.get("args") or []
         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"
@@ -627,7 +654,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 = {
@@ -749,12 +777,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:
@@ -788,6 +817,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,
@@ -796,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)

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
+                )