refactor the agent

Shixiaowei02 · Shixiaowei02 · commit 5cca4a0edf4d · 2026-01-08T13:40:51.000Z
Signed-off-by: Shixiaowei02 &lt;39303645+Shixiaowei02@users.noreply.github.com&gt;
diff --git a/tensorrt_llm/_torch/disaggregation/base/agent.py b/tensorrt_llm/_torch/disaggregation/base/agent.py
@@ -2,126 +2,77 @@
 from dataclasses import dataclass
 from typing import Enum, List, Tuple, Union
 
-from tensorrt_llm.logger import logger
-
-# Try to import C++ bindings for zero-copy performance
-try:
-    from tensorrt_llm.tensorrt_llm_transfer_agent_binding import (
-        BaseTransferAgent,
-        MemoryDesc,
-        MemoryDescs,
-        MemoryType,
-        TransferOp,
-        TransferRequest,
-        TransferStatus,
-    )
-
-    _CPP_BINDING_AVAILABLE = True
-except ImportError:
-    _CPP_BINDING_AVAILABLE = False
-    logger.warning(
-        "C++ transfer agent bindings not available. "
-        "Falling back to Python implementations which may have lower performance."
-    )
-
-
-def is_cpp_binding_available() -> bool:
-    """Check if C++ transfer agent bindings are available."""
-    return _CPP_BINDING_AVAILABLE
-
-
-# Fallback Python implementations when C++ bindings not available
-if not _CPP_BINDING_AVAILABLE:
-
-    class TransferOp(Enum):
-        READ = "READ"
-        WRITE = "WRITE"
-
-    class MemoryType(Enum):
-        DRAM = "DRAM"
-        VRAM = "VRAM"
-        BLK = "BLK"
-        OBJ = "OBJ"
-        FILE = "FILE"
-
-    @dataclass
-    class MemoryDesc:
-        ptr: int
-        size: int
-        device_id: int
-
-    @dataclass
-    class MemoryDescs:
-        type: str
-        descs: List[Union[Tuple[int, int, int], MemoryDesc]]
-
-    @dataclass
-    class TransferRequest:
-        op: TransferOp
-        src_descs: MemoryDescs
-        dst_descs: MemoryDescs
-        remote_name: str
-        sync_message: str
-
-    class TransferStatus(ABC):
-        @abstractmethod
-        def is_completed(self) -> bool: ...
-
-        @abstractmethod
-        def wait(self, timeout: float | None = None) -> None: ...
-
-    class BaseTransferAgent(ABC):
-        @abstractmethod
-        def register_memory(self, descs: MemoryDescs) -> None:
-            """Register a set of memory descriptors on the agent."""
-            ...
-
-        @abstractmethod
-        def deregister_memory(self, descs: MemoryDescs) -> None:
-            """De-register a set of memory descriptors on the agent."""
-            ...
-
-        @abstractmethod
-        def load_remote_agent(self, name: str, agent_desc: str) -> None:
-            """
-            Load information about a remote agent specified by name.
-
-            Args:
-                name (str): The remote agent's identifier.
-                agent_desc (str): A serialized description of the agent.
-            """
-            ...
-
-        @abstractmethod
-        def get_local_agent_desc(self) -> str:
-            """Return the serialized description of this agent."""
-            ...
-
-        @abstractmethod
-        def invalidate_remote_agent(self, name: str) -> None:
-            """Invalidate any cached information about the specified remote agent."""
-            ...
-
-        @abstractmethod
-        def submit_transfer_requests(self, request: TransferRequest) -> TransferStatus:
-            """Submit transfer tasks to the agent based on a request."""
-            ...
-
-        @abstractmethod
-        def notify_sync_message(self, name: str, sync_message: str) -> None:
-            """Send a synchronization message to the specified remote agent."""
-            ...
-
-        @abstractmethod
-        def check_remote_descs(self, name: str, memory_descs: List[int]) -> bool:
-            """
-            Verify the remote agent's memory descriptors.
-            """
-            ...
-
-
-# RegMemoryDescs is Python-only (used for registration with name field)
+
+# Common Enumerations
+class TransferOp(Enum):
+    READ = "READ"
+    WRITE = "WRITE"
+
+
+class MemoryType(Enum):
+    DRAM = "DRAM"
+    VRAM = "VRAM"
+    BLK = "BLK"
+    OBJ = "OBJ"
+    FILE = "FILE"
+
+
+# Common Data Structures
 @dataclass
-class RegMemoryDescs:
+class MemoryDesc:
+    ptr: int
+    size: int
+    device_id: int
+
+
+@dataclass
+class MemoryDescs:
     type: str
-    descs: List[Tuple[int, int, int, str]]
+    descs: List[Union[Tuple[int, int, int], MemoryDesc]]
+
+
+@dataclass
+class TransferRequest:
+    op: TransferOp
+    src_descs: MemoryDescs
+    dst_descs: MemoryDescs
+    remote_name: str
+    sync_message: str
+
+
+class BaseTransferStatus(ABC):
+    """Abstract base class for transfer status."""
+
+    @abstractmethod
+    def is_completed(self) -> bool: ...
+
+    @abstractmethod
+    def wait(self, timeout: float | None = None) -> None: ...
+
+
+class BaseTransferAgent(ABC):
+    """Abstract base class for transfer agents."""
+
+    @abstractmethod
+    def register_memory(self, descs: MemoryDescs) -> None: ...
+
+    @abstractmethod
+    def deregister_memory(self, descs: MemoryDescs) -> None: ...
+
+    @abstractmethod
+    def load_remote_agent(self, name: str, agent_desc: str) -> None: ...
+
+    @abstractmethod
+    def get_local_agent_desc(self) -> str: ...
+
+    @abstractmethod
+    def invalidate_remote_agent(self, name: str) -> None: ...
+
+    @abstractmethod
+    def submit_transfer_requests(self, request: TransferRequest) -> BaseTransferStatus: ...
+
+    @abstractmethod
+    def notify_sync_message(self, name: str, sync_message: str) -> None: ...
+
+    @abstractmethod
+    def check_remote_descs(self, name: str, memory_descs: List[int]) -> bool: ...
diff --git a/tensorrt_llm/_torch/disaggregation/native/messenger.py b/tensorrt_llm/_torch/disaggregation/native/messenger.py
@@ -150,7 +150,7 @@ def listener():
             poller.register(self._control_socket, zmq.POLLIN)
 
             while not self._stop_event.is_set():
-                events = dict(poller.poll())
+                events = dict(poller.poll(timeout=100))
                 try:
                     if self._control_socket in events:
                         self._stop_event.set()
