modelexpress-client 0.7.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
from abc import ABC, abstractmethod

import grpc

from . import auth
from . import envs
from . import p2p_pb2
from . import p2p_pb2_grpc

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


class MxClientBase(ABC):
    """Abstract metadata client.

    Both the central-coordinator :class:`MxClient` and the decentralized
    :class:`~.k8s_service_client.MxK8sServiceClient` implement this
    interface; callers (loaders, strategies) depend on the base class so
    backend swaps are a factory-level choice.

    ``REQUIRES_P2P_METADATA`` lets a backend opt into forcing the P2P
    path in :func:`metadata.publish_metadata_and_ready` regardless of
    the ``MX_P2P_METADATA`` env var. Decentralized backends set it to
    True because they have no central store to fall back to.
    """

    REQUIRES_P2P_METADATA: bool = False

    def worker_rpc_retry_policy(self) -> tuple[int, float]:
        """Return fresh-channel retries and backoff for worker RPCs."""
        return 0, 0.0

    @abstractmethod
    def publish_metadata(
        self,
        identity: "p2p_pb2.SourceIdentity",
        worker: "p2p_pb2.WorkerMetadata",
        worker_id: str,
    ) -> str:
        """Publish worker metadata and return the computed mx_source_id."""

    @abstractmethod
    def list_sources(
        self,
        identity: "p2p_pb2.SourceIdentity | None" = None,
        status_filter: "p2p_pb2.SourceStatus | None" = None,
    ) -> "p2p_pb2.ListSourcesResponse":
        """List candidate source workers matching the given identity."""

    @abstractmethod
    def get_metadata(
        self,
        mx_source_id: str,
        worker_id: str,
    ) -> "p2p_pb2.GetMetadataResponse":
        """Fetch full worker metadata for one specific source."""

    @abstractmethod
    def update_status(
        self,
        mx_source_id: str,
        worker_id: str,
        worker_rank: int,
        status: "p2p_pb2.SourceStatus",
        source_load: float | None = None,
    ) -> bool:
        """Update a source worker's lifecycle status."""

    @abstractmethod
    def close(self) -> None:
        """Release any resources held by the 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 (deprecated, but still takes precedence:
       the TRT-LLM live-transfer integration reads only this name)
    3. ``MX_SERVER_ADDRESS`` env var (the name ModelExpress is standardizing on)
    4. Default ``localhost:8001``
    """
    if explicit_url:
        return _parse_server_address(explicit_url)
    url = envs.MODEL_EXPRESS_URL
    if url is None:
        url = envs.MX_SERVER_ADDRESS
        if url is None:
            url = "localhost:8001"
    return _parse_server_address(url)


class MxClient(MxClientBase):
    """
    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 = auth.with_auth(
                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,
            pod_name=envs.POD_NAME,
            pod_uid=envs.POD_UID,
            pod_namespace=envs.POD_NAMESPACE,
        )
        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",
        source_load: float | None = None,
    ) -> 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,
        )
        # Leave the optional field unset when there is no reading: presence is
        # how the server and pullers tell "unknown" from a measured 0.0.
        if source_load is not None:
            request.source_load = source_load
        response = self.stub.UpdateStatus(request, timeout=30)
        if not response.success:
            logger.error("UpdateStatus failed: %s", response.message)
        return response.success