feat(kvbm): register operation-bound multipart GPU endpoints [D24]

components/src/dynamo/vllm/tests/test_gpu_offloading.py

@@ -12,13 +12,16 @@
 import sys
 import uuid
 import weakref
+from contextlib import nullcontext
 from types import SimpleNamespace
+from unittest.mock import MagicMock
 
 import pytest
 import torch
 from kvbm.vllm_integration import gpu_memory, gpu_mnnvl, gpu_offloading
 from kvbm.vllm_integration.gpu_memory import GPUCacheRegion, GPUCacheWorker, import_ipc
 from kvbm.vllm_integration.gpu_offloading import (
+    GPUKVCRSecondaryTierManager,
     GPUOffloadingConnector,
     GPURetainedLoadStoreSpec,
 )
@@ -28,7 +31,7 @@
     OffloadingConnector,
 )
 from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash
-from vllm.v1.kv_offload.base import GPULoadStoreSpec, make_offload_key
+from vllm.v1.kv_offload.base import GPULoadStoreSpec, ReqContext, make_offload_key
 
 pytestmark = [
     pytest.mark.post_merge,
@@ -38,6 +41,108 @@
 ]
 
 
+@pytest.mark.pre_merge
+@pytest.mark.gpu_0
+@pytest.mark.parametrize("tp_size", [1, 2, 3, 4, 8])
+@pytest.mark.parametrize("native_failure", [False, True])
+def test_gpu_kvcr_rank_registration_and_failed_start_preserve_bank_owners(
+    monkeypatch, tp_size, native_failure
+):
+    spec = SimpleNamespace(
+        config=SimpleNamespace(parallel=SimpleNamespace(world_size=tp_size)),
+        gpu_transport="cuda_ipc",
+        num_chunks=8,
+        row_bytes=384,
+        logical_row_bytes=384 * tp_size,
+        base_domain="test-gpu-layout-v1",
+        extra_config={"router_hint_transfer_domain": "test-gpu-layout-v1"},
+    )
+    banks = [
+        SimpleNamespace(
+            shape=(spec.num_chunks, spec.row_bytes),
+            strides=(spec.row_bytes, 1),
+            tensor=SimpleNamespace(
+                device=SimpleNamespace(index=tp_size - rank - 1),
+                data_ptr=lambda rank=rank: 4096 * (rank + 1),
+            ),
+        )
+        for rank in range(tp_size)
+    ]
+    region = (
+        banks[0]
+        if tp_size == 1
+        else SimpleNamespace(
+            shape=(spec.num_chunks, spec.logical_row_bytes),
+            strides=(spec.logical_row_bytes, 1),
+            banks=banks,
+        )
+    )
+    listener = MagicMock()
+    listener.__enter__.return_value = listener
+    listener.getsockname.return_value = ("127.0.0.1", 49152)
+    socket_factory = MagicMock(return_value=listener)
+    control = object()
+    control_factory = MagicMock(return_value=control)
+    client = object()
+    forwarded = []
+
+    def create_kvcr(config, bindings, backends):
+        listener.__exit__.assert_called_once_with(None, None, None)
+        assert config.nixl_listen_port == 49152
+        assert bindings.framework_control is control
+        forwarded.append(backends)
+        if native_failure:
+            raise RuntimeError("native startup failed after registration")
+        return client
+
+    monkeypatch.setattr(gpu_offloading.socket, "socket", socket_factory)
+    monkeypatch.setattr(gpu_offloading, "ZmqPeerControlChannel", control_factory)
+    monkeypatch.setattr(gpu_offloading, "KVCR", create_kvcr)
+    retained = []
+    monkeypatch.setattr(gpu_offloading, "_RETAINED_FAILED_GPU_MANAGERS", retained)
+    expected = (
+        pytest.raises(RuntimeError, match="native startup failed")
+        if native_failure
+        else nullcontext()
+    )
+    with expected:
+        manager = GPUKVCRSecondaryTierManager(
+            spec,
+            region,
+            {
+                "compatibility_domain": spec.base_domain,
+                "control_ports": [49153],
+                "control_host": "0.0.0.0",
+                "control_advertise_host": "gpu-cache.example",
+                "router_capabilities": ["router_hint_gpu_direct"],
+            },
+        )
+
+    socket_factory.assert_called_once_with()
+    listener.bind.assert_called_once_with(("127.0.0.1", 0))
+    listener.getsockname.assert_called_once_with()
+    control_factory.assert_called_once_with("0.0.0.0", 49153, "gpu-cache.example")
+    if native_failure:
+        assert len(retained) == 1
+        assert retained[0]._region is region
+    else:
+        assert manager._kvcr is client
+        assert retained == []
+    assert [
+        (entry.label, entry.device_Id, entry.addr, entry.count, entry.size)
+        for entry in forwarded[0].framework_regions
+    ] == [
+        (
+            "" if tp_size == 1 else f":rank{rank}",
+            tp_size - rank - 1,
+            4096 * (rank + 1),
+            spec.num_chunks,
+            spec.row_bytes,
+        )
+        for rank in range(tp_size)
+    ]
+
+
 @pytest.mark.pre_merge
 @pytest.mark.gpu_0
 def test_gpu_inventory_has_its_own_medium(monkeypatch):
