diff --git a/python/sglang/srt/observability/scheduler_metrics_mixin.py b/python/sglang/srt/observability/scheduler_metrics_mixin.py index dac256b1a..b161d9b51 100644 --- a/python/sglang/srt/observability/scheduler_metrics_mixin.py +++ b/python/sglang/srt/observability/scheduler_metrics_mixin.py @@ -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( diff --git a/test/registered/unit/observability/test_scheduler_metrics_load.py b/test/registered/unit/observability/test_scheduler_metrics_load.py new file mode 100644 index 000000000..8dd830d8c --- /dev/null +++ b/test/registered/unit/observability/test_scheduler_metrics_load.py @@ -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()