Skip to content

fix(grpc): forward scheduler load info for DP-aware load balancing #1115

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Kangyan-Zhou wants to merge 1 commit into from
Closed

fix(grpc): forward scheduler load info for DP-aware load balancing #1115

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions grpc_servicer/smg_grpc_servicer/sglang/request_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
HealthCheckOutput,
TokenizedEmbeddingReqInput,
TokenizedGenerateReqInput,
WatchLoadUpdateReq,
)
from sglang.srt.observability.req_time_stats import (
APIServerReqTimeStats,
Expand Down Expand Up @@ -679,6 +680,15 @@ async def cleanup(request_id):

cleanup_tasks.append(asyncio.create_task(cleanup(rid)))

# Forward load info to DataParallelController for token-aware balancing.
# Mirrors TokenizerManager._handle_batch_output logic: when dp_size > 1,
# each scheduler piggybacks its load (num_reqs, num_tokens) on batch output.
# Without this, DPBudget stays at zero and total_tokens/total_requests
# policies degenerate to always picking rank 0.
if self.server_args.dp_size > 1 and batch_out.load is not None:
load_update = WatchLoadUpdateReq(loads=[batch_out.load])
self.send_to_scheduler.send_pyobj(load_update)
Comment on lines +688 to +690

@gemini-code-assist gemini-code-assist Bot Apr 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To maintain consistency with the rest of the GrpcRequestManager class (see lines 368, 427, and 471), please use the _send_to_scheduler helper method instead of calling self.send_to_scheduler.send_pyobj directly. This ensures that the send operation is covered by the helper's error logging. Since _handle_batch_output is an async function, the call should be awaited to follow the established pattern in this class.

Suggested change
if self.server_args.dp_size > 1 and batch_out.load is not None:
load_update = WatchLoadUpdateReq(loads=[batch_out.load])
self.send_to_scheduler.send_pyobj(load_update)
if self.server_args.dp_size > 1 and batch_out.load is not None:
load_update = WatchLoadUpdateReq(loads=[batch_out.load])
await self._send_to_scheduler(load_update)
References
  1. If a code block is identified as duplicated across multiple functions or modules, consider refactoring to unify the logic into a shared helper function or class. This improves maintainability and reduces the chance of inconsistencies when changes are needed.

Comment on lines +688 to +690

@coderabbitai coderabbitai Bot Apr 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

鈿狅笍 Potential issue | 馃敶 Critical

Make DP load forwarding best-effort to avoid dropping batch outputs.

At Line 690, an unguarded send can raise and abort _handle_batch_output before Line 694 flushes put_tasks, which can stall/delay client-visible outputs for that batch. Handle this path as non-fatal.

Suggested fix 
-        if self.server_args.dp_size > 1 and batch_out.load is not None:
-            load_update = WatchLoadUpdateReq(loads=[batch_out.load])
-            self.send_to_scheduler.send_pyobj(load_update)
+        if self.server_args.dp_size > 1 and batch_out.load is not None:
+            try:
+                await self._send_to_scheduler(WatchLoadUpdateReq(loads=[batch_out.load]))
+            except Exception as e:
+                logger.warning(f"Failed to forward DP load update: {e}")
馃 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@grpc_servicer/smg_grpc_servicer/sglang/request_manager.py` around lines 688 -
690, The send of DP load updates (WatchLoadUpdateReq with batch_out.load) from
_handle_batch_output can raise and abort before put_tasks are flushed; make this
best-effort by catching exceptions around send_to_scheduler.send_pyobj so
failures won't propagate. Wrap the call that constructs WatchLoadUpdateReq and
calls self.send_to_scheduler.send_pyobj(...) in a try/except that logs/debugs
the error (non-fatal) and continues, ensuring the subsequent flush of put_tasks
still runs; keep the guard on self.server_args.dp_size > 1 and only apply the
try/except when batch_out.load is not None.


# Execute all queue.put() operations in parallel
if put_tasks:
await asyncio.gather(*put_tasks, return_exceptions=True)
Expand Down