@@ -263,6 +368,21 @@ def test_local_fabric_descriptor_mismatch_fails_before_importing_any_gpu_handle(
             )
 
 
+def _multipart_tier(tp_size=2):
+    tier = GPUKVCRSecondaryTierManager.__new__(GPUKVCRSecondaryTierManager)
+    tier._part_labels = tuple(f":rank{rank}" for rank in range(tp_size))
+    tier._region = SimpleNamespace(shape=(4, 384 * tp_size))
+    tier._key_adapter = gpu_offloading._VllmKeyAdapter()
+    tier._finished_jobs = []
+    tier._jobs_by_op = {}
+    calls = []
+    tier._kvcr = SimpleNamespace(
+        config=SimpleNamespace(nixl_agent_name="multipart"),
+        deliver=lambda blocks, **kwargs: calls.append((blocks, kwargs)) or 91,
+    )
+    return tier, calls
+
+
 @pytest.mark.pre_merge
 @pytest.mark.gpu_0
 def test_gpu_key_adapter_preserves_group_keys_separately_from_router_hashes():
@@ -278,3 +398,36 @@ def test_gpu_key_adapter_preserves_group_keys_separately_from_router_hashes():
     assert [adapter.encode_to_router(key) for key in canonical] == [router_key] * 2
     with pytest.raises(NotImplementedError, match="group identity"):
         adapter.decode_from_router(router_key)
+
+
+@pytest.mark.pre_merge
+@pytest.mark.gpu_0
+@pytest.mark.parametrize("tp_size", [2, 3, 4, 8])
+def test_tp_delivers_every_rank_part_in_one_logical_operation(tp_size):
+    tier, calls = _multipart_tier(tp_size)
+    keys = [make_offload_key(value, 0) for value in (b"first", b"second")]
+    tier.submit_load(
+        SimpleNamespace(
+            job_id=7, keys=keys, chunk_ids=[3, 1], req_context=ReqContext("load")
+        )
+    )
+    assert len(calls) == 1
+    blocks, kwargs = calls[0]
+    assert kwargs == {"request_id": "load"}
+    assert len(blocks) == 2
+    for key, slot in zip(keys, [3, 1], strict=True):
+        assert [(ref.label, ref.element_index) for ref in blocks[BlockKey(key)]] == [
+            (f":rank{rank}", slot) for rank in range(tp_size)
+        ]
+    assert tier._jobs_by_op == {91: (7, 2, True, set())}
+    assert tier._finished_jobs == []
+
+
+@pytest.mark.pre_merge
+@pytest.mark.gpu_0
+@pytest.mark.parametrize("slot", [-1, 4, True])
+def test_tp2_descriptor_bounds_fail_before_native_transfer(slot):
+    tier, calls = _multipart_tier()
+    with pytest.raises(ValueError, match="outside"):
+        tier._make_descriptors(slot)
+    assert calls == []