laoyao0822
ceb5345410
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.
90 lines
3.1 KiB
Python
90 lines
3.1 KiB
Python
"""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()
|