[TRTLLM-10827][feat] Add KV Cache metrics to MetricsCollector for more Prometheus metrics#11243
Conversation
📝 WalkthroughWalkthroughThis PR adds Prometheus Gauge metrics for KV cache statistics and refactors Python metrics logging with new public methods. It also renames a C++ member variable for naming consistency and updates related integration points. Changes
Sequence DiagramsequenceDiagram
actor Client
participant OpenAIServer
participant MetricsCollector
participant Prometheus
Client->>OpenAIServer: request completion
OpenAIServer->>OpenAIServer: process request
OpenAIServer->>MetricsCollector: log_request_metrics_dict(metrics_dict)
MetricsCollector->>Prometheus: emit histograms & counter
alt enable_iter_perf_stats enabled
OpenAIServer->>OpenAIServer: get_stats_async(timeout=2.0)
loop for each iteration stat
OpenAIServer->>MetricsCollector: log_iteration_stats(llm_stat)
MetricsCollector->>Prometheus: update kv_cache_hit_rate gauge
MetricsCollector->>Prometheus: update kv_cache_utilization gauge
end
end
Client->>Client: receive response with metrics
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Important Action Needed: IP Allowlist UpdateIf your organization protects your Git platform with IP whitelisting, please add the new CodeRabbit IP address to your allowlist:
Reviews will stop working after February 8, 2026 if the new IP is not added to your allowlist. 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: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (5)
cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h (1)
1-2:⚠️ Potential issue | 🟡 MinorUpdate the copyright year to reflect the 2026 modification.
🔧 Suggested header update
- * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved.As per coding guidelines: All TensorRT-LLM source files (.cpp, .h, .cu, .py, and other source files) should contain an NVIDIA copyright header with the year of latest meaningful modification.
cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp (1)
1-3:⚠️ Potential issue | 🟡 MinorUpdate the copyright year to include 2026 changes.
🔧 Suggested header update
- * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.As per coding guidelines: All TensorRT-LLM source files (.cpp, .h, .cu, .py, and other source files) should contain an NVIDIA copyright header with the year of latest meaningful modification.
tensorrt_llm/serve/openai_server.py (1)
1-2:⚠️ Potential issue | 🟡 MinorAdd NVIDIA copyright header for 2026 modifications.
🔧 Suggested header addition
#!/usr/bin/env python +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# SPDX-License-Identifier: Apache-2.0As per coding guidelines: All TensorRT-LLM source files (.cpp, .h, .cu, .py, and other source files) should contain an NVIDIA copyright header with the year of latest meaningful modification.
tests/unittest/llmapi/apps/_test_openai_prometheus.py (2)
1-1:⚠️ Potential issue | 🟠 MajorAdd the NVIDIA copyright header.
This file is missing the required NVIDIA copyright header with the year of latest meaningful modification (2026). Please add the standard header used in this repo at the top of the file.
As per coding guidelines, "All TensorRT-LLM source files (.cpp, .h, .cu, .py, and other source files) should contain an NVIDIA copyright header with the year of latest meaningful modification."
53-68:⚠️ Potential issue | 🟡 MinorUse equality check for HTTP status.
Line 68 uses
isfor integer comparison, which is identity-based. Use==for value comparison instead.Proposed fix
- assert response.status is 200 + assert response.status == 200
🤖 Fix all issues with AI agents
In `@tensorrt_llm/metrics/collector.py`:
- Line 1: Add the required NVIDIA copyright header (with the latest modification
year) at the very top of this file, above the existing module docstring (the
string """Utilities for Prometheus Metrics Collection."""), so the header
precedes the docstring and contains the standard NVIDIA copyright notice and
year of last meaningful change.
- Around line 116-139: Update the docstring of log_request_metrics_dict to
Google-style: add an Args section describing metrics_dict and list expected keys
(MetricsCollector.labelname_finish_reason as finish reason string,
MetricNames.E2E, MetricNames.TTFT, MetricNames.TPOT,
MetricNames.REQUEST_QUEUE_TIME as floats in seconds), add a Returns section
stating None (metrics are logged to Prometheus), and include a short Note that
metrics are only recorded when MetricsCollector.labelname_finish_reason is
present; keep the existing high-level description and mention that
histograms/counters (histogram_e2e_time_request, histogram_time_to_first_token,
histogram_time_per_output_token, histogram_queue_time_request,
counter_request_success) are updated.
- Around line 89-96: Rename the misspelled metric attribute kv_cahce_hit_rate to
kv_cache_hit_rate wherever defined/used (e.g., in the Gauge creation and any
references), and change truthy checks on cacheHitRate to explicit existence
checks (e.g., cacheHitRate is not None) so zero values are allowed;
additionally, before computing utilization using maxNumBlocks, guard the
division by ensuring maxNumBlocks is truthy/non-zero (or skip/update metric when
maxNumBlocks == 0) when setting kv_cache_utilization. Ensure you update all
references to these symbols (kv_cahce_hit_rate, kv_cache_hit_rate, cacheHitRate,
maxNumBlocks, kv_cache_utilization) consistently.
🧹 Nitpick comments (2)
tensorrt_llm/metrics/collector.py (1)
4-4: Keep typing imports namespaced.♻️ Suggested refactor
-import time -from typing import Dict, Union +import time +import typing @@ - def __init__(self, labels: Dict[str, str]) -> None: + def __init__(self, labels: typing.Dict[str, str]) -> None: @@ - def _label_merge(self, labels: Dict[str, str]) -> Dict[str, str]: + def _label_merge(self, labels: typing.Dict[str, str]) -> typing.Dict[str, str]: @@ - def _log_counter(self, counter, labels: Dict[str, str], - data: Union[int, float]) -> None: + def _log_counter(self, counter, labels: typing.Dict[str, str], + data: typing.Union[int, float]) -> None: @@ - def _log_histogram(self, histogram, data: Union[int, float]) -> None: + def _log_histogram(self, histogram, data: typing.Union[int, float]) -> None: @@ - def _log_gauge(self, gauge, data: Union[int, float]) -> None: + def _log_gauge(self, gauge, data: typing.Union[int, float]) -> None:As per coding guidelines: Always maintain the namespace when importing Python modules, even if only one class or function from a module is used.
tensorrt_llm/serve/openai_server.py (1)
468-475: Avoid blocking request finalization while draining iteration stats.
get_stats_async(timeout=2.0)can add up to the timeout in latency per request. Consider draining in a background task so response completion isn’t delayed.♻️ Suggested refactor
if self.metrics_collector: self.metrics_collector.log_request_metrics_dict(res.metrics_dict) if self.llm.args.enable_iter_perf_stats: - async for llm_stat in self.llm.get_stats_async(timeout=2.0): - self.metrics_collector.log_iteration_stats(llm_stat) + async def _drain_iter_stats(): + async for llm_stat in self.llm.get_stats_async(timeout=2.0): + self.metrics_collector.log_iteration_stats(llm_stat) + asyncio.create_task(_drain_iter_stats())
yijingl-nvidia
force-pushed
the
more_metrics
branch
from
9850e86 to
29f07d5
Compare
|
Addressing comments from coderabbit |
yijingl-nvidia
force-pushed
the
more_metrics
branch
from
3f541e0 to
337d84f
Compare
|
/bot run |
|
PR_Github #35717 [ run ] triggered by Bot. Commit: |
|
PR_Github #35717 [ run ] completed with state
|
|
/bot run |
|
PR_Github #35832 [ run ] triggered by Bot. Commit: |
|
PR_Github #35832 [ run ] completed with state
|
yijingl-nvidia
force-pushed
the
more_metrics
branch
from
337d84f to
0410f59
Compare
|
/bot run |
|
PR_Github #35846 [ run ] triggered by Bot. Commit: |
|
PR_Github #35846 [ run ] completed with state
|
|
/bot run |
|
PR_Github #35923 [ run ] triggered by Bot. Commit: |
|
PR_Github #35923 [ run ] completed with state
|
yijingl-nvidia
force-pushed
the
more_metrics
branch
from
0410f59 to
5f3b274
Compare
|
/bot run |
|
PR_Github #35952 [ run ] triggered by Bot. Commit: |
|
PR_Github #35952 [ run ] completed with state
|
yijingl-nvidia
force-pushed
the
more_metrics
branch
from
5f3b274 to
d369966
Compare
|
/bot run |
|
PR_Github #36056 [ run ] triggered by Bot. Commit: |
|
PR_Github #36056 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #36095 [ run ] triggered by Bot. Commit: |
|
PR_Github #36095 [ run ] completed with state
|
Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Add kv cache stats in MetricsCollector. Add feature in OpenAIServer to collect iteration stats and upload them to Prometheus using MetricsCollector. Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
yijingl-nvidia
force-pushed
the
more_metrics
branch
from
d369966 to
c08a603
Compare
|
/bot run |
|
PR_Github #36122 [ run ] triggered by Bot. Commit: |
|
PR_Github #36122 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #36154 [ run ] triggered by Bot. Commit: |
|
PR_Github #36154 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #36175 [ run ] triggered by Bot. Commit: |
|
PR_Github #36175 [ run ] completed with state |
…e Prometheus metrics (NVIDIA#11243) Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com> Signed-off-by: Yuanjing Xue <197832395+yuanjingx87@users.noreply.github.com>
…e Prometheus metrics (NVIDIA#11243) Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com> Signed-off-by: Yuanjing Xue <197832395+yuanjingx87@users.noreply.github.com>
| request_queue_time) | ||
| self.last_log_time = time.time() | ||
|
|
||
| def log_iteration_stats(self, iteration_stats: dict) -> None: |
11 participants
Description
Current TRTLLM exposes too few Prometheus metrics. This PR adds KV cache utilization and hit rate metrics to Prometheus metrics. It will benefit Dynamo in monitoring TRTLLM.
Minor clean-up:
Test Coverage
Modified tests/unittest/llmapi/apps/_test_openai_prometheus.py to cover the change
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.
Summary by CodeRabbit
Release Notes
New Features
Improvements