update tests

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

Author: Shixiaowei02

tensorrt_llm/_torch/disaggregation/base/agent.py

@@ -1,23 +1,28 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import List, Tuple, Union
+from typing import Enum, List, Tuple, Union
+
+from tensorrt_llm.logger import logger
 
 # Try to import C++ bindings for zero-copy performance
-_CPP_BINDING_AVAILABLE = False
 try:
-    import tensorrt_llm.tensorrt_llm_transfer_agent_binding as _cpp_binding
+    from tensorrt_llm.tensorrt_llm_transfer_agent_binding import (
+        BaseTransferAgent,
+        MemoryDesc,
+        MemoryDescs,
+        MemoryType,
+        TransferOp,
+        TransferRequest,
+        TransferStatus,
+    )
 
     _CPP_BINDING_AVAILABLE = True
-    # Use C++ types directly when available
-    MemoryType = _cpp_binding.MemoryType
-    TransferOp = _cpp_binding.TransferOp
-    MemoryDesc = _cpp_binding.MemoryDesc
-    MemoryDescs = _cpp_binding.MemoryDescs
-    TransferRequest = _cpp_binding.TransferRequest
-    TransferStatus = _cpp_binding.TransferStatus
-    BaseTransferAgent = _cpp_binding.BaseTransferAgent
 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:
@@ -28,11 +33,11 @@ def is_cpp_binding_available() -> bool:
 # Fallback Python implementations when C++ bindings not available
 if not _CPP_BINDING_AVAILABLE:
 
-    class TransferOp:
+    class TransferOp(Enum):
         READ = "READ"
         WRITE = "WRITE"
 
-    class MemoryType:
+    class MemoryType(Enum):
         DRAM = "DRAM"
         VRAM = "VRAM"
         BLK = "BLK"
@@ -67,28 +72,52 @@ def wait(self, timeout: float | None = None) -> None: ...
 
     class BaseTransferAgent(ABC):
         @abstractmethod
-        def register_memory(self, descs: MemoryDescs) -> None: ...
+        def register_memory(self, descs: MemoryDescs) -> None:
+            """Register a set of memory descriptors on the agent."""
+            ...
 
         @abstractmethod
-        def deregister_memory(self, descs: MemoryDescs) -> None: ...
+        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: ...
+        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: ...
+        def get_local_agent_desc(self) -> str:
+            """Return the serialized description of this agent."""
+            ...
 
         @abstractmethod
-        def invalidate_remote_agent(self, name: str) -> None: ...
+        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: ...
+        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: ...
+        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: ...
+        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)

tensorrt_llm/_torch/disaggregation/base/kv_transfer.py

@@ -21,6 +21,8 @@ class KVSlice:
 
 
 class SessionStatus(Enum):
+    """Status of a transfer session."""
+
     INIT = "INIT"
     READY = "READY"
     TRANSFERRING = "TRANSFERRING"
@@ -36,63 +38,107 @@ class SessionStatus(Enum):
 
 @dataclass
 class SessionState:
+    """State of a transfer session."""
+
     status: SessionStatus
     finished_tasks: List[TaskIdType]
 
 
 @dataclass
 class SessionArgsBase:
-    request_id: int
+    """Base arguments for transfer sessions."""
+
     params: DisaggregatedParams
 
 
-class SenderBase(ABC): ...
+class SenderBase(ABC):
+    """Base class for sending KV cache data."""
+
+    ...
 
 
-class ReceiverBase(ABC): ...
+class ReceiverBase(ABC):
+    """Base class for receiving KV cache data."""
+
+    ...
 
 
 class TxSessionBase(ABC):
     def __init__(self, sender: SenderBase, args: SessionArgsBase):
+        """
+        Initializes the transmission session.
+        :param sender: The sender instance responsible for sending data.
+        :param args: The session arguments.
+        """
         self._base_args = args
 
     @property
     @abstractmethod
-    def state(self) -> SessionState: ...
+    def state(self) -> SessionState:
+        """
+        Returns the current state of the session.
+        """
+        ...
 
     @abstractmethod
-    def poll_task(self, id: TaskIdType) -> SessionStatus: ...
+    def poll_task(self, id: TaskIdType) -> SessionStatus:
+        """
+        Polls the status of a specific task by its ID.
+        :param id: The task ID to poll.
+        """
+        ...
 
     @abstractmethod
-    def send(self, slice: KVSlice) -> TaskIdType: ...
-
-    """
-    Async send slice to the peer. return the task id. Task state can be polled by poll_task().
-    """
+    def send(self, slice: KVSlice) -> TaskIdType:
+        """
+        Sends a slice of KV cache data and returns the task ID.
+        :param slice: The KV slice to send.
+        """
+        ...
 
     @property
     @abstractmethod
-    def exception(self) -> Optional[Exception]: ...
+    def exception(self) -> Optional[Exception]:
+        """
+        Returns any exception that occurred during the session.
+        """
+        ...
 
 
 class RxSessionBase(ABC):
     def __init__(self, receiver: ReceiverBase, args: SessionArgsBase):
+        """
+        Initializes the reception session.
+        :param receiver: The receiver instance responsible for receiving data.
+        """
         self._base_args = args
 
     @property
     @abstractmethod
-    def state(self) -> SessionState: ...
+    def state(self) -> SessionState:
+        """
+        Returns the current state of the session.
+        """
+        ...
 
     @abstractmethod
-    def poll_task(self, id: TaskIdType) -> SessionStatus: ...
+    def poll_task(self, id: TaskIdType) -> SessionStatus:
+        """
+        Polls the status of a specific task by its ID.
+        :param id: The task ID to poll.
+        """
+        ...
 
     @abstractmethod
-    def receive(self, slice: KVSlice) -> TaskIdType: ...
-
-    """
-    Async receive slice from the peer. return the task id. Task state can be polled by poll_task().
-    """
+    def receive(self, slice: KVSlice) -> TaskIdType:
+        """
+        Receives a slice of KV cache data and returns the task ID.
+        :param slice: The KV slice to receive.
+        """
+        ...
 
     @property
     @abstractmethod
-    def exception(self) -> Optional[Exception]: ...
+    def exception(self) -> Optional[Exception]:
+        """Returns any exception that occurred during the session."""
+        ...