Feature/offloading manager stats

[Srinivasoo7]

Purpose

This PR adds support for telemetry emissions from the OffloadingManager interface inside the KV Connector.

As part of isolating the changes from the block-reuse frequency tracking PR, the OffloadingManager interface requires a standardized method to emit statistics up to the scheduler's KVConnector interfaces. This change:

1. Adds get_stats() -> dict[str, Any] to the base OffloadingManager abstract class.
2. Implements stats aggregation inside OffloadingConnectorStats.reduce to gracefully support flat scalar statistics (ints, floats) alongside its standard list metrics.
3. Aggregates metrics polled from the currently active self.connector_scheduler.manager within OffloadingConnector.get_kv_connector_stats().
4. Exposes stores_skipped directly from the FilteredOffloadingManager (formerly StoreReusedOffloadingManager), tracking exactly how many block hashes failed the LRU reuse threshold.

Test Plan

• Ran pytest tests/v1/kv_offload to ensure manager abstraction and logic remain structurally sound.
• Ran ruff format and ruff check on the modified files to ensure compliance with the repository layout.

Test Result

• Pytest and Ruff linters pass cleanly across the modified abstract and connector files.
• Emitted OffloadingConnectorStats successfully aggregate scalar dict values without validation errors.

[gemini-code-assist]

Code Review

This pull request effectively adds telemetry support for the OffloadingManager and introduces a FilteredOffloadingManager to gate offloading based on block reuse frequency. The implementation is well-structured, particularly the use of a decorator pattern for the filtering logic and the inclusion of comprehensive unit tests. I have identified a couple of areas for improvement to enhance configurability and robustness.

[gemini-code-assist]

The FilteredOffloadingManager is initialized without a configurable max_tracker_size. This means the LRU tracker for block reuse frequency will always use the default size. This parameter can have a significant impact on memory usage and filtering effectiveness, and should be configurable for performance tuning. I suggest reading max_tracker_size from self.extra_config, similar to how store_threshold is handled.

            store_threshold = int(self.extra_config.get("store_threshold", 0))
            if store_threshold > 1:
                max_tracker_size = int(
                    self.extra_config.get("max_tracker_size", 64_000)
                )
                self._manager = FilteredOffloadingManager(
                    backing=self._manager,
                    store_threshold=store_threshold,
                    max_tracker_size=max_tracker_size,
                )

[gemini-code-assist]

The get_stats method modifies the dictionary returned by self._backing.get_stats() in-place. While this is currently safe because the existing backing managers return a new empty dictionary, this pattern is fragile. If a backing manager's implementation changes in the future to return a shared or cached dictionary, this could lead to unintended side effects. It's safer to create a new dictionary to avoid mutating the object returned by the backing manager.

Suggested change

	stats = self._backing.get_stats()
	stats["stores_skipped"] = self.stores_skipped
	return stats
	return {**self._backing.get_stats(), "stores_skipped": self.stores_skipped}

[mergify]

Hi @Srinivasoo7, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[Srinivasoo7]

Hi @orozery
Since we merged PR #35342, can you look into the supporting stats for this functionality?

Thanks

[orozery]

@Srinivasoo7 can you please rebase?

[orozery]

The stats dictionary returned by the connector is later used to feed OffloadingConnectorStats.
OffloadingConnectorStats currently assumes the a hard-code structure of transfer_type -> ops_list.
The changed proposed here will break this assumption.
Specifically, I expect the OffloadingConnectorStats.aggregate and OffloadingConnectorStats.reduce to fail.

[Srinivasoo7]

Hi @orozery
You are right, in the original OffloadingConnectorStats, I assumed a hard-coded transfer_type -> ops_list structure, but the branch already updates both aggregate() and reduce() to handle scalar values alongside list values.
Hence, the stores_skipped flows through both methods correctly without breaking the existing transfer-type logic.

[mergify]

Hi @Srinivasoo7, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Srinivasoo7.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[orozery]

Please rebase :)

[Srinivasoo7]

Done @orozery

[Srinivasoo7]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces telemetry for the offloading manager, which is a valuable addition. The changes to support a new statistics structure with both transfers and gauges are well-implemented in metrics.py and offloading_connector.py. However, I've identified a critical bug in reuse_manager.py that could lead to a runtime error, and another high-risk issue related to in-place data mutation. I have provided suggestions to address both of these concerns.

[gemini-code-assist]

