Routed per-subset state dict gets into NIXL CUDA destinations

tests/routing/test_routing.py

@@ -6,6 +6,8 @@
 
 from __future__ import annotations
 
+from unittest.mock import Mock
+
 import pytest
 import torch
 from torch.distributed.device_mesh import DeviceMesh
@@ -234,6 +236,80 @@ async def test_transfer_dtype_preserves_keys() -> None:
     torch.testing.assert_close(result["buffer"], source["buffer"])
 
 
+@pytest.mark.asyncio
+async def test_strict_read_rejects_key_skipped_by_non_strict_read() -> None:
+    source = torch.arange(2)
+    destination = {"weight": torch.zeros_like(source), "local_only": torch.zeros(2)}
+    clients = await routing_clients(
+        publishers={"publisher/0": {"weight": source}},
+        requesters={"requester/0/0": destination},
+    )
+    await put_state_dict(clients["publisher/0"], {"weight": source}, "model")
+    requester = clients["requester/0/0"]
+    await get_state_dict(requester, "model", destination, strict=False)
+
+    with pytest.raises(KeyError, match="no publisher"):
+        await get_state_dict(
+            requester, "model", {"local_only": destination["local_only"]}
+        )
+
+
+@pytest.mark.asyncio
+async def test_non_strict_reads_plan_and_fetch_one_subset_per_call() -> None:
+    """Read one subset of the state dict per call, as a per-layer load does."""
+    sources = {"first": torch.arange(2), "second": torch.arange(2, 4)}
+    sources["third"] = torch.arange(4, 6)
+    clients = await routing_clients(
+        publishers={
+            "publisher/0": {"first": sources["first"]},
+            "publisher/1": {"second": sources["second"], "third": sources["third"]},
+        },
+        requesters={"requester/0/0": {}},
+    )
+    await put_state_dict(clients["publisher/0"], {"first": sources["first"]}, "model")
+    await put_state_dict(
+        clients["publisher/1"],
+        {"second": sources["second"], "third": sources["third"]},
+        "model",
+    )
+    requester = clients["requester/0/0"]
+    requester._control_plane = Mock(wraps=requester._control_plane)
+    get_batch = Mock(wraps=requester.get_batch)
+    requester.get_batch = get_batch
+
+    # A failed plan leaves its keys to be planned by the next get.
+    with pytest.raises(ValueError, match="inconsistent global shape"):
+        await get_state_dict(
+            requester, "model", {"first": torch.zeros(3)}, strict=False
+        )
+    # Subsets: new, new, already planned, mixed, then planned after a get with
+    # no published keys.
+    for names, local_only in (
+        (["first"], False),
+        (["second"], False),
+        (["first"], False),
+        (["second", "third"], False),
+        ([], True),
+        (["first", "third"], False),
+    ):
+        get_batch.reset_mock()
+        destination = {name: torch.zeros_like(sources[name]) for name in names}
+        if local_only:
+            destination["local_only"] = torch.zeros(2)
+        result = await get_state_dict(requester, "model", destination, strict=False)
+
+        assert set(result) == set(names)
+        for name in names:
+            assert result[name] is destination[name]
+            torch.testing.assert_close(result[name], sources[name])
+        if names:
+            ((batch,), _kwargs) = get_batch.call_args
+            assert set(batch) == {f"model/{name}" for name in names}
+        else:
+            get_batch.assert_not_called()
+    requester._control_plane.get_layouts.assert_called_once_with("model")
+
+
 # Shard dim of each tensor on two publishers, or None for a whole copy on each.
 # Column shards of out.weight are strided in every requester's buffer, so reads
 # copy them in rather than writing in place.

tests/test_nixl_transport.py

@@ -108,12 +108,21 @@ def __init__(self, name, config) -> None:
             setattr(self, method_name, Mock(wraps=getattr(self, method_name)))
         FakeNixlAgent.instances.append(self)
 
+    def create_backend(self, backend, initParams) -> None:
+        self.backend_params = {backend: initParams}
+
     def register_memory(self, tensor, *, backends):
         assert tensor.numel() > 0
