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

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

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

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

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

autoinference_utils-0.2.11/src/autoinference_utils/cluster/__init__.py

@@ -0,0 +1 @@
+"""Utilities for app-local inference clusters."""

autoinference_utils-0.2.11/src/autoinference_utils/cluster/_net.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+import socket
+
+
+def i6pn_address() -> str:
+    return socket.getaddrinfo("i6pn.modal.local", None, socket.AF_INET6)[0][4][0]
+
+
+def endpoint(host: str, port: int) -> str:
+    return f"[{host}]:{port}" if ":" in host else f"{host}:{port}"
+
+
+def dialable(address: str) -> str:
+    """Bracket bare IPv6 endpoints for gRPC."""
+    host, _, port = address.rpartition(":")
+    if not host or not port.isdigit():
+        return address
+    return endpoint(host.strip("[]"), int(port))

autoinference_utils-0.2.11/src/autoinference_utils/cluster/config.py

@@ -0,0 +1,11 @@
+"""Shared cluster discovery configuration."""
+
+from __future__ import annotations
+
+DICT_NAME_ENV = "AUTOINFERENCE_CLUSTER_DICT"
+SCOPE_ENV = "AUTOINFERENCE_CLUSTER_SCOPE"
+TTL_ENV = "AUTOINFERENCE_CLUSTER_TTL"
+
+DEFAULT_DICT_NAME = "autoinf-cluster"
+
+DEFAULT_TTL_SECONDS = 120.0

autoinference_utils-0.2.11/src/autoinference_utils/cluster/image.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+import modal
+
+_ABSL_REPOSITORY = "https://github.com/abseil/abseil-cpp.git"
+_UCX_REPOSITORY = "https://github.com/openucx/ucx.git"
+_LIBFABRIC_REPOSITORY = "https://github.com/ofiwg/libfabric.git"
+
+
+NIXL_REPOSITORY = "https://github.com/ai-dynamo/nixl.git"
+NIXL_COMMIT = "26eabeb5667b386642298af948811ef77355d1d8"
+# Upstream's version; Ubuntu's libabsl-dev lacks absl_log.
+ABSL_TAG = "lts_2025_08_14"
+LIBFABRIC_TAG = "v1.21.0"
+# Match the UCX ABI already loaded by the engine images.
+UCX_COMMIT = "5288e74cd40be622e109489de7bb80d821341bd5"
+
+_NIXL_SOURCE_PATH = "/opt/nixl"
+
+# UCX serves mlx5 hosts; libfabric serves EFA hosts.
+_NIXL_MESON_ARGS = (
+    "-Csetup-args=-Denable_plugins=UCX,LIBFABRIC "
+    "-Csetup-args=-Ducx_path=/usr/local/ucx "
+    "-Csetup-args=-Dlibfabric_path=/usr/local "
+    "-Csetup-args=-Dbuild_tests=false "
+    "-Csetup-args=-Dbuild_examples=false "
+    "-Csetup-args=-Dbuild_docs=false"
+)
+
+# Catch stock libraries left behind under a different distribution name.
+_VERIFY_NIXL_IPV6 = (
+    'python -c "'
+    "import pathlib, nixl; "
+    "libs = sorted(pathlib.Path(nixl.__file__).parent.parent.rglob('libnixl*.so*')); "
+    "stock = [str(p) for p in libs if b'inet_pton failed for ip_addr' in p.read_bytes()]; "
+    "ipv6 = [str(p) for p in libs if b'Invalid IPv4 or IPv6 address: ' in p.read_bytes()]; "
+    "assert not stock, ('stock nixl still present: ' + str(stock)); "
+    "assert ipv6, ('NIXL lacks IPv6 support: ' + str([str(p) for p in libs])); "
+    "print('IPv6 NIXL:', ipv6)\""
+)
+
+
+def install_nixl(image: modal.Image) -> modal.Image:
+    """Build and install the pinned NIXL revision."""
+    return (
+        image.apt_install(
+            "git",
+            "cmake",
+            "pkg-config",
+            "ninja-build",
+            "autoconf",
+            "automake",
+            "libtool",
+            "libnuma-dev",
+            "libibverbs-dev",
+            "librdmacm-dev",
+            "ibverbs-providers",
+        )
+        .run_commands(
+            'uv pip install --python "$(command -v python)" --no-cache '
+            "meson meson-python pybind11 patchelf build"
+        )
+        .run_commands(
+            "apt-get remove -y libabsl-dev || true",
+            "rm -rf /opt/absl && git init /opt/absl",
+            f"git -C /opt/absl remote add origin {_ABSL_REPOSITORY}",
+            f"git -C /opt/absl fetch --depth 1 origin {ABSL_TAG}",
+            "git -C /opt/absl checkout --detach FETCH_HEAD",
+            "cmake -S /opt/absl -B /opt/absl/build -GNinja "
+            "-DCMAKE_INSTALL_PREFIX=/usr/local -DCMAKE_BUILD_TYPE=Release "
+            "-DCMAKE_CXX_STANDARD=20 -DBUILD_SHARED_LIBS=ON "
+            "-DABSL_PROPAGATE_CXX_STD=ON -DABSL_ENABLE_INSTALL=ON",
+            "ninja -C /opt/absl/build && ninja -C /opt/absl/build install && ldconfig",
+            "rm -rf /opt/absl",
+        )
+        .run_commands(
+            "rm -rf /opt/ucx && git init /opt/ucx",
+            f"cd /opt/ucx && git remote add origin {_UCX_REPOSITORY} "
+            f"&& git fetch --depth 1 origin {UCX_COMMIT} "
+            "&& git checkout --detach FETCH_HEAD",
+            f'test "$(git -C /opt/ucx rev-parse HEAD)" = "{UCX_COMMIT}"',
+            # gVisor returns EINVAL for UCX's invalid-fd DMA-BUF probe.
+            "sed -i 's/errno == EBADF/errno == EBADF || errno == EINVAL/' "
+            "/opt/ucx/src/uct/ib/base/ib_md.c",
+            "grep -Fq 'errno == EBADF || errno == EINVAL' "
+            "/opt/ucx/src/uct/ib/base/ib_md.c",
+            "cd /opt/ucx && ./autogen.sh && ./contrib/configure-release-mt "
+            "--prefix=/usr/local/ucx --enable-shared --disable-static "
+            "--disable-doxygen-doc --enable-optimizations --enable-cma "
+            "--enable-devel-headers --with-cuda=/usr/local/cuda --with-verbs "
+            "--with-dm --without-gdrcopy",
+            "cd /opt/ucx && make -j$(nproc) && make install && ldconfig",
+            "rm -rf /opt/ucx",
+        )
+        .run_commands(
+            f"git clone --depth 1 -b {LIBFABRIC_TAG} {_LIBFABRIC_REPOSITORY} /opt/libfabric",
+            "cd /opt/libfabric && ./autogen.sh && ./configure --prefix=/usr/local "
+            "--disable-verbs --disable-psm3 --disable-opx --disable-usnic "
+            "--disable-rstream --enable-efa --with-cuda=/usr/local/cuda "
+            "--enable-cuda-dlopen --without-gdrcopy",
+            "cd /opt/libfabric && make -j$(nproc) && make install && ldconfig",
+            "rm -rf /opt/libfabric",
+        )
+        .run_commands(
+            f"rm -rf {_NIXL_SOURCE_PATH} && git init {_NIXL_SOURCE_PATH}",
+            f"cd {_NIXL_SOURCE_PATH} && git remote add origin {NIXL_REPOSITORY} "
+            f"&& git fetch --depth 1 origin {NIXL_COMMIT} "
+            "&& git checkout --detach FETCH_HEAD",
+            f'test "$(git -C {_NIXL_SOURCE_PATH} rev-parse HEAD)" = "{NIXL_COMMIT}"',
+            'uv pip install --python "$(command -v python)" --no-cache tomlkit',
+            f"cd {_NIXL_SOURCE_PATH} && "
+            "nixl_cuda_major=$(nvcc --version | sed -n 's/.*release \\([0-9]*\\)\\..*/\\1/p') "
+            '&& python contrib/tomlutil.py --wheel-name "nixl-cu${nixl_cuda_major}" pyproject.toml',
+            'uv pip uninstall --python "$(command -v python)" nixl nixl-cu12 nixl-cu13',
+            f'cd {_NIXL_SOURCE_PATH} && uv pip install --python "$(command -v python)" '
+            f"--no-cache --no-deps --no-build-isolation -Cbuild-dir={_NIXL_SOURCE_PATH}/build "
+            f"{_NIXL_MESON_ARGS} .",
+            'uv pip install --python "$(command -v python)" --no-deps '
+            f"{_NIXL_SOURCE_PATH}/build/src/bindings/python/nixl-meta/nixl-*-py3-none-any.whl",
+            _VERIFY_NIXL_IPV6,
+            'python -c "from nixl._api import nixl_agent, nixl_agent_config; '
+            "agent = nixl_agent('build-check', nixl_agent_config(backends=[])); "
+            "plugins = agent.get_plugin_list(); "
+            "assert {'UCX', 'LIBFABRIC'} <= set(plugins), plugins\"",
+        )
+    )

