Skip to content

Conversation

@ixlmar
Copy link
Collaborator

@ixlmar ixlmar commented Jun 12, 2025

[TRTLLM-5508] feat: check input tokens + improve error handling

Description

Fixes https://nvbugspro.nvidia.com/bug/5307858 and implements necessary error handling improvements in PyTorch backend and trtllm-bench.

Test Coverage

Tests covering invalid requests still need to be designed and implemented. Tests should best cover the LLM API as well as trtllm-bench.

Added test cases:

  • tests/unittest/llmapi/test_llm_pytorch.py::test_llm_invalid_input_token_async, tests/unittest/llmapi/test_llm_pytorch.py::test_llm_invalid_input_token
  • tests/integration/defs/test_e2e.py::test_trtllm_bench_invalid_token_pytorch

GitHub Bot Help

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

run [--disable-fail-fast --skip-test --stage-list "A10-1, xxx" --gpu-type "A30, H100_PCIe" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-[Post-Merge]-1, xxx"]

Launch build/test pipelines. All previously running jobs will be killed.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-1, xxx". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests. Will also run L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-[Post-Merge]-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-[Post-Merge]-1, xxx".

For guidance on mapping tests to stage names, see docs/source/reference/ci-overview.md.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

Summary by CodeRabbit

  • New Features

    • Added validation to ensure input token IDs are within the valid range, providing immediate error feedback for invalid tokens.
    • Improved error handling and reporting in benchmarking and asynchronous task management, ensuring clear messages and graceful shutdowns.
    • Exposed token ID range checking in both C++ and Python APIs.
  • Bug Fixes

    • Enhanced request queue logic to handle invalid tokens and prevent processing of erroneous requests.
    • Executor now skips stats dispatch when no results are available, avoiding unnecessary errors.
  • Tests

    • Introduced new unit and integration tests to verify correct handling and reporting of invalid token IDs.
    • Improved test output visibility by explicitly logging exceptions.
  • Refactor

    • Refactored request queue and asynchronous manager logic for better error aggregation, resource cleanup, and type safety.
  • Style

    • Added and updated type annotations for improved code clarity.

@ixlmar
Copy link
Collaborator Author

ixlmar commented Jun 12, 2025

/bot run

1 similar comment
@ixlmar
Copy link
Collaborator Author

ixlmar commented Jun 13, 2025

/bot run

@tensorrt-cicd
Copy link
Collaborator

PR_Github #8771 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #8771 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #6369 completed with status: 'FAILURE'

@ixlmar ixlmar force-pushed the feat/nvbug-5307858 branch from 108db04 to b923131 Compare June 13, 2025 11:59
@ixlmar
Copy link
Collaborator Author

ixlmar commented Jun 13, 2025

/bot run --stage-list "A100X-PyTorch-1"

@tensorrt-cicd
Copy link
Collaborator

PR_Github #8808 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #8808 [ run ] completed with state FAILURE
/LLM/main/L0_MergeRequest_PR pipeline #6396 (Partly Tested) completed with status: 'FAILURE'

@ixlmar ixlmar force-pushed the feat/nvbug-5307858 branch from b923131 to 1e7daaa Compare June 13, 2025 13:46
@ixlmar
Copy link
Collaborator Author

ixlmar commented Jun 13, 2025

/bot run --stage-list "A100X-PyTorch-1"

@tensorrt-cicd
Copy link
Collaborator

PR_Github #8815 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #8815 [ run ] completed with state FAILURE
/LLM/main/L0_MergeRequest_PR pipeline #6403 (Partly Tested) completed with status: 'FAILURE'

@ixlmar ixlmar force-pushed the feat/nvbug-5307858 branch from 1e7daaa to 23df114 Compare June 13, 2025 15:33
@ixlmar
Copy link
Collaborator Author

ixlmar commented Jun 13, 2025

/bot run --stage-list "A100X-PyTorch-1"

@tensorrt-cicd
Copy link
Collaborator

PR_Github #8829 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #8829 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #6416 (Partly Tested) completed with status: 'FAILURE'

@ixlmar ixlmar force-pushed the feat/nvbug-5307858 branch from 23df114 to e3fb4f4 Compare July 17, 2025 13:55
@coderabbitai
Copy link
Contributor

coderabbitai bot commented Jul 17, 2025

📝 Walkthrough

Walkthrough

This update introduces token ID range validation for requests in both the Python and C++ components of the system. It adds new methods for checking token validity, enhances error handling and reporting in the benchmarking and executor logic, and extends the test suite with cases for invalid token handling. Several methods are refactored for clarity, and type annotations are improved throughout.

Changes

Cohort / File(s) Change Summary
Token ID Validation (C++ and Python Bindings)
cpp/include/tensorrt_llm/batch_manager/llmRequest.h, cpp/tensorrt_llm/batch_manager/llmRequest.cpp, cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp, cpp/tensorrt_llm/pybind/batch_manager/bindings.cpp
Added checkTokenIdRange method to LlmRequest (C++), with Python bindings via nanobind and pybind. Enables checking if request tokens are within the vocabulary range.
Request Validation and Error Handling (PyTorch Executor)
tensorrt_llm/_torch/pyexecutor/py_executor.py, tensorrt_llm/_torch/pyexecutor/executor_request_queue.py, tensorrt_llm/_torch/pyexecutor/model_engine.py
Introduced token validation during request fetch, improved error handling for invalid requests, refactored enqueue logic, and updated method signatures with type annotations.
Benchmarking Error Handling and Shutdown
tensorrt_llm/bench/benchmark/throughput.py, tensorrt_llm/bench/benchmark/utils/asynchronous.py
Improved error handling in benchmarking: explicit process termination on errors, aggregated asynchronous task errors, graceful shutdown, and improved worker loop logic.
Executor and Proxy Robustness
tensorrt_llm/executor/proxy.py, tensorrt_llm/executor/worker.py
Added guard clause in stats dispatch to handle missing results, and type annotation for iteration result task.
Unit and Integration Testing for Invalid Tokens
tests/unittest/llmapi/test_llm_pytorch.py, tests/integration/defs/test_e2e.py, tests/integration/test_lists/test-db/l0_a10.yml
Added new unit and integration tests for invalid token handling, including async/sync PyTorch LLM tests and a benchmark integration test. Updated test list YAML to include new test.
Test Exception Reporting and Subprocess Timeout Handling
tests/unittest/conftest.py, tests/integration/defs/trt_test_alternative.py
Enhanced pytest to print exceptions and tracebacks for better visibility; improved subprocess timeout error propagation.

Sequence Diagram(s)

sequenceDiagram
    participant User
    participant PyExecutor
    participant ExecutorRequestQueue
    participant LlmRequest

    User->>PyExecutor: Submit requests
    PyExecutor->>ExecutorRequestQueue: Enqueue requests
    ExecutorRequestQueue-->>PyExecutor: Request IDs

    loop On fetch/activation
        PyExecutor->>ExecutorRequestQueue: Fetch new requests
        ExecutorRequestQueue-->>PyExecutor: LlmRequest(s)
        PyExecutor->>LlmRequest: check_token_id_range(vocab_size)
        alt Token ID valid
            PyExecutor->>PyExecutor: Add to active_requests
        else Token ID invalid
            PyExecutor->>PyExecutor: _handle_errors (fail specific requests)
        end
    end

    Note over PyExecutor: Only validated requests proceed to generation
Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Suggested labels

Community want to contribute

Suggested reviewers

  • FrankD412
  • pcastonguay
  • litaotju
  • yilin-void

Note

⚡️ Unit Test Generation is now available in beta!

Learn more here, or try it out under "Finishing Touches" below.


📜 Recent review details

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

📥 Commits

Reviewing files that changed from the base of the PR and between f20658b and 99093d3.

📒 Files selected for processing (16)
  • cpp/include/tensorrt_llm/batch_manager/llmRequest.h (1 hunks)
  • cpp/tensorrt_llm/batch_manager/llmRequest.cpp (1 hunks)
  • cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp (1 hunks)
  • cpp/tensorrt_llm/pybind/batch_manager/bindings.cpp (1 hunks)
  • tensorrt_llm/_torch/pyexecutor/executor_request_queue.py (7 hunks)
  • tensorrt_llm/_torch/pyexecutor/model_engine.py (1 hunks)
  • tensorrt_llm/_torch/pyexecutor/py_executor.py (9 hunks)
  • tensorrt_llm/bench/benchmark/throughput.py (2 hunks)
  • tensorrt_llm/bench/benchmark/utils/asynchronous.py (5 hunks)
  • tensorrt_llm/executor/proxy.py (1 hunks)
  • tensorrt_llm/executor/worker.py (1 hunks)
  • tests/integration/defs/test_e2e.py (2 hunks)
  • tests/integration/defs/trt_test_alternative.py (1 hunks)
  • tests/integration/test_lists/test-db/l0_a10.yml (1 hunks)
  • tests/unittest/conftest.py (2 hunks)
  • tests/unittest/llmapi/test_llm_pytorch.py (2 hunks)
✅ Files skipped from review due to trivial changes (1)
  • cpp/include/tensorrt_llm/batch_manager/llmRequest.h
🚧 Files skipped from review as they are similar to previous changes (13)
  • tests/integration/defs/trt_test_alternative.py
  • cpp/tensorrt_llm/pybind/batch_manager/bindings.cpp
  • cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp
  • tests/unittest/conftest.py
  • tensorrt_llm/bench/benchmark/throughput.py
  • cpp/tensorrt_llm/batch_manager/llmRequest.cpp
  • tests/integration/test_lists/test-db/l0_a10.yml
  • tensorrt_llm/executor/worker.py
  • tensorrt_llm/executor/proxy.py
  • tests/integration/defs/test_e2e.py
  • tensorrt_llm/_torch/pyexecutor/model_engine.py
  • tests/unittest/llmapi/test_llm_pytorch.py
  • tensorrt_llm/_torch/pyexecutor/executor_request_queue.py
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py

📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)

