modelexpress-client 0.5.1

# 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):
    return SimpleNamespace(
        global_rank=0,
        worker_rank=0,
        worker_id="target-0",
        nixl_manager=manager,
        accelerator_backend=SimpleNamespace(name="cuda", synchronize=MagicMock()),
        adapter=MagicMock(),
        identity=SimpleNamespace(model_name="m"),
    )


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):
    strategy = RdmaStrategy()
    with (
        patch("modelexpress.load_strategy.rdma_strategy.register_tensors"),
        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"
        )


class TestReleaseOnTheLoadPath:
    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")

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

        assert order == ["transfer", "sync", "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_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