tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py (2)

70-78: Consider validating all descriptor elements, not just the first.

The check only validates descs.descs[0]. If the list contains heterogeneous tuples (e.g., first has 4 elements, later ones have 3), malformed entries will pass validation and fail deeper in get_reg_descs.

♻️ Proposed fix to validate all elements

-        if isinstance(descs.descs[0], tuple) and len(descs.descs[0]) != 4:
-            raise ValueError(
-                f"Expected 4 elements per desc, got {len(descs.descs[0])}: {descs.descs[0]}"
-            )
+        for i, desc in enumerate(descs.descs):
+            if isinstance(desc, tuple) and len(desc) != 4:
+                raise ValueError(
+                    f"Expected 4 elements per desc at index {i}, got {len(desc)}: {desc}"
+                )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py` around lines 70 - 78,
The validation currently only checks descs.descs[0], which allows other
malformed descriptor tuples to slip through; update the validation to iterate
over all entries in descs.descs and ensure each is a tuple of length 4 (raise
ValueError including the offending index and value if any fail), before calling
self.agent.get_reg_descs(descs.descs, descs.type) so get_reg_descs always
receives uniformly validated descriptors.

---

81-93: Consider extracting duplicated validation into a helper method.

The validation logic in deregister_memory (lines 81-93) is nearly identical to register_memory (lines 67-79). A private helper could reduce duplication.

♻️ Example helper extraction

def _get_validated_reg_descs(self, descs: RegMemoryDescs):
    """Validate and return registration descriptors."""
    if not descs.descs:
        raise ValueError("descs.descs must not be empty")
    for i, desc in enumerate(descs.descs):
        if isinstance(desc, tuple) and len(desc) != 4:
            raise ValueError(
                f"Expected 4 elements per desc at index {i}, got {len(desc)}: {desc}"
            )
    reg_descs = self.agent.get_reg_descs(descs.descs, descs.type)
    if reg_descs is None:
        raise RuntimeError(
            f"nixl get_reg_descs returned None for type={descs.type}, count={len(descs.descs)}"
        )
    return reg_descs

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py` around lines 81 - 93,
Extract the duplicated validation and conversion logic used in register_memory
and deregister_memory into a private helper (e.g., _get_validated_reg_descs)
that accepts a RegMemoryDescs, validates descs.descs is non-empty, iterates
through each desc and ensures tuple entries have length 4 (including index in
error messages), calls self.agent.get_reg_descs(descs.descs, descs.type) and
raises the same RuntimeError if None, and returns reg_descs; then replace the
validation block in both register_memory and deregister_memory to call this
helper and use its return value when calling self.agent.register_memory or
self.agent.deregister_memory.

tensorrt_llm/_torch/disaggregation/nixl/agent.py (1)

36-42: Optional: Consider cleaner f-string formatting for the error message.

The conditional concatenation works but could be slightly more readable.

♻️ Alternative formatting

     if _py_agent is None:
-        raise ImportError(
-            "Failed to load pure Python NIXL Transfer Agent."
-            + (f" Caused by: {_py_agent_err}" if _py_agent_err else "")
-        )
+        cause = f" Caused by: {_py_agent_err}" if _py_agent_err else ""
+        raise ImportError(f"Failed to load pure Python NIXL Transfer Agent.{cause}")

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/disaggregation/nixl/agent.py` around lines 36 - 42, The
ImportError construction using string concatenation around _py_agent and
_py_agent_err is readable but can be simplified to a single f-string; update the
raise ImportError to use one f-string that conditionally includes the "Caused
by: {_py_agent_err}" segment (e.g. f"Failed to load pure Python NIXL Transfer
Agent{f' Caused by: {_py_agent_err}' if _py_agent_err else ''}"), keeping the
same variables (_py_agent, _py_agent_err) and preserving the current behavior
and spacing, then leave the subsequent references to NixlTransferStatus and
NixlTransferAgent unchanged.