[KVConnector][Metrics] Aggregate scheduler-side KVConnectorStats

[QierLi]

Purpose

#22188 introduced generic metrics reporting for KVConnector. Building on that, this PR extends metrics to include stats from the Scheduler, in addition to those from Workers.

The Scheduler-side KVConnector should export several important metrics, such as prefix matching for CPU KV (get_num_new_matched_tokens - duration, token counts), KV Cache event's (take_events) etc, while Workers' export KV transfers'.

Test Plan

Existed tests.
The changes are a no-op if no KVConnectorStats are reported from the Scheduler.

Test Result

No new breakages.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[QierLi]

@NickLucche a small extension on top of your #24786 :)

[gemini-code-assist]

Code Review

This pull request extends metrics reporting for KVConnector to include stats from the Scheduler. The change correctly identifies the place to aggregate scheduler-side stats with worker-side stats. However, there is a potential issue in how the aggregation is performed. The aggregate method on KVConnectorStats might return a new stats object rather than modifying it in-place. The current implementation discards the return value, which could lead to scheduler-side stats not being reported. I've suggested a fix to address this.

[gemini-code-assist]

The aggregate method of KVConnectorStats might return a new instance rather than modifying the object in-place. Other parts of the codebase, like KVConnectorLogging, reassign the result of aggregate. To ensure correctness and prevent potential loss of the aggregated stats, you should assign the result of aggregate back to kv_connector_stats.

Suggested change

	kv_connector_stats.aggregate(stats)
	kv_connector_stats = kv_connector_stats.aggregate(stats)

[NickLucche]

I think gemini's concern is valid given the interface is pretty flexible here

[NickLucche]

Hey @QierLi thanks a lot for adding this feature!

Although the change is very small due to the flexibility of the existing infra, I think it would be great if you could add a simple unit test showcasing scheduler<>worker metrics merging!

[NickLucche]

I think gemini's concern is valid given the interface is pretty flexible here

[mergify]

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

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

[QierLi]

Hey @QierLi thanks a lot for adding this feature!

Although the change is very small due to the flexibility of the existing infra, I think it would be great if you could add a simple unit test showcasing scheduler<>worker metrics merging!

Added a test to guard the scheduler stat aggregation via update_from_output path. : )

[NickLucche]

LGTM

[xuechendi]

I found this PR will crash:
bash vllm/tests/v1/kv_connector/nixl_integratio/run_accuracy_test.sh

with error

 File "/workspace/vllm/vllm/v1/engine/core.py", line 328, in step
engine_core_outputs = self.scheduler.update_from_output(
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/workspace/vllm/vllm/v1/core/sched/scheduler.py", line 926, in update_from_output
stats = self.connector.get_kv_connector_stats()
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/workspace/vllm/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py", line 244, in get_kv_connector_stats
assert self.connector_worker is not None

It calls get_kv_connector_stats() in scheduler, connector_worker for SCHEDULER role is None

Would like to know should we enable something else to init connector_worker for SCHEDULER ?

[QierLi]

I found this PR will crash:

bash vllm/tests/v1/kv_connector/nixl_integratio/run_accuracy_test.sh

with error

 File "/workspace/vllm/vllm/v1/engine/core.py", line 328, in step

engine_core_outputs = self.scheduler.update_from_output(

 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "/workspace/vllm/vllm/v1/core/sched/scheduler.py", line 926, in update_from_output

stats = self.connector.get_kv_connector_stats()

^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "/workspace/vllm/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py", line 244, in get_kv_connector_stats

assert self.connector_worker is not None

It calls get_kv_connector_stats() in scheduler, connector_worker for SCHEDULER role is None

Would like to know should we enable something else to init connector_worker for SCHEDULER ?

Thanks for catching this - I think returning None on get_kv_connector_stats() if the ROLE is Scheduler should simply fix this. Do you want to apply a PR to patch it? I can also file one later today.

[xuechendi]

@QierLi , thanks for quick reply, I initiated a quick fix,
When will we initiate connector_worker for Scheduler?

[njhill]

I'm wondering whether this could introduce some confusion. Currently the connector interface methods are either scheduler-side or worker-side. We're now saying one of them will be called on both sides.

I'm thinking we should even separate the KVConnector interface (abstract class) into two separate ones for the worker side and scheduler side.

And maybe there should be a separate method for querying the scheduler-side stats.

cc @NickLucche @ApostaC