tensorrt_llm/executor/worker.py (1)

375-386: Fix TOCTOU race around LoRA CPU cache check; gate weights only for PyTorch backend

There’s a known race when using is_adapter_in_cpu_cache() to decide whether to omit weights (see our past learning). You compute adapter_in_cache before attempting a load, then unconditionally call _load_lora_adapter(). Re-evaluating cache status after the load reduces the window. Also, only omit weights for the PyTorch backend where a CPU peft cache is effectively shared; for TRT backend keep sending weights.

Proposed adjustment:

-            adapter_in_cache = self._lora_manager.is_adapter_in_cpu_cache(
-                request.lora_request.adapter_id)
-            self._load_lora_adapter(request.lora_request)
+            # Load the adapter from checkpoint if missing; no-op if already loaded.
+            self._load_lora_adapter(request.lora_request)
+            # Check cache status after the load to minimize TOCTOU race.
+            adapter_in_cache = self._lora_manager.is_adapter_in_cpu_cache(
+                request.lora_request.adapter_id)
             uid = str(request.lora_request.adapter_id)
-            lora_config = tllm.LoraConfig(
-                task_id=request.lora_request.adapter_id,
-                weights=self._lora_manager.cpp_lora_weights[uid]
-                if not adapter_in_cache else None,
-                config=self._lora_manager.cpp_lora_config[uid])
+            # For TRT backend, always send weights. For PyTorch backend, omit weights if present in CPU cache.
+            send_weights = (not self._is_pytorch_backend) or (not adapter_in_cache)
+            lora_config = tllm.LoraConfig(
+                task_id=request.lora_request.adapter_id,
+                weights=self._lora_manager.cpp_lora_weights[uid] if send_weights else None,
+                config=self._lora_manager.cpp_lora_config[uid])
             py_lora_path = request.lora_request.lora_path

This keeps TRT safe and narrows the race for PyTorch.

tensorrt_llm/_torch/pyexecutor/_util.py (1)

503-509: PeftCacheManager signature validated; lora_config propagation is correct.

Nit: duplicate ModelConfig imports in tensorrt_llm/_torch/pyexecutor/_util.py—remove or alias one to avoid ambiguity.

• Line 10: from tensorrt_llm._torch.model_config import ModelConfig
• Line 21: from ..model_config import ModelConfig

cpp/tensorrt_llm/pybind/batch_manager/bindings.cpp (1)

97-97: Build breaker: py::classh should be py::class_.

These occurrences will fail to compile. Replace with py::class_.

-    py::classh<GenLlmReq>(m, "GenericLlmRequest")
+    py::class_<GenLlmReq>(m, "GenericLlmRequest")
...
-    py::classh<tb::LlmRequest, GenLlmReq>(m, "LlmRequest", pybind11::dynamic_attr())
+    py::class_<tb::LlmRequest, GenLlmReq>(m, "LlmRequest", pybind11::dynamic_attr())

Also applies to: 258-258

cpp/tests/unit_tests/executor/loraConfigTest.cpp (1)

117-117: Test bug: writing to the wrong buffer (weightsData instead of configData).

This invalidates the serialization test for config.

-            weightsData[i * configDim1 + j] = 3 * (i * configDim1 + j);
+            configData[i * configDim1 + j] = 3 * (i * configDim1 + j);

cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp (1)

338-342: Prefer using LlmRequest::removeLoraTensors to clear LoRA data.

Simplify and centralize cleanup with the new API.

-            // free memory associated with lora weights in llmRequest
-            req->clearLoraWeights();
-            req->clearLoraConfig();
+            // free memory associated with LoRA tensors in the request
+            req->removeLoraTensors();

tests/unittest/_torch/test_resource_manager.py (1)

244-249: Enforce LoRA config dtype to int32 when converting from NumPy

The executor expects LoRA config to be INT32. Make the dtype explicit to avoid accidental int64 loads on some platforms.

-        if lora_config is not None:
-            lora_config = torch.from_numpy(lora_config)
+        if lora_config is not None:
+            lora_config = torch.from_numpy(lora_config).to(torch.int32)

tensorrt_llm/_torch/pyexecutor/resource_manager.py (1)

1237-1242: Fix ensure_batch return annotation

Tests treat ensure_batch’s return as a mapping (request_id -> peft table). The annotation says List[LlmRequest], which is misleading.

-    def ensure_batch(self,
+    def ensure_batch(self,
                      context_batch: List[LlmRequest],
                      generation_batch: List[LlmRequest],
-                     reset_gpu_cache: bool = False) -> List[LlmRequest]:
+                     reset_gpu_cache: bool = False) -> dict:
         return self.impl.ensure_batch(context_batch, generation_batch,
                                       reset_gpu_cache)

If there’s a concrete type for the peft table entries, prefer annotating that mapping precisely.

tensorrt_llm/_torch/pyexecutor/llm_request.py (1)

288-289: Carry py_lora_path in LlmRequest: LGTM; please document the new field

The addition is consistent with other py_* attributes and child-copying behavior. Please add a short docstring/comment in the constructor documenting py_lora_path semantics (expected to be an on-disk LoRA ckpt path, optional), per guidelines.

tensorrt_llm/_utils.py (1)

198-202: Type annotate and document binding_to_str_dtype

Minor polishing for clarity and Sphinx docs.

-def binding_to_str_dtype(binding_dtype) -> str:
-    ret = _binding_to_str_dtype.get(binding_dtype)
+def binding_to_str_dtype(binding_dtype: DataType) -> str:
+    """Return the string dtype name corresponding to a bindings DataType."""
+    ret = _binding_to_str_dtype.get(binding_dtype)
     assert ret is not None, f'Unsupported binding dtype: {binding_dtype}'
     return ret

