chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Account decode handoff queues in DP load snapshots

Browse Source 
Decode DP dispatch was collapsing onto a few ranks because the controller only randomizes among exact minimum load pairs. The load snapshot undercounted decode handoff work: pending prefill-info requests were absent, and DecodeRequest wrappers in prealloc/transfer queues were skipped because their rid lives on .req.

This makes scheduler load accounting unwrap DecodeRequest items and include pending decode requests, so TOTAL_TOKENS sees queued handoff backlog instead of repeatedly treating busy ranks as empty.

Constraint: Do not mask imbalance with synthetic per-request token penalties; dispatch should be driven by accurate observed load.

Rejected: Add req*4000 or other queue penalties | heuristic, workload-dependent, and hides the accounting bug.

Confidence: medium

Scope-risk: moderate

Directive: Any new decode handoff queue must be included in get_load() or DP routing can regress to stale/min-load collapse.

Tested: Remote cjy-glm5-new: PYTHONPATH=python python -m pytest -q test/registered/unit/observability/test_scheduler_metrics_load.py test/registered/unit/managers/test_prefill_adder.py -> 27 passed.

Not-tested: Fresh decode ETE distribution after service restart.
This commit is contained in:
laoyao0822
2026-06-13 02:35:10 +08:00
parent 69ca7045ea
commit ceb5345410
2 changed files with 136 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

58
python/sglang/srt/observability/scheduler_metrics_mixin.py
View File

