modelexpress 0.4.0__tar.gz

modelexpress-0.4.0/PKG-INFO

@@ -0,0 +1,157 @@
+Metadata-Version: 2.4
+Name: modelexpress
+Version: 0.4.0
+Summary: Python client for ModelExpress P2P GPU transfer service
+Author-email: NVIDIA Corporation <sw-dl-dynamo@nvidia.com>
+License-Expression: Apache-2.0
+Project-URL: Homepage, https://github.com/ai-dynamo/modelexpress
+Project-URL: Repository, https://github.com/ai-dynamo/modelexpress.git
+Project-URL: Documentation, https://github.com/ai-dynamo/modelexpress/tree/main/docs
+Keywords: llm,gpu,transfer,rdma,nixl,vllm
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: grpcio>=1.66.2
+Requires-Dist: huggingface_hub>=0.20.0
+Requires-Dist: nixl[cu12]; sys_platform == "linux"
+Requires-Dist: numpy>=1.24.0
+Requires-Dist: protobuf<6.0.0,>=5.27.0
+Requires-Dist: pydantic>=2.0.0
+Requires-Dist: torch>=2.6.0
+Requires-Dist: runai-model-streamer[azure,gcs,s3]; sys_platform == "linux"
+Provides-Extra: dev
+Requires-Dist: grpcio-tools<=1.66.2,>=1.60.0; extra == "dev"
+Requires-Dist: opentelemetry-api>=1.41.1; extra == "dev"
+Requires-Dist: opentelemetry-sdk>=1.41.1; extra == "dev"
+Requires-Dist: pytest>=7.0.0; extra == "dev"
+Requires-Dist: pytest-asyncio>=0.21.0; extra == "dev"
+Provides-Extra: otel
+Requires-Dist: opentelemetry-api>=1.20.0; extra == "otel"
+Provides-Extra: vmm
+Requires-Dist: cuda-python>=12.0; extra == "vmm"
+
+# ModelExpress Python Client
+
+Python client for ModelExpress -- high-performance GPU-to-GPU model weight transfers using NVIDIA NIXL over RDMA/InfiniBand.
+
+Instead of each vLLM instance loading model weights from storage, one "source" instance loads the model and transfers weights directly to "target" instances via GPUDirect RDMA, bypassing the CPU entirely.
+
+## Installation
+
+```bash
+# From PyPI (coming soon)
+pip install modelexpress
+
+# Editable install from source
+pip install -e .
+
+# With dev dependencies (pytest, grpcio-tools)
+pip install -e ".[dev]"
+```
+
+### Requirements
+
+- Python >= 3.10
+- NVIDIA GPUs with RDMA/InfiniBand support
+- [NIXL](https://github.com/ai-dynamo/nixl) (NVIDIA Interconnect eXchange Library)
+- A running [ModelExpress server](https://github.com/ai-dynamo/modelexpress/tree/main/modelexpress_server) (Rust gRPC service backed by Redis)
+
+## Quick Start with vLLM
+
+ModelExpress integrates with vLLM via custom model loaders. vLLM can discover the package through its `vllm.general_plugins` entrypoint; set `VLLM_PLUGINS=modelexpress` if your vLLM deployment requires explicit plugin selection. For manual registration, call `register_modelexpress_loaders()` in your code.
+
+```bash
+export MX_SERVER_ADDRESS="modelexpress-server:8001"
+
+vllm serve deepseek-ai/DeepSeek-V3 \
+    --load-format modelexpress \
+    --tensor-parallel-size 8
+```
+
+Starting the vLLM engine with the `modelexpress` load format on the source worker will load the weights from disk and register/publish the NIXL and tensor metadata to the MX server. The `mx` load format is kept as a backward-compatible alias.
+And on the target worker, it will retrieve these metadata from MX serverand stream weights over RDMA from GPU to GPU.
+
+## Programmatic Usage
+
+### MxClient
+
+`MxClient` is a lightweight gRPC client for communicating with the ModelExpress server:
+
+```python
+from modelexpress import MxClient
+
+client = MxClient(server_url="modelexpress-server:8001")
+
+# Query for a source model
+response = client.get_metadata("deepseek-ai/DeepSeek-V3")
+if response.found:
+    for worker in response.workers:
+        print(f"Worker rank {worker.worker_rank}: {len(worker.tensors)} tensors")
+
+# Wait for source readiness (blocks until ready or timeout)
+success, session_id, metadata_hash = client.wait_for_ready(
+    model_name="deepseek-ai/DeepSeek-V3",
+    worker_id=0,
+    timeout_seconds=7200,
+)
+
+client.close()
+```
+
+### Registering Loaders Manually
+
+```python
+from modelexpress import register_modelexpress_loaders
+
+register_modelexpress_loaders()
+# Now vLLM recognizes --load-format modelexpress and mx
+```
+
+## Environment Variables
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `MX_SERVER_ADDRESS` | `localhost:8001` | ModelExpress gRPC server address (recommended) |
+| `MODEL_EXPRESS_URL` | `localhost:8001` | Deprecated, pending removal in a future release. Still read by all client paths and takes precedence when both are set; keep setting it during the transition. |
+| `MX_EXPECTED_WORKERS` | Auto-detected from TP size | Number of GPU workers to coordinate |
+| `MX_SYNC_PUBLISH` | `0` | Source: wait for all workers before publishing metadata |
+| `MX_SYNC_START` | `1` | Target: wait for all source workers before transferring |
+| `MX_POOL_REG` | `0` | Allocation-level NIXL registration (registers cudaMalloc blocks instead of individual tensors) |
+
+### UCX/NIXL Tuning
+
+| Variable | Recommended | Description |
+|----------|-------------|-------------|
+| `UCX_RNDV_SCHEME` | `get_zcopy` | Zero-copy RDMA reads |
+| `UCX_RNDV_THRESH` | `0` | Force rendezvous for all transfers |
+| `NIXL_LOG_LEVEL` | `INFO` | NIXL logging level |
+
+## Package Structure
+
+| Module | Description |
+|--------|-------------|
+| `modelexpress.client` | `MxClient` -- gRPC client for the ModelExpress server |
+| `modelexpress.metadata` | Metadata clients, source identity, heartbeat, and worker manifest serving |
+| `modelexpress.engines.vllm.loader` | `MxModelLoader` -- vLLM integration |
+| `modelexpress.vllm_loader` | Compatibility shim for the vLLM loader |
+| `modelexpress.nixl_transfer` | `NixlTransferManager` -- NIXL agent lifecycle and RDMA transfers |
+| `modelexpress.types` | `TensorDescriptor`, `WorkerMetadata` -- core data types |
+| `modelexpress.vllm_worker` | Compatibility worker extension for older manual-registration workflows |
+
+## How It Works
+
+1. **Source** loads weights from disk, registers raw tensors with NIXL *before* FP8 processing, and publishes metadata to the ModelExpress server.
+2. **Target** creates dummy weights, waits for the source ready flag, then pulls raw tensors via RDMA read.
+3. Both source and target run `process_weights_after_loading()` independently, producing identical FP8-transformed weights.
+
+This pre-processing transfer strategy is critical for FP8 models (e.g., DeepSeek-V3) where `weight_scale_inv` tensors are renamed and transformed during processing.
+
+## License
+
+Apache-2.0

modelexpress-0.4.0/README.md

@@ -0,0 +1,119 @@
+# ModelExpress Python Client
+
+Python client for ModelExpress -- high-performance GPU-to-GPU model weight transfers using NVIDIA NIXL over RDMA/InfiniBand.
+
+Instead of each vLLM instance loading model weights from storage, one "source" instance loads the model and transfers weights directly to "target" instances via GPUDirect RDMA, bypassing the CPU entirely.
+
+## Installation
+
+```bash
+# From PyPI (coming soon)
+pip install modelexpress
+
+# Editable install from source
+pip install -e .
+
+# With dev dependencies (pytest, grpcio-tools)
+pip install -e ".[dev]"
+```
+
+### Requirements
+
+- Python >= 3.10
+- NVIDIA GPUs with RDMA/InfiniBand support
+- [NIXL](https://github.com/ai-dynamo/nixl) (NVIDIA Interconnect eXchange Library)
+- A running [ModelExpress server](https://github.com/ai-dynamo/modelexpress/tree/main/modelexpress_server) (Rust gRPC service backed by Redis)
+
+## Quick Start with vLLM
+
+ModelExpress integrates with vLLM via custom model loaders. vLLM can discover the package through its `vllm.general_plugins` entrypoint; set `VLLM_PLUGINS=modelexpress` if your vLLM deployment requires explicit plugin selection. For manual registration, call `register_modelexpress_loaders()` in your code.
+
+```bash
+export MX_SERVER_ADDRESS="modelexpress-server:8001"
+
+vllm serve deepseek-ai/DeepSeek-V3 \
+    --load-format modelexpress \
+    --tensor-parallel-size 8
+```
+
+Starting the vLLM engine with the `modelexpress` load format on the source worker will load the weights from disk and register/publish the NIXL and tensor metadata to the MX server. The `mx` load format is kept as a backward-compatible alias.
+And on the target worker, it will retrieve these metadata from MX serverand stream weights over RDMA from GPU to GPU.
+
+## Programmatic Usage
+
+### MxClient
+
+`MxClient` is a lightweight gRPC client for communicating with the ModelExpress server:
+
+```python
+from modelexpress import MxClient
+
+client = MxClient(server_url="modelexpress-server:8001")
+
+# Query for a source model
+response = client.get_metadata("deepseek-ai/DeepSeek-V3")
+if response.found:
+    for worker in response.workers:
+        print(f"Worker rank {worker.worker_rank}: {len(worker.tensors)} tensors")
+
+# Wait for source readiness (blocks until ready or timeout)
+success, session_id, metadata_hash = client.wait_for_ready(
+    model_name="deepseek-ai/DeepSeek-V3",
+    worker_id=0,
+    timeout_seconds=7200,
+)
+
+client.close()
+```
+
+### Registering Loaders Manually
+
+```python
+from modelexpress import register_modelexpress_loaders
+
+register_modelexpress_loaders()
+# Now vLLM recognizes --load-format modelexpress and mx
+```
+
+## Environment Variables
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `MX_SERVER_ADDRESS` | `localhost:8001` | ModelExpress gRPC server address (recommended) |
+| `MODEL_EXPRESS_URL` | `localhost:8001` | Deprecated, pending removal in a future release. Still read by all client paths and takes precedence when both are set; keep setting it during the transition. |
+| `MX_EXPECTED_WORKERS` | Auto-detected from TP size | Number of GPU workers to coordinate |
+| `MX_SYNC_PUBLISH` | `0` | Source: wait for all workers before publishing metadata |
+| `MX_SYNC_START` | `1` | Target: wait for all source workers before transferring |
+| `MX_POOL_REG` | `0` | Allocation-level NIXL registration (registers cudaMalloc blocks instead of individual tensors) |
+
+### UCX/NIXL Tuning
+
+| Variable | Recommended | Description |
+|----------|-------------|-------------|
+| `UCX_RNDV_SCHEME` | `get_zcopy` | Zero-copy RDMA reads |
+| `UCX_RNDV_THRESH` | `0` | Force rendezvous for all transfers |
+| `NIXL_LOG_LEVEL` | `INFO` | NIXL logging level |
+
+## Package Structure
+
+| Module | Description |
+|--------|-------------|
+| `modelexpress.client` | `MxClient` -- gRPC client for the ModelExpress server |
+| `modelexpress.metadata` | Metadata clients, source identity, heartbeat, and worker manifest serving |
+| `modelexpress.engines.vllm.loader` | `MxModelLoader` -- vLLM integration |
+| `modelexpress.vllm_loader` | Compatibility shim for the vLLM loader |
+| `modelexpress.nixl_transfer` | `NixlTransferManager` -- NIXL agent lifecycle and RDMA transfers |
+| `modelexpress.types` | `TensorDescriptor`, `WorkerMetadata` -- core data types |
+| `modelexpress.vllm_worker` | Compatibility worker extension for older manual-registration workflows |
+
+## How It Works
+
+1. **Source** loads weights from disk, registers raw tensors with NIXL *before* FP8 processing, and publishes metadata to the ModelExpress server.
+2. **Target** creates dummy weights, waits for the source ready flag, then pulls raw tensors via RDMA read.
+3. Both source and target run `process_weights_after_loading()` independently, producing identical FP8-transformed weights.
+
+This pre-processing transfer strategy is critical for FP8 models (e.g., DeepSeek-V3) where `weight_scale_inv` tensors are renamed and transformed during processing.
+
+## License
+
+Apache-2.0

modelexpress-0.4.0/modelexpress/__init__.py

@@ -0,0 +1,84 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+ModelExpress - High-performance GPU-to-GPU model weight transfers.
+
+This package provides:
+- NIXL-based RDMA transfers for GPU tensors
+- GPUDirect Storage (GDS) for direct file-to-GPU loading
+- vLLM worker extension for serving model weights
+- Custom model loaders for FP8 model support (DeepSeek-V3, etc.)
+
+Quick Start (vLLM):
+    from modelexpress import register_modelexpress_loaders
+    register_modelexpress_loaders()
+
+    # vllm serve model --load-format modelexpress
+    # Auto-detects: RDMA -> GDS -> disk
+"""
+
+import logging
+import os
+
+_logger = logging.getLogger(__name__)
+_loaders_registered = False
+
+
+def configure_vllm_logging():
+    """Ensure modelexpress loggers are visible in vLLM worker subprocesses.
+
+    vLLM only attaches log handlers to the "vllm" namespace. Without this,
+    all "modelexpress.*" output is silently dropped in EngineCore worker
+    processes. Copies vLLM's handlers onto the "modelexpress" parent logger
+    so every child inherits them via propagation. Idempotent.
+    """
+    mx_root = logging.getLogger("modelexpress")
+    if mx_root.handlers:
+        return
+    vllm_logger = logging.getLogger("vllm")
+    for handler in vllm_logger.handlers:
+        mx_root.addHandler(handler)
+    mx_level = os.environ.get("MODEL_EXPRESS_LOG_LEVEL", "").upper()
+    if mx_level and hasattr(logging, mx_level):
+        mx_root.setLevel(getattr(logging, mx_level))
+    elif vllm_logger.level != logging.NOTSET:
+        mx_root.setLevel(vllm_logger.level)
+
+
+def register_modelexpress_loaders():
+    """
+    Register ModelExpress loaders with vLLM.
+
+    This function ensures loaders are registered exactly once. It can be called
+    multiple times safely (idempotent).
+
+    Enables:
+        --load-format modelexpress  (auto-detect: RDMA -> GDS -> disk)
+        --load-format mx            (backward-compatible alias)
+    """
+    global _loaders_registered
+    if _loaders_registered:
+        return
+
+    from .engines.vllm import register_modelexpress_loaders as register_vllm_loaders
+
+    register_vllm_loaders()
+
+    _loaders_registered = True
+    _logger.debug("ModelExpress loaders registered")
+
+
+from .client import MxClient  # noqa: F401
+from .gds_loader import MxGdsLoader  # noqa: F401
+from .gds_transfer import GdsTransferManager  # noqa: F401
+from .metadata.heartbeat import HeartbeatThread  # noqa: F401
+
+__all__ = [
+    "GdsTransferManager",
+    "HeartbeatThread",
+    "MxClient",
+    "MxGdsLoader",
+    "configure_vllm_logging",
+    "register_modelexpress_loaders",
+]

modelexpress-0.4.0/modelexpress/adapter.py

@@ -0,0 +1,176 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Engine adapter contract for ModelExpress loading strategies."""
+
+from __future__ import annotations
+
+import functools
+import os
+from typing import TYPE_CHECKING, Iterator
+
+import torch
+
+from . import p2p_pb2
+
+if TYPE_CHECKING:
+    from .load_strategy.context import LoadResult
+
+
+class UnsupportedCapability(NotImplementedError):
+    """Raised by default bodies for adapter capabilities."""
+
+
+class StrategyFailed(RuntimeError):
+    """Expected strategy miss that lets the chain try the next strategy.
+
+    Set mutated=True when the current strategy may have changed model weights
+    or model structure before failing. The chain will run rollback() for
+    strategy-owned cleanup, then ask the adapter to re-initialize the model.
+    """
+
+    def __init__(self, message: str, *, mutated: bool = False):
+        super().__init__(message)
+        self.mutated = mutated
+
+
+def gated_capability(method):
+    """Create an optional adapter method that engines must override to support it.
+
+    Strategies compare their required EngineAdapter methods by object identity.
+    If an engine inherits this default method, the strategy is not eligible; if
+    the engine overrides it, the strategy may call the engine implementation.
+    Calling the inherited default still raises UnsupportedCapability.
+    """
+
+    @functools.wraps(method)
+    def default(self, *args, **kwargs):
+        raise UnsupportedCapability(method.__name__)
+
+    default.__gated_capability__ = True
+    return default
+
+
+class EngineAdapter:
+    """Engine-specific boundary used by shared loading strategies.
+
+    Load strategies own ModelExpress policy: fallback order, metadata
+    publishing, RDMA transfer, and retry decisions. Engine adapters own the
+    operations that depend on a framework's model object, device mapping, and
+    post-load processing rules.
+
+    Methods decorated with @gated_capability are optional capabilities. A
+    strategy lists the exact adapter methods it needs in its `requires` tuple;
+    the strategy is eligible only when the concrete engine overrides those
+    methods. Plain lifecycle hooks below are no-op by default because they are
+    safe extension points around an already-selected strategy.
+    """
+
+    @gated_capability
+    def build_identity(self) -> p2p_pb2.SourceIdentity:
+        """Return the stable identity used to match compatible source workers."""
+        ...
+
+    @gated_capability
+    def get_worker_rank(self) -> int:
+        """Return the model-shard key used to pair source and target workers.
+
+        Data-parallel replicas should normally return the same key because
+        their weights are interchangeable. Tensor, pipeline, and expert
+        parallel shards should return distinct keys when they own distinct
+        model weights.
+        """
+        ...
+
+    @gated_capability
+    def get_global_rank(self) -> int:
+        """Return the engine's global distributed rank for logging and metadata."""
+        ...
+
+    @gated_capability
+    def get_device_id(self) -> int:
+        """Return the local CUDA device id owned by this adapter instance."""
+        ...
+
+    @gated_capability
+    def discover_tensors(self, result: LoadResult) -> dict[str, torch.Tensor]:
+        """Return publishable tensors from the loaded engine model.
+
+        Strategies call this after weights are ready. Implementations may run
+        engine-specific tensor adoption or normalization before collecting the
+        tensors, but should not change model weights.
+        """
+        ...
+
+    @gated_capability
+    def apply_weight_iter(
+        self,
+        result: LoadResult,
+        weights_iter: Iterator[tuple[str, torch.Tensor]],
+    ) -> LoadResult:
+        """Apply a stream of named tensors to the engine model.
+
+        This hook may mutate model weights. If a strategy catches an exception
+        from this method, it should raise StrategyFailed(mutated=True) so the
+        chain reinitializes the model before trying another strategy.
+        """
+        ...
+
+    @gated_capability
+    def build_model_streamer_weight_iter(
+        self,
+        model_uri: str,
+        model: torch.nn.Module | None = None,
+    ) -> Iterator[tuple[str, torch.Tensor]]:
+        """Return the engine-native ModelStreamer weight iterator.
+
+        The shared strategy owns fallback and registration policy. Engines own
+        the concrete ModelStreamer integration because native loader APIs and
+        distributed-device selection are framework-specific. Some engine-native
+        loaders need the initialized model to discover secondary weight sources.
+        """
+        ...
+
+    @gated_capability
+    def load_via_native(self, result: LoadResult) -> LoadResult:
+        """Load the model using the engine's native disk/checkpoint loader."""
+        ...
+
+    @gated_capability
+    def reinit_for_retry(self, result: LoadResult) -> LoadResult:
+        """Replace a possibly-mutated model with a fresh engine model instance."""
+        ...
+
+    def get_unique_id(self) -> str:
+        """Return a best-effort unique id for engines without custom identity."""
+        if torch.distributed.is_available() and torch.distributed.is_initialized():
+            return f"rank-{torch.distributed.get_rank()}"
+        return f"pid-{os.getpid()}"
+
+    def get_target_device(self) -> torch.device:
+        """Return the torch device where target model state should live."""
+        return torch.device(f"cuda:{self.get_device_id()}")
+
+    def is_cuda_alike(self) -> bool:
+        """Return whether this engine is running on a CUDA-like platform."""
+        return False
+
+    def prepare_rdma_target(self, result: LoadResult) -> LoadResult:
+        """Prepare target-side model storage before receiving RDMA weights."""
+        return result
+
+    def before_rdma_receive(self, result: LoadResult) -> LoadResult:
+        """Run engine post-processing needed before RDMA writes into tensors."""
+        return result
+
+    def after_rdma_receive(self, result: LoadResult) -> LoadResult:
+        """Run engine post-processing after RDMA weights have been received."""
+        return result
+
+    def after_weight_iter_load(self, result: LoadResult) -> LoadResult:
+        """Run engine post-processing after apply_weight_iter() succeeds."""
+        return result
+
+    def after_native_load(self, result: LoadResult) -> LoadResult:
+        """Run engine post-processing after load_via_native() succeeds."""
+        return result