The block_hashes parameter is an Iterable. Calling len(block_hashes) after it has been potentially consumed by the list comprehension on the preceding lines can lead to incorrect behavior or a TypeError if block_hashes is a generator. To ensure correctness, you should first convert the iterable to a list and use that list for both operations.

        block_hashes_list = list(block_hashes)
        eligible = [
            bh for bh in block_hashes_list
            if self.counts.get(bh, 0) >= self.store_threshold
        ]

        self.stores_skipped += len(block_hashes_list) - len(eligible)

[gemini-code-assist]

The current implementation modifies the dictionary returned by self._backing.get_stats() in-place. This can cause unexpected side effects if the backing manager's returned dictionary is not meant to be mutated. To prevent this, you should operate on a copy. The safest approach is to use copy.deepcopy.

You will need to add import copy to the file.

    def get_stats(self) -> dict[str, Any]:
        stats = copy.deepcopy(self._backing.get_stats())
        stats.setdefault("gauges", {})["stores_skipped"] = self.stores_skipped
        return stats

[Srinivasoo7]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces telemetry for the OffloadingManager, allowing statistics to be collected and emitted. The changes include adding a get_stats method to the OffloadingManager interface, implementing it in FilterReusedOffloadingManager to report stores_skipped, and updating the metrics aggregation logic to handle these new scalar statistics (gauges) alongside existing transfer metrics. The OffloadingConnector is also modified to collect and aggregate these stats from the scheduler's manager.

My main concern is with the new test file tests/v1/kv_offload/test_reuse_manager.py. It appears to be written for a different version of the code, referencing classes (FilteredOffloadingManager, BlockReuseTracker) that do not exist in vllm/v1/kv_offload/reuse_manager.py. This is a critical issue as it means the new functionality is not being tested correctly. Please see my detailed comment on this file.

[gemini-code-assist]

This test file attempts to import BlockReuseTracker and FilteredOffloadingManager from the reuse_manager module. However, the implementation in vllm/v1/kv_offload/reuse_manager.py defines a class named FilterReusedOffloadingManager and does not contain a BlockReuseTracker class. This will cause an AttributeError at runtime, and these tests will fail to run.

Please ensure the test file is updated to use the correct class names from the module under test. FilteredOffloadingManager should likely be FilterReusedOffloadingManager. The logic for BlockReuseTracker seems to be part of FilterReusedOffloadingManager now, so the tests might need significant refactoring to target the public API of FilterReusedOffloadingManager.

[orozery]

I don't think we need aggregation, as it is not possible for both self.connector_scheduler and self.connector_worker to be not None.

[orozery]

Why do we need this change?

[orozery]

Let's try and make it clearer by:

1. Defining "transfers" and "gauges" as string constants, e..g TRANSFERS_KEY = "transfers", GAUGUES_KEY = "gauges".
2. Add a docstring to OffloadingConnectorStats describing its expected structure.

[orozery]

Can we create and use a KVConnectorStats.set_gauge(gauge_name, gauge_value) method?

[orozery]

Let's define it as:

Suggested change

	def get_stats(self) -> dict[str, Any]:
	def get_stats(self) -> dict[str, Any] | None:

to save the dictionary allocation if stats are unused.

[orozery]

Looks like this is an old file?
We already have a unit test in test_cpu_manager.py
You should instead add a test that verifies the stores_skipped gauge.

[Srinivasoo7]

Hi @orozery
Addressed the changes in the latest commit.

Thanks

[mergify]

Hi @Srinivasoo7, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[orozery]

@Srinivasoo7 Also need to fix 2 more things:

1. pre-commit
2. assert self._connector_stats is None which fails some of the existing tests (see CI logs).

[Srinivasoo7]

Yes boss @orozery, we'll fix it asap!

[orozery]

Thanks @Srinivasoo7 for the hard work! (and sorry for all of the nit picking :) )
I added an e2e test and noticed an issue that metric metadata is not serialized.
I changed the stats structs to next the metadata under the self.data dictionary, which is serialized.
Also another issue was we were missing empty worker side stats (until #43877 lands).
I pushed the fixes to your branch.

[Srinivasoo7]

Gotcha @orozery.
This was a good stuff we worked on, though it took significant time, it was a great learning experience all together.
Yup I was trying to land both offloading metrics and worker side stats at once.

Also, from this PR we branched out to metrics redesign (#44008) request to help with your views there to start the redesign PR against the RFC.

Looking forward to more such contributions!