**/*.py: Python code should conform to Python 3.8+.
Indent Python code with 4 spaces. Do not use tabs.
Always maintain the namespace when importing in Python, even if only one class or function from a module is used.
Python filenames should use snake_case (e.g., some_file.py).
Python classes should use PascalCase (e.g., class SomeClass).
Python functions and methods should use snake_case (e.g., def my_awesome_function():).
Python local variables should use snake_case, and prefix k for variable names that start with a number (e.g., k_99th_percentile = ...).
Python global variables should use upper snake_case and prefix G (e.g., G_MY_GLOBAL = ...).
Python constants should use upper snake_case (e.g., MY_CONSTANT = ...).
Avoid shadowing variables declared in an outer scope in Python.
Initialize all externally visible members of a Python class in the constructor.
For interfaces that may be used outside a Python file, prefer docstrings over comments.
Comments in Python should be reserved for code within a function, or interfaces that are local to a file.
Use Google style docstrings for Python classes and functions, which can be parsed by Sphinx.
Attributes and variables in Python can be documented inline; attribute docstrings will be rendered under the docstring for the class.
Avoid using reflection in Python when functionality can be easily achieved without reflection.
When using try-except blocks in Python, limit the except to the smallest set of errors possible.
When using try-except blocks to handle multiple possible variable types in Python, keep the body of the try as small as possible, using the else block to implement the logic.

Files:

  • tensorrt_llm/_torch/pyexecutor/py_executor.py
  • tensorrt_llm/bench/benchmark/utils/asynchronous.py
**/*.{cpp,h,cu,py,cc,cxx,hpp}

📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)

All TensorRT-LLM Open Source Software code should contain an NVIDIA copyright header that includes the current year. This includes .cpp, .h, .cu, .py, and any other source files which are compiled or interpreted.

Files:

  • tensorrt_llm/_torch/pyexecutor/py_executor.py
  • tensorrt_llm/bench/benchmark/utils/asynchronous.py
🧠 Learnings (3)
📓 Common learnings
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.
📚 Learning: in tensorrt_llm/executor/worker.py, the lora adapter cache optimization logic that checks `is_adapte...
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.

Applied to files:

  • tensorrt_llm/_torch/pyexecutor/py_executor.py
📚 Learning: in tensorrt-llm's multimodal processing pipeline, shared tensor recovery using `from_shared_tensor()...
Learnt from: yechank-nvidia
PR: NVIDIA/TensorRT-LLM#6254
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:1201-1204
Timestamp: 2025-07-22T09:22:14.726Z
Learning: In TensorRT-LLM's multimodal processing pipeline, shared tensor recovery using `from_shared_tensor()` is only needed during the context phase. Generation requests reuse the already-recovered tensor data and only need to call `strip_for_generation()` to remove unnecessary multimodal data while preserving the recovered tensors. This avoids redundant tensor recovery operations during generation.

Applied to files:

  • tensorrt_llm/_torch/pyexecutor/py_executor.py
🔇 Additional comments (10)
tensorrt_llm/_torch/pyexecutor/py_executor.py (5)

11-11: LGTM! Clean type annotation improvements.

These changes enhance code clarity by adding proper type hints without affecting functionality. The Iterable import supports the flexible parameter type in _enqueue_responses, and the return type annotations accurately reflect the methods' actual return types.

Also applies to: 32-32, 288-288, 380-380


1119-1147: Well-structured refactoring with clean error handling.

The method refactoring effectively separates validation concerns and implements fail-fast behavior for invalid requests. The nested _respond_if_invalid function provides clear separation of concerns, and the list comprehension approach is clean and readable.

The error handling properly isolates invalid requests and fails them individually without affecting valid requests in the batch.


1422-1444: Excellent enhancement enabling granular error handling.

This refactoring adds valuable functionality while maintaining backward compatibility. Key strengths:

  • Keyword-only parameter: Prevents accidental positional argument usage
  • Selective vs. global handling: Supports both targeted request failures and global error scenarios
  • Clean list comprehension: Efficiently removes only the failed requests from active_requests
  • Backward compatibility: Existing calls continue to work unchanged

The design properly handles the different error scenarios without breaking existing functionality.


1472-1472: Good type flexibility improvement.

Changing the parameter type from List to Iterable provides better interface flexibility and supports the usage pattern in _handle_errors where dict.items() is passed. This follows Python best practices of accepting the most general interface that meets the method's needs.


660-660: Consistent method name updates.

The calls have been properly updated to use the new _fetch_and_activate_new_requests method name, maintaining consistency with the refactored implementation.

Also applies to: 848-848

tensorrt_llm/bench/benchmark/utils/asynchronous.py (5)

7-7: Well-designed error aggregation system.

The new error collection mechanism effectively aggregates failures from multiple asynchronous tasks and provides informative error summaries. Key strengths:

  • Comprehensive collection: Gathers all task errors rather than failing on first error
  • Useful reporting: Groups identical errors with counts for easier debugging
  • Clean implementation: Proper type hints and clear logic flow
  • Actionable feedback: The aggregated error messages help identify patterns in failures

This approach provides much better visibility into batch processing failures compared to immediate termination.

Also applies to: 39-41, 94-105


107-113: Good callback refactoring supporting error aggregation.

The simplified callback now properly focuses on task bookkeeping and error collection, allowing the main error handling logic to aggregate and report errors comprehensively. The approach of continuing processing despite individual task failures is appropriate for batch operations.


55-92: Excellent async refactoring with proper resource management.

The rewritten process_request method demonstrates excellent async programming practices:

  • Clean resource management: Proper use of semaphore_guard context manager
  • Clear control flow: Well-structured handling of streaming vs non-streaming cases
  • Comprehensive metrics: Thorough performance tracking including first-token timing
  • Proper error propagation: Allows exceptions to bubble up for aggregation
  • Readable structure: Logical flow from request initiation through response processing

This refactoring significantly improves the code quality and maintainability.


194-201: Proper async conversion for clean shutdown.

Converting stop() to async enables proper awaiting of background tasks instead of abrupt cancellation. The sequential awaiting of iteration log task followed by backend task ensures orderly shutdown.


205-206: Improved busy state detection and task management.

The enhanced busy property now accurately reflects both queued and actively running tasks, providing better visibility into the manager's state. The conditional task assignment in run() is also cleaner.

Also applies to: 211-211

✨ Finishing Touches
  • 📝 Generate Docstrings
🧪 Generate unit tests
  • Create PR with unit tests
  • Post copyable unit tests in a comment

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share
🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

  • Review comments: Directly reply to a review comment made by CodeRabbit. Example:
    • I pushed a fix in commit <commit_id>, please review it.
    • Explain this complex logic.
    • Open a follow-up GitHub issue for this discussion.
  • Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
    • @coderabbitai explain this code block.
  • PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
    • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
    • @coderabbitai read src/utils.ts and explain its main purpose.
    • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

CodeRabbit Commands (Invoked using PR comments)

  • @coderabbitai pause to pause the reviews on a PR.
  • @coderabbitai resume to resume the paused reviews.
  • @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
  • @coderabbitai full review to do a full review from scratch and review all the files again.
  • @coderabbitai summary to regenerate the summary of the PR.
  • @coderabbitai generate docstrings to generate docstrings for this PR.
  • @coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
  • @coderabbitai generate unit tests to generate unit tests for this PR.
  • @coderabbitai resolve resolve all the CodeRabbit review comments.
  • @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
  • @coderabbitai help to get help.

Other keywords and placeholders

  • Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
  • Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
  • Add @coderabbitai or @coderabbitai title anywhere in the PR title to generate the title automatically.

Documentation and Community

  • Visit our Documentation for detailed information on how to use CodeRabbit.
  • Join our Discord Community to get help, request features, and share feedback.
  • Follow us on X/Twitter for updates and announcements.

@ixlmar
Copy link
Collaborator Author

ixlmar commented Jul 17, 2025

/bot run --stage-list "A100X-PyTorch-1"

Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 0

🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/model_engine.py (1)

930-938: Return type annotation narrows the contract—verify all overrides/stubs

Adding -> DecoderModelForCausalLM improves readability, but _load_model is occasionally monkey-patched in tests (see comments around line 299) and might return stub models that are not DecoderModelForCausalLM.
If any such override now returns a different subtype (e.g., a lightweight torch.nn.Module), static checkers will emit errors. Confirm that:

  1. All overrides/stubs indeed inherit from DecoderModelForCausalLM, or
  2. You relax the hint to a wider type (e.g., torch.nn.Module | DecoderModelForCausalLM) to keep tests/builds warning-free.

No functional issue—just a typing concern.

📜 Review details

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

📥 Commits

Reviewing files that changed from the base of the PR and between 58d22a7 and e3fb4f4.

📒 Files selected for processing (4)
  • tensorrt_llm/_torch/pyexecutor/model_engine.py (1 hunks)
  • tensorrt_llm/_torch/pyexecutor/py_executor.py (11 hunks)
  • tensorrt_llm/bench/benchmark/throughput.py (2 hunks)
  • tensorrt_llm/bench/benchmark/utils/asynchronous.py (5 hunks)
🧰 Additional context used
🧠 Learnings (1)
tensorrt_llm/_torch/pyexecutor/py_executor.py (1)
Learnt from: amitz-nv
PR: NVIDIA/TensorRT-LLM#5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.374Z
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.
🧬 Code Graph Analysis (1)
tensorrt_llm/bench/benchmark/throughput.py (1)
tensorrt_llm/logger.py (1)
  • error (125-126)
