[Router] Fix k8s Service Discovery #6766

YouNeedCryDear · 2025-05-30T07:31:07Z

Motivation

The current k8s service discovery directly update the worker-urls of the Router instance. This might cause the worker tree and router status not correctly updated. The motivation is to use the public function add_worker and remove_worker to update the worker-urls. Thus all the other attributes of the router and status of the router would be correctly updated.

Modifications

Use Arc Router instead of Arc worker-urls in service_discovery
Call add_worker and remove_worker to update the worker-urls and other attributes
Update the server to use Arc Router
Minor improvement on logging

Unit tests and benchmarking result will come later

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.

This reverts commit f6f34c25df78fcd20661c1e59d81ef2add25817b.

gemini-code-assist

Hello @YouNeedCryDear, 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 team, gemini-code-assist here to provide a summary of this pull request. This PR addresses an issue in the Kubernetes service discovery logic where it was directly modifying the router's internal worker URL list. This direct manipulation could lead to inconsistencies in the router's state, such as the worker tree or overall status not being correctly updated. The core change is to refactor the service discovery to hold an Arc<Router> and use the router's public add_worker and remove_worker methods. This ensures that all necessary internal state updates within the Router are handled correctly when workers are added or removed based on Kubernetes pod events.

Highlights

Improved Router State Management: The service discovery component now holds a shared reference (Arc) to the Router instance itself, allowing it to interact with the router through its defined public interface.
Utilizing Router API for Worker Updates: Instead of directly modifying the worker URL list, the service discovery now calls router.add_worker() when a healthy pod is found and router.remove_worker() when a pod is deleted or becomes unhealthy. This delegates the state update responsibility to the Router.
Server Integration Update: The main server startup logic (startup function) has been updated to create the Router wrapped in an Arc and pass this Arc<Router> to the service discovery initialization function.
Logging Improvements: Minor adjustments to logging messages have been made for better clarity on service discovery actions.

Changelog

Click here to see the changelog

sgl-router/src/server.rs
- Updated AppState struct to hold router: Arc<Router> instead of router: Router (DIFF 1).
- Modified AppState::new to wrap the newly created Router in an Arc (DIFF 2).
- Refactored the startup function to correctly initialize AppState with Arc<Router> and pass a clone of this Arc to start_service_discovery (DIFF 3).
- Adjusted the info log message during startup to read the worker URLs from the router's state via the Arc (DIFF 4).
sgl-router/src/service_discovery.rs
- Imported the Router type and removed the direct RwLock import (DIFF 1).
- Changed the signature of start_service_discovery to accept router: Arc<Router> instead of worker_urls: Arc<RwLock<Vec<String>>> (DIFF 2).
- Removed the cloning of the old worker_urls_clone variable as it's no longer needed (DIFF 3).
- Updated the pod event processing loop (try_for_each) to clone the router Arc for inner closures (DIFF 4).
- Changed the signatures of handle_pod_event and handle_pod_deletion to accept router: Arc<Router> (DIFF 5, DIFF 6).
- Replaced direct modification of the worker URL vector in handle_pod_event with a call to router.add_worker() (DIFF 6).
- Replaced direct modification of the worker URL vector in handle_pod_deletion with a call to router.remove_worker() (DIFF 6).
- Updated info log messages in event handlers for better clarity (DIFF 6).
- Added a debug log message in handle_pod_deletion for events related to untracked pods (DIFF 6).

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.

State held in Arc,
Shared safely, a vital spark.
Workers come and go.

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 refactors the Kubernetes service discovery to use the Router's add_worker and remove_worker methods, which is a significant improvement for maintaining consistent state. The core changes are well-implemented.

One main issue has been identified in handle_pod_event concerning the logic for tracking pods after an add_worker attempt, which could lead to healthy workers being missed if their initial addition to the router fails.

Overall, this is a valuable change. Addressing the identified issue and adding comprehensive unit tests will make it robust.

Summary of Findings