cpp/include/tensorrt_llm/batch_manager/llmRequest.h (1)

2351-2353: Prefer Doxygen-style comment for new public API

Use Doxygen single-line comments for headers to match project guidelines.

-    // Remove LoRA weights and LoRA config tensors
-    void removeLoraTensors();
+    //! Remove LoRA weights and LoRA config tensors.
+    void removeLoraTensors();

tensorrt_llm/_torch/pyexecutor/_util.py (1)

8-13: Duplicate ModelConfig imports.

ModelConfig is imported twice from different paths; keep a single authoritative import to avoid ambiguity.

-from tensorrt_llm._torch.model_config import ModelConfig
-...
-from ..model_config import ModelConfig
+# Prefer one consistent import; for example:
+from tensorrt_llm._torch.model_config import ModelConfig

Also applies to: 21-24

cpp/tests/unit_tests/executor/loraConfigTest.cpp (1)

88-90: Misleading comment (optional).

Comment says "Wrong memory type" but this case checks mismatched shapes. Update for clarity.

-        // Wrong memory type
+        // Shapes not matching

cpp/tensorrt_llm/executor/loraConfig.cpp (1)

41-56: Tighten weights branch usability: small refactor and message polish

• Consider binding value() once to make intent and safety clearer and avoid repeated value() calls.
• Error messages mix “lora”/“LoRA”. Prefer consistent “LoRA” for clarity.
• Optional: if there’s a required dtype for weights (e.g., FP16/FP32), validate it here; otherwise, please confirm weights dtype is intentionally unconstrained.

Apply this minimal refactor for readability and consistency:

-    if (mWeights.has_value())
+    if (mWeights.has_value())
     {
-        SizeType32 constexpr expectedWeightsDims = 2;
-        TLLM_CHECK_WITH_INFO(
-            mConfig.has_value(), "Request for LoRA inference with lora weights must also have lora config");
+        SizeType32 constexpr expectedWeightsDims = 2;
+        TLLM_CHECK_WITH_INFO(
+            mConfig.has_value(), "Request for LoRA inference with LoRA weights must also have LoRA config");
 
-        TLLM_CHECK_WITH_INFO(
-            mWeights.value().getShape().size() == expectedWeightsDims, "Expected weights tensor to have 2 dimensions");
+        auto const& config = mConfig.value();
+        auto const& weights = mWeights.value();
+        TLLM_CHECK_WITH_INFO(
+            weights.getShape().size() == expectedWeightsDims, "Expected weights tensor to have 2 dimensions");
 
-        TLLM_CHECK_WITH_INFO(mWeights.value().getMemoryType() != MemoryType::kGPU
-                && mWeights.value().getMemoryType() != MemoryType::kUNKNOWN,
-            "Expected lora weights to be in CPU memory");
+        TLLM_CHECK_WITH_INFO(weights.getMemoryType() != MemoryType::kGPU
+                && weights.getMemoryType() != MemoryType::kUNKNOWN,
+            "Expected LoRA weights to be in CPU memory");
 
-        TLLM_CHECK_WITH_INFO(mConfig.value().getShape()[0] == mWeights.value().getShape()[0],
-            "Expected dim 0 of lora weights and lora config to have the same size");
+        TLLM_CHECK_WITH_INFO(config.getShape()[0] == weights.getShape()[0],
+            "Expected dim 0 of LoRA weights and LoRA config to have the same size");
     }

tests/unittest/_torch/test_resource_manager.py (3)

12-12: Follow Python import namespace guideline

Per project guidelines: always maintain the namespace when importing in Python. Prefer importing the module, then referencing the class.

-from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest
+from tensorrt_llm._torch.pyexecutor import llm_request as py_llm_request

And below (Line 251), construct via:

-        request = LlmRequest(
+        request = py_llm_request.LlmRequest(

---

251-260: Consider adding py_lora_path coverage

You now support lazy loading via py_lora_path. To exercise that path, add a test variant that sets request.py_lora_path and omits lora_weights, ensuring add_request_peft loads from checkpoint as intended.

I can draft such a test if helpful.

---

265-269: Docstring nit: clarify expected shapes and dtypes

Add that weights are float16 and config is int32, both 2D, to make expectations explicit.

-        """Create mock LoRA weights and config.
+        """Create mock LoRA weights and config.
+
+        Weights: float16, 2D CPU tensor.
+        Config: int32, 2D CPU tensor.

tests/unittest/llmapi/test_llm.py (1)

1462-1480: Avoid duplicating the helper across test suites

This helper mirrors logic added in the PyTorch test file. Consider centralizing it in tests/unittest/llmapi/lora_test_utils.py and reusing from both places to reduce drift.

tests/unittest/llmapi/test_llm_pytorch.py (1)

314-333: Mirror helper exists in TRT tests—consider centralizing

This helper duplicates the variant in tests/unittest/llmapi/test_llm.py. To reduce drift, move it to llmapi/lora_test_utils.py and import from both test files.

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

Learnt from: amitz-nv
PR: NVIDIA/TensorRT-LLM#5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.402Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks `is_adapter_in_cpu_cache()` and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.

Learnt from: moraxu
PR: NVIDIA/TensorRT-LLM#6303
File: tests/integration/test_lists/qa/examples_test_list.txt:494-494
Timestamp: 2025-07-28T17:06:08.621Z
Learning: In TensorRT-LLM testing, it's common to have both CLI flow tests (test_cli_flow.py) and PyTorch API tests (test_llm_api_pytorch.py) for the same model. These serve different purposes: CLI flow tests validate the traditional command-line workflow, while PyTorch API tests validate the newer LLM API backend. Both are legitimate and should coexist.

Learnt from: amitz-nv
PR: NVIDIA/TensorRT-LLM#5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.402Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks `is_adapter_in_cpu_cache()` and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.