test and fix for agent

Signed-off-by: Shixiaowei02 <[email protected]>

Author: Shixiaowei02

tensorrt_llm/_torch/disaggregation/base/agent.py

@@ -81,7 +81,7 @@ def check_remote_descs(self, name: str, memory_descs: List[int]) -> bool: ...
 @dataclass
 class RegMemoryDescs:
     type: str
-    descs: List[Tuple[int, int, int, str]]
+    descs: List[Tuple[int, int, int, str]]  # (ptr, size, device_id, name)
 
 
 def _force_py_nixl_kv_transfer() -> bool:
@@ -123,4 +123,6 @@ def _try_load_cpp_binding():
     BaseTransferAgent = _cpp_binding.BaseTransferAgent
     logger.info("Using Pybind transfer agent binding for Transfer Agent implementation.")
 else:
-    logger.info("Failed to import Pybind transfer agent binding, using pure Python implementation.")
+    logger.warning(
+        "Failed to import Pybind transfer agent binding, using pure Python implementation."
+    )

tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py

@@ -37,19 +37,33 @@ def wait(self):
 class NixlTransferAgent(BaseTransferAgent):
     """NixlTransferAgent using Python nixl library."""
 
-    def __init__(self, name: str, use_prog_thread: bool, num_workers: int = 1):
+    def __init__(self, name: str, use_prog_thread: bool = True, num_workers: int = 1):
+        """
+        Initialize NixlTransferAgent.
+        :param name: Name of the agent.
+        :param use_prog_thread: Whether to enable the progress thread, if available.
+        :param num_workers: Specify number of threads for the supported multi-threaded backends.
+        """
         self.name = name
+        self.backends = ["UCX"]
         agent_config = nixl_agent_config(
-            enable_prog_thread=use_prog_thread, backends=["UCX"], num_threads=num_workers
+            enable_prog_thread=use_prog_thread, backends=self.backends, num_threads=num_workers
         )
         self.agent = nixl_agent(name, agent_config)
 
     def register_memory(self, descs: RegMemoryDescs):
+        if isinstance(descs.descs[0], tuple):
+            assert len(descs.descs[0]) == 4, f"Expected 4 elements per desc, got {descs.descs[0]}"
         reg_descs = self.agent.get_reg_descs(descs.descs, descs.type)
+        assert reg_descs is not None, "Failed to get reg_descs"
         self.agent.register_memory(reg_descs)
 
     def deregister_memory(self, descs: RegMemoryDescs):
-        self.agent.deregister_memory(descs.descs, descs.type)
+        if isinstance(descs.descs[0], tuple):
+            assert len(descs.descs[0]) == 4, f"Expected 4 elements per desc, got {descs.descs[0]}"
+        reg_descs = self.agent.get_reg_descs(descs.descs, descs.type)
+        assert reg_descs is not None, "Failed to get reg_descs"
+        self.agent.deregister_memory(reg_descs)
 
     def load_remote_agent(self, name: str, agent_desc: bytes):
         self.agent.add_remote_agent(agent_desc)
@@ -70,6 +84,8 @@ def notify_sync_message(self, name: str, sync_message: str):
     def submit_transfer_requests(self, request: TransferRequest) -> TransferStatus:
         src_xfer_descs = self.agent.get_xfer_descs(request.src_descs.descs, request.src_descs.type)
         dst_xfer_descs = self.agent.get_xfer_descs(request.dst_descs.descs, request.dst_descs.type)
