[router] introduce dp worker abstraction

[slin1237]

Motivation

The current SGLang router implementation has DP-aware (Data Parallel aware) logic throughout the Router code. To enable DP Aware routing in PD mode. A abstraction is required. The DP-aware functionality allows workers to handle data-parallel routing by injecting data_parallel_rank into requests, but this logic is currently hardcoded in the router layer.

This PR introduces a clean Worker abstraction that encapsulates DP-aware functionality, allowing:

• Any worker type (Regular, Prefill, Decode) to be DP-aware, not just Regular workers
• Cleaner separation of concerns between routing logic and worker management
• Delegating DP-specific logic to worker implementations

Modifications

1. Worker Trait

• Added DP-aware methods to the Worker trait with default implementations:
  • is_dp_aware() - Check if worker supports data-parallel routing
  • base_url() - Get base URL without DP rank suffix
  • dp_rank() / dp_size() - Get DP configuration
  • prepare_request() - Transform requests for DP-aware routing
  • endpoint_url() - Get the actual endpoint URL for requests

2. New DPAwareWorker Implementation

• Created DPAwareWorker struct that wraps any worker type with DP-aware capabilities
• Supports all worker types: Regular, Prefill, and Decode
• Automatically injects data_parallel_rank into requests
• Uses base URL for health checks and actual requests while maintaining @dp_rank suffix for identification

3. WorkerFactory

• Added async methods for creating DP-aware workers:
  • create_dp_aware() - Create individual DP-aware worker of any type
  • create_dp_aware_regular_workers() - For regular router
  • create_dp_aware_prefill_workers() - For PD router prefill workers
  • create_dp_aware_decode_workers() - For PD router decode workers
• Modified create_workers() to use the new DP-aware factory methods

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]

Summary of Changes

Hello @slin1237, 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!

This pull request significantly refactors the SGLang router's worker management by introducing a dedicated Data Parallel (DP) worker abstraction. This change aims to decouple DP-aware logic from the core router, making the system more modular and extensible. It allows any worker type to be DP-aware and centralizes the handling of data parallel ranks within the new DPAwareWorker and updated WorkerFactory.

Highlights

• Enhanced Worker Trait: The Worker trait has been extended with new methods (is_dp_aware, base_url, dp_rank, dp_size, prepare_request, endpoint_url, can_handle) to support Data Parallel (DP) awareness, providing default implementations for non-DP-aware workers.
• New DPAwareWorker Implementation: A new DPAwareWorker struct has been introduced, which wraps existing worker types (Regular, Prefill, Decode) and automatically injects data_parallel_rank into requests, centralizing DP-specific logic.
• WorkerFactory Enhancements: The WorkerFactory now includes asynchronous methods (create_dp_aware, create_dp_aware_regular_workers, create_dp_aware_prefill_workers, create_dp_aware_decode_workers) for creating DP-aware workers, and the general create_workers method has been updated to utilize these new factory methods.

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 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 or fill out our survey 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

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 is an excellent refactoring that introduces a DPAwareWorker abstraction to encapsulate data-parallel logic, which is a great improvement for modularity and maintainability. The changes are well-structured and the addition of comprehensive unit tests is fantastic.

My review focuses on a few areas for further improvement, mainly around performance and code duplication. I've suggested parallelizing worker creation, reusing HTTP clients to improve efficiency, and refactoring some repetitive logic. Overall, great work!

[gemini-code-assist]

While the current check for dp_size overflow is correct, using usize::try_from is a more idiomatic and robust way to handle fallible numeric conversions in Rust. It clearly expresses the intent of a conversion that might fail and is safer across different architectures (e.g., 32-bit vs. 64-bit).

usize::try_from(dp_size).map_err(|_| WorkerError::InvalidConfiguration {
    message: format!("dp_size {} is too large for this architecture", dp_size),
})

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[qiang-w-ruijie]

Dose it support kv cache aware？