"""Unit tests for prefill bootstrap queue compaction.""" import unittest from types import SimpleNamespace from unittest.mock import MagicMock, patch import torch import torch.distributed as dist from sglang.srt.disaggregation.base import KVPoll from sglang.srt.disaggregation.prefill import PrefillBootstrapQueue from sglang.srt.disaggregation.utils import poll_and_all_reduce_attn_cp_tp_group 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") def poll(self): return KVPoll.Success 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)]) def test_poll_consensus_debug_fails_before_shape_mismatch_hang(self): reduce_calls = [] def fake_all_reduce(tensor, op, group): reduce_calls.append((tensor.dtype, int(tensor.numel()), op, group)) # The new debug guard uses an int64 [queue_len, queue_hash] scalar # vector before building the uint8 per-request poll tensor. Simulate # another rank having a different queue length and assert that we # fail fast before reaching the old variable-length uint8 all_reduce. if tensor.dtype == torch.int64 and int(tensor.numel()) == 2: if op == dist.ReduceOp.MIN: tensor[0] = 1 elif op == dist.ReduceOp.MAX: tensor[0] = 2 return raise AssertionError( "poll tensor all_reduce should not run after queue mismatch" ) with ( patch( "sglang.srt.disaggregation.utils.envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.get", return_value=True, ), patch("sglang.srt.disaggregation.utils.dist.all_reduce", fake_all_reduce), ): with self.assertRaisesRegex(RuntimeError, "poll_queue.*inflight"): poll_and_all_reduce_attn_cp_tp_group( [FakeSender(), FakeSender()], MagicMock(name="cp_group"), MagicMock(name="tp_group"), debug_label="inflight", debug_ids=["rid-a", "rid-b"], ) self.assertTrue(reduce_calls) def test_poll_consensus_debug_disabled_preserves_old_collective_shape(self): reduce_calls = [] def fake_all_reduce(tensor, op, group): reduce_calls.append((tensor.dtype, int(tensor.numel()), op, group)) with ( patch( "sglang.srt.disaggregation.utils.envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.get", return_value=False, ), patch("sglang.srt.disaggregation.utils.dist.all_reduce", fake_all_reduce), ): polls = poll_and_all_reduce_attn_cp_tp_group( [FakeSender(), FakeSender()], MagicMock(name="cp_group"), MagicMock(name="tp_group"), debug_label="inflight", debug_ids=["rid-a", "rid-b"], ) self.assertEqual(polls, [KVPoll.Success, KVPoll.Success]) self.assertEqual( [(dtype, size) for dtype, size, _op, _group in reduce_calls], [(torch.uint8, 2), (torch.uint8, 2)], ) if __name__ == "__main__": unittest.main()