+        start, end = tensor.data_ptr(), tensor.data_ptr() + tensor.nbytes
+        # NIXL accepts this, but peers then fail to load the agent's metadata.
+        assert all(
+            (start, end) != (other_start, other_end)
+            for other_start, other_end, _ in self.regions.values()
+        ), "same memory registered twice"
         descriptor = id(tensor)
         self.regions[descriptor] = (
-            tensor.data_ptr(),
-            tensor.data_ptr() + tensor.nbytes,
+            start,
+            end,
             # Weak, so registration does not keep the storage alive.
             weakref.ref(tensor.untyped_storage()),
         )
@@ -707,6 +716,212 @@ async def test_objects_stay_on_rpc_path_without_creating_nixl_agent(ref, ctx):
     assert FakeNixlAgent.instances == []
 
 
+@pytest.mark.asyncio
+async def test_uncached_registration_lasts_until_last_drop(ref):
+    """Buffers of one get share a storage's registration until both drop."""
+    first = NixlTransportBuffer(ref)
+    second = NixlTransportBuffer(ref)
+    destination = torch.zeros(4, 16)
+
+    first._register_client_tensor(destination[0], cached=False)
+    first._register_client_tensor(destination[1], cached=False)
+    second._register_client_tensor(destination[2], cached=False)
+    cache = first._client_cache()
+    agent = cache.agent
+    ((descriptors, holders),) = cache._leases.values()
+
+    assert holders == 2
+    agent.register_memory.assert_called_once()
+    assert cache._registrations == {}
+    assert _storage_volume_copy(first)._leases == set()
+    await first.drop()
+    agent.deregister_memory.assert_not_called()
+    await second.drop()
+    agent.deregister_memory.assert_called_once_with(descriptors)
+    assert cache._leases == {}
+
+
+async def _cancel_get_during_volume_call(ref, late_error=False):
+    """Cancel a CUDA-style get while the volume is answering it.
+
+    Returns once the get waits for the volume, or has finished without waiting,
+    with a weak reference to the destination.
+    """
+    client_buffer = NixlTransportBuffer(ref)
+    destinations = []
+    volume_called = asyncio.Event()
+    answer = asyncio.Event()
+    waiting = asyncio.Event()
+
+    async def pre_get_hook(requests):
+        destination = torch.zeros(4)
+        destinations.append(weakref.ref(destination))
+        client_buffer._contexts = [
+            client_buffer._register_client_tensor(destination, cached=False)
+        ]
+
+    async def volume_get(transport_buffer, meta_requests):
+        volume_called.set()
+        await answer.wait()
+        if late_error:
+            raise RuntimeError("late transfer failed")
+        return transport_buffer
+
+    wait_for_volume = client_buffer._wait_for_volume
+
+    async def signal_and_wait_for_volume(call):
+        waiting.set()
+        await wait_for_volume(call)
+
+    client_buffer._pre_get_hook = pre_get_hook
+    client_buffer._wait_for_volume = signal_and_wait_for_volume
+    ref.volume.get = SimpleNamespace(call_one=volume_get)
+    get = asyncio.create_task(
+        client_buffer._get_requests([Request.from_any("weight", None)])
+    )
+    await asyncio.wait_for(volume_called.wait(), timeout=10)
+    get.cancel()
+    signal = asyncio.create_task(waiting.wait())
+    await asyncio.wait([get, signal], timeout=10, return_when=asyncio.FIRST_COMPLETED)
+    signal.cancel()
+    return client_buffer, get, answer, destinations[0]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("late_error", [False, True])
+async def test_cancelled_get_waits_for_volume_answer(ref, late_error):
+    """The volume may still write after the caller stops waiting."""
+    client_buffer, get, answer, _ = await _cancel_get_during_volume_call(
+        ref, late_error
+    )
+    cache = client_buffer._client_cache()
+
+    assert not get.done()
+    cache.agent.deregister_memory.assert_not_called()
+    answer.set()
+    with pytest.raises(asyncio.CancelledError):
+        await asyncio.wait_for(get, timeout=10)
+    cache.agent.deregister_memory.assert_called_once()
+    assert cache._leases == {}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("give_up", ["cancel", "timeout"])
+async def test_abandoned_get_releases_registration_when_volume_answers(
+    ref, monkeypatch, give_up
+):
+    """A second cancel or a timeout returns without deregistering.
+
+    Until the volume answers, the destination stays referenced and another get
+    cannot register its storage.
+    """
+    if give_up == "timeout":
+        monkeypatch.setattr(nixl_transport, "ENV_TORCHSTORE_NIXL_TIMEOUT_S", 0)
+    client_buffer, get, answer, destination_ref = await _cancel_get_during_volume_call(
+        ref
+    )
+    cache = client_buffer._client_cache()
+    if give_up == "cancel":
+        get.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await asyncio.wait_for(get, timeout=10)
+
+    cache.agent.deregister_memory.assert_not_called()
+    assert len(cache._leases) == 1
+    with pytest.raises(RuntimeError, match="has not answered"):
+        cache.acquire(destination_ref())
+    gc.collect()
+    assert destination_ref() is not None
+    released = asyncio.Event()
+    release = cache.release
+
+    def release_and_signal(key):
+        release(key)
+        released.set()
+
+    cache.release = release_and_signal
+    answer.set()
+    await asyncio.wait_for(released.wait(), timeout=10)
+    cache.agent.deregister_memory.assert_called_once()
+    assert cache._leases == {}
+    assert not cache._abandoned
+    gc.collect()
+    assert destination_ref() is None
+
+
+@pytest.mark.asyncio
+async def test_get_cancelled_before_volume_call_releases_registration(ref):
+    client_buffer = NixlTransportBuffer(ref)
+    destination = torch.zeros(4)
+    registered = asyncio.Event()
+
+    async def pre_get_hook(requests):
+        client_buffer._contexts = [
+            client_buffer._register_client_tensor(destination, cached=False)
+        ]
+        registered.set()
+        await asyncio.Event().wait()
+
+    client_buffer._pre_get_hook = pre_get_hook
+    get = asyncio.create_task(
+        client_buffer._get_requests([Request.from_any("weight", None)])
+    )
+    await registered.wait()
+    get.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await get
+
+    cache = client_buffer._client_cache()
+    cache.agent.deregister_memory.assert_called_once()
+    assert cache._leases == {}
+
+
+def test_lease_replaces_cached_registration(agent_cache):
+    """A CUDA put source can later be a get destination."""
+    tensor = torch.zeros(16)
+    cached = agent_cache.register(tensor)
+
+    key = agent_cache.acquire(tensor)
+
+    agent_cache.agent.deregister_memory.assert_called_once_with(cached)
+    assert agent_cache._registrations == {}
+    agent_cache.release(key)
+    agent_cache.register(tensor)
+    assert agent_cache.agent.register_memory.call_count == 3
+
+
+def test_register_reuses_lease(agent_cache):
+    tensor = torch.zeros(16)
+    key = agent_cache.acquire(tensor)
+
+    assert agent_cache.register(tensor[1:]) == agent_cache._leases[key][0]
+    assert agent_cache._registrations == {}
+    agent_cache.agent.register_memory.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_drop_releases_every_lease_when_one_fails(ref):
+    client_buffer = NixlTransportBuffer(ref)
+    first, second = torch.zeros(4), torch.zeros(4)
+    client_buffer._register_client_tensor(first, cached=False)
+    client_buffer._register_client_tensor(second, cached=False)
+    cache = client_buffer._client_cache()
+    cache.agent.deregister_memory.side_effect = [RuntimeError("failed"), None]
+
+    with pytest.raises(RuntimeError, match="failed"):
+        await client_buffer.drop()
+
+    assert cache.agent.deregister_memory.call_count == 2
+    assert cache._leases == {}
+
+
+def test_agent_disables_ucx_registration_reuse(agent_cache):
+    assert agent_cache.agent.config.kwargs["backends"] == []
+    assert agent_cache.agent.backend_params == {
+        "UCX": {"engine_config": "CUDA_COPY_REG_WHOLE_ALLOC=off,RCACHE_ENABLE=n"}
+    }
+
+
 def test_memory_registration_is_reused_for_same_tensor(agent_cache):
     tensor = torch.zeros(16)