🔇 Additional comments (10)
tensorrt_llm/bench/benchmark/throughput.py (1)

5-5: Improved error handling with explicit process termination.

The addition of explicit exception handlers for KeyboardInterrupt and generic exceptions, along with sys.exit(-1) calls, ensures proper non-zero exit status on errors. This is a good practice for benchmarking tools where exit codes are important for CI/CD pipelines.

Also applies to: 454-459

tensorrt_llm/bench/benchmark/utils/asynchronous.py (4)

39-40: Excellent refactoring of task lifecycle management and error handling.

The worker method improvements are well-designed:

  • Non-blocking queue operations with get_nowait() + sleep(1) prevent deadlocks
  • Error aggregation provides detailed failure counts for better debugging
  • Proper task cleanup ensures no resource leaks
  • Clear separation of concerns with the _raise_for_failed helper

Also applies to: 97-145


198-205: Proper async shutdown implementation.

Converting stop() to async and awaiting task completion ensures graceful shutdown without race conditions. This is the correct pattern for async task management.


209-210: Fixed critical typo and improved busy state detection.

  • The busy property now accurately reflects ongoing work by checking both queue state and active task completion
  • The typo fix ensures _iteration_log_task is properly assigned for tracking

Also applies to: 215-215


302-302: Proper async cleanup in benchmark flow.

Adding await for both enqueue_task cancellation and backend.stop() ensures clean resource management and prevents task leakage.

Also applies to: 304-304

tensorrt_llm/_torch/pyexecutor/py_executor.py (5)

353-383: Well-designed enqueue refactoring following DRY principles.

The unified _enqueue_impl method elegantly handles both single and batch enqueueing with optional query support. The use of itertools.repeat(None) for the batch case is a clean solution.


1360-1384: Excellent separation of fetch, validate, and activate phases.

The _fetch_and_activate_new_requests method properly validates requests before activation, immediately failing invalid ones. This prevents invalid requests from consuming resources and provides fast feedback to users.


1981-2002: Improved granular error handling allows selective request failure.

The enhanced _handle_errors method now supports failing specific requests without affecting others, which is crucial for maintaining service availability when individual requests have issues (e.g., invalid tokens).


1210-1219: Clean extraction of request counting logic.

Extracting _count_active_requests_all_ranks improves code organization and testability while maintaining the same functionality.


1350-1359: Token validation logic has an off-by-one error.

The upper bound check should use >= self.model_engine.model.lm_head.num_embeddings instead of > ... - 1 for clarity and correctness.

-    if any(
-            min(tokens) < 0 or max(tokens) >
-            self.model_engine.model.lm_head.num_embeddings - 1
-            for tokens in token_lists):
+    if any(
+            min(tokens) < 0 or max(tokens) >= 
+            self.model_engine.model.lm_head.num_embeddings
+            for tokens in token_lists):

Likely an incorrect or invalid review comment.

@ixlmar ixlmar force-pushed the feat/nvbug-5307858 branch from e3fb4f4 to 10536f6 Compare July 17, 2025 14:00
@ixlmar
Copy link
Collaborator Author

ixlmar commented Jul 17, 2025

/bot run --stage-list "A100X-PyTorch-1"

@tensorrt-cicd
Copy link
Collaborator

PR_Github #12209 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #12209 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #9065 (Partly Tested) completed with status: 'FAILURE'

@ixlmar ixlmar force-pushed the feat/nvbug-5307858 branch 2 times, most recently from 072deea to 51d8ac0 Compare July 22, 2025 09:40
@ixlmar
Copy link
Collaborator Author

ixlmar commented Jul 22, 2025

/bot run

Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (4)
tensorrt_llm/bench/benchmark/throughput.py (2)

454-456: Consider using exit code 130 for KeyboardInterrupt.

While the implementation correctly handles KeyboardInterrupt and ensures clean termination, the standard Unix convention for keyboard interrupt (SIGINT) is exit code 130, not -1.

-        sys.exit(-1)
+        sys.exit(130)

457-459: Improve error logging and consider exit code 1.

The generic exception handler has two areas for improvement:

  1. The error message could be more informative by including exception type
  2. Exit code 1 is more conventional than -1 for general errors
-    except Exception as e:
-        logger.error(f"Error during benchmarking: {e}")
-        sys.exit(-1)
+    except Exception as e:
+        logger.error(f"Error during benchmarking: {type(e).__name__}: {e}")
+        sys.exit(1)
tensorrt_llm/_torch/pyexecutor/executor_request_queue.py (2)

122-122: Consider replacing assertion with proper exception

While the assertion provides clear error messaging, assertions can be disabled with Python's -O flag. Consider using a proper exception for more reliable error handling in production.

-            assert self.active, "PyExecutor has already been shutdown."
+            if not self.active:
+                raise RuntimeError("PyExecutor has already been shutdown.")

140-146: Good backward-compatible signature enhancement