Incorrect Pod Tracking on add_worker Failure: In handle_pod_event, a pod is added to the local tracked_pods set even if the router.add_worker() call fails for reasons other than the worker already existing. This can prevent future attempts to add a healthy worker if its initial addition to the router failed, as service discovery will incorrectly believe the pod is already handled.
Unit Testing: The PR description mentions that unit tests will come later. Given the changes to the service discovery logic and its interaction with the router, comprehensive unit tests are crucial to ensure correctness and prevent regressions. This was noted as 'not commented on due to review settings' as it's a general process point rather than a specific code line issue in the diff.

Merge Readiness

The pull request introduces valuable improvements to the service discovery mechanism. However, there is a high-severity issue identified in handle_pod_event related to pod tracking logic that should be addressed before merging. Additionally, as noted in the PR checklist, unit tests for these changes are pending and will be important for ensuring long-term stability.

I recommend that the identified issue be fixed and unit tests be added/updated before this PR is merged. As a reviewer, I am not authorized to approve pull requests, but based on this review, changes are requested.

sgl-router/src/service_discovery.rs

zhyncs · 2025-06-02T19:22:32Z

Hi @YouNeedCryDear May you help fix the CI check

YouNeedCryDear · 2025-06-02T19:49:16Z

@zhyncs Just followed steps in pr-test-rust workflow. Could you trigger it one more time?

Co-authored-by: Simo Lin <linsimo.mark@gmail.com>

