modelexpress-client 0.7.0

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

"""RDMA P2P loading strategy: receive weights from an existing source via NIXL."""

from __future__ import annotations

import logging
import math
import os
import time

from .. import envs, p2p_pb2
from ..adapter import EngineAdapter, StrategyFailed, StrategyRecoveryError
from ..metadata.payload import (
    accelerators_compatible,
    worker_tensor_count,
    worker_tensor_descriptors,
)
from ..metrics import metrics as selection_metrics
from ..nixl_transfer import is_nixl_available
from ..source_selection import (
    configured_policy_label,
    get_configured_selector,
)
from ..transfer_safety import check_transfer_allowed
from ..types import TensorDescriptor
from .base import (
    LoadContext,
    LoadStrategy,
    SourceTransferError,
    _as_load_result,
    clear_exception_tracebacks,
    register_tensors,
)
from .context import LoadResult

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

MAX_SOURCE_RETRIES = 3

# Fallback when MX_TRANSFER_TIMEOUT is unset or unusable. Matches the value this
# path used when the timeout was hard-coded, so behaviour is unchanged by default.
DEFAULT_RDMA_TRANSFER_TIMEOUT_S = 300.0


def _transfer_timeout_seconds() -> float:
    """Per-source RDMA receive budget, in seconds.

    Every candidate costs this long when a source is wedged, because a wedged QP
    produces neither a completion nor an error status - the timeout is the only
    thing that ends the wait. With MAX_SOURCE_RETRIES candidates the worst case is
    that multiplied by three before the target falls back to disk, so operators
    need to be able to size it against their model rather than accept a constant.

    Honours MX_TRANSFER_TIMEOUT only when it is *explicitly set*, and reads
    os.environ directly to tell that apart from the default. This is deliberate:
    that variable defaults to 900, so simply adopting ``envs.MX_TRANSFER_TIMEOUT``
    would triple this path's budget from 300s to 900s and make the stall being
    fixed three times more expensive for anyone who never configured it.

    A non-positive or unparseable value falls back rather than being honoured. In
    particular 0 must not become "no timeout": ``receive_from_source`` treats None
    as wait-forever, which is the failure mode this exists to bound.
    """
    raw = os.environ.get("MX_TRANSFER_TIMEOUT")
    if raw is None:
        return DEFAULT_RDMA_TRANSFER_TIMEOUT_S
    try:
        configured = float(raw)
    except (TypeError, ValueError):
        logger.warning(
            "MX_TRANSFER_TIMEOUT=%r is not a number; using %.1fs for the RDMA "
            "receive budget",
            raw,
            DEFAULT_RDMA_TRANSFER_TIMEOUT_S,
        )
        return DEFAULT_RDMA_TRANSFER_TIMEOUT_S
    if not math.isfinite(configured) or configured <= 0:
        logger.warning(
            "MX_TRANSFER_TIMEOUT=%r is not finite and positive; using %.1fs for the RDMA "
            "receive budget rather than waiting indefinitely",
            raw,
            DEFAULT_RDMA_TRANSFER_TIMEOUT_S,
        )
        return DEFAULT_RDMA_TRANSFER_TIMEOUT_S
    return configured


class RdmaStrategy(LoadStrategy):
    """Load weights via RDMA P2P transfer from an existing source.

    Overrides load() entirely since RDMA has a fundamentally different flow:
    prepare target storage -> RDMA receive -> register + publish.
    """

    name = "rdma"
    requires = (EngineAdapter.discover_tensors,)
    requires_tensor_read_lease = False

    def rollback(self, ctx: LoadContext) -> None:
        """Clean up NIXL state from a failed RDMA target attempt."""
        if ctx.nixl_manager is not None:
            ctx.nixl_manager.shutdown()
        ctx.tensors = {}
        ctx.nixl_manager = None

    def is_available(self, ctx: LoadContext) -> bool:
        if not ctx.p2p_enabled:
            return False
        if not super().is_available(ctx):
            return False
        if not is_nixl_available():
            return False

        if not ctx.accelerator_backend.supports_rdma_p2p():
            logger.info(
                f"[Worker {ctx.global_rank}] Backend "
                f"{ctx.accelerator_backend.name} does not support RDMA P2P, skipping"
            )
            return False

        # Decentralized backends (k8s-service) serve their own
        # metadata; skip the central-server precondition for them.
        # Strict `is True` check so MagicMock's auto-attribute doesn't
        # masquerade as the flag in tests.
        server_addr = (
            ctx.mx_server_url or envs.MODEL_EXPRESS_URL or envs.MX_SERVER_ADDRESS
        )
        requires_p2p = getattr(ctx.mx_client, "REQUIRES_P2P_METADATA", False) is True
        if not server_addr and not requires_p2p:
            logger.info(
                f"[Worker {ctx.global_rank}] No MX server configured, skipping RDMA"
            )
            return False

        allowed, reason = check_transfer_allowed(ctx.model_config)
        if not allowed:
            logger.info(f"[Worker {ctx.global_rank}] RDMA transfer disabled: {reason}")
            return False

        return True

    def load(self, result: LoadResult, ctx: LoadContext) -> LoadResult:
        """Load from a READY source or raise StrategyFailed for fallback.

        Source discovery and metadata misses do not mutate the target model.
        A failed transfer retries the next candidate. If it may have mutated
        the model, the adapter first replaces the model with a fresh instance.
        """
        result = _as_load_result(result)
        candidates = self._find_source_instances(ctx)
        if not candidates:
            logger.info(
                f"[Worker {ctx.global_rank}] No RDMA source available, skipping"
            )
            raise StrategyFailed("No RDMA source available", mutated=False)

        attempts = candidates[:MAX_SOURCE_RETRIES]
        policy = configured_policy_label()
        for attempt_index, instance in enumerate(attempts):
            mx_source_id = instance.mx_source_id
            worker_id = instance.worker_id
            logger.info(
                f"[Worker {ctx.global_rank}] Source attempt: "
                f"source_attempt_index={attempt_index} "
                f"source_worker_id={worker_id} mx_source_id={mx_source_id}"
            )

            try:
                with selection_metrics.time_source_attempt_phase(policy, "metadata"):
                    source_worker = self._fetch_worker_metadata(
                        ctx,
                        mx_source_id,
                        worker_id,
                    )
            except Exception as e:
                logger.warning(
                    f"[Worker {ctx.global_rank}] Failed to fetch metadata for worker {worker_id}: {e}. "
                    f"Trying next candidate."
                )
                selection_metrics.record_metadata_failure(policy)
                selection_metrics.record_attempt(policy, "metadata_miss")
                continue

            if source_worker is None:
                selection_metrics.record_attempt(policy, "metadata_miss")
                continue

            if not self._accelerator_compatible(ctx, source_worker, worker_id):
                continue

            logger.info(
                f"[Worker {ctx.global_rank}] Trying source worker {worker_id} "
                f"({worker_tensor_count(source_worker)} tensors)"
            )

            # The peer id is always passed; the collector drops it unless
            # MX_METRICS_SOURCE_ID_LABEL=1, so the cardinality decision lives in
            # one place rather than being duplicated at every call site.
            selection_metrics.record_selection(policy, worker_id)
            transfer_start = time.perf_counter()
            try:
                out = self._load_as_target(
                    result,
                    ctx,
                    source_worker,
                    mx_source_id,
                    worker_id,
                )
            except StrategyFailed as e:
                has_next_candidate = attempt_index + 1 < len(attempts)
                selection_metrics.observe_transfer_seconds(
                    policy,
                    "retry" if has_next_candidate else "fallback",
                    time.perf_counter() - transfer_start,
                )
                selection_metrics.record_attempt(
                    policy,
                    "transfer_retry" if has_next_candidate else "transfer_fallback",
                )
                if not has_next_candidate:
                    raise

                logger.warning(
                    f"[Worker {ctx.global_rank}] RDMA source worker {worker_id} "
                    f"failed: {e}. Trying next candidate."
                )
                try:
                    self.rollback(ctx)
                except Exception as cleanup_error:
                    raise StrategyFailed(
                        f"Failed to clean up target after source worker "
                        f"{worker_id} failed: {cleanup_error}",
                        mutated=True,
                    ) from cleanup_error
                if e.mutated:
                    try:
                        clear_exception_tracebacks(e)
                        reinitialized = ctx.adapter.reinit_for_retry(result)
                        # LoadResult is the stable envelope shared with the
                        # outer strategy chain. Some adapters return a new
                        # envelope, so copy its restored state back rather than
                        # leaving the outer owner with the cleared pre-retry
                        # object if all later candidates miss.
                        if reinitialized is not result:
                            vars(result).update(vars(reinitialized))
                    except Exception as reinit_error:
                        raise StrategyRecoveryError(
                            f"Failed to reinitialize target after source worker "
                            f"{worker_id} failed: {reinit_error}",
                        ) from reinit_error
                continue
            except BaseException:
                selection_metrics.observe_transfer_seconds(
                    policy, "fallback", time.perf_counter() - transfer_start
                )
                selection_metrics.record_attempt(policy, "transfer_fallback")
                raise
            selection_metrics.observe_transfer_seconds(
                policy, "success", time.perf_counter() - transfer_start
            )
            selection_metrics.record_attempt(policy, "success")
            return out

        tried = min(len(candidates), MAX_SOURCE_RETRIES)
        logger.warning(
            f"[Worker {ctx.global_rank}] Tried {tried} of {len(candidates)} source workers "
            f"(max retries={MAX_SOURCE_RETRIES}), falling through"
        )
        raise StrategyFailed(
            "No RDMA source succeeded",
            mutated=False,
        )

    def _find_source_instances(
        self,
        ctx: LoadContext,
    ) -> list[p2p_pb2.SourceInstanceRef]:
        """Return READY source instances ranked by the configured selector.

        Filters listed instances to the target's worker_rank, then delegates
        ordering to the policy named by MX_P2P_SOURCE_SELECTOR (default
        ``random``). The retry slice (MAX_SOURCE_RETRIES) is applied by the
        caller in load(), so the selector controls ordering only.
        """
        policy = configured_policy_label()
        try:
            list_resp = ctx.mx_client.list_sources(
                identity=ctx.identity,
                status_filter=p2p_pb2.SOURCE_STATUS_READY,
            )
            if not list_resp.instances:
                logger.debug(
                    f"[Worker {ctx.global_rank}] No ready source instances found"
                )
                # Record the empty funnel before returning. Without this,
                # ``stage="listed"`` could never observe zero -- the one bucket
                # that distinguishes "no peers published" from "peers listed but
                # every one filtered out" was unreachable, so the two looked
                # identical on a dashboard.
                selection_metrics.record_list_sources(policy, "empty")
                for stage in ("listed", "rank_matched", "accelerator_matched"):
                    selection_metrics.observe_candidates(policy, stage, 0)
                return []

            rank_matched = [
                inst
                for inst in list_resp.instances
                if inst.worker_rank == ctx.worker_rank
            ]

            # Drop accelerator-incompatible sources before ordering and the
            # retry-cap slice, so the selector only ranks eligible peers and a
            # compatible source can never be pushed past MAX_SOURCE_RETRIES by
            # incompatible ones. The post-GetMetadata check in load() stays as
            # defense-in-depth (empty refs, stale records, metadata drift).
            target_accelerator = ctx.accelerator_backend.name
            candidates = [
                inst
                for inst in rank_matched
                if accelerators_compatible(
                    target_accelerator,
                    inst.accelerator,
                    mx_source_type=ctx.identity.mx_source_type,
                    quantization=ctx.identity.quantization,
                    dtype=ctx.identity.dtype,
                )
            ]

            selector = get_configured_selector()
            select_start = time.perf_counter()
            ordered = selector.order(candidates, ctx)
            select_seconds = time.perf_counter() - select_start

            selection_metrics.record_list_sources(selector.name, "ok")
            selection_metrics.observe_candidates(
                selector.name, "listed", len(list_resp.instances)
            )
            selection_metrics.observe_candidates(
                selector.name, "rank_matched", len(rank_matched)
            )
            selection_metrics.observe_candidates(
                selector.name, "accelerator_matched", len(candidates)
            )
            selection_metrics.observe_selection_seconds(selector.name, select_seconds)

            # Surface the source-published load the client saw, so a dashboard can
            # compare how load_aware steers vs random/rendezvous. The collector
            # decides what one series means when the per-peer label is off.
            selection_metrics.observe_candidate_loads(ordered)

            logger.info(
                f"[Worker {ctx.global_rank}] Source selection: "
                f"source_selector={selector.name} "
                f"source_candidates_total={len(list_resp.instances)} "
                f"source_candidates_rank_matched={len(rank_matched)} "
                f"source_candidates_accelerator_matched={len(candidates)}"
            )
            if ordered:
                logger.debug(
                    f"[Worker {ctx.global_rank}] Ranked source workers: "
                    f"{[inst.worker_id for inst in ordered]}"
                )
            return ordered

        except Exception as e:
            logger.warning(
                f"[Worker {ctx.global_rank}] Error listing sources, falling through: {e}"
            )
            # A ListSources RPC failure used to record nothing at all, so a
            # complete backend outage was indistinguishable from a cluster with
            # no peers -- both showed up as an absence. This counter is what
            # separates them.
            selection_metrics.record_list_sources(policy, "error")
            return []

    def _accelerator_compatible(
        self,
        ctx: LoadContext,
        source_worker: p2p_pb2.WorkerMetadata,
        worker_id: str,
    ) -> bool:
        """Return whether source and target accelerator metadata are compatible.

        Defense-in-depth re-check on the authoritative ``WorkerMetadata`` after
        GetMetadata: the pre-slice filter in ``_find_source_instances`` uses the
        lightweight ``SourceInstanceRef.accelerator``, which may be empty on old
        servers or drift between list and fetch.
        """
        target_accelerator = ctx.accelerator_backend.name
        source_accelerator = source_worker.accelerator
        if accelerators_compatible(
            target_accelerator,
            source_accelerator,
            mx_source_type=ctx.identity.mx_source_type,
            quantization=ctx.identity.quantization,
            dtype=ctx.identity.dtype,
            defer_unknown=False,
        ):
            return True

        logger.info(
            f"[Worker {ctx.global_rank}] Skipping source worker {worker_id}: "
            f"accelerator mismatch source={source_accelerator!r}, "
            f"target={target_accelerator!r}"
        )
        return False

    def _fetch_worker_metadata(
        self,
        ctx: LoadContext,
        mx_source_id: str,
        worker_id: str,
    ) -> p2p_pb2.WorkerMetadata | None:
        """Fetch tensor metadata for one worker."""
        fetch_start = time.perf_counter()
        metadata_resp = ctx.mx_client.get_metadata(
            mx_source_id=mx_source_id,
            worker_id=worker_id,
        )
        if not metadata_resp.found:
            logger.debug(
                f"[Worker {ctx.global_rank}] Metadata not found for worker {worker_id}, skipping"
            )
            return None
        worker = metadata_resp.worker
        has_tensor_descriptors = bool(worker_tensor_descriptors(worker))
        if not has_tensor_descriptors and not worker.worker_grpc_endpoint:
            logger.debug(
                f"[Worker {ctx.global_rank}] Worker {worker_id} has no tensors "
                f"and no P2P endpoint, skipping"
            )
            return None
        # A worker ID requires endpoint validation. Without one, fetch only
        # when the metadata response did not already include the manifest.
        needs_manifest_prefetch = worker_id != "" or not has_tensor_descriptors
        if worker.worker_grpc_endpoint and needs_manifest_prefetch:
            self._prefetch_tensor_manifest(ctx, worker, mx_source_id, worker_id)
        fetch_time = time.perf_counter() - fetch_start
        mode = "P2P (lightweight)" if worker.worker_grpc_endpoint else "centralized"
        tensor_count = worker_tensor_count(worker)
        logger.info(
            f"[Worker {ctx.global_rank}] [TIMING] GetMetadata ({mode}): "
            f"{fetch_time:.3f}s, {tensor_count} tensors"
        )
        return worker

    def _prefetch_tensor_manifest(
        self,
        ctx: LoadContext,
        source_worker: p2p_pb2.WorkerMetadata,
        mx_source_id: str,
        worker_id: str,
    ) -> None:
        """Fetch once before target preparation and retain it for the transfer."""
        from ..metadata.worker_server import fetch_tensor_manifest

        manifest_start = time.perf_counter()
        tensor_protos, manifest_bytes = fetch_tensor_manifest(
            endpoint=source_worker.worker_grpc_endpoint,
            mx_source_id=mx_source_id,
            worker_id=worker_id,
        )
        # Store the validated descriptors on the metadata object so the
        # transfer can reuse them without a second manifest RPC.
        source_worker.tensor_source.ClearField("tensors")
        source_worker.tensor_source.tensors.extend(tensor_protos)
        manifest_time = time.perf_counter() - manifest_start
        logger.info(
            f"[Worker {ctx.global_rank}] [TIMING] P2P tensor manifest: "
            f"{manifest_time:.3f}s ({len(tensor_protos)} tensors, "
            f"{manifest_bytes} bytes)"
        )

    @staticmethod
    def _use_prefetched_manifest(source_worker):
        return (
            None,
            worker_tensor_descriptors(source_worker),
            source_worker.metadata_endpoint,
            source_worker.agent_name,
        )

    def _prepare_leased_manifest(
        self,
        ctx: LoadContext,
        source_worker,
        mx_source_id: str,
        source_worker_id: str,
    ):
        from ..metadata.worker_server import prepare_tensor_read

        max_retries, retry_backoff = ctx.mx_client.worker_rpc_retry_policy()
        tensor_read, _ = prepare_tensor_read(
            source_worker.worker_grpc_endpoint,
            mx_source_id,
            worker_id=source_worker_id,
            timeout=_transfer_timeout_seconds(),
            max_retries=max_retries,
            retry_backoff_seconds=retry_backoff,
        )
        manifest = tensor_read.manifest
        return (
            tensor_read,
            manifest.tensors,
            manifest.metadata_endpoint,
            manifest.agent_name,
        )

    def _load_as_target(
        self,
        result: LoadResult,
        ctx: LoadContext,
        source_worker,
        mx_source_id: str,
        source_worker_id: str,
    ) -> LoadResult:
        """Receive fully-processed weights via RDMA from an existing source."""
        # The source-selection policy (random, rendezvous_hash, load_aware, ...)
        # that chose this candidate. Every P2P client family carries it, so a
        # phase duration can be read against the policy that picked the peer.
        policy = configured_policy_label()
        try:
            with selection_metrics.time_source_attempt_phase(policy, "prepare"):
                result = ctx.adapter.prepare_rdma_target(result)
                result = ctx.adapter.before_rdma_receive(result)
            self._receive_from_peer(
                result,
                ctx,
                source_worker,
                mx_source_id,
                source_worker_id,
            )
            with selection_metrics.time_source_attempt_phase(policy, "finalize"):
                return ctx.adapter.after_rdma_receive(result)
        except StrategyFailed:
            raise
        except Exception as e:
            raise StrategyFailed(str(e), mutated=True) from e

    def _receive_from_peer(
        self,
        result: LoadResult,
        ctx: LoadContext,
        source_worker,
        mx_source_id: str,
        source_worker_id: str,
    ) -> None:
        """Receive fully-processed tensors via RDMA from the detected source."""
        receive_start = time.perf_counter()
        policy = configured_policy_label()
        with selection_metrics.time_source_attempt_phase(policy, "register"):
            register_tensors(result, ctx)

        is_p2p = bool(source_worker.worker_grpc_endpoint)
        remote_agent_name = None

        tensor_read = None
        try:
            if is_p2p:
                if self.requires_tensor_read_lease:
                    tensor_read, tensor_protos, ep, agent_name = (
                        self._prepare_leased_manifest(
                            ctx,
                            source_worker,
                            mx_source_id,
                            source_worker_id,
                        )
                    )
                else:
                    tensor_read, tensor_protos, ep, agent_name = (
                        self._use_prefetched_manifest(source_worker)
                    )
                source_tensors = [
                    TensorDescriptor(
                        name=t.name,
                        addr=t.addr,
                        size=t.size,
                        device_id=t.device_id,
                        dtype=t.dtype,
                    )
                    for t in tensor_protos
                ]
                host, port_str = ep.rsplit(":", 1)
                # Claimed before the dial so a fetch that fails part-way, leaving
                # metadata that lands later, is still released.
                remote_agent_name = agent_name
                # One handshake phase, two transports: here the peer's NIXL
                # metadata is fetched from the peer itself; the centralized
                # branch below loads the copy the MX server holds. Exactly one
                # of the two runs per attempt.
                try:
                    with selection_metrics.time_source_attempt_phase(
                        policy, "handshake"
                    ) as handshake:
                        ctx.nixl_manager.fetch_remote_and_wait(
                            remote_agent_name=agent_name,
                            ip=host,
                            port=int(port_str),
                        )
                except BaseException:
                    ctx.nixl_manager.remove_remote_agent(remote_agent_name)
                    remote_agent_name = None
                    raise
                logger.info(
                    f"[Worker {ctx.global_rank}] [TIMING] P2P NIXL metadata fetch: "
                    f"{handshake.seconds:.3f}s"
                )
            else:
                source_tensors = [
                    TensorDescriptor(
                        name=t.name,
                        addr=t.addr,
                        size=t.size,
                        device_id=t.device_id,
                        dtype=t.dtype,
                    )
                    for t in worker_tensor_descriptors(source_worker)
                ]
                # Loaded here rather than inside receive_from_source so this method
                # holds the name it is responsible for releasing.
                with selection_metrics.time_source_attempt_phase(
                    policy, "handshake"
                ) as handshake:
                    remote_agent_name = ctx.nixl_manager.add_remote_agent(
                        source_worker.nixl_metadata
                    )
                logger.info(
                    f"[Worker {ctx.global_rank}] [TIMING] add_remote_agent: "
                    f"{handshake.seconds:.3f}s (agent={remote_agent_name})"
                )

            try:
                logger.info(
                    f"[Worker {ctx.global_rank}] Receiving "
                    f"{len(source_tensors)} tensors from source"
                    f"{' (P2P)' if is_p2p else ''}"
                )

                # Cross-family (heterogeneous) transfers must name the exact same tensor
                # set on both sides: a name diff can mean vendor-specific hidden/derived
                # tensors, which would leave part of the target at dummy values while
                # RDMA reports success. Same-family transfers tolerate subset transfers.
                target_accelerator = ctx.accelerator_backend.name
                source_accelerator = source_worker.accelerator
                require_exact_match = (
                    ctx.adapter.requires_exact_tensor_catalog()
                    or bool(
                        target_accelerator
                        and source_accelerator
                        and target_accelerator != source_accelerator
                    )
                )

                try:
                    with selection_metrics.time_source_attempt_phase(
                        policy, "receive"
                    ) as received:
                        (
                            bytes_transferred,
                            tensor_count,
                            _,
                        ) = ctx.nixl_manager.receive_from_source(
                            source_metadata=source_worker.nixl_metadata,
                            source_tensors=source_tensors,
                            timeout_seconds=_transfer_timeout_seconds(),
                            remote_agent_name=remote_agent_name,
                            require_exact_match=require_exact_match,
                        )
                except Exception as e:
                    raise SourceTransferError(f"RDMA receive failed: {e}") from e
                transfer_time = received.seconds

                bandwidth_gbps = (
                    (bytes_transferred * 8) / (transfer_time * 1e9)
                    if transfer_time > 0
                    else 0
                )
                logger.info(
                    f"[Worker {ctx.global_rank}] [TIMING] RDMA transfer complete: "
                    f"{tensor_count} tensors, {bytes_transferred / 1e9:.2f} GB, "
                    f"{transfer_time:.3f}s, {bandwidth_gbps:.1f} Gbps"
                )

                ctx.accelerator_backend.synchronize()
            finally:
                # A weight load is one-shot, so release the source here instead of at
                # process exit: the engine process is torn down without running atexit
                # hooks, and in P2P only the reader holds a record of the peer to
                # invalidate. None means acquisition failed with nothing to release.
                if remote_agent_name is not None:
                    with selection_metrics.time_source_attempt_phase(
                        policy, "release"
                    ):
                        ctx.nixl_manager.remove_remote_agent(remote_agent_name)
        finally:
            if tensor_read is not None:
                tensor_read.close()

        total_time = time.perf_counter() - receive_start
        logger.info(
            f"[Worker {ctx.global_rank}] [TIMING] Total receive time: {total_time:.2f}s"
        )