The addition of the optional query parameter maintains backward compatibility while supporting StarAttention functionality. The method correctly delegates to the new unified implementation.

Consider using a more specific type than Optional[List] if the query structure is well-defined:

query: Optional[List[int]] = None  # if it's token IDs
📜 Review details

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

📥 Commits

Reviewing files that changed from the base of the PR and between 072deea and 51d8ac0.

📒 Files selected for processing (6)
  • tensorrt_llm/_torch/pyexecutor/executor_request_queue.py (7 hunks)
  • tensorrt_llm/_torch/pyexecutor/model_engine.py (1 hunks)
  • tensorrt_llm/_torch/pyexecutor/py_executor.py (8 hunks)
  • tensorrt_llm/bench/benchmark/throughput.py (2 hunks)
  • tensorrt_llm/bench/benchmark/utils/asynchronous.py (5 hunks)
  • tests/unittest/conftest.py (2 hunks)
🚧 Files skipped from review as they are similar to previous changes (4)
  • tests/unittest/conftest.py
  • tensorrt_llm/_torch/pyexecutor/model_engine.py
  • tensorrt_llm/_torch/pyexecutor/py_executor.py
  • tensorrt_llm/bench/benchmark/utils/asynchronous.py
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
  • GitHub Check: Pre-commit Check
🔇 Additional comments (6)
tensorrt_llm/bench/benchmark/throughput.py (1)

5-5: LGTM: Import added for process termination.

The sys import is correctly added to support the new sys.exit(-1) calls in the exception handlers.

tensorrt_llm/_torch/pyexecutor/executor_request_queue.py (5)

8-9: LGTM: Import improvements enhance type safety

The addition of repeat from itertools and comprehensive type imports support the refactored enqueue methods and improved type annotations throughout the file.

Also applies to: 17-18


116-147: Excellent refactoring: Unified enqueue implementation

The introduction of _enqueue_impl as a private helper eliminates code duplication while maintaining the same functionality. The use of zip(requests, repeat(None)) for batch operations is elegant and efficient.


121-121: Excellent modernization: Context manager for lock safety

Converting to context managers (with self.enqueue_lock:) ensures locks are always properly released, even during exceptions. This eliminates potential deadlock scenarios and follows modern Python concurrency best practices.

Also applies to: 149-149, 154-154, 159-160


202-202: LGTM: Comprehensive type annotations

The addition of explicit return type annotations (-> List[LlmRequest]) and modern parameter type syntax improves type safety and IDE support without changing functionality.

Also applies to: 210-210, 224-224, 429-430, 448-449


116-147: Well-coordinated refactoring supports broader improvements

This refactoring provides a solid foundation for the request validation and error handling improvements mentioned in the PR objectives, while maintaining backward compatibility and improving code maintainability.

@tensorrt-cicd
Copy link
Collaborator

PR_Github #12553 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #12553 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #9334 completed with status: 'FAILURE'

@tensorrt-cicd
Copy link
Collaborator

PR_Github #14013 [ run ] completed with state ABORTED

@ixlmar ixlmar force-pushed the feat/nvbug-5307858 branch from 6af3918 to f20658b Compare August 4, 2025 17:31
@ixlmar
Copy link
Collaborator Author

ixlmar commented Aug 4, 2025

/bot run --stage-list="A100X-PyTorch-1"

@tensorrt-cicd
Copy link
Collaborator

PR_Github #14023 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #14019 [ run ] completed with state ABORTED

Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

📜 Review details

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

📥 Commits

Reviewing files that changed from the base of the PR and between 6af3918 and f20658b.

📒 Files selected for processing (16)
  • cpp/include/tensorrt_llm/batch_manager/llmRequest.h (1 hunks)
  • cpp/tensorrt_llm/batch_manager/llmRequest.cpp (1 hunks)
  • cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp (1 hunks)
  • cpp/tensorrt_llm/pybind/batch_manager/bindings.cpp (1 hunks)
  • tensorrt_llm/_torch/pyexecutor/executor_request_queue.py (7 hunks)
  • tensorrt_llm/_torch/pyexecutor/model_engine.py (1 hunks)
  • tensorrt_llm/_torch/pyexecutor/py_executor.py (9 hunks)
  • tensorrt_llm/bench/benchmark/throughput.py (2 hunks)
  • tensorrt_llm/bench/benchmark/utils/asynchronous.py (6 hunks)
  • tensorrt_llm/executor/proxy.py (1 hunks)
  • tensorrt_llm/executor/worker.py (1 hunks)
  • tests/integration/defs/test_e2e.py (2 hunks)
  • tests/integration/defs/trt_test_alternative.py (1 hunks)
  • tests/integration/test_lists/test-db/l0_a10.yml (1 hunks)
  • tests/unittest/conftest.py (2 hunks)
  • tests/unittest/llmapi/test_llm_pytorch.py (2 hunks)
