modelexpress-client 0.7.0

# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""
ModelExpress Custom Model Loader for vLLM.

This loader hooks into vLLM's weight loading pipeline to perform RDMA transfers
of fully-processed model tensors. Registration happens AFTER
process_weights_after_loading() so that all final tensors are captured.
Tensor discovery uses named_parameters() and named_buffers(); bare tensor
attributes created during post-processing (e.g. FP8 scales, MLA projections)
are auto-promoted to non-persistent buffers via capture_tensor_attrs().

Uses LoadStrategyChain to auto-detect the best loading strategy:
    1. RDMA (P2P GPU transfer via NIXL) - if a source is already serving
    2. ServerCache (stream weights from ModelExpress Server) - set MODEL_EXPRESS_NO_SHARED_STORAGE=1
    3. InstantTensor (fast local safetensors, direct I/O + GDS) - set MX_INSTANT_TENSOR=0 to disable
    4. ModelStreamer (S3/GCS/Azure/local via runai-model-streamer) - set MX_MODEL_URI
    5. GDS (GPUDirect Storage) - direct file-to-GPU, bypassing CPU
    6. Default (vLLM DefaultModelLoader) - standard CPU-staged loading

Usage:
    --load-format modelexpress
    --load-format mx  (backward-compatible alias)
"""

from __future__ import annotations

import logging
import time

import torch
import torch.nn as nn

from ... import configure_vllm_logging, envs, model_prefetch
from ...load_strategy import (
    LoadContext,
    drain_tensor_readers,
    publish_metadata,
    run_load_strategy_chain,
    unpublish_metadata,
)
from ...metrics import enable_metrics, metrics
from ...nixl_transfer import NixlTransferManager
from ...vmm.runtime import log_arena_post_load, maybe_enter_vmm_arena
from .adapter import _is_speculative_draft, build_vllm_load_context
from .artifacts import (
    _vllm_health_ready,
    install_vllm_cache_artifacts,
    schedule_vllm_cache_artifact_publish,
)

from vllm.config import ModelConfig, VllmConfig
from vllm.config.load import LoadConfig
from vllm.model_executor.model_loader.base_loader import BaseModelLoader
from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
from vllm.model_executor.model_loader.utils import initialize_model
from vllm.utils.torch_utils import set_default_torch_dtype

logger = logging.getLogger(__name__)


# Global storage for tensor metadata, keyed by device_id (local CUDA ordinal).
_tensor_registry: dict[int, dict[str, torch.Tensor]] = {}
_nixl_managers: dict[int, NixlTransferManager] = {}
_loader_registry: dict[int, MxModelLoader] = {}


def get_model_loader(device_id: int) -> MxModelLoader | None:
    """Return the ModelExpress loader that completed this device's main load."""
    return _loader_registry.get(device_id)