Merge branch 'sgl_20250610_sync_tag047 of git@code.alipay.com:Theta/SGLang.git into main https://code.alipay.com/Theta/SGLang/pull_requests/52 Reviewed-by: 剑川 <jianchuan.gys@antgroup.com> * [Bugfix] Fix slice operation when chunk size mismatch (sgl-project#6697) * [Bugfix] Fix ChatCompletion endpoint of mini_lb when stream is set (sgl-project#6703) * [CI] Fix setup of disaggregation with different tp (sgl-project#6706) * [PD] Remove Unnecessary Exception Handling for FastQueue.get() (sgl-project#6712) * Fuse routed_scaling_factor in DeepSeek (sgl-project#6710) * Overlap two kernels in DeepSeek with communication (sgl-project#6711) * Minor refactor two-batch overlap (sgl-project#6682) * Speed up when having padding tokens two-batch overlap (sgl-project#6668) * [Feature] Support Flashinfer fp8 blockwise GEMM kernel on Blackwell (sgl-project#6479) * Fix LoRA bench (sgl-project#6719) * temp * Fix PP for Qwen3 MoE (sgl-project#6709) * [feat] triton kernel for get_last_loc (sgl-project#6676) * [fix] more mem for draft_extend cuda_graph (sgl-project#6726) * [PD] bug fix: Update status if nixl receiver send a a dummy req. (sgl-project#6720) * Tune memory arguments on B200 (sgl-project#6718) * Add DeepSeek-R1-0528 function call chat template (sgl-project#6725) * refactor(tool call): Fix BaseFormatDetector tool_index issue and refactor `parse_streaming_increment` (sgl-project#6715) * Add draft extend CUDA graph for Triton backend (sgl-project#6705) * refactor apply_w8a8_block_fp8_linear in fp (sgl-project#6545) * [PD] Support completion endpoint (sgl-project#6729) * PD Rust LB (PO2) (sgl-project#6437) * Super tiny enable sole usage of expert distribution metrics and update doc (sgl-project#6680) * Support picking variants of EPLB algorithms (sgl-project#6728) * Support tuning DeepEP configs (sgl-project#6742) * [test] add ut and bm for get_last_loc (sgl-project#6746) * Fix mem_fraction_static for AMD CI (sgl-project#6748) * [fix][RL] Fix DeepSeekV3ForCausalLM.post_load_weights for multiple update weight (sgl-project#6265) * Improve EPLB logical to physical dispatch map (sgl-project#6727) * Update DeepSeek-R1-0528 function call chat template (sgl-project#6765) * [PD] Optimize time out logic and add env var doc for mooncake (sgl-project#6761) * Fix aiohttp 'Chunk too big' in bench_serving (sgl-project#6737) * Support sliding window in triton backend (sgl-project#6509) * Fix shared experts fusion error (sgl-project#6289) * Fix one bug in the grouped-gemm triton kernel (sgl-project#6772) * update llama4 chat template and pythonic parser (sgl-project#6679) * feat(tool call): Enhance Llama32Detector for improved JSON parsing in non-stream (sgl-project#6784) * Support token-level quantization for EP MoE (sgl-project#6782) * Temporarily lower mmlu threshold for triton sliding window backend (sgl-project#6785) * ci: relax test_function_call_required (sgl-project#6786) * Add intel_amx backend for Radix Attention for CPU (sgl-project#6408) * Fix incorrect LoRA weight loading for fused gate_up_proj (sgl-project#6734) * fix(PD-disaggregation): Can not get local ip (sgl-project#6792) * [FIX] mmmu bench serving result display error (sgl-project#6525) (sgl-project#6791) * Bump torch to 2.7.0 (sgl-project#6788) * chore: bump sgl-kernel v0.1.5 (sgl-project#6794) * Improve profiler and integrate profiler in bench_one_batch_server (sgl-project#6787) * chore: upgrade sgl-kernel v0.1.5 (sgl-project#6795) * [Minor] Always append newline after image token when parsing chat message (sgl-project#6797) * Update CI tests for Llama4 models (sgl-project#6421) * [Feat] Enable PDL automatically on Hopper architecture (sgl-project#5981) * chore: update blackwell docker (sgl-project#6800) * misc: cache is_hopper_arch (sgl-project#6799) * Remove contiguous before Flashinfer groupwise fp8 gemm (sgl-project#6804) * Correctly abort the failed grammar requests & Improve the handling of abort (sgl-project#6803) * [EP] Add cuda kernel for moe_ep_pre_reorder (sgl-project#6699) * Add draft extend CUDA graph for flashinfer backend (sgl-project#6805) * Refactor CustomOp to avoid confusing bugs (sgl-project#5382) * Tiny log prefill time (sgl-project#6780) * Tiny fix EPLB assertion about rebalancing period and recorder window size (sgl-project#6813) * Add simple utility to dump tensors for debugging (sgl-project#6815) * Fix profiles do not have consistent names (sgl-project#6811) * Speed up rebalancing when using non-static dispatch algorithms (sgl-project#6812) * [1/2] Add Kernel support for Cutlass based Fused FP4 MoE (sgl-project#6093) * [Router] Fix k8s Service Discovery (sgl-project#6766) * Add CPU optimized kernels for topk and rope fusions (sgl-project#6456) * fix new_page_count_next_decode (sgl-project#6671) * Fix wrong weight reference in dynamic EPLB (sgl-project#6818) * Minor add metrics to expert location updater (sgl-project#6816) * [Refactor] Rename `n_share_experts_fusion` as `num_fused_shared_experts` (sgl-project#6735) * [FEAT] Add transformers backend support (sgl-project#5929) * [fix] recover auto-dispatch for rmsnorm and rope (sgl-project#6745) * fix ep_moe_reorder kernel bugs (sgl-project#6858) * [Refactor] Multimodal data processing for VLM (sgl-project#6659) * Decoder-only Scoring API (sgl-project#6460) * feat: add dp-rank to KV events (sgl-project#6852) * Set `num_fused_shared_experts` as `num_shared_experts` when shared_experts fusion is not disabled (sgl-project#6736) * Fix one missing arg in DeepEP (sgl-project#6878) * Support LoRA in TestOpenAIVisionServer and fix fused kv_proj loading bug. (sgl-project#6861) * support 1 shot allreduce in 1-node and 2-node using mscclpp (sgl-project#6277) * Fix Qwen3MoE missing token padding optimization (sgl-project#6820) * Tiny update error hints (sgl-project#6846) * Support layerwise rebalancing experts (sgl-project#6851) * Tiny allow profiler API to auto create directory (sgl-project#6865) * Support Blackwell DeepEP docker images (sgl-project#6868) * [EP] Add cuda kernel for moe_ep_post_reorder (sgl-project#6837) * [theta]merge 0605 * oai: fix openAI client error with single request via batch api (sgl-project#6170) * [PD] Fix potential perf spike caused by tracker gc and optimize doc (sgl-project#6764) * Use deepgemm instead of triton for fused_qkv_a_proj_with_mqa (sgl-project#6890) * [CUTLASS-FP4-MOE] Introduce CutlassMoEParams class for easy initialization of Cutlass Grouped Gems Metadata (sgl-project#6887) * bugfix(OAI): Fix image_data processing for jinja chat templates (sgl-project#6877) * [CPU] enable CI for PRs, add Dockerfile and auto build task (sgl-project#6458) * AITER backend extension and workload optimizations (sgl-project#6838) * [theta]merge * [theta]merge * [Feature] Support Flashinfer fmha on Blackwell (sgl-project#6930) * Fix a bug in abort & Improve docstrings for abort (sgl-project#6931) * Tiny support customize DeepEP max dispatch tokens per rank (sgl-project#6934) * Sync the changes on cuda graph runners (sgl-project#6932) * [PD] Optimize transfer queue forward logic for dummy rank (sgl-project#6922) * [Refactor] image data process in bench_serving (sgl-project#6879) * [fix] logical_to_all_physical_map index 256 is out of bounds in EP parallel. (sgl-project#6767) * Add triton fused moe kernel config for E=257 on B200 (sgl-project#6939) * [sgl-kernel] update deepgemm (sgl-project#6942) * chore: bump sgl-kernel v0.1.6 (sgl-project#6943) * Minor compile fused topk (sgl-project#6944) * [Bugfix] pipeline parallelism and Eagle Qwen2 (sgl-project#6910) * Tiny re-introduce profile id logging (sgl-project#6912) * Add triton version as a fused_moe_triton config search key to avoid performace decrease in different Triton version (sgl-project#5955) * reduce torch.zeros overhead in moe align block size kernel (sgl-project#6369) * chore: upgrade sgl-kernel v0.1.6 (sgl-project#6945) * add fbgemm moe grouped gemm kernel benchmark (sgl-project#6924) * [Docker] Add docker file for SGL Router (sgl-project#6915) * Disabling mixed chunked prefill when eagle is enabled (sgl-project#6874) * Add canary for EPLB rebalancing (sgl-project#6895) * Refactor global_server_args_dict (sgl-project#6866) * Fuse routed scaling factor in topk_reduce kernel (sgl-project#6220) * Update server timeout time in AMD CI. (sgl-project#6953) * [misc] add is_cpu() (sgl-project#6950) * Add H20 fused MoE kernel tuning configs for DeepSeek-R1/V3 (sgl-project#6885) * Add a CUDA kernel for fusing mapping and weighted sum for MoE. (sgl-project#6916) * chore: bump sgl-kernel v0.1.6.post1 (sgl-project#6955) * chore: upgrade sgl-kernel v0.1.6.post1 (sgl-project#6957) * [DeepseekR1-FP4] Add Support for nvidia/DeepSeekR1-FP4 model (sgl-project#6853) * Revert "Fuse routed scaling factor in topk_reduce kernel (sgl-project#6220)" (sgl-project#6968) * [AMD] Add more tests to per-commit-amd (sgl-project#6926) * chore: bump sgl-kernel v0.1.7 (sgl-project#6963) * Slightly improve the sampler to skip unnecessary steps (sgl-project#6956) * rebase h20 fused_moe config (sgl-project#6966) * Fix CI and triton moe Configs (sgl-project#6974) * Remove unnecessary kernels of num_token_non_padded (sgl-project#6965) * Extend cuda graph capture bs for B200 (sgl-project#6937) * Fuse routed scaling factor in deepseek (sgl-project#6970) * Sync cuda graph runners (sgl-project#6976) * Fix draft extend ut stability with flush cache (sgl-project#6979) * Fix triton sliding window test case (sgl-project#6981) * Fix expert distribution dumping causes OOM (sgl-project#6967) * Minor remove one kernel for DeepSeek (sgl-project#6977) * [perf][sgl-kernel] extend cutlass_mla_decode to support num_head < 128 (sgl-project#6929) * Enable more unit tests for AMD CI. (sgl-project#6983) * Use torch.compile to fuse flash attention decode metadata preparation (sgl-project#6973) * Eliminate stream sync to speed up LoRA batch init (sgl-project#6960) * support qwen3 emebedding (sgl-project#6990) * Fix torch profiler bugs for bench_offline_throughput.py (sgl-project#6557) * chore: upgrade flashinfer v0.2.6.post1 jit (sgl-project#6958) * cleanup tmp dir (sgl-project#7007) * chore: update pr test xeon (sgl-project#7008) * Fix cutlass MLA gets almost zero accuracy (sgl-project#6998) * Update amd nightly models CI. (sgl-project#6992) * feat: add direct routing strategy to DP worker (sgl-project#6884) * Fallback to lower triton version for unfound fused moe configs (sgl-project#7013) * Fix torchvision version for Blackwell (sgl-project#7015) * Simplify prepare_extend_after_decode (sgl-project#6987) * Migrate to assertEqual (sgl-project#6741) * Fix torch version in blackwell dockerfile (sgl-project#7017) * chore: update pr test xeon (sgl-project#7018) * Update default settings for blackwell (sgl-project#7023) * Support both approximate and exact expert distribution collection (sgl-project#6964) * Add decode req pool (sgl-project#6980) * [theta]merge 0610 * [theta]merge 0610 * [CI] Add CI workflow for sgl-router docker build (sgl-project#7027) * Fix fused_moe triton configs (sgl-project#7029) * CPU: map changes from developing branch in sgl-kernel (sgl-project#6833) * chore: bump v0.4.7 (sgl-project#7038) * Update README.md (sgl-project#7040)