autoinference_utils-0.2.11/src/autoinference_utils/cluster/modelexpress/__init__.py

@@ -0,0 +1,18 @@
+"""ModelExpress weight-loading integration."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import modal
+
+BACKEND_NAME = "modal-dict"
+
+
+def install_cluster_deps(
+    image: "modal.Image", *, engine: str, enabled: bool = True
+) -> "modal.Image":
+    from .image import install_cluster_deps
+
+    return install_cluster_deps(image, engine=engine, enabled=enabled)

autoinference_utils-0.2.11/src/autoinference_utils/cluster/modelexpress/dict_client.py

@@ -0,0 +1,182 @@
+"""ModelExpress worker discovery through Modal Dict."""
+
+from __future__ import annotations
+
+from dataclasses import replace
+import logging
+import os
+import threading
+import time
+
+from modelexpress import p2p_pb2
+from modelexpress.client import MxClientBase
+from modelexpress.metadata.source_id import compute_mx_source_id
+
+from .._net import dialable
+from ..config import (
+    DEFAULT_DICT_NAME,
+    DEFAULT_TTL_SECONDS,
+    DICT_NAME_ENV,
+    SCOPE_ENV,
+    TTL_ENV,
+)
+from .peers import PeerRecord, PeerStatus, PeerTable
+
+logger = logging.getLogger("autoinference.cluster.modelexpress")
+
+
+def _source_status(value: int) -> PeerStatus:
+    return PeerStatus[p2p_pb2.SourceStatus.Name(value).removeprefix("SOURCE_STATUS_")]
+
+
+class MxModalDictClient(MxClientBase):
+    # Makes ModelExpress start WorkerGrpcServer for peer manifest retrieval.
+    REQUIRES_P2P_METADATA = True
+
+    def __init__(self, table: PeerTable):
+        self._table = table
+        self._published: dict[tuple[str, str], PeerRecord] = {}
+        self._lock = threading.Lock()
+        self._closed = False
+
+    def close(self) -> None:
+        with self._lock:
+            self._closed = True
+            for mx_source_id, worker_id in tuple(self._published):
+                try:
+                    self._table.remove(mx_source_id, worker_id)
+                except Exception:
+                    logger.warning(
+                        "failed to remove source %s/%s",
+                        mx_source_id,
+                        worker_id,
+                        exc_info=True,
+                    )
+                else:
+                    self._published.pop((mx_source_id, worker_id))
+
+    def publish_metadata(
+        self,
+        identity: "p2p_pb2.SourceIdentity",
+        worker: "p2p_pb2.WorkerMetadata",
+        worker_id: str,
+    ) -> str:
+        with self._lock:
+            if self._closed:
+                raise RuntimeError("cannot publish metadata after closing the client")
+            mx_source_id = compute_mx_source_id(identity)
+            record = PeerRecord(
+                group_id=mx_source_id,
+                worker_id=worker_id,
+                worker_rank=worker.worker_rank,
+                accelerator=worker.accelerator,
+                status=_source_status(
+                    worker.status or p2p_pb2.SOURCE_STATUS_INITIALIZING
+                ),
+                grpc_endpoint=dialable(worker.worker_grpc_endpoint),
+                metadata_endpoint=worker.metadata_endpoint,
+                agent_name=worker.agent_name,
+                updated_at=time.time(),
+            )
+            self._table.put(record)
+            self._published[mx_source_id, worker_id] = record
+        logger.info(
+            "published source %s rank=%d at %s",
+            mx_source_id,
+            worker.worker_rank,
+            worker.worker_grpc_endpoint,
+        )
+        return mx_source_id
+
+    def list_sources(
+        self,
+        identity: "p2p_pb2.SourceIdentity | None" = None,
+        status_filter: "p2p_pb2.SourceStatus | None" = None,
+    ) -> "p2p_pb2.ListSourcesResponse":
+        if identity is None:
+            raise ValueError(
+                "list_sources requires an identity so mx_source_id can be "
+                "computed locally without a central coordinator"
+            )
+        mx_source_id = compute_mx_source_id(identity)
+        expected_status = (
+            _source_status(status_filter) if status_filter is not None else None
+        )
+        instances = [
+            p2p_pb2.SourceInstanceRef(
+                mx_source_id=record.group_id,
+                worker_id=record.worker_id,
+                model_name=identity.model_name,
+                worker_rank=record.worker_rank,
+                accelerator=record.accelerator,
+            )
+            for record in self._table.list(mx_source_id)
+            if record.grpc_endpoint
+            and (expected_status is None or record.status == expected_status)
+        ]
+        logger.info("listed %d source(s) for %s", len(instances), mx_source_id)
+        return p2p_pb2.ListSourcesResponse(instances=instances)
+
+    def get_metadata(
+        self,
+        mx_source_id: str,
+        worker_id: str,
+    ) -> "p2p_pb2.GetMetadataResponse":
+        record = self._table.get(mx_source_id, worker_id)
+        if record is None or not record.grpc_endpoint:
+            logger.warning("no live record for %s/%s", mx_source_id, worker_id)
+            return p2p_pb2.GetMetadataResponse(
+                found=False, mx_source_id=mx_source_id, worker_id=worker_id
+            )
+
+        return p2p_pb2.GetMetadataResponse(
+            found=True,
+            mx_source_id=mx_source_id,
+            worker_id=worker_id,
+            worker=p2p_pb2.WorkerMetadata(
+                worker_rank=record.worker_rank,
+                accelerator=record.accelerator,
+                status=p2p_pb2.SourceStatus.Value(
+                    f"SOURCE_STATUS_{record.status.name}"
+                ),
+                worker_grpc_endpoint=record.grpc_endpoint,
+                metadata_endpoint=record.metadata_endpoint,
+                agent_name=record.agent_name,
+            ),
+        )
+
+    def update_status(
+        self,
+        mx_source_id: str,
+        worker_id: str,
+        worker_rank: int,
+        status: "p2p_pb2.SourceStatus",
+        source_load: float | None = None,
+    ) -> bool:
+        with self._lock:
+            record = self._published.get((mx_source_id, worker_id))
+            if record is None or self._closed:
+                logger.warning(
+                    "cannot update status for unknown %s/%s", mx_source_id, worker_id
+                )
+                return False
+            record = replace(
+                record, status=_source_status(status), updated_at=time.time()
+            )
+            self._table.put(record)
+            self._published[mx_source_id, worker_id] = record
+        return True
+
+
+def build_dict_client() -> MxModalDictClient:
+    scope = os.environ.get(SCOPE_ENV, "").strip()
+    if not scope:
+        raise RuntimeError(
+            f"{SCOPE_ENV} is unset; the cluster scope must reach the engine"
+        )
+    table = PeerTable(
+        scope=scope,
+        dict_name=os.environ.get(DICT_NAME_ENV) or DEFAULT_DICT_NAME,
+        ttl=float(os.environ.get(TTL_ENV) or DEFAULT_TTL_SECONDS),
+    )
+    return MxModalDictClient(table)

autoinference_utils-0.2.11/src/autoinference_utils/cluster/modelexpress/engine.py

@@ -0,0 +1,130 @@
+"""ModelExpress engine configuration for peer weight loading."""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Mapping, Optional
+
+from .._net import i6pn_address
+from ..config import (
+    DEFAULT_DICT_NAME,
+    DEFAULT_TTL_SECONDS,
+    DICT_NAME_ENV,
+    SCOPE_ENV,
+    TTL_ENV,
+)
+
+from . import BACKEND_NAME
+
+# Bound upstream's transfer and publish waits for cold deployments.
+DEFAULT_TRANSFER_TIMEOUT = 120
+DEFAULT_PUBLISH_TIMEOUT = 120
+
+# UCX auto-detection includes the CUDA transports needed to register weights.
+DEFAULT_NIXL_UCX_TLS = None
+
+# SGLang and ModelExpress require a non-empty URL; the Dict client ignores it.
+MX_PLACEHOLDER_URL = "modal-dict://peers"
+MX_CONFIG = '{"transport":"nixl","url":"%s"}' % MX_PLACEHOLDER_URL
+
+_SGLANG_ARGS = {
+    "--load-format": "remote_instance",
+    "--remote-instance-weight-loader-backend": "modelexpress",
+    "--modelexpress-config": MX_CONFIG,
+    # RemoteInstanceModelLoader rejects SGLangEndpoint's default extra config.
+    "--model-loader-extra-config": "{}",
+}
+_VLLM_ARGS = {"--load-format": "modelexpress"}
+
+
+@dataclass(frozen=True)
+class _EngineOverrides:
+    """Launch overrides for one engine process. Falsy when the cluster is off."""
+
+    scope: Optional[str] = None
+    server_args: Mapping[str, str] = field(default_factory=dict)
+    env: Mapping[str, str] = field(default_factory=dict)
+
+    def __bool__(self) -> bool:
+        return self.scope is not None
+
+
+def _engine_overrides(
+    *,
+    app_name: str,
+    app_id: str,
+    engine: str,
+    revision: str = "",
+    dict_name: Optional[str] = None,
+    ttl: float = DEFAULT_TTL_SECONDS,
+    nixl_ucx_tls: Optional[str] = DEFAULT_NIXL_UCX_TLS,
+    transfer_timeout: int = DEFAULT_TRANSFER_TIMEOUT,
+    publish_timeout: int = DEFAULT_PUBLISH_TIMEOUT,
+) -> _EngineOverrides:
+    """Render the overrides for an engine launched from this container."""
+    if engine not in ("sglang", "vllm"):
+        raise ValueError(f"unsupported engine {engine!r}")
+    if not app_id:
+        raise ValueError(
+            "app_id is unset; configure the cluster from a running Modal app"
+        )
+
+    region = os.environ.get("MODAL_REGION", "")
+    if not region:
+        return _off("MODAL_REGION is unset")
+    try:
+        worker_host = i6pn_address()
+    except OSError as error:
+        return _off(f"no i6pn address: {error}")
+
+    scope = f"weights/{app_id}/{region}"
+    # Modal limits Dict names to 64 characters.
+    dict_prefix = f"{DEFAULT_DICT_NAME}-{app_name}"
+    dict_name = dict_name or f"{dict_prefix[:55]}-{app_id[-8:]}"
+    env = {
+        "MX_METADATA_BACKEND": BACKEND_NAME,
+        "MX_SERVER_ADDRESS": MX_PLACEHOLDER_URL,
+        SCOPE_ENV: scope,
+        DICT_NAME_ENV: dict_name,
+        TTL_ENV: str(ttl),
+        # NIXL needs a bare host; the ModelExpress adapter brackets gRPC endpoints.
+        "MX_WORKER_HOST": worker_host,
+        "MX_NIXL_BACKEND": _nixl_backend(),
+        "MX_P2P_METADATA": "1",
+        "MX_P2P_SOURCE_SELECTOR": "random",
+        # i6pn is IPv6; container IPv4 addresses are not peer-routable.
+        "UCX_TCP_AF_PRIO": "inet6",
+        "MX_NIXL_METADATA_TIMEOUT": "5",
+        "MX_TRANSFER_TIMEOUT": str(transfer_timeout),
+        "MX_PUBLISH_TIMEOUT_SECS": str(publish_timeout),
+        # SGLang's multimodal workers fork after gRPC has started.
+        "GRPC_ENABLE_FORK_SUPPORT": "1",
+        "GRPC_POLL_STRATEGY": "poll",
+    }
+    if revision:
+        env["MX_MODEL_REVISION"] = revision
+    if nixl_ucx_tls:
+        env["NIXL_UCX_TLS"] = nixl_ucx_tls
+
+    print(f"[cluster] enabled, scope={scope} host={worker_host}", flush=True)
+    return _EngineOverrides(
+        scope=scope,
+        server_args=dict(_SGLANG_ARGS if engine == "sglang" else _VLLM_ARGS),
+        env=env,
+    )
+
+
+def _nixl_backend() -> str:
+    """Use libfabric for mounted EFA devices; UCX handles verbs and TCP."""
+    for device in Path("/dev/infiniband").glob("uverbs*"):
+        driver = Path("/sys/class/infiniband_verbs") / device.name / "device" / "driver"
+        if driver.is_symlink() and os.path.basename(os.readlink(driver)) == "efa":
+            return "LIBFABRIC"
+    return "UCX"
+
+
+def _off(reason: str) -> _EngineOverrides:
+    print(f"[cluster] disabled, loading weights normally: {reason}", flush=True)
+    return _EngineOverrides()

autoinference_utils-0.2.11/src/autoinference_utils/cluster/modelexpress/image.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+import base64
+
+import modal
+
+from ..image import install_nixl
+
+MODELEXPRESS_REPOSITORY = "https://github.com/modal-projects/modelexpress.git"
+MODELEXPRESS_COMMIT = "7fc747444459423029a37bb279b9bc6d401d9d42"
+MODELEXPRESS_VERSION = "0.7.0"
+_SOURCE_PATH = "/opt/modelexpress"
+
+# Run after other .pth files have adjusted sys.path.
+_PTH_NAME = "zz-autoinference-cluster.pth"
+
+# .pth files execute single lines starting with "import ". Guard optional imports.
+_PTH_LINE = (
+    r"import sys; exec('try:\n"
+    r"    import autoinference_utils.cluster.modelexpress.shim\n"
+    r"except Exception:\n"
+    r"    pass\n')"
+)
+_PTH_B64 = base64.b64encode(f"{_PTH_LINE}\n".encode()).decode()
+
+# Co-locate with the package: pip and uv can select different site directories.
+_WRITE_PTH = (
+    'python -c "import base64, pathlib, autoinference_utils; '
+    "pathlib.Path(autoinference_utils.__file__).parent.parent"
+    f".joinpath('{_PTH_NAME}').write_bytes(base64.b64decode('{_PTH_B64}'))\""
+)
+
+# The startup hook must work without Modal's injected PYTHONPATH.
+_VERIFY_SHIM = (
+    'env -u PYTHONPATH MX_METADATA_BACKEND=modal-dict python -c "'
+    "import sys; "
+    "names = [type(f).__name__ for f in sys.meta_path]; "
+    "assert '_PatchOnImport' in names, names\""
+)
+
+_VERIFY_SHIM_IS_SILENT = (
+    "test -z \"$(env -u PYTHONPATH -u MX_METADATA_BACKEND python -c 'pass' 2>&1)\""
+)
+
+# SGLang #24723 replaced inline ModelExpress loading with package delegation.
+SGLANG_MIN_COMMIT = "435ea41cf0"
+SGLANG_MIN_DATE = "2026-05-16"
+
+_ENGINE_DELEGATION_CHECKS = {
+    "sglang": (
+        'python -c "'
+        "import inspect; from sglang.srt.model_loader import loader; "
+        "assert 'modelexpress.engines.sglang.loader' in inspect.getsource(loader), "
+        "'This SGLang build predates #24723 (%s, %s): it builds its own "
+        "MxClient and ignores the modal-dict backend. Use a newer image.'\""
+        % (SGLANG_MIN_COMMIT, SGLANG_MIN_DATE)
+    ),
+    "vllm": (
+        'python -c "'
+        "from vllm.model_executor.model_loader import _LOAD_FORMAT_TO_MODEL_LOADER as m; "
+        "assert 'modelexpress' in m, sorted(m)\""
+    ),
+}
+
+_VERIFY_MODELEXPRESS = (
+    'python -c "from importlib.metadata import version; '
+    f"assert version('modelexpress') == '{MODELEXPRESS_VERSION}'; "
+    "import modelexpress.nixl_transfer; "
+    'from modelexpress.metadata.client_factory import create_metadata_client"'
+)
+
+
+def install_cluster_deps(
+    image: modal.Image,
+    *,
+    engine: str,
+    enabled: bool = True,
+) -> modal.Image:
+    """Add NIXL and ModelExpress peer loading to an engine image.
+
+    autoinference-utils and runtime dependencies must already be installed.
+    Use --no-deps to preserve the engine image's torch installation.
+    """
+    if not enabled:
+        return image
+    if engine not in _ENGINE_DELEGATION_CHECKS:
+        raise ValueError(f"unsupported engine {engine!r}")
+    return install_nixl(image).run_commands(
+        f"rm -rf {_SOURCE_PATH} && git init {_SOURCE_PATH}",
+        f"git -C {_SOURCE_PATH} remote add origin {MODELEXPRESS_REPOSITORY}",
+        f"git -C {_SOURCE_PATH} fetch --depth 1 origin {MODELEXPRESS_COMMIT}",
+        f"git -C {_SOURCE_PATH} checkout --detach FETCH_HEAD",
+        f'test "$(git -C {_SOURCE_PATH} rev-parse HEAD)" = "{MODELEXPRESS_COMMIT}"',
+        # --system can select a different interpreter in engine images.
+        'uv pip install --python "$(command -v python)" --no-cache --no-deps '
+        f"-e {_SOURCE_PATH}/modelexpress_client/python",
+        _VERIFY_MODELEXPRESS,
+        _WRITE_PTH,
+        _VERIFY_SHIM,
+        _VERIFY_SHIM_IS_SILENT,
+        _ENGINE_DELEGATION_CHECKS[engine],
+    )

autoinference_utils-0.2.11/src/autoinference_utils/cluster/modelexpress/peers.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Optional
+
+import modal
+
+from ..config import DEFAULT_DICT_NAME, DEFAULT_TTL_SECONDS
+
+
+class PeerStatus(str, Enum):
+    UNKNOWN = "unknown"
+    INITIALIZING = "initializing"
+    READY = "ready"
+    STALE = "stale"
+
+
+@dataclass(frozen=True)
+class PeerRecord:
+    group_id: str
+    worker_id: str
+    worker_rank: int
+    accelerator: str
+    status: PeerStatus
+    grpc_endpoint: str
+    metadata_endpoint: str
+    agent_name: str
+    updated_at: float
+
+    def as_payload(self) -> dict[str, Any]:
+        return {
+            "worker_rank": self.worker_rank,
+            "accelerator": self.accelerator,
+            "status": self.status.value,
+            "grpc_endpoint": self.grpc_endpoint,
+            "metadata_endpoint": self.metadata_endpoint,
+            "agent_name": self.agent_name,
+            "updated_at": self.updated_at,
+        }
+
+    @classmethod
+    def from_payload(
+        cls, raw: Any, *, group_id: str, worker_id: str
+    ) -> Optional["PeerRecord"]:
+        if not isinstance(raw, dict):
+            return None
+        try:
+            return cls(
+                group_id=group_id,
+                worker_id=worker_id,
+                worker_rank=int(raw["worker_rank"]),
+                accelerator=str(raw["accelerator"]),
+                status=PeerStatus(raw["status"]),
+                grpc_endpoint=str(raw["grpc_endpoint"]),
+                metadata_endpoint=str(raw["metadata_endpoint"]),
+                agent_name=str(raw["agent_name"]),
+                updated_at=float(raw["updated_at"]),
+            )
+        except (KeyError, TypeError, ValueError):
+            return None
+
+
+class PeerTable:
+    def __init__(
+        self,
+        *,
+        scope: str,
+        dict_name: str = DEFAULT_DICT_NAME,
+        ttl: float = DEFAULT_TTL_SECONDS,
+        registry: Optional[modal.Dict] = None,
+    ):
+        if not scope:
+            raise ValueError("a peer table needs a scope")
+        self.scope = scope
+        self.ttl = ttl
+        self._registry = registry or modal.Dict.from_name(
+            dict_name, create_if_missing=True
+        )
+
+    def key(self, group_id: str, worker_id: str) -> str:
+        if not worker_id or "/" in worker_id:
+            raise ValueError("worker_id must be non-empty and cannot contain '/'")
+        return f"{self.scope}/{group_id}/{worker_id}"
+
+    def put(self, record: PeerRecord) -> None:
+        self._registry.put(
+            self.key(record.group_id, record.worker_id),
+            record.as_payload(),
+        )
+
+    def get(self, group_id: str, worker_id: str) -> Optional[PeerRecord]:
+        raw = self._registry.get(self.key(group_id, worker_id))
+        if raw is None:
+            return None
+        record = PeerRecord.from_payload(raw, group_id=group_id, worker_id=worker_id)
+        cutoff = time.time() - self.ttl
+        if record is None or record.updated_at < cutoff:
+            self._reap(self.key(group_id, worker_id), group_id, worker_id, cutoff)
+            return None
+        return record
+
+    def list(self, group_id: str) -> list[PeerRecord]:
+        prefix = f"{self.scope}/"
+        cutoff = time.time() - self.ttl
+        live: list[PeerRecord] = []
+        for key, raw in self._registry.items():
+            if not isinstance(key, str) or not key.startswith(prefix):
+                continue
+            group, _, worker_id = key[len(prefix) :].partition("/")
+            if not group or not worker_id or "/" in worker_id:
+                continue
+            record = PeerRecord.from_payload(raw, group_id=group, worker_id=worker_id)
+            if record is None or record.updated_at < cutoff:
+                self._reap(key, group, worker_id, cutoff)
+            elif group == group_id:
+                live.append(record)
+        return live
+
+    def _reap(self, key: str, group_id: str, worker_id: str, cutoff: float) -> None:
+        raw = self._registry.pop(key, None)
+        record = PeerRecord.from_payload(raw, group_id=group_id, worker_id=worker_id)
+        if record is not None and record.updated_at >= cutoff:
+            self._registry.put(key, raw, skip_if_exists=True)
+
+    def remove(self, group_id: str, worker_id: str) -> None:
+        self._registry.pop(self.key(group_id, worker_id), None)

autoinference_utils-0.2.11/src/autoinference_utils/cluster/modelexpress/shim.py

@@ -0,0 +1,81 @@
+"""Register the Dict backend through a .pth import hook.
+
+ModelExpress has no backend registration API. Defer its import until the engine
+loads it so the startup hook does not import torch.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import os
+import sys
+from importlib.abc import Loader, MetaPathFinder
+from typing import Any, Optional
+
+from . import BACKEND_NAME
+
+TARGET_MODULE = "modelexpress.metadata.client_factory"
+
+
+def install() -> None:
+    if os.environ.get("MX_METADATA_BACKEND", "").strip().lower() != BACKEND_NAME:
+        return
+    if any(isinstance(finder, _PatchOnImport) for finder in sys.meta_path):
+        return
+    sys.meta_path.insert(0, _PatchOnImport())
+
+
+class _PatchOnImport(MetaPathFinder):
+    def find_spec(self, fullname: str, path: Any = None, target: Any = None) -> Any:
+        if fullname != TARGET_MODULE:
+            return None
+        sys.meta_path.remove(self)
+        try:
+            spec = importlib.util.find_spec(fullname)
+        except Exception:
+            spec = None
+        finally:
+            sys.meta_path.insert(0, self)
+        if spec is None or spec.loader is None:
+            return None
+        spec.loader = _PatchingLoader(spec.loader)
+        return spec
+
+
+class _PatchingLoader(Loader):
+    def __init__(self, inner: Loader):
+        self._inner = inner
+
+    def create_module(self, spec: Any) -> Any:
+        return self._inner.create_module(spec)  # type: ignore[attr-defined]
+
+    def exec_module(self, module: Any) -> None:
+        self._inner.exec_module(module)  # type: ignore[attr-defined]
+        patch_factory(module)
+
+
+def patch_factory(module: Any) -> None:
+    original = getattr(module, "create_metadata_client", None)
+    if original is None or getattr(original, "_cluster_patched", False):
+        return
+
+    def create_metadata_client(
+        worker_rank: Optional[int] = None,
+        server_url: Optional[str] = None,
+    ) -> Any:
+        backend = os.environ.get("MX_METADATA_BACKEND", "").strip().lower()
+        if backend != BACKEND_NAME:
+            return original(worker_rank, server_url)
+        from .dict_client import build_dict_client
+
+        return build_dict_client()
+
+    create_metadata_client._cluster_patched = True  # type: ignore[attr-defined]
+    module.create_metadata_client = create_metadata_client
+
+
+try:
+    install()
+except Exception:  # noqa: BLE001
+    # .pth errors would be printed by every interpreter using this image.
+    pass

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

@@ -90,6 +90,59 @@ class Endpoint(ABC):
 
     def __init__(self, base_url: str):
         self.base_url = base_url.rstrip("/")
+        self._cluster_env: dict[str, str] = {}
+
+    def with_cluster(self, app, *, revision: str = "", enabled: bool = True):
+        """Configure peer weight loading before starting the endpoint."""
+        if not enabled:
+            return self
+        from .cluster.modelexpress.engine import _engine_overrides
+
+        if isinstance(self, SGLangEndpoint):
+            engine = "sglang"
+        elif isinstance(self, VLLMEndpoint):
+            engine = "vllm"
+        else:
+            raise ValueError(f"cluster is unsupported for {type(self).__name__}")
+        if self._proc is not None:
+            raise RuntimeError("configure the cluster before starting the endpoint")
+        overrides = _engine_overrides(
+            app_name=app.name, app_id=app.app_id, engine=engine, revision=revision
+        )
+        if not overrides:
+            return self
+        loader_args = {}
+        if isinstance(self, SGLangEndpoint) and self.load_format is not None:
+            loader_args["--load-format"] = self.load_format
+        loader_args = _merge_server_args(loader_args, self.extra_server_args)
+        for key, value in loader_args.items():
+            flag, attached, inline_value = key.partition("=")
+            value = inline_value if attached else value
+            if flag == "--load-format":
+                supported = value in ("auto", overrides.server_args[flag])
+            elif flag == "--model-loader-extra-config":
+                supported_configs = [{}]
+                if isinstance(self, SGLangEndpoint):
+                    supported_configs.append(
+                        json.loads(self.DEFAULT_OPERATIONAL_ARGS[flag])
+                    )
+                try:
+                    supported = json.loads(value) in supported_configs
+                except ValueError:
+                    supported = False
+            elif flag in overrides.server_args:
+                supported = value == overrides.server_args[flag]
+            else:
+                continue
+            if not supported:
+                raise ValueError(f"cluster does not support {flag}={value!r}")
+        self.extra_server_args = _merge_server_args(
+            self.extra_server_args, overrides.server_args
+        )
+        self._cluster_env = dict(overrides.env)
+        if isinstance(self, SGLangEndpoint):
+            self.load_format = None
+        return self
 
     def __enter__(self):
         try:
@@ -300,21 +353,25 @@ class SGLangEndpoint(Endpoint):
     def start(self):
         cmd = self._build_cmd()
         print(f"[endpoint] starting: {shlex.join(cmd)}")
-        self._proc = subprocess.Popen(cmd)
-        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,
-        )
-        if self.bench_mode:
-            self._bench_server = start_bench_proxy(
-                listen_port=self.listen_port,
-                upstream_port=self.worker_port,
-                upstream_base_url=self.base_url,
+        self._proc = subprocess.Popen(cmd, env=os.environ | self._cluster_env)
+        try:
+            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,
             )
+            if self.bench_mode:
+                self._bench_server = start_bench_proxy(
+                    listen_port=self.listen_port,
+                    upstream_port=self.worker_port,
+                    upstream_base_url=self.base_url,
+                )
+        except BaseException:
+            self.stop()
+            raise
 
     def stop(self):
         if self._bench_server is not None:
@@ -378,19 +435,23 @@ class VLLMEndpoint(Endpoint):
     def start(self):
         cmd = self._build_cmd()
         print(f"[vllm] starting: {shlex.join(cmd)}")
-        self._proc = subprocess.Popen(cmd)
-        wait_ready(
-            self._proc,
-            port=self.worker_port,
-            timeout=self.health_timeout,
-            poll_interval=self.health_poll_interval,
-            request_timeout=self.health_request_timeout,
-        )
-        if self.bench_mode:
-            self._bench_server = start_bench_proxy(
-                listen_port=self.listen_port,
-                upstream_port=self.worker_port,
+        self._proc = subprocess.Popen(cmd, env=os.environ | self._cluster_env)
+        try:
+            wait_ready(
+                self._proc,
+                port=self.worker_port,
+                timeout=self.health_timeout,
+                poll_interval=self.health_poll_interval,
+                request_timeout=self.health_request_timeout,
             )
+            if self.bench_mode:
+                self._bench_server = start_bench_proxy(
+                    listen_port=self.listen_port,
+                    upstream_port=self.worker_port,
+                )
+        except BaseException:
+            self.stop()
+            raise
 
     def stop(self):
         if self._bench_server is not None:

{autoinference_utils-0.2.9 → autoinference_utils-0.2.11}/src/autoinference_utils/pd.py

@@ -5,6 +5,8 @@ from __future__ import annotations
 import asyncio
 import math
 import os
+import socket
+import struct
 import subprocess
 import threading
 import time
@@ -40,6 +42,9 @@ _MANAGED_ROUTER_FLAGS = frozenset(
     }
 )
 