+        assert src_xfer_descs is not None, "Failed to get src_xfer_descs"
+        assert dst_xfer_descs is not None, "Failed to get dst_xfer_descs"
         sync_message = "" if request.sync_message is None else request.sync_message
         handle = self.agent.initialize_xfer(
             request.op,

tests/unittest/disaggregated/test_agent.py

@@ -0,0 +1,156 @@
+from dataclasses import dataclass, field
+
+import pytest
+import torch
+
+from tensorrt_llm import logger
+from tensorrt_llm._torch.disaggregation.base.agent import (
+    MemoryDescs,
+    MemoryType,
+    RegMemoryDescs,
+    TransferOp,
+    TransferRequest,
+)
+from tensorrt_llm._torch.disaggregation.nixl.agent import NixlTransferAgent
+
+
+def _convert_to_memory_descs(reg_descs: RegMemoryDescs) -> MemoryDescs:
+    tuples = [(ptr, size, device_id) for (ptr, size, device_id, _) in reg_descs.descs]
+
+    def _convert_memory_type(py_type: str) -> MemoryType:
+        """Convert Python memory type string to C++ MemoryType."""
+        type_map = {
+            "DRAM": MemoryType.DRAM,
+            "VRAM": MemoryType.VRAM,
+            "GPU": MemoryType.VRAM,
+            "BLK": MemoryType.BLK,
+            "OBJ": MemoryType.OBJ,
+            "FILE": MemoryType.FILE,
+        }
+        return type_map.get(py_type.upper(), MemoryType.VRAM)
+
+    return MemoryDescs(_convert_memory_type(reg_descs.type), tuples)
+
+
+@dataclass
+class MemoryManager:
+    allocated_memory: list[torch.Tensor] = field(default_factory=list)
+
+    def allocate_memory(
+        self, size: int, name: str, memory_type=MemoryType.VRAM, device_id: int = 0
+    ) -> RegMemoryDescs:
+        device = torch.device(f"cuda:{device_id}" if memory_type == MemoryType.VRAM else "cpu")
+
+        # Allocate memory block using torch.Tensor and track it
+        block = torch.zeros(size, dtype=torch.uint8, device=device)
+        self.allocated_memory.append(block)
+
+        # Return RegMemoryDescs with position arguments
+        memory_descs = RegMemoryDescs(
+            type=memory_type, descs=[(block.data_ptr(), block.numel(), device_id, name)]
+        )
+        return memory_descs
+
+    def clear_memory(self):
+        """Clear all tracked memory blocks."""
+        self.allocated_memory.clear()
+
+
+@pytest.fixture
+def memory_manager():
+    return MemoryManager()
+
+
+@pytest.fixture(params=[256, 512])
+def memory_size(request):
+    return request.param
+
+
+@pytest.fixture(params=["DRAM", "VRAM"])
+def memory_type(request):
+    return request.param
+
+
+@pytest.fixture
+def alloc(memory_manager, memory_size, memory_type):
+    """Allocate memory for source and destination, based on the memory_size and memory_type parameters."""
+    assert memory_size > 0, "Memory size must be a positive integer."
+    src_descs = memory_manager.allocate_memory(
+        size=memory_size, name="src_mem", memory_type=memory_type
+    )
+    dst_descs = memory_manager.allocate_memory(
+        size=memory_size, name="dst_mem", memory_type=memory_type
+    )
+    return src_descs, dst_descs
+
+
+@pytest.fixture
+def transfer_agent_src():
+    return NixlTransferAgent(name="src_agent")
+
+
+@pytest.fixture
+def transfer_agent_dst():
+    return NixlTransferAgent(name="dst_agent")
+
+
+def test_transfer_between_agents(
+    transfer_agent_src,
+    transfer_agent_dst,
+    memory_manager,
+    alloc,
+    memory_size,
+    memory_type,
+):
+    """End-to-end test of data transfer between two agents with parameterized memory sizes and types."""
+    # Debug log the parameters being tested
+    logger.info(f"Testing with memory_size={memory_size}, memory_type={memory_type}")
+
+    # Unpack source and destination memory descriptions
+    memory_descs_src, memory_descs_dst = alloc
+
+    # Fill source memory with sequential data for validation
+    src_data = memory_manager.allocated_memory[0]
+    assert memory_size > 0, "Memory size must be positive."
+    tensor = torch.arange(memory_size, dtype=torch.uint8) % 10
+    src_data.copy_(tensor)
+
+    # Register memory with source and destination agents
+    transfer_agent_src.register_memory(memory_descs_src)
+    transfer_agent_dst.register_memory(memory_descs_dst)
+
+    src_agent_desc = transfer_agent_src.get_local_agent_desc()
+    transfer_agent_dst.load_remote_agent("src_agent", src_agent_desc)
+
+    dst_agent_desc = transfer_agent_dst.get_local_agent_desc()
+    transfer_agent_src.load_remote_agent("dst_agent", dst_agent_desc)
+
+    # Create and submit the transfer request
+    transfer_request = TransferRequest(
+        op=TransferOp.WRITE,
+        src_descs=_convert_to_memory_descs(memory_descs_src),
+        dst_descs=_convert_to_memory_descs(memory_descs_dst),
+        remote_name="dst_agent",
+        sync_message=None,
+    )
+    transfer_status = transfer_agent_src.submit_transfer_requests(transfer_request)
+    transfer_status.wait()
+
+    # Validate transfer completion
+    assert transfer_status.is_completed(), "Transfer did not complete successfully."
+
+    # Validate that the destination data matches the source data
+    dst_data = memory_manager.allocated_memory[1]
+    assert torch.equal(dst_data, src_data), "Destination data does not match source data."
+
+    # Clean up by deregistering memory and clearing allocations
+    transfer_agent_src.deregister_memory(memory_descs_src)
+    transfer_agent_dst.deregister_memory(memory_descs_dst)
+    memory_manager.clear_memory()
+
+    transfer_agent_src.invalidate_remote_agent("dst_agent")
+    transfer_agent_dst.invalidate_remote_agent("src_agent")
+
+
+if __name__ == "__main__":
+    pytest.main()