perf(disaggregation): compact prefill bootstrap queue in one pass

2026-04-19 04:23:48 +08:00
parent 09acec1908
commit 9470d6d8dc
2 changed files with 135 additions and 7 deletions
@@ -0,0 +1,129 @@
+"""Unit tests for prefill bootstrap queue compaction."""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+from sglang.srt.disaggregation.base import KVPoll
+from sglang.srt.disaggregation.prefill import PrefillBootstrapQueue
+from sglang.srt.disaggregation.utils import ReqToMetadataIdxAllocator
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=6, suite="stage-a-test-cpu")
+
+
+class FakeSender:
+    def __init__(self, should_fail=False):
+        self.should_fail = should_fail
+        self.init_calls = []
+
+    def init(self, num_pages, metadata_buffer_index):
+        self.init_calls.append((num_pages, metadata_buffer_index))
+
+    def failure_exception(self):
+        if self.should_fail:
+            raise RuntimeError("boom")
+
+
+class TestPrefillBootstrapQueue(CustomTestCase):
+    def _make_req(self, rid, bootstrap_room, origin_input_ids, sender):
+        req = MagicMock()
+        req.rid = rid
+        req.bootstrap_room = bootstrap_room
+        req.origin_input_ids = origin_input_ids
+        req.disagg_kv_sender = sender
+        req.return_logprob = False
+        req.metadata_buffer_index = -1
+        req.time_stats = SimpleNamespace(
+            set_bootstrap_done_time=MagicMock(),
+            set_wait_queue_entry_time=MagicMock(),
+            trace_ctx=SimpleNamespace(abort=MagicMock()),
+        )
+        return req
+
+    def _build_queue(self, allocator_size=4):
+        queue = PrefillBootstrapQueue.__new__(PrefillBootstrapQueue)
+        queue.req_to_metadata_buffer_idx_allocator = ReqToMetadataIdxAllocator(
+            size=allocator_size
+        )
+        queue.token_to_kv_pool = SimpleNamespace(page_size=4)
+        queue.scheduler = SimpleNamespace(
+            stream_output=MagicMock(),
+            enable_metrics=False,
+            enable_hicache_storage=False,
+            attn_cp_cpu_group=MagicMock(),
+            attn_tp_cpu_group=MagicMock(),
+        )
+        queue.tp_rank = 0
+        queue.queue = []
+        return queue
+
+    def test_pop_bootstrapped_compacts_queue_in_one_pass(self):
+        queue = self._build_queue(allocator_size=2)
+
+        waiting = self._make_req("wait", 1, list(range(8)), FakeSender())
+        ready = self._make_req("ready", 2, list(range(10)), FakeSender())
+        failed = self._make_req("failed", 3, list(range(6)), FakeSender(True))
+        queue.queue = [waiting, ready, failed]
+
+        with patch(
+            "sglang.srt.disaggregation.prefill.poll_and_all_reduce_attn_cp_tp_group",
+            return_value=[
+                KVPoll.Bootstrapping,
+                KVPoll.WaitingForInput,
+                KVPoll.Failed,
+            ],
+        ):
+            bootstrapped, failed_reqs = queue.pop_bootstrapped(return_failed_reqs=True)
+
+        self.assertEqual(bootstrapped, [ready])
+        self.assertEqual(failed_reqs, [failed])
+        self.assertEqual(queue.queue, [waiting])
+        self.assertEqual(ready.metadata_buffer_index, 0)
+        self.assertEqual(ready.disagg_kv_sender.init_calls, [(3, 0)])
+
+    def test_pop_bootstrapped_stops_on_metadata_exhaustion(self):
+        queue = self._build_queue(allocator_size=1)
+
+        first_ready = self._make_req("ready-1", 1, list(range(8)), FakeSender())
+        second_ready = self._make_req("ready-2", 2, list(range(12)), FakeSender())
+        trailing = self._make_req("trailing", 3, list(range(4)), FakeSender())
+        queue.queue = [first_ready, second_ready, trailing]
+
+        with patch(
+            "sglang.srt.disaggregation.prefill.poll_and_all_reduce_attn_cp_tp_group",
+            return_value=[
+                KVPoll.WaitingForInput,
+                KVPoll.WaitingForInput,
+                KVPoll.Bootstrapping,
+            ],
+        ):
+            bootstrapped = queue.pop_bootstrapped()
+
+        self.assertEqual(bootstrapped, [first_ready])
+        self.assertEqual(queue.queue, [second_ready, trailing])
+        self.assertEqual(first_ready.disagg_kv_sender.init_calls, [(2, 0)])
+        self.assertEqual(second_ready.disagg_kv_sender.init_calls, [])
+
+    def test_pop_bootstrapped_preserves_unchecked_entries(self):
+        queue = self._build_queue(allocator_size=2)
+
+        skipped = self._make_req("skip", 1, list(range(8)), FakeSender())
+        checked = self._make_req("check", 2, list(range(10)), FakeSender())
+        queue.queue = [skipped, checked]
+
+        with patch(
+            "sglang.srt.disaggregation.prefill.poll_and_all_reduce_attn_cp_tp_group",
+            return_value=[KVPoll.WaitingForInput, KVPoll.WaitingForInput],
+        ):
+            bootstrapped = queue.pop_bootstrapped(rids_to_check=["check"])
+
+        self.assertEqual(bootstrapped, [checked])
+        self.assertEqual(queue.queue, [skipped])
+        self.assertEqual(skipped.disagg_kv_sender.init_calls, [])
+        self.assertEqual(checked.disagg_kv_sender.init_calls, [(3, 0)])
+
+
+if __name__ == "__main__":
+    unittest.main()