✅ Files skipped from review due to trivial changes (1)
  • cpp/include/tensorrt_llm/batch_manager/llmRequest.h
🚧 Files skipped from review as they are similar to previous changes (14)
  • tests/integration/defs/trt_test_alternative.py
  • cpp/tensorrt_llm/batch_manager/llmRequest.cpp
  • tensorrt_llm/executor/proxy.py
  • tensorrt_llm/bench/benchmark/throughput.py
  • tests/unittest/conftest.py
  • cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp
  • tensorrt_llm/executor/worker.py
  • tensorrt_llm/_torch/pyexecutor/model_engine.py
  • tests/integration/test_lists/test-db/l0_a10.yml
  • tests/integration/defs/test_e2e.py
  • cpp/tensorrt_llm/pybind/batch_manager/bindings.cpp
  • tensorrt_llm/_torch/pyexecutor/py_executor.py
  • tensorrt_llm/_torch/pyexecutor/executor_request_queue.py
  • tests/unittest/llmapi/test_llm_pytorch.py
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py

📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)

**/*.py: The code developed for TensorRT-LLM should conform to Python 3.8+.
Indent Python code with 4 spaces. Do not use tabs.
Always maintain the namespace when importing in Python, even if only one class or function from a module is used.
Python filenames should use snake_case (e.g., some_file.py).
Python classes should use PascalCase (e.g., class SomeClass).
Python functions and methods should use snake_case (e.g., def my_awesome_function():).
Python local variables should use snake_case. Prefix k for variable names that start with a number (e.g., k_99th_percentile = ...).
Python global variables should use upper snake_case and prefix G (e.g., G_MY_GLOBAL = ...).
Python constants should use upper snake_case (e.g., MY_CONSTANT = ...).
Avoid shadowing variables declared in an outer scope in Python.
Initialize all externally visible members of a Python class in the constructor.
For interfaces that may be used outside a file, prefer docstrings over comments in Python.
Comments in Python should be reserved for code within a function, or interfaces that are local to a file.
Use Google style docstrings for classes and functions in Python, which can be parsed by Sphinx.
Attributes and variables in Python can be documented inline; attribute docstrings will be rendered under the docstring for the class.
Avoid using reflection in Python when functionality can be easily achieved without it.
When using try-except blocks in Python, limit the except to the smallest set of errors possible.
When using try-except blocks to handle multiple possible variable types in Python, keep the body of the try as small as possible, using the else block to implement the logic.

Files:

  • tensorrt_llm/bench/benchmark/utils/asynchronous.py
**/*.{cpp,h,hpp,cc,cxx,cu,py}

📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)

All TensorRT-LLM Open Source Software code should contain an NVIDIA copyright header that includes the current year. This includes .cpp, .h, .cu, .py, and any other source files which are compiled or interpreted.

Files:

  • tensorrt_llm/bench/benchmark/utils/asynchronous.py
🧠 Learnings (1)
📓 Common learnings
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.
🔇 Additional comments (9)
tensorrt_llm/bench/benchmark/utils/asynchronous.py (9)

7-7: LGTM: Import addition supports new type hints.

The addition of Set to the typing imports is necessary for the new self._tasks: Set[asyncio.Task] type annotation.


39-41: Good: Enhanced error tracking and task management.

The new instance variables improve error handling by:

  • _task_errors: Collecting all task exceptions instead of failing fast
  • _backend_task and _iteration_log_task: Proper type hints for better async task management

80-80: Minor improvement: More efficient token extraction.

Converting the list comprehension to a generator expression with chain is more memory-efficient for large token sequences.


98-116: Excellent: Comprehensive error aggregation and reporting.

The new error handling approach provides several benefits:

  • Collects all task errors instead of stopping on first failure
  • Groups identical errors with counts for cleaner reporting
  • Provides detailed error summary in the exception message

This aligns well with the --disable-fail-fast behavior mentioned in the PR comments.


138-144: Good: Proper cleanup with task cancellation and error propagation.

The cleanup logic in the finally block correctly:

  • Cancels all remaining tasks
  • Waits for their completion
  • Raises any accumulated errors

This ensures no tasks are left hanging and all errors are properly reported.


197-204: LGTM: Proper async conversion with task awaiting.

Converting stop() to async and awaiting the tasks instead of cancelling them ensures proper cleanup. The assertion on _backend_task is appropriate since it should always be set after run() is called.


208-209: Good: More accurate busy state detection.

The updated busy property now correctly checks both queue state and running tasks, providing a more accurate representation of the manager's activity state.


214-214: LGTM: Type hint consistency.

Adding the explicit type hint for _iteration_log_task maintains consistency with the new type annotations.


117-144: Critical: Ensure non-blocking worker loop doesn’t spin the CPU

While switching from blocking Queue.get() to get_nowait() plus await asyncio.sleep(0) yields better task interleaving, an always‐empty inbox could still cause a tight yield loop under high concurrency. We do have benchmarks in place—

  • tests/unittest/scaffolding/test_bench.py: async_scaffolding_benchmark
  • tensorrt_llm/bench/benchmark/throughput.py & low_latency.py invoking async_benchmark

