modelexpress-client 0.7.0

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

"""A reader must release the source as part of the load, not at process exit.

Regression cover for nvbug 6519532, which reproduced on 0.5.0-rc.3 *after* the
peer-disconnect machinery had shipped: it was reachable only from an ``atexit``
hook, and the NIXL agent lives in a process that dies without running one. The
disconnect existed and never ran, leaving the source with a half-open QP that
stalled every later reader.

These drive ``_receive_from_peer`` rather than calling teardown directly, so
that reachability is asserted rather than assumed.
"""

from __future__ import annotations

from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import pytest

from modelexpress.load_strategy.base import SourceTransferError
from modelexpress.load_strategy.rdma_strategy import RdmaStrategy

P2P_AGENT = "mx-auto-worker0-0d6a01a9"
BLOB_AGENT = "agent-from-blob"


def _descriptor(name="w0"):
    return SimpleNamespace(name=name, addr=4096, size=128, device_id=0, dtype="torch.bfloat16")


def _manager():
    mgr = MagicMock()
    mgr.receive_from_source.return_value = (128, 1, 0.01)
    mgr.add_remote_agent.return_value = BLOB_AGENT
    mgr.remove_remote_agent.return_value = True
    return mgr


def _ctx(manager):
    # Matching accelerators and an adapter that does not demand an exact catalog
    # keep require_exact_match off, so these tests exercise peer lifecycle rather
    # than manifest strictness. Pinned rather than left to MagicMock, whose auto
    # attributes are truthy and would quietly turn strictness on.
    adapter = MagicMock()
    adapter.requires_exact_tensor_catalog.return_value = False
    return SimpleNamespace(
        global_rank=0,
        worker_rank=0,
        worker_id="target-0",
        nixl_manager=manager,
        accelerator_backend=SimpleNamespace(name="cuda", synchronize=MagicMock()),
        adapter=adapter,
        identity=SimpleNamespace(model_name="m"),
        mx_client=SimpleNamespace(
            worker_rpc_retry_policy=lambda: (0, 0.0),
        ),
    )


def _source_worker(p2p: bool):
    """A source the strategy will read from, in P2P or centralized mode."""
    return SimpleNamespace(
        agent_name=P2P_AGENT,
        # Presence of this endpoint is what selects the P2P branch.
        worker_grpc_endpoint="10.0.18.37:5556" if p2p else "",
        metadata_endpoint="10.0.18.37:5555",
        nixl_metadata=b"blob",
        accelerator="cuda",
    )


def _receive(manager, p2p=True, ctx=None, leased=True):
    strategy = RdmaStrategy()
    strategy.requires_tensor_read_lease = leased

    class _Lease:
        manifest = SimpleNamespace(
            tensors=[_descriptor()],
            metadata_endpoint="10.0.18.37:5555",
            agent_name=P2P_AGENT,
        )

        def __enter__(self):
            return self

        def close(self):
            callback = getattr(manager, "_tensor_lease_release", None)
            if callback is not None:
                callback()

        def __exit__(self, *_args):
            self.close()
            return None

    with (
        patch("modelexpress.load_strategy.rdma_strategy.register_tensors"),
        patch(
            "modelexpress.metadata.worker_server.prepare_tensor_read",
            return_value=(_Lease(), 0),
        ) as prepare_read,
        patch(
            "modelexpress.load_strategy.rdma_strategy.worker_tensor_descriptors",
            return_value=[_descriptor()],
        ),
    ):
        strategy._receive_from_peer(
            MagicMock(),
            ctx or _ctx(manager),
            _source_worker(p2p),
            "src-1",
            "source-worker-1",
        )
    return prepare_read