class MxModelLoader(BaseModelLoader):
    """
    Auto-detecting model loader for ModelExpress.

    Uses LoadStrategyChain to find the best available loading strategy
    (RDMA P2P, GDS, or default disk loading), then registers tensors
    with NIXL and publishes metadata so future nodes can discover this
    one as a source.
    """

    def __init__(self, load_config: LoadConfig):
        super().__init__(load_config)
        configure_vllm_logging()
        # Unconditionally, and off the load path: a run that skips P2P and falls
        # back to a local or HuggingFace path must still bring the exporter up,
        # or it produces output byte-identical to MX_METRICS_ENABLED=0 -- the run
        # you most need to diagnose. No-op unless enabled; never raises.
        enable_metrics()
        self._ctx: LoadContext | None = None

    def load_model(
        self,
        vllm_config: VllmConfig,
        model_config: ModelConfig,
        prefix: str = "",
    ) -> nn.Module:
        """Load model, auto-detecting the best loading strategy.

        `prefix` is vLLM's BaseModelLoader.load_model argument for initializing
        a model subtree. ModelExpress does not interpret it; it is passed through
        to vLLM's initialize_model().
        """
        load_start = time.perf_counter()

        is_speculative_draft = _is_speculative_draft(vllm_config, model_config)
        if is_speculative_draft and envs.MX_LOAD_STRATEGY_CHAIN == "RL":
            raise ValueError(
                "RL initial loading does not support speculative draft models"
            )

        ctx = build_vllm_load_context(vllm_config, model_config)
        ctx.p2p_enabled = not is_speculative_draft
        if envs.MX_ARTIFACT_READY_URL.strip():
            ctx.source_ready_fn = lambda: _vllm_health_ready(ctx)
        if ctx.p2p_enabled:
            self._ctx = ctx

        logger.info(
            f"[Worker {ctx.global_rank}] MxModelLoader starting "
            f"(model={ctx.identity.model_name}, p2p_enabled={ctx.p2p_enabled})"
        )

        # A speculative draft loads through this same path and finishes far
        # sooner than the model the user asked for, so timing them together
        # makes the p99 of neither meaningful. Asked directly rather than
        # inferred from p2p_enabled: that flag happens to agree today, but it is
        # a capability switch and any future reason to clear it would silently
        # relabel real loads as drafts.
        model_role = "draft" if is_speculative_draft else "main"

        # L0 wraps everything below, and the four L1 phases inside it are
        # disjoint, so their sum is bounded by the total by construction. The
        # timers only bracket existing calls; nothing here changes load order.
        model_id = ctx.identity.model_name
        with metrics.time_load("vllm", model_id, model_role):
            with maybe_enter_vmm_arena(ctx):
                if ctx.p2p_enabled:
                    with metrics.time_load_phase("vllm", model_id, "artifact_install"):
                        install_vllm_cache_artifacts(ctx)
                with set_default_torch_dtype(model_config.dtype):
                    with ctx.target_device:
                        with metrics.time_load_phase("vllm", model_id, "model_init"):
                            model = initialize_model(
                                vllm_config=vllm_config,
                                model_config=model_config,
                                prefix=prefix,
                            )

                    with metrics.time_load_phase("vllm", model_id, "chain"):
                        model = run_load_strategy_chain(model, ctx)

                    if ctx.p2p_enabled:
                        _loader_registry[ctx.device_id] = self
                        _tensor_registry[ctx.device_id] = ctx.tensors
                        if ctx.nixl_manager is not None:
                            _nixl_managers[ctx.device_id] = ctx.nixl_manager
                        else:
                            _nixl_managers.pop(ctx.device_id, None)

                        # Scheduling the publish, not completing it: the work is
                        # handed to a background thread, so this phase measures
                        # the handoff and never the upload.
                        with metrics.time_load_phase("vllm", model_id, "publish"):
                            schedule_vllm_cache_artifact_publish(ctx)

            log_arena_post_load(ctx)

        total_time = time.perf_counter() - load_start
        logger.info(
            f"[Worker {ctx.global_rank}] MxModelLoader.load_model() COMPLETE "
            f"in {total_time:.2f}s"
        )
        return model.eval()

    def download_model(self, model_config: ModelConfig) -> None:
        """Download the model so it can be loaded immediately."""
        if envs.MX_LOAD_STRATEGY_CHAIN == "RL":
            logger.info(
                "RL initial load is selected; leaving weight acquisition to "
                "the RL strategy chain"
            )
            return

        if model_prefetch.is_enabled():
            # Without shared storage this would pull the full weight set from
            # Hugging Face before any strategy runs, defeating P2P-first and
            # failing outright when the worker is offline. The strategy chain
            # decides where the weights come from.
            logger.info(
                "MODEL_EXPRESS_NO_SHARED_STORAGE is set; leaving weight "
                "acquisition to the ModelExpress strategy chain"
            )
            return

        import copy

        disk_config = copy.copy(self.load_config)
        try:
            disk_config.load_format = "auto"
        except AttributeError:
            object.__setattr__(disk_config, "load_format", "auto")
        DefaultModelLoader(disk_config).download_model(model_config)

    def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
        """Load weights into an already-initialized model (standalone API)."""
        import copy

        disk_config = copy.copy(self.load_config)
        try:
            disk_config.load_format = "auto"
        except AttributeError:
            object.__setattr__(disk_config, "load_format", "auto")
        DefaultModelLoader(disk_config).load_weights(model, model_config)

    @property
    def nixl_manager(self) -> NixlTransferManager | None:
        """Access the NIXL manager for external use."""
        if self._ctx is not None:
            return self._ctx.nixl_manager
        return None

    @property
    def tensors(self) -> dict[str, torch.Tensor]:
        """Access the registered tensor dict."""
        if self._ctx is not None:
            return self._ctx.tensors
        return {}

    @property
    def worker_id(self) -> str | None:
        """Return the inference P2P worker ID after a completed main load."""
        if self._ctx is not None:
            return self._ctx.worker_id
        return None

    def unpublish_runtime_tensors(self) -> None:
        """Withdraw this loader's runtime tensors before an active refit."""
        if self._ctx is not None:
            drain_tensor_readers(
                self._ctx,
                timeout=envs.MX_TRANSFER_TIMEOUT,
            )
            unpublish_metadata(self._ctx)

    def publish_runtime_tensors(self, version_id: str) -> None:
        """Publish this loader's installed runtime tensors at an exact version."""
        if self._ctx is not None:
            self._ctx.identity.revision = version_id
            publish_metadata(self._ctx)