modelexpress-client 0.6.0

# 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 -> InstantTensor -> ModelStreamer -> GDS -> disk
"""

import logging

from . import envs

_logger = logging.getLogger(__name__)
_loaders_registered = False


def _configure_engine_logging(engine_logger_name: str) -> None:
    """Attach an engine's Python log handlers to the ModelExpress namespace."""
    mx_root = logging.getLogger("modelexpress")
    if mx_root.handlers:
        return
    engine_logger = logging.getLogger(engine_logger_name)
    for handler in engine_logger.handlers:
        mx_root.addHandler(handler)
    mx_level = envs.MODEL_EXPRESS_LOG_LEVEL
    if mx_level and hasattr(logging, mx_level):
        mx_root.setLevel(getattr(logging, mx_level))
    elif engine_logger.level != logging.NOTSET:
        mx_root.setLevel(engine_logger.level)


def configure_vllm_logging() -> None:
    """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.
    """
    _configure_engine_logging("vllm")


def configure_trtllm_logging() -> None:
    """Ensure modelexpress loggers are visible in TRT-LLM worker processes."""
    _configure_engine_logging("TRT-LLM")


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 -> InstantTensor -> ModelStreamer -> 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.publisher import PublisherThread  # noqa: F401
from .model_client import ModelCacheClient  # noqa: F401

__all__ = [
    "GdsTransferManager",
    "ModelCacheClient",
    "MxClient",
    "MxGdsLoader",
    "PublisherThread",
    "configure_trtllm_logging",
    "configure_vllm_logging",
    "register_modelexpress_loaders",
]