diff --git a/tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py b/tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py
@@ -0,0 +1,108 @@
+from tensorrt_llm._utils import nvtx_range
+from tensorrt_llm.tensorrt_llm_transfer_agent_binding import (
+    AgentDesc,
+    BaseAgentConfig,
+    MemoryDescs,
+    MemoryType,
+    TransferState,
+)
+from tensorrt_llm.tensorrt_llm_transfer_agent_binding import (
+    NixlTransferAgent as CppNixlTransferAgent,
+)
+from tensorrt_llm.tensorrt_llm_transfer_agent_binding import (
+    NixlTransferStatus as CppNixlTransferStatus,
+)
+
+from ..base.agent import BaseTransferAgent, BaseTransferStatus, RegMemoryDescs, TransferRequest
+
+
+class NixlTransferStatus(BaseTransferStatus):
+    def __init__(self, cpp_status: CppNixlTransferStatus):
+        self._cpp_status = cpp_status
+
+    def is_completed(self) -> bool:
+        """Check if transfer is completed (releases GIL)."""
+        return self._cpp_status.is_completed()
+
+    @nvtx_range("NixlTransferStatus.wait")
+    def wait(self, timeout: float = None) -> bool:
+        """Wait for transfer to complete (releases GIL)."""
+        return self._cpp_status.wait() == TransferState.SUCCESS
+
+
+class NixlTransferAgent(BaseTransferAgent):
+    """NixlTransferAgent using C++ bindings with GIL release support.
+
+    This implementation uses the standalone nixl_bindings C++ module which releases
+    the GIL during blocking operations like wait().
+
+    The nixl_bindings module is independent from the main trtllm bindings,
+    so trtllm can function normally even without NIXL.
+    """
+
+    def __init__(self, name: str, use_prog_thread: bool = True, num_workers: int = 1):
+        config = BaseAgentConfig(
+            name=name,
+            use_prog_thread=use_prog_thread,
+            multi_thread=False,
+            use_listen_thread=False,
+            num_workers=num_workers,
+        )
+        self._cpp_agent = CppNixlTransferAgent(config)
+        self.name = name
+
+    def register_memory(self, descs: RegMemoryDescs):
+        cpp_descs = self._convert_reg_memory_descs(descs)
+        self._cpp_agent.register_memory(cpp_descs)
+
+    def deregister_memory(self, descs: RegMemoryDescs):
+        cpp_descs = self._convert_reg_memory_descs(descs)
+        self._cpp_agent.deregister_memory(cpp_descs)
+
+    def load_remote_agent(self, name: str, agent_desc: bytes):
+        desc_str = agent_desc if isinstance(agent_desc, bytes) else agent_desc.encode()
+        cpp_desc = AgentDesc(desc_str)
+        self._cpp_agent.load_remote_agent(name, cpp_desc)
+
+    def load_remote_agent_by_connection(self, name: str, connection_info: str):
+        self._cpp_agent.load_remote_agent_by_connection(name, connection_info)
+
+    def get_local_agent_desc(self) -> bytes:
+        agent_desc = self._cpp_agent.get_local_agent_desc()
+        return agent_desc.backend_agent_desc
+
+    def get_local_connection_info(self) -> str:
+        return self._cpp_agent.get_local_connection_info()
+
+    def invalidate_remote_agent(self, name: str):
+        self._cpp_agent.invalidate_remote_agent(name)
+
+    def check_remote_descs(self, name: str, memory_descs: MemoryDescs) -> bool:
+        return self._cpp_agent.check_remote_descs(name, memory_descs)
+
+    def notify_sync_message(self, name: str, sync_message: str):
+        self._cpp_agent.notify_sync_message(name, sync_message)
+
+    def get_notified_sync_messages(self):
+        return self._cpp_agent.get_notified_sync_messages()
+
+    @nvtx_range("BindingsNixlTransferAgent.submit_transfer_requests")
+    def submit_transfer_requests(self, request: TransferRequest) -> BaseTransferStatus:
+        cpp_status = self._cpp_agent.submit_transfer_requests(request)
+        return NixlTransferStatus(cpp_status)
+
+    def _convert_reg_memory_descs(self, descs: RegMemoryDescs) -> "MemoryDescs":
+        mem_type = self._convert_memory_type(descs.type)
+        tuples = [(d[0], d[1], d[2]) for d in descs.descs]  # Extract (ptr, size, device_id)
+        return MemoryDescs(mem_type, tuples)
+
+    def _convert_memory_type(self, py_type: str) -> "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)
diff --git a/tensorrt_llm/_torch/disaggregation/nixl/_agent_python.py b/tensorrt_llm/_torch/disaggregation/nixl/_agent_python.py
@@ -0,0 +1,81 @@
+import time
+
+from nixl import nixl_agent, nixl_agent_config, nixl_xfer_handle
+
+from tensorrt_llm._utils import nvtx_range
+
+from ..base.agent import BaseTransferAgent, BaseTransferStatus, RegMemoryDescs, TransferRequest
+
+
+class NixlTransferStatus(BaseTransferStatus):
+    """TransferStatus using the Python NIXL library."""
+
+    def __init__(self, agent: nixl_agent, handle: nixl_xfer_handle):
+        self.agent = agent
+        self.handle = handle
+
+    def is_completed(self) -> bool:
+        status = self.agent.check_xfer_state(self.handle)
+        return status == "DONE"
+
+    def wait(self) -> bool:
+        status = "PROC"
+        sleep_time = 0.0001  # 0.1ms
+        max_sleep_time = 0.01  # 10ms
+        while status == "PROC":
+            status = self.agent.check_xfer_state(self.handle)
+            if status == "ERROR":
+                return False  # Transfer failed
+            time.sleep(sleep_time)  # Sleep to release GIL
+            sleep_time = min(sleep_time * 2, max_sleep_time)
+        return status == "DONE"
+
+
+class NixlTransferAgent(BaseTransferAgent):
+    """Python-based TransferAgent using the NIXL library."""
+
+    def __init__(self, name: str, use_prog_thread: bool, num_workers: int = 1):
+        self.name = name
+        agent_config = nixl_agent_config(
+            enable_prog_thread=use_prog_thread,
+            backends=["UCX"],
+            num_threads=num_workers,
+        )
+        self.agent = nixl_agent(name, agent_config)
+
+    def register_memory(self, descs: RegMemoryDescs):
+        reg_descs = self.agent.get_reg_descs(descs.descs, descs.type)
+        self.agent.register_memory(reg_descs)
+
+    def deregister_memory(self, descs: RegMemoryDescs):
+        self.agent.deregister_memory(descs.descs, descs.type)
+
+    def load_remote_agent(self, name: str, agent_desc: bytes):
+        self.agent.add_remote_agent(agent_desc)
+
+    def get_local_agent_desc(self) -> bytes:
+        return self.agent.get_agent_metadata()
+
+    def invalidate_remote_agent(self, name: str):
+        self.agent.remove_remote_agent(name)
+
+    def check_remote_descs(self, name: str, memory_descs: list[int]) -> bool:
+        raise NotImplementedError("check_remote_descs is not implemented.")
+
+    def notify_sync_message(self, name: str, sync_message: str):
+        raise NotImplementedError("notify_sync_message is not implemented.")
+
+    @nvtx_range("NixlTransferAgent.submit_transfer_requests")
+    def submit_transfer_requests(self, request: TransferRequest) -> BaseTransferStatus:
+        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)
+        handle = self.agent.initialize_xfer(
+            request.op,
+            src_xfer_descs,
+            dst_xfer_descs,
+            request.remote_name,
+            request.sync_message,
+        )
+        status = self.agent.transfer(handle)
+        assert status != "ERROR", "Transfer failed during initialization."
+        return NixlTransferStatus(self.agent, handle)
diff --git a/tensorrt_llm/_torch/disaggregation/nixl/agent.py b/tensorrt_llm/_torch/disaggregation/nixl/agent.py
