[PD] Log KVTransfer metrics #20443
[PD] Log KVTransfer metrics #20443
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @NickLucche, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request introduces initial observability for KV cache transfers by implementing a mechanism to collect, aggregate, and log transfer-related metrics. The primary goal is to provide insights into the performance of KV cache transfers, addressing the challenge of gathering worker-level metrics and propagating them to the scheduler for centralized logging.
Highlights
- KV Transfer Metrics Definition: Introduced new dataclasses
KVTransferAggregatedStatsandKVTransferStatsto define and track key performance metrics for KV cache transfers, along with aKVTransferLoggingutility for periodic output. - Metric Propagation: Extended core data flow objects such as
EngineCoreOutput,ModelRunnerOutput, andIterationStatsto include and propagateKVTransferAggregatedStatsthroughout the system. - Connector Interface Updates: Modified the
KVConnectorBase_V1interface and its implementations (e.g.,lmcache_connector,multi_connector,nixl_connector,p2p_nccl_connector,p2p_nccl_engine) to return and aggregate KV transfer metrics from worker processes to the scheduler. - Nixl Connector Implementation: Implemented the actual collection and aggregation of KV transfer metrics within the
nixl_connector, including tracking transfer durations, bytes, and block counts, and passing these up the communication chain. - Logging Integration: Integrated the new KV transfer metrics into the existing
loggers.pysystem, enabling periodic output of aggregated transfer statistics tostd.out.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
There was a problem hiding this comment.
Code Review
This pull request is a great first step towards adding observability for KVCache transfers. The changes to plumb the new metrics through the system are well-structured. My review focuses on the new metrics collection logic, where I've found a few critical issues that could lead to runtime errors, as well as some incomplete implementations that would report incorrect metrics. I've provided detailed feedback and suggestions to address these points. Once these are addressed, this will be a solid and valuable contribution.
| call to this method (this call or a prior one). | ||
| """ | ||
| return self._lmcache_engine.get_finished(finished_req_ids) | ||
| return self._lmcache_engine.get_finished(finished_req_ids), None |
There was a problem hiding this comment.
The return value of this function does not match the updated type hint. self._lmcache_engine.get_finished(finished_req_ids) returns a tuple of two elements, so this function is currently returning ((elem1, elem2), None), which will cause a TypeError at runtime. The new signature is tuple[Optional[set[str]], Optional[set[str]], Optional[KVTransferAggregatedStats]].
You should unpack the tuple returned by _lmcache_engine.get_finished and add None as the third element for the stats.
| return self._lmcache_engine.get_finished(finished_req_ids), None | |
| return *self._lmcache_engine.get_finished(finished_req_ids), None |
| class KVTransferStats: | ||
| """Container for transfer performance metrics""" | ||
| transfer_durations: list[float] # Transfer durations in seconds | ||
| bytes_transferred: list[int] # Bytes transferred per transfer | ||
| num_blocks_transferred: list[int] # Number of blocks per transfer | ||
| last_log_time: float # Last time we logged metrics | ||
| num_transfers: int = 0 |
There was a problem hiding this comment.
The KVTransferStats dataclass has fields (transfer_durations, bytes_transferred, num_blocks_transferred, last_log_time) that are not given default values. Instantiating this dataclass with KVTransferStats() will raise a TypeError at runtime.
You should provide default values for these fields. For the lists, field(default_factory=list) is appropriate. For last_log_time, you could use field(default_factory=time.time).
Note that you'll also need to import field from dataclasses at the top of the file: from dataclasses import dataclass, field.
| class KVTransferStats: | |
| """Container for transfer performance metrics""" | |
| transfer_durations: list[float] # Transfer durations in seconds | |
| bytes_transferred: list[int] # Bytes transferred per transfer | |
| num_blocks_transferred: list[int] # Number of blocks per transfer | |
| last_log_time: float # Last time we logged metrics | |
| num_transfers: int = 0 | |
| @dataclass | |
| class KVTransferStats: | |
| """Container for transfer performance metrics""" | |
| transfer_durations: list[float] = field(default_factory=list) # Transfer durations in seconds | |
| bytes_transferred: list[int] = field(default_factory=list) # Bytes transferred per transfer | |
| num_blocks_transferred: list[int] = field(default_factory=list) # Number of blocks per transfer | |
| last_log_time: float = field(default_factory=time.time) # Last time we logged metrics | |
| num_transfers: int = 0 |
| def reduce_and_reset(self) -> KVTransferAggregatedStats: | ||
| # NOTE (NickLucche): to have statistical significance, we assume the | ||
| # size of the measurements groups to be the same. This allows to bound | ||
| # the size of the messages. | ||
| # TODO finish this | ||
| stats = KVTransferAggregatedStats( | ||
| avg_transfer_durations=11.0, | ||
| avg_bytes_transferred=0.0, | ||
| num_blocks_transferred=0, | ||
| num_successful_transfers=self.num_transfers) | ||
| self.num_transfers = 0 | ||
| return stats |
There was a problem hiding this comment.
The reduce_and_reset method is incomplete and returns hardcoded values (e.g., avg_transfer_durations=11.0). This will lead to incorrect metrics being reported. The method should compute the aggregated stats from the collected lists (transfer_durations, bytes_transferred, etc.) before resetting them.
The TODO comment indicates this is a work in progress, but as it stands, it's a correctness bug.
| def get_latency_stats(self) -> tuple[float, float, float]: | ||
| """Get transfer latency statistics""" | ||
| # TODO possible use | ||
| import numpy as np |
There was a problem hiding this comment.
The numpy import is inside the get_latency_stats method. It's generally better practice to have all imports at the top of the file for clarity and to avoid repeated import overhead if the method is called multiple times. Please move import numpy as np to the top of the file.
|
This pull request has merge conflicts that must be resolved before it can be |
|
cc @njhill |
|
closing in favor of #22188 |
1 participant
First attempt at addressing observability for KVCache transfers, starting from outputting to std.out.
The main concern here is the fact that these metrics have to be gathered at the worker level -and in the worker process-while having to go all the way to the scheduler process.
This results in a couple perhaps-not-so-pleasant interface changes that I'd like to discuss in this thread.
How does it look: