[None][fix] Prevent memory leaks with max_iteration_result_size #9257
Conversation
Signed-off-by: zhanghaotong <[email protected]>
📝 WalkthroughWalkthroughA new Changes
Sequence Diagram(s)sequenceDiagram
participant User as User Config
participant LlmArgs as BaseLlmArgs
participant GenExec as GenerationExecutor
participant PostConfig as PostprocWorkerConfig
participant IterResult as IterationResult
participant Queue as AsyncQueue/Queue
User->>LlmArgs: max_iteration_result_size: 16384
LlmArgs->>GenExec: __init__(max_iteration_result_size)
GenExec->>PostConfig: PostprocWorkerConfig(max_iteration_result_size)
GenExec->>GenExec: super().__init__(max_iteration_result_size)
GenExec->>IterResult: IterationResult(maxsize)
IterResult->>Queue: AsyncQueue(maxsize) or Queue(maxsize)
Queue->>Queue: collections.deque with bounded capacity
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Areas requiring extra attention:
Pre-merge checks and finishing touches❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Tip 📝 Customizable high-level summaries are now available in beta!You can now customize how CodeRabbit generates the high-level summary in your pull requests — including its content, structure, tone, and formatting.
Example instruction:
Note: This feature is currently in beta for Pro-tier users, and pricing will be announced later. 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: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/executor/executor.py (1)
243-257: Applymax_iteration_result_sizelimit to_iter_kv_events_resultfor consistency and OOM prevention.The code initializes
_iter_stats_resultwithpostproc_config.max_iteration_result_size(line 248), but_iter_kv_events_resultis created without any size limit (line 254). SinceIterationResultaccepts an optionalmaxsizeparameter and both results follow the same pattern of accumulating data when the API is not called, the KV events result should also be bounded to prevent potential memory buildup.Consider updating line 254 to:
self._iter_kv_events_result = IterationResult(self.postproc_config.max_iteration_result_size)If KV cache events require different handling than stats, document the rationale for unbounded initialization.
🧹 Nitpick comments (3)
tensorrt_llm/llmapi/llm_args.py (1)
1814-1818: Clarify and validatemax_iteration_result_sizesemanticsThe new knob is useful, but its semantics aren’t fully specified, especially given
PostprocWorkerConfig.max_iteration_result_sizedefaults to0.Consider:
- Explicitly documenting what values mean (e.g.,
> 0= bounded queue,<= 0= unbounded/disabled) to avoid confusion across backends/configs.- Adding a small
field_validator("max_iteration_result_size")to reject negatives or normalize<= 0to the “unbounded” sentinel before it reaches the queue implementation.This keeps behavior predictable and prevents invalid values from propagating into the runtime.
tensorrt_llm/executor/postproc_worker.py (1)
41-47: Exposemax_iteration_result_sizesemantics in config docAdding
max_iteration_result_sizetoPostprocWorkerConfigis consistent with the rest of the plumbing, but its behavior isn’t documented here.It would help future readers if the class/field docstring clarified:
- What
max_iteration_result_sizecontrols (bounded size of iteration/KV event result queues).- How
0is interpreted (e.g., “use executor default/unbounded”).This keeps the config self-describing and aligned with
BaseLlmArgs.max_iteration_result_size.tensorrt_llm/llmapi/llm.py (1)
1071-1075: Consider wiringmax_iteration_result_sizefor the TRT backend as wellHere
_TorchLLMcorrectly passesself.args.max_iteration_result_sizeintoPostprocWorkerConfig, so the PyTorch path respects the new bounded-queue knob.In
_TrtLLM._build_model, thePostprocWorkerConfigconstruction still only sets:PostprocWorkerConfig( num_postprocess_workers=self.args.num_postprocess_workers, postprocess_tokenizer_dir=self.args.postprocess_tokenizer_dir, )which means the new
BaseLlmArgs.max_iteration_result_sizeis effectively ignored for the TRT backend.For consistency (and to have the leak fix apply uniformly), it’s worth mirroring this addition in the TRT path, e.g.:
- self._executor = self._executor_cls.create( + self._executor = self._executor_cls.create( self._engine_dir, executor_config=self._executor_config, batched_logits_processor=self.args.batched_logits_processor, model_world_size=self.args.parallel_config.world_size, mpi_session=self.mpi_session, reuse_mpi_comm=external_mpi_comm_available( self.args.parallel_config.world_size), return_logits=return_logits, - postproc_worker_config=PostprocWorkerConfig( - num_postprocess_workers=self.args.num_postprocess_workers, - postprocess_tokenizer_dir=self.args.postprocess_tokenizer_dir, - ), + postproc_worker_config=PostprocWorkerConfig( + num_postprocess_workers=self.args.num_postprocess_workers, + postprocess_tokenizer_dir=self.args.postprocess_tokenizer_dir, + max_iteration_result_size=self.args.max_iteration_result_size, + ), is_llm_executor=True)so that
max_iteration_result_sizebehaves the same regardless of backend.
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (9)
tensorrt_llm/executor/base_worker.py(1 hunks)tensorrt_llm/executor/executor.py(2 hunks)tensorrt_llm/executor/postproc_worker.py(1 hunks)tensorrt_llm/executor/proxy.py(1 hunks)tensorrt_llm/executor/result.py(1 hunks)tensorrt_llm/executor/rpc_proxy.py(1 hunks)tensorrt_llm/llmapi/llm.py(1 hunks)tensorrt_llm/llmapi/llm_args.py(1 hunks)tensorrt_llm/llmapi/utils.py(1 hunks)
🧰 Additional context used
🧠 Learnings (5)
📚 Learning: 2025-08-26T09:37:10.463Z
Learnt from: jiaganc
Repo: NVIDIA/TensorRT-LLM PR: 7031
File: tensorrt_llm/bench/dataclasses/configuration.py:90-104
Timestamp: 2025-08-26T09:37:10.463Z
Learning: In TensorRT-LLM, the `get_pytorch_perf_config()` method returns `self.pytorch_config` which can contain default `cuda_graph_config` values, so `llm_args` may already have this config before the extra options processing.
Applied to files:
tensorrt_llm/llmapi/llm.py
📚 Learning: 2025-09-23T15:12:38.312Z
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/thop/allreduceOp.cpp:352-446
Timestamp: 2025-09-23T15:12:38.312Z
Learning: In TensorRT-LLM NCCL device implementation, NCCL version 2.28+ requirements are handled at runtime in the nccl_device/config layer rather than with compile-time guards. This allows the allreduceOp to remain version-agnostic and delegates version compatibility validation to the appropriate lower-level components that can gracefully handle unsupported configurations.
Applied to files:
tensorrt_llm/llmapi/llm.py
📚 Learning: 2025-08-21T00:16:56.457Z
Learnt from: farshadghodsian
Repo: NVIDIA/TensorRT-LLM PR: 7101
File: docs/source/blogs/tech_blog/blog9_Deploying_GPT_OSS_on_TRTLLM.md:36-36
Timestamp: 2025-08-21T00:16:56.457Z
Learning: TensorRT-LLM container release tags in documentation should only reference published NGC container images. The README badge version may be ahead of the actual published container versions.
Applied to files:
tensorrt_llm/llmapi/llm.py
📚 Learning: 2025-08-14T15:38:01.771Z
Learnt from: MatthiasKohl
Repo: NVIDIA/TensorRT-LLM PR: 6904
File: cpp/tensorrt_llm/pybind/thop/bindings.cpp:55-57
Timestamp: 2025-08-14T15:38:01.771Z
Learning: In TensorRT-LLM Python bindings, tensor parameter collections like mla_tensor_params and spec_decoding_tensor_params are kept as required parameters without defaults to maintain API consistency, even when it might affect backward compatibility.
Applied to files:
tensorrt_llm/llmapi/llm.py
📚 Learning: 2025-08-06T13:58:07.506Z
Learnt from: galagam
Repo: NVIDIA/TensorRT-LLM PR: 6487
File: tests/unittest/_torch/auto_deploy/unit/singlegpu/test_ad_trtllm_bench.py:1-12
Timestamp: 2025-08-06T13:58:07.506Z
Learning: In TensorRT-LLM, test files (files under tests/ directories) do not require NVIDIA copyright headers, unlike production source code files. Test files typically start directly with imports, docstrings, or code.
Applied to files:
tensorrt_llm/llmapi/llm.py
🧬 Code graph analysis (3)
tensorrt_llm/llmapi/llm_args.py (1)
tensorrt_llm/builder.py (1)
default(45-50)
tensorrt_llm/executor/result.py (2)
tensorrt_llm/executor/utils.py (1)
has_event_loop(68-73)tensorrt_llm/llmapi/utils.py (2)
AsyncQueue(361-445)sync_q(380-381)
tensorrt_llm/executor/executor.py (2)
tensorrt_llm/executor/postproc_worker.py (1)
PostprocWorkerConfig(42-50)tensorrt_llm/executor/result.py (1)
IterationResult(951-1001)
🪛 Ruff (0.14.5)
tensorrt_llm/llmapi/utils.py
373-373: PEP 484 prohibits implicit Optional
Convert to T | None
(RUF013)
⏰ 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 (4)
tensorrt_llm/executor/base_worker.py (1)
85-91: Correctly forwardingmax_iteration_result_sizeintoGenerationExecutorWiring
postproc_config.max_iteration_result_sizeinto the baseGenerationExecutorensures the in-process worker path respects the bounded-queue setting; the defaultPostprocWorkerConfig()fallback keeps behavior well-defined when no config is provided.tensorrt_llm/executor/proxy.py (1)
48-58: Proxy now respectsmax_iteration_result_sizein iteration/KV queuesPropagating
postproc_worker_config.max_iteration_result_sizeinto the baseGenerationExecutorhere ensures the proxy path uses the same bounded-queue semantics as the in-process worker. Combined with thequeue.full()/queue.get()logic in_iteration_result_task, this will naturally drop the oldest stats/events once the configured capacity is reached.tensorrt_llm/executor/rpc_proxy.py (1)
304-312: LGTM! Clean parameter forwarding.The
max_iteration_result_sizeparameter is correctly forwarded from thepostproc_worker_configto the base class constructor, following the same pattern as the other configuration parameters.tensorrt_llm/executor/executor.py (1)
80-89: LGTM! Parameter properly integrated.The
max_iteration_result_sizeparameter is correctly added to theGenerationExecutor.__init__method and properly stored in thePostprocWorkerConfig. The default value of 0 (unbounded) maintains backward compatibility.
svc-trtllm-gh-bot
added
the
Community want to contribute
Signed-off-by: zhanghaotong <[email protected]>
|
All bot issues have been fixed. |
|
@hchings could you please review? Thanks. |
| description="The path to the tokenizer directory for postprocessing.", | ||
| status="prototype") | ||
|
|
||
| max_iteration_result_size: int = Field( |
There was a problem hiding this comment.
We will need at least a unit test to verify that this is working as expected. See tests/unittest/llmapi/test_llm_pytorch.py for examples.
| """ | ||
|
|
||
| def __init__(self): | ||
| def __init__(self, maxsize: int = 0): |
There was a problem hiding this comment.
Use max_size here and in other places.
3 participants
Summary by CodeRabbit
Release Notes
Description
We discovered that when TRT runs without ever calling the /metrics API (which we do not invoke in production), the
_iter_stats_resultqueue in the executor keeps growing indefinitely until it causes an OOM error. To address this, we introduced amax_iteration_result_sizeparameter to cap the queue size. When the queue reaches its maximum capacity, the oldest results are automatically discarded.Before the fix, memory utilization kept increasing.
After the fix, the results are still under stress testing, coming soon~
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.
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.