diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 7d97f123b..1a0a3ce2c 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -260,7 +260,7 @@ class PrefillBootstrapQueue: bootstrapped_reqs = [] failed_reqs = [] - indices_to_remove = set() + remaining_queue = [] if len(self.queue) == 0: if return_failed_reqs is False: @@ -278,9 +278,11 @@ class PrefillBootstrapQueue: if rids_to_check is not None: # if req not in reqs_info_to_check, skip if req.rid not in rids_to_check: + remaining_queue.append(req) continue if poll == KVPoll.Bootstrapping: + remaining_queue.append(req) continue elif poll == KVPoll.Failed: error_message = f"Prefill bootstrap failed for request rank={self.tp_rank} {req.rid=} {req.bootstrap_room=}" @@ -294,7 +296,6 @@ class PrefillBootstrapQueue: req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR ) self.scheduler.stream_output([req], req.return_logprob) - indices_to_remove.add(i) failed_reqs.append(req) if self.scheduler.enable_metrics: self.scheduler.metrics_collector.increment_bootstrap_failed_reqs() @@ -307,6 +308,8 @@ class PrefillBootstrapQueue: req.time_stats.set_bootstrap_done_time() num_kv_indices = len(req.origin_input_ids) if self.req_to_metadata_buffer_idx_allocator.available_size() == 0: + remaining_queue.append(req) + remaining_queue.extend(self.queue[i + 1 :]) break req.metadata_buffer_index = ( @@ -318,12 +321,9 @@ class PrefillBootstrapQueue: req.disagg_kv_sender.init(num_pages, req.metadata_buffer_index) bootstrapped_reqs.append(req) - indices_to_remove.add(i) req.time_stats.set_wait_queue_entry_time() - self.queue = [ - entry for i, entry in enumerate(self.queue) if i not in indices_to_remove - ] + self.queue = remaining_queue if return_failed_reqs is False: return bootstrapped_reqs @@ -563,7 +563,6 @@ class SchedulerDisaggregationPrefillMixin: undone_reqs: List[Req] = [] # Check .poll() for the reqs in disagg_prefill_inflight_queue. If Success, respond to the client and remove it from the queue for req, poll in zip(self.disagg_prefill_inflight_queue, polls): - if rids_to_check is not None: if req.rid not in rids_to_check: undone_reqs.append(req) diff --git a/test/registered/unit/disaggregation/test_prefill_bootstrap_queue.py b/test/registered/unit/disaggregation/test_prefill_bootstrap_queue.py new file mode 100644 index 000000000..25c2f8e0e --- /dev/null +++ b/test/registered/unit/disaggregation/test_prefill_bootstrap_queue.py @@ -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()