+# Leave headroom inside Modal's 30-second exit-handler deadline.
+_SHUTDOWN_TIMEOUT = 25
+
 
 def _managed_router_tokens(router_args: Mapping[str, str]) -> list[str]:
     """Rendered override tokens that would reach a managed flag.
@@ -177,6 +182,7 @@ class PDEndpoint(Endpoint):
         health_failure_timeout: float = 30,
         router_args: Mapping[str, str] | None = None,
         allow_custom_router: bool = False,
+        shutdown_port: int | None = None,
     ):
         from modal.experimental import get_cluster_info
 
@@ -217,6 +223,13 @@ class PDEndpoint(Endpoint):
                 f"P/D ratio {ratio} requires a {len(roles)}-container Modal cluster"
             )
         self._host_ip = hosts[cluster.rank]
+        self._hosts = hosts
+        self._rank = cluster.rank
+        self._shutdown_port = (
+            router_port + 1 if shutdown_port is None else shutdown_port
+        )
+        if not 1 <= self._shutdown_port <= 65535:
+            raise ValueError("shutdown_port must be between 1 and 65535")
         self.role = roles[cluster.rank]
         super().__init__(_url(hosts[0], router_port))
         self.engine = engine
@@ -266,6 +279,8 @@ class PDEndpoint(Endpoint):
         self._stopped = threading.Event()
         self._lock = threading.Lock()
         self._started = False
+        self._serving = False
+        self._shutdown_listener: socket.socket | None = None
         self._threads: list[threading.Thread] = []
 
     def start(self, *, warmup: Callable[[], None] | None = None) -> None:
@@ -280,6 +295,13 @@ class PDEndpoint(Endpoint):
             SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION="1",
         )
         try:
+            if self._rank == 0:
+                self._shutdown_listener = socket.socket(socket.AF_INET6)
+                self._shutdown_listener.setsockopt(
+                    socket.SOL_SOCKET, socket.SO_REUSEADDR, 1
+                )
+                self._shutdown_listener.bind(("::", self._shutdown_port))
+                self._shutdown_listener.listen(len(self._hosts) - 1)
             for endpoint in (self.engine, self.router):
                 if endpoint is not None:
                     endpoint.start()
@@ -291,6 +313,7 @@ class PDEndpoint(Endpoint):
                 for _, host in self.router.pd_config:
                     self._watch_health(_url(host, self.engine.worker_port) + "/health")
                 self._watch_health(self.router.base_url + "/health")
+            self._serving = True
         except BaseException:
             self.stop()
             raise
@@ -352,7 +375,7 @@ class PDEndpoint(Endpoint):
             if self._stopped.is_set():
                 return
             self._stopped.set()
-        deadline = time.monotonic() + self.drain_timeout + 30
+        deadline = time.monotonic() + _SHUTDOWN_TIMEOUT
         try:
             for endpoint in (self.router, self.engine):
                 if endpoint is None:
@@ -364,9 +387,58 @@ class PDEndpoint(Endpoint):
                         process.wait(timeout=max(0, deadline - time.monotonic()))
                     except subprocess.TimeoutExpired:
                         process.kill()
+                        process.wait()
                 endpoint.stop()
         finally:
             for thread in self._threads:
-                thread.join(
-                    timeout=self.engine.health_request_timeout + self.health_interval
+                thread.join(timeout=max(0, deadline - time.monotonic()))
+            if self._serving:
+                if self._rank == 0:
+                    self._wait_for_followers(deadline)
+                else:
+                    self._notify_leader(deadline)
+            if self._shutdown_listener is not None:
+                self._shutdown_listener.close()
+                self._shutdown_listener = None
+
+    def _notify_leader(self, deadline: float) -> None:
+        while True:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                print("[pd] Failed to notify rank zero of shutdown", flush=True)
+                return
+            try:
+                with socket.create_connection(
+                    (self._hosts[0], self._shutdown_port), timeout=min(remaining, 1)
+                ) as connection:
+                    connection.sendall(struct.pack("!I", self._rank))
+                return
+            except OSError:
+                time.sleep(min(0.1, max(0, deadline - time.monotonic())))
+
+    def _wait_for_followers(self, deadline: float) -> None:
+        listener = self._shutdown_listener
+        if listener is None:
+            return
+        pending = set(range(1, len(self._hosts)))
+        while pending:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                print(
+                    f"[pd] Timed out waiting for ranks {sorted(pending)} to stop",
+                    flush=True,
                 )
+                return
+            listener.settimeout(min(remaining, 1))
+            try:
+                connection, _ = listener.accept()
+            except TimeoutError:
+                continue
+            with connection:
+                connection.settimeout(min(remaining, 1))
+                try:
+                    payload = connection.recv(4, socket.MSG_WAITALL)
+                except TimeoutError:
+                    continue
+            if len(payload) == 4:
+                pending.discard(struct.unpack("!I", payload)[0])