modelexpress-client 0.3.0

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

"""
ModelExpress Client for P2P GPU Weight Transfers.

Orchestrates NIXL/RDMA transfers between vLLM workers. The client fetches
NIXL metadata from workers via ZMQ, queries the ModelExpress server for
existing sources, and instructs workers to receive weights if found.

NIXL agents live in vLLM workers (not here) because GPU memory must be
registered by the owning process for GPUDirect RDMA.
"""

import logging
import os

import grpc

from . import p2p_pb2
from . import p2p_pb2_grpc

logger = logging.getLogger("modelexpress.client")


def _parse_server_address(address: str) -> str:
    """Strip http:// or https:// prefix from server address for gRPC."""
    if address.startswith("http://"):
        return address[7:]
    elif address.startswith("https://"):
        return address[8:]
    return address


def _get_server_url(explicit_url: str | None = None) -> str:
    """
    Resolve the ModelExpress server URL.

    Priority:
    1. Explicit ``server_url`` argument
    2. ``MODEL_EXPRESS_URL`` env var (Dynamo-consistent)
    3. ``MX_SERVER_ADDRESS`` env var (backward compat)
    4. Default ``localhost:8001``
    """
    if explicit_url:
        return _parse_server_address(explicit_url)
    url = os.environ.get(
        "MODEL_EXPRESS_URL",
        os.environ.get("MX_SERVER_ADDRESS", "localhost:8001"),
    )
    return _parse_server_address(url)


class MxClient:
    """
    Lightweight gRPC client for ModelExpress server communication.

    Provides typed methods for every P2P RPC (``PublishMetadata``,
    ``ListSources``, ``GetMetadata``, ``UpdateStatus``) so that callers
    (loaders, coordinators) never need to create gRPC channels or
    stubs directly.

    The connection is created lazily on first use.

    Args:
        server_url: Explicit server address (``host:port``).  When
            *None* the address is resolved via ``MODEL_EXPRESS_URL``
            or ``MX_SERVER_ADDRESS`` env vars, falling back to
            ``localhost:8001``.
        max_message_size: Max send/receive message size in bytes.
    """

    def __init__(
        self,
        server_url: str | None = None,
        max_message_size: int = 100 * 1024 * 1024,  # 100 MB
    ):
        self.server_url = _get_server_url(server_url)
        self._max_message_size = max_message_size
        self._channel: grpc.Channel | None = None
        self._stub: p2p_pb2_grpc.P2pServiceStub | None = None

    # -- connection management ------------------------------------------------

    @property
    def stub(self) -> p2p_pb2_grpc.P2pServiceStub:
        """Return (and lazily create) the gRPC stub."""
        if self._channel is None:
            options = [
                ("grpc.max_send_message_length", self._max_message_size),
                ("grpc.max_receive_message_length", self._max_message_size),
            ]
            self._channel = grpc.insecure_channel(self.server_url, options=options)
            self._stub = p2p_pb2_grpc.P2pServiceStub(self._channel)
            logger.debug("MxClient connected to %s", self.server_url)
        return self._stub

    def close(self) -> None:
        """Close the underlying gRPC channel."""
        if self._channel is not None:
            self._channel.close()
            self._channel = None
            self._stub = None

    # -- RPC wrappers ---------------------------------------------------------

    def publish_metadata(
        self,
        identity: "p2p_pb2.SourceIdentity",
        worker: "p2p_pb2.WorkerMetadata",
        worker_id: str,
    ) -> str:
        """Publish metadata for one worker so targets can discover this source.

        Returns the *mx_source_id* (16-char hex) on success, raises on failure.
        """
        request = p2p_pb2.PublishMetadataRequest(
            identity=identity,
            worker=worker,
            worker_id=worker_id,
        )
        response = self.stub.PublishMetadata(request, timeout=30)
        if not response.success:
            raise RuntimeError(f"PublishMetadata failed: {response.message}")
        return response.mx_source_id

    def list_sources(
        self,
        identity: "p2p_pb2.SourceIdentity | None" = None,
        status_filter: "p2p_pb2.SourceStatus | None" = None,
    ) -> "p2p_pb2.ListSourcesResponse":
        """List available source workers, optionally filtered by identity and status."""
        request = p2p_pb2.ListSourcesRequest(
            identity=identity,
            status_filter=status_filter,
        )
        return self.stub.ListSources(request, timeout=30)

    def get_metadata(
        self,
        mx_source_id: str,
        worker_id: str,
    ) -> "p2p_pb2.GetMetadataResponse":
        """Fetch full tensor metadata for one specific worker."""
        request = p2p_pb2.GetMetadataRequest(
            mx_source_id=mx_source_id,
            worker_id=worker_id,
        )
        return self.stub.GetMetadata(request, timeout=30)

    def update_status(
        self,
        mx_source_id: str,
        worker_id: str,
        worker_rank: int,
        status: "p2p_pb2.SourceStatus",
    ) -> bool:
        """Update worker status.  Returns *True* on success."""
        request = p2p_pb2.UpdateStatusRequest(
            mx_source_id=mx_source_id,
            worker_id=worker_id,
            worker_rank=worker_rank,
            status=status,
        )
        response = self.stub.UpdateStatus(request, timeout=30)
        if not response.success:
            logger.error("UpdateStatus failed: %s", response.message)
        return response.success