[TRTLLM-10185][feat] AutoTuner Cache: Support cache file lock and merge all ranks into one #10336
Conversation
…ge all ranks into one Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthroughThe changes introduce rank-aware cache operations with file-locking in the autotuner, evolve the cache file format from single-blob to per-rank keyed JSON with metadata sections, refactor serialization/deserialization methods, and update public API signatures for profiling cache methods to include rank parameters. Changes
Sequence Diagram(s)sequenceDiagram
participant Rank0 as Rank 0 Worker
participant RankN as Rank N Worker
participant AutoTuner as AutoTuner
participant Dist as Distributed Backend<br/>(MPIDist/TorchDist)
participant Cache as Cache File<br/>(Rank-keyed JSON)
Rank0->>Dist: setup_distributed_state(mapping, dist)
RankN->>Dist: setup_distributed_state(mapping, dist)
rect rgba(0, 100, 200, 0.1)
note over Rank0,RankN: Cache Path Synchronization
Rank0->>Rank0: create temp cache_path
Rank0->>Dist: broadcast cache_path
Dist->>RankN: receive cache_path
end
rect rgba(100, 150, 100, 0.1)
note over Rank0,RankN: Distributed Tuning
Rank0->>AutoTuner: autotune(cache_path=cache_path)
RankN->>AutoTuner: autotune(cache_path=cache_path)
AutoTuner->>AutoTuner: profile kernels
end
rect rgba(150, 100, 100, 0.1)
note over Rank0,RankN: Cache Persistence & Reload
Rank0->>Cache: save_cache(cache_path, rank=0)<br/>writes rank_0 entry with metadata
RankN->>Cache: save_cache(cache_path, rank=N)<br/>writes rank_N entry
Rank0->>Cache: load_cache(cache_path, rank=0)<br/>reads rank_0 entry
RankN->>Cache: load_cache(cache_path, rank=N)<br/>reads rank_N entry
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Pre-merge checks and finishing touches❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✨ Finishing touches
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/autotuner.py (1)
583-593: PotentialUnboundLocalErrorwhen tactic deserialization fails.If
ast.literal_eval(value["tactic"])fails at line 584 and raises aValueErrororTypeError, thetacticvariable is never assigned, but the code continues (nocontinuestatement) and tries to usetacticat line 593. This will cause anUnboundLocalError.🔎 Proposed fix to skip entry on deserialization failure
try: tactic = ast.literal_eval(value["tactic"]) except (ValueError, TypeError): logger.warning_once( f"[AutoTuner] Could not deserialize tactic: {value['tactic']} for cache key {key_str}", key=value["tactic"]) + continue runner_id = value["runner_id"] min_time = value["min_time"] cache[key] = (runner_id, tactic, min_time)
🧹 Nitpick comments (2)
tensorrt_llm/_torch/autotuner.py (2)
455-476: Consider atomic write pattern to prevent file corruption.The current read-modify-write pattern with truncate could leave the file in a corrupted state if an exception occurs between
f.truncate()(line 468) andjson.dump()(line 470). Consider writing to a temporary file first, then atomically renaming it.Also, while the lock is implicitly released when the file is closed, explicit release with
fcntl.flock(f, fcntl.LOCK_UN)in afinallyblock would make the intent clearer.🔎 Proposed atomic write pattern
try: serialized_rank_cache_data = self._serialize_cache_data() - with open(file_path, 'a+') as f: - fcntl.flock(f, fcntl.LOCK_EX) - f.seek(0) - content = f.read() - if content.strip(): - current_cache = json.loads(content) - else: - current_cache = { - "metadata": self._serialize_metadata(), - } - f.seek(0) - f.truncate() - current_cache[f"rank_{rank}"] = serialized_rank_cache_data - json.dump(current_cache, f, indent=2, default=str) + with open(file_path, 'a+') as f: + fcntl.flock(f, fcntl.LOCK_EX) + try: + f.seek(0) + content = f.read() + if content.strip(): + current_cache = json.loads(content) + else: + current_cache = { + "metadata": self._serialize_metadata(), + } + current_cache[f"rank_{rank}"] = serialized_rank_cache_data + # Write to temp file first for atomicity + temp_path = file_path.with_suffix('.tmp') + with open(temp_path, 'w') as tmp_f: + json.dump(current_cache, tmp_f, indent=2, default=str) + temp_path.replace(file_path) + finally: + fcntl.flock(f, fcntl.LOCK_UN)
496-509: Lock scope issue: cache deserialization occurs after lock release.The shared lock is held only within the
with open(file_path, 'r') as f:block (lines 497-501), but the cache deserialization (self._deserialize_cache_data) on line 502-503 occurs after the file is closed and the lock is released. While this is functionally correct sincecurrent_cache_contentsis already loaded into memory, consider keeping the lock until deserialization completes for consistency with save_cache, or move the deserialization inside the with block.Additionally, consider explicit lock release in a finally block for clarity.
🔎 Proposed fix to keep lock scope consistent
try: with open(file_path, 'r') as f: fcntl.flock(f, fcntl.LOCK_SH) - current_cache_contents = json.load(f) - self._deserialize_metadata(current_cache_contents["metadata"]) - assert f"rank_{rank}" in current_cache_contents, f"Rank {rank} cache not found in {file_path}" - self.cache = self._deserialize_cache_data( - current_cache_contents[f'rank_{rank}']) + try: + current_cache_contents = json.load(f) + self._deserialize_metadata(current_cache_contents["metadata"]) + assert f"rank_{rank}" in current_cache_contents, f"Rank {rank} cache not found in {file_path}" + self.cache = self._deserialize_cache_data( + current_cache_contents[f'rank_{rank}']) + finally: + fcntl.flock(f, fcntl.LOCK_UN)
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
tensorrt_llm/_torch/autotuner.pytests/integration/test_lists/test-db/l0_dgx_b200.ymltests/unittest/_torch/misc/test_autotuner.py
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
**/*.py: 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 files should use snake_case naming:some_file.py
Python classes should use PascalCase naming:class SomeClass
Python functions and methods should use snake_case naming:def my_awesome_function():
Python local variables should use snake_case naming:my_variable = ...
Python variable names that start with a number should be prefixed with 'k':k_99th_percentile = ...
Python global variables should use upper snake_case with prefix 'G':G_MY_GLOBAL = ...
Python constants should use upper snake_case naming: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 Python interfaces that may be used outside a file, prefer docstrings over comments
Python comments should be reserved for code within a function, or interfaces that are local to a file
Use Google style docstrings in Python for classes and functions, which can be parsed by Sphinx
Python attributes and variables can be documented inline with type and description
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 in Python to handle multiple possible variable types (duck-typing), keep the body of the try as small as possible, using the else block for logic
Files:
tests/unittest/_torch/misc/test_autotuner.pytensorrt_llm/_torch/autotuner.py
**/*.{cpp,h,cu,cuh,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
All TensorRT-LLM Open Source Software code should contain an NVIDIA copyright header that includes the year of its latest meaningful modification
Files:
tests/unittest/_torch/misc/test_autotuner.pytensorrt_llm/_torch/autotuner.py
🧠 Learnings (5)
📓 Common learnings
Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6768
File: cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h:0-0
Timestamp: 2025-08-20T06:48:45.368Z
Learning: There is a planned refactoring to move cache block bookkeeping utilities from BlockManager/WindowBlockManager into the GenerationRequest class itself to improve code organization and make responsibilities clearer.
Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6768
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:2010-2045
Timestamp: 2025-08-21T09:41:49.347Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, updateSequenceCacheBlockOffsets is specifically for updating bookkeeping when blocks are added during the context phase, not for refreshing offsets after detach operations. During detach operations, GenerationRequest::removeFrontBlock handles the necessary cache block bookkeeping internally.
📚 Learning: 2025-09-09T09:40:45.658Z
Learnt from: fredricz-20070104
Repo: NVIDIA/TensorRT-LLM PR: 7645
File: tests/integration/test_lists/qa/llm_function_core.txt:648-648
Timestamp: 2025-09-09T09:40:45.658Z
Learning: In TensorRT-LLM test lists, it's common and intentional for the same test to appear in multiple test list files when they serve different purposes (e.g., llm_function_core.txt for comprehensive core functionality testing and llm_function_core_sanity.txt for quick sanity checks). This duplication allows tests to be run in different testing contexts.
Applied to files:
tests/integration/test_lists/test-db/l0_dgx_b200.yml
📚 Learning: 2025-08-26T09:49:04.956Z
Learnt from: pengbowang-nv
Repo: NVIDIA/TensorRT-LLM PR: 7192
File: tests/integration/test_lists/test-db/l0_dgx_b200.yml:56-72
Timestamp: 2025-08-26T09:49:04.956Z
Learning: In TensorRT-LLM test configuration files, the test scheduling system handles wildcard matching with special rules that prevent duplicate test execution even when the same tests appear in multiple yaml files with overlapping GPU wildcards (e.g., "*b200*" and "*gb200*").
Applied to files:
tests/integration/test_lists/test-db/l0_dgx_b200.yml
📚 Learning: 2025-07-28T17:06:08.621Z
Learnt from: moraxu
Repo: NVIDIA/TensorRT-LLM PR: 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.
Applied to files:
tests/integration/test_lists/test-db/l0_dgx_b200.yml
📚 Learning: 2025-09-02T13:42:44.885Z
Learnt from: pcastonguay
Repo: NVIDIA/TensorRT-LLM PR: 7455
File: tensorrt_llm/_torch/pyexecutor/py_executor.py:1852-1860
Timestamp: 2025-09-02T13:42:44.885Z
Learning: In MPI communication within TensorRT-LLM pipeline parallelism, different communication types (tokens, logits, termination sync) must use disjoint tag namespaces to avoid message routing collisions when using the same source/destination patterns.
Applied to files:
tests/unittest/_torch/misc/test_autotuner.py
🧬 Code graph analysis (2)
tests/unittest/_torch/misc/test_autotuner.py (3)
tensorrt_llm/_torch/distributed/communicator.py (2)
MPIDist(340-428)TorchDist(448-782)tensorrt_llm/_utils.py (1)
mpi_disabled(533-535)tensorrt_llm/_torch/autotuner.py (6)
autotune(256-294)AutoTuner(598-1574)get(640-643)load_cache(478-509)choose_one(731-886)clear(379-380)
tensorrt_llm/_torch/autotuner.py (1)
tensorrt_llm/logger.py (2)
info(138-139)error(126-127)
⏰ 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 (9)
tests/integration/test_lists/test-db/l0_dgx_b200.yml (1)
18-18: LGTM!The new test entry for
test_autotuner_distributed_strategyis correctly placed in the MPI pre-merge test group with 4 GPU requirement, which aligns with the test's@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True)decorator and GPU count skip condition.tests/unittest/_torch/misc/test_autotuner.py (5)
20-23: LGTM!The new imports for
MPIDist,TorchDist, andmpi_disabledare correctly added to support the runtime distributed backend selection in the distributed worker function.
328-342: LGTM!The cache path handling is correctly updated to use the new per-rank cache API with
cache_pathparameter inautotune()and explicitrank=0inload_cache().
432-443: LGTM!Consistent cache path handling update matching the new per-rank cache API.
651-693: LGTM!The distributed worker function correctly implements:
- Runtime backend selection using
mpi_disabled()to choose betweenTorchDistandMPIDist- Cache path broadcast from rank 0 to ensure all ranks write to the same file
- Proper verification of single-file cache creation and rank-aware cache reload
One minor note: The
temp_diris only held by rank 0, so ensure the test completes before cleanup. This should be fine since MPI barriers are implicit in the broadcast/gather operations.
732-732: LGTM!The test function rename to
test_autotuner_distributed_strategyis descriptive and consistent with the test entry inl0_dgx_b200.yml.tensorrt_llm/_torch/autotuner.py (3)
271-294: LGTM!The
autotune()context manager correctly integrates the per-rank cache loading and saving with the new API signatures.
511-524: LGTM!The metadata serialization/deserialization methods are correctly implemented. Consider adding a version compatibility check in
_deserialize_metadatato warn users when loading a cache from a different library version.
525-557: LGTM!The
_serialize_cache_datamethod correctly converts cache entries to a JSON-serializable flat dictionary format, maintaining backward compatibility with the tactic serialization pattern.
| import contextlib | ||
| import copy | ||
| import enum | ||
| import fcntl |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Search for fcntl imports and usage in the codebase
echo "=== Searching for fcntl imports ==="
rg "import fcntl|from fcntl" --type=py -n
echo -e "\n=== Checking for Windows-related documentation or platform checks ==="
rg -i "windows|win32|platform" --type=py -l | head -20
echo -e "\n=== Checking README or setup files for platform requirements ==="
fd "README|setup|pyproject" -e md -e py -e txt | head -10Repository: NVIDIA/TensorRT-LLM
Length of output: 1566
🏁 Script executed:
#!/bin/bash
# Check how fcntl is actually being used in autotuner.py
echo "=== Checking fcntl usage in autotuner.py ==="
rg "fcntl\." tensorrt_llm/_torch/autotuner.py -n
# Check if there are any existing platform-specific imports or guards in the file
echo -e "\n=== Checking for platform-specific patterns in autotuner.py ==="
rg "sys\.platform|os\.name|platform\." tensorrt_llm/_torch/autotuner.py -nRepository: NVIDIA/TensorRT-LLM
Length of output: 272
🏁 Script executed:
#!/bin/bash
# Check if there's any Windows exclusion or platform documentation in the project
echo "=== Checking project configuration for platform requirements ==="
cat pyproject.toml 2>/dev/null | grep -A 10 "platform" || echo "No platform config found in pyproject.toml"
echo -e "\n=== Checking for filelock or alternative locking mechanisms ==="
rg "filelock|FileLock|threading\.Lock|multiprocessing" --type=py -l | head -10
echo -e "\n=== Check if this is a torch-specific module (which may have platform constraints) ==="
ls -la tensorrt_llm/_torch/ 2>/dev/null | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 1915
🏁 Script executed:
#!/bin/bash
# Check README for platform/OS requirements
echo "=== Checking README.md for platform documentation ==="
head -100 README.md
echo -e "\n=== Check _torch module documentation ==="
head -50 tensorrt_llm/_torch/__init__.py
echo -e "\n=== Check if filelock is already a dependency ==="
rg "filelock|FileLock" pyproject.toml setup.py setup.cfg requirements.txt 2>/dev/null | head -10Repository: NVIDIA/TensorRT-LLM
Length of output: 7965
🏁 Script executed:
#!/bin/bash
# Check the test file that explicitly marks fcntl as Linux-only
echo "=== Context around fcntl import in test file ==="
sed -n '75,85p' tests/unittest/_torch/auto_deploy/unit/multigpu/test_ad_allreduce_strategies.py
# Check if there are other locking mechanisms used in autotuner.py
echo -e "\n=== Checking broader context around fcntl usage in autotuner.py ==="
sed -n '450,465p' tensorrt_llm/_torch/autotuner.py
sed -n '490,505p' tensorrt_llm/_torch/autotuner.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1925
Add platform-specific handling for fcntl module to prevent ImportError on Windows.
The fcntl module is Unix-only and will cause an import failure on non-Unix systems. The codebase already demonstrates the correct pattern for handling this in tests/unittest/_torch/auto_deploy/unit/multigpu/test_ad_allreduce_strategies.py with a try-except block. Apply the same approach:
try:
import fcntl
except ImportError:
fcntl = NoneThen wrap the fcntl.flock() calls at lines 458 and 498 with checks for fcntl is not None to gracefully handle platforms where the module is unavailable.
🤖 Prompt for AI Agents
In tensorrt_llm/_torch/autotuner.py at top around line 5 and for usages at lines
458 and 498, the direct import of fcntl is Unix-only and will raise on Windows;
change the import to a safe try/except that assigns fcntl = None on ImportError,
and then guard the two fcntl.flock() calls by checking if fcntl is not None
before calling flock so code runs gracefully on platforms without fcntl.
|
PR_Github #30152 [ run ] triggered by Bot. Commit: |
|
PR_Github #30152 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #30213 [ run ] triggered by Bot. Commit: |
2 participants
Summary by CodeRabbit
New Features
API Changes
save_cache()andload_cache()methods to require a rank parameter for proper cache isolation.Tests
✏️ Tip: You can customize this high-level summary in your review settings.
Description
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
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.
Details
run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]Launch build/test pipelines. All previously running jobs will be killed.
--reuse-test (optional)pipeline-id(OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.--disable-reuse-test(OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.--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-PyTorch-1, xxx"(OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-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.--test-backend "pytorch, cpp"(OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline 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 in addition to running 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-TensorRT-Post-Merge-1, xxx"(OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".--detailed-log(OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.--debug(OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in thestage-listparameter to access the appropriate container environment. Note: Does NOT update GitHub check status.For guidance on mapping tests to stage names, see
docs/source/reference/ci-overview.mdand the
scripts/test_to_stage_mapping.pyhelper.kill
killKill all running builds associated with pull request.
skip
skip --comment COMMENTSkip 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-pipelineReuse 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.