class TestReleaseOnTheLoadPath:
    def test_standard_p2p_does_not_prepare_a_tensor_read_lease(self):
        mgr = _manager()

        assert RdmaStrategy.requires_tensor_read_lease is False
        prepare_read = _receive(mgr, p2p=True, leased=False)

        prepare_read.assert_not_called()

    def test_p2p_uses_metadata_backend_retry_policy(self):
        mgr = _manager()
        ctx = _ctx(mgr)
        ctx.mx_client = SimpleNamespace(
            worker_rpc_retry_policy=lambda: (4, 0.25),
        )

        prepare_read = _receive(mgr, p2p=True, ctx=ctx)

        assert prepare_read.call_args.kwargs["max_retries"] == 4
        assert prepare_read.call_args.kwargs["retry_backoff_seconds"] == 0.25

    def test_p2p_source_is_released_after_a_successful_load(self):
        """The case rc.3 disproved: cleanup must not wait for interpreter exit."""
        mgr = _manager()
        _receive(mgr, p2p=True)
        mgr.remove_remote_agent.assert_called_once_with(P2P_AGENT)

    def test_the_released_peer_is_the_one_that_was_fetched(self):
        """Releasing a different name would leave the real QP half-open."""
        mgr = _manager()
        _receive(mgr, p2p=True)
        fetched = mgr.fetch_remote_and_wait.call_args.kwargs["remote_agent_name"]
        assert mgr.remove_remote_agent.call_args.args[0] == fetched

    def test_the_peer_outlives_the_transfer_and_the_device_sync(self):
        """Release has to be last. Disconnecting before the READ completes would
        tear down the QP under it, and doing so before the device sync would rely
        on receive_from_source's internal sync staying where it is - a silent
        corruption if that ever moves, rather than a loud failure."""
        mgr = _manager()
        order = []
        mgr.receive_from_source.side_effect = lambda **_: (
            order.append("transfer") or (128, 1, 0.01)
        )
        mgr.remove_remote_agent.side_effect = lambda *_: order.append("release") or True
        ctx = _ctx(mgr)
        ctx.accelerator_backend.synchronize.side_effect = lambda: order.append("sync")
        mgr._tensor_lease_release = lambda: order.append("lease_release")

        _receive(mgr, p2p=True, ctx=ctx)

        assert order == ["transfer", "sync", "release", "lease_release"]


class TestReleaseOnFailure:
    def test_a_failed_transfer_still_releases_the_source(self):
        """The wedged-source case. A reader that times out and moves to the next
        candidate must not leave connection state behind on the way out."""
        mgr = _manager()
        mgr.receive_from_source.side_effect = TimeoutError("Transfer timed out")

        with pytest.raises(SourceTransferError):
            _receive(mgr, p2p=True)

        mgr.remove_remote_agent.assert_called_once_with(P2P_AGENT)

    def test_the_transfer_error_reaches_the_caller(self):
        """Retry and fallback both depend on the original failure propagating.

        The release cannot get in the way: remove_remote_agent logs and swallows
        its own failures by contract, so it has nothing to mask this with.
        """
        mgr = _manager()
        mgr.receive_from_source.side_effect = TimeoutError("Transfer timed out")

        with pytest.raises(SourceTransferError) as raised:
            _receive(mgr, p2p=True)

        assert isinstance(raised.value.__cause__, TimeoutError)

    def test_a_failure_before_the_transfer_still_releases_the_peer(self):
        """The peer is acquired before the transfer starts, so anything raising in
        between would leak it - and atexit is exactly what cannot be relied on to
        collect it."""
        mgr = _manager()

        class ExplodingBackend:
            @property
            def name(self):
                raise RuntimeError("accelerator backend gone")

        ctx = _ctx(mgr)
        ctx.accelerator_backend = ExplodingBackend()

        with pytest.raises(RuntimeError, match="accelerator backend gone"):
            _receive(mgr, p2p=True, ctx=ctx)

        mgr.remove_remote_agent.assert_called_once_with(P2P_AGENT)

    def test_a_half_finished_p2p_fetch_is_released(self):
        """fetch_remote_and_wait dials the source and only then waits, so a
        timeout can leave metadata that arrives afterwards. Claiming the name
        before the dial is what gives that late arrival something to release it."""
        mgr = _manager()
        mgr.fetch_remote_and_wait.side_effect = TimeoutError(
            "Timed out waiting for remote metadata"
        )

        with pytest.raises(TimeoutError):
            _receive(mgr, p2p=True)

        mgr.remove_remote_agent.assert_called_once_with(P2P_AGENT)

    def test_a_failed_blob_load_releases_nothing(self):
        """The centralized name only exists once add_remote_agent returns, so
        there is no peer to release if it raises. Releasing a name we never got
        would mean calling the manager with None."""
        mgr = _manager()
        mgr.add_remote_agent.side_effect = RuntimeError("bad metadata blob")

        with pytest.raises(RuntimeError, match="bad metadata blob"):
            _receive(mgr, p2p=False)

        mgr.remove_remote_agent.assert_not_called()


class TestCentralizedMode:
    def test_centralized_load_releases_the_peer_it_loaded(self):
        """Centralized readers load the source from a blob, so they own the same
        teardown duty - the asymmetry that wedges the source is identical."""
        mgr = _manager()
        _receive(mgr, p2p=False)
        mgr.add_remote_agent.assert_called_once_with(b"blob")
        mgr.remove_remote_agent.assert_called_once_with(BLOB_AGENT)

    def test_centralized_load_does_not_double_load_the_peer(self):
        """The strategy loads the peer so it holds the name; passing it through
        keeps receive_from_source from loading a second, untracked copy."""
        mgr = _manager()
        _receive(mgr, p2p=False)
        assert mgr.receive_from_source.call_args.kwargs["remote_agent_name"] == BLOB_AGENT