From 7858ac94fb5c13ecea47be98968016b957eae7b3 Mon Sep 17 00:00:00 2001 From: Kangyan Zhou Date: Mon, 13 Apr 2026 15:25:06 -0700 Subject: [PATCH] fix(grpc): forward scheduler load info for DP-aware load balancing In gRPC mode, the GrpcRequestManager receives BatchTokenIDOutput from the scheduler but does not forward the piggybacked load info (num_reqs, num_tokens per DP rank) back to the DataParallelController. This causes total_tokens and total_requests load balance policies to see all-zero budgets and degenerate to always picking rank 0. Mirror the TokenizerManager pattern: when dp_size > 1, send a WatchLoadUpdateReq with the load info back to the scheduler input socket so the DataParallelController can make informed dispatch decisions. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../smg_grpc_servicer/sglang/request_manager.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/grpc_servicer/smg_grpc_servicer/sglang/request_manager.py b/grpc_servicer/smg_grpc_servicer/sglang/request_manager.py index 8814f7f702..673d76bcc0 100644 --- a/grpc_servicer/smg_grpc_servicer/sglang/request_manager.py +++ b/grpc_servicer/smg_grpc_servicer/sglang/request_manager.py @@ -28,6 +28,7 @@ HealthCheckOutput, TokenizedEmbeddingReqInput, TokenizedGenerateReqInput, + WatchLoadUpdateReq, ) from sglang.srt.observability.req_time_stats import ( APIServerReqTimeStats, @@ -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) + # Execute all queue.put() operations in parallel if put_tasks: await asyncio.gather(*put_tasks, return_exceptions=True)