YouNeedCryDear added 6 commits May 30, 2025 00:51

temp fix for router toctou worker urls

a629994

Revert "temp fix for router toctou worker urls"

7a6212a

This reverts commit f6f34c25df78fcd20661c1e59d81ef2add25817b.

use add_worker and remove_worker in service discovery

5648832

change server behaviour with arc router in service discovery

122dfc1

format

8ddeb8b

remove extra clone

fc0a52a

YouNeedCryDear requested review from ByronHsu, Ying1123 and slin1237 as code owners May 30, 2025 07:31

gemini-code-assist bot reviewed May 30, 2025

View reviewed changes

Merge branch 'main' into fix/router-toc-tou-fix-service-discovery

a11fae7

slin1237 approved these changes May 30, 2025

View reviewed changes

gemini-code-assist bot suggested changes May 30, 2025

View reviewed changes

sgl-router/src/service_discovery.rs Outdated Show resolved Hide resolved

YouNeedCryDear added 3 commits June 2, 2025 18:17

move tracker inside OK condition in add_worker

c1059f6

Merge branch 'main' into fix/router-toc-tou-fix-service-discovery

6f06729

add unit tests for service_discovery

d5bab7e

YouNeedCryDear requested a review from slin1237 June 2, 2025 19:17

run cargo fmt

e43095b

slin1237 merged commit ff91474 into sgl-project:main Jun 2, 2025
4 checks passed

YouNeedCryDear deleted the fix/router-toc-tou-fix-service-discovery branch June 3, 2025 00:54

Layssy pushed a commit to Layssy/sglang-iaas that referenced this pull request Jun 9, 2025

[Router] Fix k8s Service Discovery (sgl-project#6766)

19be43f

Co-authored-by: Simo Lin <linsimo.mark@gmail.com>

xwu-intel pushed a commit to xwu-intel/sglang that referenced this pull request Jun 17, 2025

[Router] Fix k8s Service Discovery (sgl-project#6766)

b963c16

Co-authored-by: Simo Lin <linsimo.mark@gmail.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Router] Fix k8s Service Discovery #6766

[Router] Fix k8s Service Discovery #6766

Uh oh!

YouNeedCryDear commented May 30, 2025 •

edited

Loading

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

zhyncs commented Jun 2, 2025

Uh oh!

YouNeedCryDear commented Jun 2, 2025

Uh oh!

Uh oh!

Uh oh!

[Router] Fix k8s Service Discovery #6766

[Router] Fix k8s Service Discovery #6766

Uh oh!

Conversation

YouNeedCryDear commented May 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Checklist

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Changelog

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Summary of Findings

Merge Readiness

Uh oh!

Uh oh!

zhyncs commented Jun 2, 2025

Uh oh!

YouNeedCryDear commented Jun 2, 2025

Uh oh!

Uh oh!

Uh oh!

YouNeedCryDear commented May 30, 2025 •

edited

Loading