@@ -679,7 +679,33 @@ class SchedulerMetricsMixin:
/ self.stats.max_running_requests_under_SLO,
self.stats.token_usage / 0.9,
)
def _unwrap_load_req(self, item):
# Decode disaggregation queues store DecodeRequest wrappers, while
# normal waiting/retracted queues store Req directly. Load accounting
# must see the underlying request in both cases; otherwise DP dispatch
# loses queued decode work from reqs_list and can repeatedly choose the
# same few ranks under bursty PD traffic.
return getattr(item, "req", item)
def _get_load_item_num_tokens(self, item):
num_tokens = getattr(item, "seqlen", None)
if num_tokens is not None:
return num_tokens
req = self._unwrap_load_req(item)
num_tokens = getattr(req, "seqlen", None)
if num_tokens is not None:
return num_tokens
origin_input_ids = getattr(req, "origin_input_ids", None)
output_ids = getattr(req, "output_ids", None)
prompt_tokens = len(origin_input_ids) if origin_input_ids is not None else 0
completion_tokens = len(output_ids) if output_ids is not None else 0
return prompt_tokens + completion_tokens
def _build_request_info(self, req):
req = self._unwrap_load_req(req)
prompt_tokens = 0
if getattr(req, "origin_input_ids_unpadded", None) is not None:
prompt_tokens = len(req.origin_input_ids_unpadded)
@@ -703,16 +729,17 @@ class SchedulerMetricsMixin:
completion_tokens=completion_tokens,
timestamp=time.perf_counter(),
)
def _collect_request_infos(self, queues,request_infos):
def _collect_request_infos(self, queues, request_infos):
# all queues
for queue in queues:
for req in queue:
rid = getattr(req, "rid", None)
if rid is None:
# logger.info("rid is None, skipping request")
continue
request_infos.append(self._build_request_info(req))
for item in queue:
req = self._unwrap_load_req(item)
rid = getattr(req, "rid", None)
if rid is None:
# logger.info("rid is None, skipping request")
continue
request_infos.append(self._build_request_info(req))
def get_load(self: Scheduler, _: GetLoadReqInput = None) -> GetLoadReqOutput:
if self.is_hybrid_swa:
@@ -728,13 +755,22 @@ class SchedulerMetricsMixin:
if self.disaggregation_mode == DisaggregationMode.PREFILL:
waiting_queues.append(self.disagg_prefill_bootstrap_queue.queue)
elif self.disaggregation_mode == DisaggregationMode.DECODE:
# Include every decode handoff stage. pending_reqs are not yet
# wrapped in DecodeRequest and were previously invisible to the DP
# controller, so burst routing could keep dispatching to a rank with
# a large unresolved handoff backlog.
waiting_queues.append(self.disagg_decode_prealloc_queue.pending_reqs)
waiting_queues.append(self.disagg_decode_prealloc_queue.queue)
waiting_queues.append(self.disagg_decode_transfer_queue.queue)
waiting_queues.append(self.disagg_decode_prealloc_queue.retracted_queue)
num_tokens += sum(req.seqlen for queue in waiting_queues for req in queue)
num_tokens += sum(
self._get_load_item_num_tokens(item)
for queue in waiting_queues
for item in queue
)
num_waiting_reqs = sum(len(queue) for queue in waiting_queues)
request_infos=[]
self._collect_request_infos(waiting_queues,request_infos)
request_infos = []
self._collect_request_infos(waiting_queues, request_infos)
# logger.info("current dp rank %d num tokens %d waiting queues %d",self.dp_rank,num_tokens,num_waiting_reqs)
return GetLoadReqOutput(

89
test/registered/unit/observability/test_scheduler_metrics_load.py Normal file
View File

@@ -0,0 +1,89 @@
"""Tests for scheduler load accounting used by DP dispatch."""
from types import SimpleNamespace
# Import scheduler first to mirror production import order and avoid the
# scheduler <-> metrics circular import during isolated unit collection.
import sglang.srt.managers.scheduler # noqa: F401
from sglang.srt.disaggregation.decode import DecodeRequest
from sglang.srt.disaggregation.utils import DisaggregationMode
from sglang.srt.observability.scheduler_metrics_mixin import SchedulerMetricsMixin
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
class FakeReq:
def __init__(self, rid: str, prompt_tokens: int, output_tokens: int = 0):
self.rid = rid
self.origin_input_ids = list(range(prompt_tokens))
self.origin_input_ids_unpadded = self.origin_input_ids
self.output_ids = list(range(output_tokens))
self.cached_tokens = 0
self.dp_rank = -1
@property
def seqlen(self) -> int:
return len(self.origin_input_ids) + len(self.output_ids)
class FakeScheduler(SchedulerMetricsMixin):
def __init__(self):
self.is_hybrid_swa = False
self.is_hybrid_ssm = False
self.waiting_queue = []
self.disaggregation_mode = DisaggregationMode.DECODE
self.disagg_decode_prealloc_queue = SimpleNamespace(
pending_reqs=[], queue=[], retracted_queue=[]
)
self.disagg_decode_transfer_queue = SimpleNamespace(queue=[])
self.running_batch = SimpleNamespace(reqs=[])
self.dp_rank = 3
def _get_token_info(self):
return (100, None, None)
class TestSchedulerLoadAccounting(CustomTestCase):
def test_decode_load_counts_pending_and_wrapped_handoff_reqs(self):
scheduler = FakeScheduler()
scheduler.running_batch.reqs = [FakeReq("running", 11)]
scheduler.waiting_queue = [FakeReq("waiting", 7)]
scheduler.disagg_decode_prealloc_queue.pending_reqs = [
FakeReq("pending", 101)
]
scheduler.disagg_decode_prealloc_queue.queue = [
DecodeRequest(req=FakeReq("prealloc", 103), kv_receiver=None)
]
scheduler.disagg_decode_transfer_queue.queue = [
DecodeRequest(req=FakeReq("transfer", 107), kv_receiver=None)
]
scheduler.disagg_decode_prealloc_queue.retracted_queue = [
FakeReq("retracted", 109)
]
load = scheduler.get_load()
self.assertEqual(load.dp_rank, 3)
self.assertEqual(load.pool_tokens, 100)
self.assertEqual(load.num_waiting_reqs, 5)
self.assertEqual(load.num_reqs, 6)
self.assertEqual(load.num_tokens, 100 + 7 + 101 + 103 + 107 + 109)
self.assertEqual(
{info.rid: info.num_tokens for info in load.reqs_list},
{
"waiting": 7,
"pending": 101,
"prealloc": 103,
"transfer": 107,
"retracted": 109,
},
)
if __name__ == "__main__":
import unittest
unittest.main()