—but none explicitly measure CPU utilization when the queue is empty. Please add CPU‐profiling or utilization metrics to your existing benchmarks to confirm there’s no busy‐waiting regression.

• tensorrt_llm/bench/benchmark/utils/asynchronous.py (worker: lines 117–144)
• tests/unittest/scaffolding/test_bench.py (async_scaffolding_benchmark)
• tensorrt_llm/bench/benchmark/throughput.py & low_latency.py (async_benchmark callers)

@tensorrt-cicd
Copy link
Collaborator

PR_Github #14023 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #10572 (Partly Tested) completed with status: 'SUCCESS'

@ixlmar ixlmar force-pushed the feat/nvbug-5307858 branch from f20658b to 99093d3 Compare August 5, 2025 06:43
@ixlmar
Copy link
Collaborator Author

ixlmar commented Aug 5, 2025

/bot run --stage-list="A100X-PyTorch-1"

@tensorrt-cicd
Copy link
Collaborator

PR_Github #14103 [ run ] triggered by Bot

@ixlmar
Copy link
Collaborator Author

ixlmar commented Aug 5, 2025

@FrankD412 Did a short sanity-check using:

$ cat /tmp/input.yml
cuda_graph_config:
  enable_padding: true
  batch_sizes:
  - 1
  - 2
  - 4
  - 8
  - 16
  - 32
  - 64
  - 128
  - 256
  - 384

$ trtllm-bench \
    --model TinyLlama-1.1B-Chat-v1.0 \
    --model_path /models/llama-models-v2/TinyLlama-1.1B-Chat-v1.0 \
    throughput \
    --dataset tiny_input1_output1_requests10000_good.json \
    --backend pytorch \
    --extra_llm_api_options /tmp/input.yml \
    --concurrency 1000

Results with new [old] code:

===========================================================
= REQUEST DETAILS
===========================================================
Number of concurrent requests:  887.5591 [892.2088]
Average Input Length (tokens):  1.0000 [1.0000]
Average Output Length (tokens): 1.0000 [1.0000]
===========================================================
= WORLD + RUNTIME INFORMATION
===========================================================
Max Runtime Batch Size: 4096 [4096]
Max Runtime Tokens:     8192 [8192]
Scheduling Policy:      GUARANTEED_NO_EVICT [GUARANTEED_NO_EVICT]
KV Memory Percentage:   90.00% [90.00%]
Issue Rate (req/sec):   5.6360E+14 [5.8349E+14]

===========================================================
= PERFORMANCE OVERVIEW
===========================================================
Request Throughput (req/sec):                     1581.2645 [1530.5127]
Total Output Throughput (tokens/sec):             1581.2645 [1530.5127]
Total Token Throughput (tokens/sec):              3162.5290 [3061.0254]
Total Latency (ms):                               6324.0526 [6533.7583]
Average request latency (ms):                     561.2971 [582.9477]
Per User Output Throughput [w/ ctx] (tps/user):   1.8354 [1.7774]
Per GPU Output Throughput (tps/gpu):              1581.2645 [1530.5127]

-- Request Latency Breakdown (ms) -----------------------

[Latency] P50    : 549.1097 [589.3429]
[Latency] P90    : 696.3812 [718.5348]
[Latency] P95    : 719.9421 [736.1621]
[Latency] P99    : 757.0753 [761.9933]
[Latency] MINIMUM: 143.7844 [128.6150]
[Latency] MAXIMUM: 779.3706 [784.5101]
[Latency] AVERAGE: 561.2971 [582.9477]

Re-running with the old code suggests that the observed differences are within the spread.

@tensorrt-cicd
Copy link
Collaborator

PR_Github #14103 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #10640 (Partly Tested) completed with status: 'SUCCESS'

@ixlmar
Copy link
Collaborator Author

ixlmar commented Aug 5, 2025

/bot run

@tensorrt-cicd
Copy link
Collaborator

PR_Github #14117 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #14117 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #10651 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

@DomBrown DomBrown enabled auto-merge (squash) August 5, 2025 17:26
@DomBrown DomBrown merged commit 1ebceb7 into NVIDIA:main Aug 5, 2025
4 checks passed
lancelly pushed a commit to lancelly/TensorRT-LLM that referenced this pull request Aug 6, 2025
@ixlmar ixlmar deleted the feat/nvbug-5307858 branch August 6, 2025 06:07
@ixlmar ixlmar restored the feat/nvbug-5307858 branch August 6, 2025 06:23
@ixlmar ixlmar deleted the feat/nvbug-5307858 branch August 6, 2025 06:24
jain-ria pushed a commit to jain-ria/TensorRT-LLM that referenced this pull request Aug 7, 2025
@coderabbitai coderabbitai bot mentioned this pull request Aug 26, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants