feat: send kvmetrics from sglang scheduler

[zixuanzhang226]

Motivation

This PR introduces the first step toward integrating the Dynamo Planner with SGLANG. It enables the scheduler to send KV-related metrics—required by the planner to make scaling decisions for prefill/decode workers—via a ZeroMQ (ZMQ) socket.

Background:
Dynamo Planner is a service designed to monitor the state of the inference system and dynamically scale prefill and decode workers based on metrics such as KV cache usage and prefill queue size.

Reference Issue Link: #6163

Completed end-to-end testing and confirmed that KV metrics are successfully sent from the SGLang scheduler and received by the Dynamo worker.

The receiver-side implementation (on the Dynamo Planner worker) PR link: ai-dynamo/dynamo#1465

Modifications

The scheduler now gathers relevant KV metrics and sends them through a ZMQ socket. These metrics will be consumed by the Dynamo Planner for making scale-up/scale-down decisions.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[gemini-code-assist]

Hello @zixuanzhang226, I'm Gemini Code Assist¹! 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!

Summary of Changes

Hello! Gemini here, providing a summary of this pull request to help everyone quickly understand the changes and context.

This PR represents the initial step towards integrating the SGLANG scheduler with the external Dynamo Planner service. The core goal is to enable the scheduler to report key-value (KV) cache related metrics, which the Dynamo Planner will consume via ZeroMQ (ZMQ) to make informed decisions about scaling prefill and decode workers. This PR focuses solely on the scheduler's side of sending these metrics; the receiver implementation in the Dynamo Planner will be handled in a subsequent PR.

Highlights

• KV Metrics Structure: Introduced a new KvMetrics class in scheduler.py to encapsulate various KV cache and request-related statistics that are relevant for external monitoring and scaling decisions.
• ZMQ Socket for Metrics: Added a new ZeroMQ PUSH socket (send_metrics_from_scheduler) to the Scheduler class, specifically dedicated to sending the collected metrics.
• Metrics Collection and Sending: Modified the log_prefill_stats and log_decode_stats methods in the scheduler to gather the defined KV metrics, populate the KvMetrics object, and send it over the newly created ZMQ socket.
• Configuration for Metrics Socket: Added a new field metrics_ipc_name to the PortArgs class in server_args.py to allow configuration of the address for the metrics ZMQ socket, initializing it with a temporary IPC file path for both single and MPI distribution modes.

Changelog

• python/sglang/srt/managers/scheduler.py
  • Added KvMetrics class definition (around line 167).
  • Initialized a new ZMQ PUSH socket self.send_metrics_from_scheduler using port_args.metrics_ipc_name (around line 233).
  • Set self.send_metrics_from_scheduler to None in the non-server initialization path (around line 253).
  • Added logic in log_prefill_stats to create, populate, and send a KvMetrics object via the ZMQ socket (around line 1184).
  • Added logic in log_decode_stats to create, populate, and send a KvMetrics object via the ZMQ socket (around line 1250).
• python/sglang/srt/server_args.py
  • Added metrics_ipc_name: str field to the PortArgs class (around line 1341).
  • Initialized metrics_ipc_name with a temporary IPC file path in PortArgs.init_new for single distribution mode (around line 1363).
  • Initialized metrics_ipc_name with a temporary IPC file path in PortArgs.init_new for MPI distribution mode (around line 1394).

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.

---

Metrics flow like streams,
Through sockets, data gleams.
Planner watches close,
As KV usage grows,
Scaling up our inference dreams.

Footnotes

1. 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. ↩

[gemini-code-assist]

Code Review

This pull request introduces an important first step towards integrating the Dynamo Planner by enabling the scheduler to send KV-related metrics. The overall approach of using a ZMQ socket for this purpose is sound.

I've identified a few areas for improvement, primarily concerning the calculation and naming of one of the metrics, code duplication, management of temporary IPC files, and consistency of IPC mechanisms in distributed setups. Additionally, adding unit tests for this new functionality would be beneficial.

Great work on laying this foundation!

Summary of Findings

• Metric Calculation Clarity: The calculation for kv_active_blocks in scheduler.py might be misinterpreted if 'block' doesn't mean 'token slot'. Clarification or renaming could improve understanding.
• Code Duplication: The logic for creating and sending KvMetrics is duplicated in log_prefill_stats and log_decode_stats in scheduler.py. Refactoring into a helper method is recommended.
• IPC Temporary File Management: Temporary files created for IPC sockets in server_args.py are not cleaned up, potentially leading to file system clutter. A cleanup strategy should be considered.
• IPC Mechanism Consistency in DP Mode: In server_args.py, metrics_ipc_name uses file-based IPC in DP mode, while other IPCs use TCP. This might be problematic if the metrics consumer is remote.
• Unit Testing: The new metrics sending functionality lacks unit tests, which are important for ensuring correctness and preventing regressions.
• Docstrings for KvMetrics: The KvMetrics class in scheduler.py could benefit from a docstring explaining its purpose and attributes. (Low severity, not commented directly)
• Docstrings for PortArgs.metrics_ipc_name: The new metrics_ipc_name field in PortArgs in server_args.py could use a more detailed docstring. (Low severity, not commented directly)
• KvMetrics class structure: The KvMetrics class in scheduler.py could be defined using @dataclass for better conciseness and to enable type hinting more easily. (Low severity, not commented directly)

Merge Readiness

This PR introduces valuable functionality for metrics collection. However, there are a few medium-severity issues related to metric clarity, code duplication, IPC file management, and IPC consistency in distributed setups that should be addressed. Additionally, adding unit tests for this new feature is highly recommended. I am unable to approve the pull request, and recommend that these changes be made and reviewed by others before merging.

[trevor-m]

Overall looks good to me.

[xiezhq-hermann]

Thank you for the PR and I agree these are useful logs for monitoring. That being said, can you jusitify why introducing KvMetrics instead of reusing existing observability interfaces like https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/metrics/collector.py?

[zixuanzhang226]

Thank you for the PR and I agree these are useful logs for monitoring. That being said, can you jusitify why introducing KvMetrics instead of reusing existing observability interfaces like https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/metrics/collector.py?

Hello @xiezhq-hermann, thank you very much for your review. We introduced the KvMetrics class primarily for clarity and ease of use when sending and receiving KV metrics over the ZMQ socket. Additionally, not all KV metrics originate from SchedulerMetricsCollector (which primarily logs metrics from SchedulerStats); some are obtained directly within the Scheduler class itself.

[zhyncs]

@zixuanzhang226 please fix the conflicts

[zixuanzhang226]

@zixuanzhang226 please fix the conflicts

Hi @zhyncs, I resolved the conflicts. Please take a look. Thank you!

[zixuanzhang226]

Hello @zhyncs, just a quick note — I updated test/srt/test_server_args.py to ensure the unit test passes.

[ishandhanani]

Kicking off CI @zixuanzhang226

[zixuanzhang226]

Hi @ByronHsu, we believe the current unit test failures are not related to the code changes in this PR. Could you please take a look? Thanks!

[faradawn]

The KV metric schema looks good. Good idea to separate it into a emit event helper function.

[zhyncs]

please fix the lint issue pre-commit run --all-files

[zixuanzhang226]

please fix the lint issue pre-commit run --all-files

Hi @zhyncs I fixed it. Thanks!

[ishandhanani]

Running CI @zixuanzhang226