130 lines
4.9 KiB
Python
130 lines
4.9 KiB
Python
"""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()
|