313 lines
11 KiB
Python
313 lines
11 KiB
Python
import unittest
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from sglang.srt.managers.schedule_batch import Req
|
|
from sglang.srt.managers.schedule_policy import PrefillAdder
|
|
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
|
from sglang.test.test_utils import CustomTestCase
|
|
|
|
register_cuda_ci(est_time=1, suite="stage-b-test-small-1-gpu")
|
|
register_amd_ci(est_time=2, suite="stage-b-test-small-1-gpu-amd")
|
|
|
|
|
|
class TestPrefillAdder(CustomTestCase):
|
|
def setUp(self):
|
|
self.mock_tree_cache = self.create_tree_cache()
|
|
self.mock_token_allocator = self.create_token_allocator()
|
|
patcher = patch(
|
|
"sglang.srt.managers.schedule_policy.is_nsa_prefill_cp_in_seq_split",
|
|
return_value=False,
|
|
)
|
|
self.mock_is_nsa = patcher.start()
|
|
self.addCleanup(patcher.stop)
|
|
|
|
def create_tree_cache(
|
|
self,
|
|
*,
|
|
full_evictable_size: int = 0,
|
|
swa_evictable_size: int = 0,
|
|
evictable_size: int = 0,
|
|
) -> MagicMock:
|
|
tree_cache = MagicMock()
|
|
tree_cache.full_evictable_size.return_value = full_evictable_size
|
|
tree_cache.swa_evictable_size.return_value = swa_evictable_size
|
|
tree_cache.evictable_size.return_value = evictable_size
|
|
tree_cache.disable = False
|
|
tree_cache.inc_lock_ref.return_value = None
|
|
tree_cache.dec_lock_ref.return_value = None
|
|
return tree_cache
|
|
|
|
def create_token_allocator(
|
|
self,
|
|
*,
|
|
full_available_size: int = 0,
|
|
swa_available_size: int = 0,
|
|
available_size: int = 0,
|
|
) -> MagicMock:
|
|
allocator = MagicMock()
|
|
allocator.full_available_size.return_value = full_available_size
|
|
allocator.swa_available_size.return_value = swa_available_size
|
|
allocator.available_size.return_value = available_size
|
|
return allocator
|
|
|
|
def create_running_batch(self, reqs=None) -> MagicMock:
|
|
batch = MagicMock()
|
|
batch.reqs = list(reqs or [])
|
|
batch.release_req.return_value = None
|
|
batch.filter_batch.return_value = None
|
|
return batch
|
|
|
|
def create_server_args(
|
|
self, *, schedule_low_priority_values_first: bool
|
|
) -> MagicMock:
|
|
server_args = MagicMock()
|
|
server_args.schedule_low_priority_values_first = (
|
|
schedule_low_priority_values_first
|
|
)
|
|
return server_args
|
|
|
|
def create_mock_req(self, rid, priority, max_new_tokens, output_len=0, wait_time=0):
|
|
req = MagicMock(spec=Req)
|
|
req.rid = str(rid)
|
|
req.priority = priority
|
|
req.extend_input_len = 0
|
|
req.extend_logprob_start_len = 0
|
|
req.output_ids = [0] * output_len
|
|
req.sampling_params = SimpleNamespace(max_new_tokens=max_new_tokens)
|
|
req.time_stats = SimpleNamespace(wait_queue_entry_time=wait_time)
|
|
req.finished.return_value = False
|
|
return req
|
|
|
|
def create_adder(self, running_batch):
|
|
return PrefillAdder(
|
|
page_size=1,
|
|
tree_cache=self.mock_tree_cache,
|
|
token_to_kv_pool_allocator=self.mock_token_allocator,
|
|
running_batch=running_batch,
|
|
new_token_ratio=1.0,
|
|
rem_input_tokens=10000,
|
|
rem_chunk_tokens=None,
|
|
mixed_with_decode_tokens=0,
|
|
priority_scheduling_preemption_threshold=0,
|
|
)
|
|
|
|
def test_preempt_success_high_priority_values_first(self):
|
|
params = [
|
|
("run1", 0, 50),
|
|
("run2", 1, 75),
|
|
("run3", 2, 100),
|
|
]
|
|
running_reqs = [
|
|
self.create_mock_req(rid, priority, max_new_tokens)
|
|
for rid, priority, max_new_tokens in params
|
|
]
|
|
mock_server_args = self.create_server_args(
|
|
schedule_low_priority_values_first=False
|
|
)
|
|
running_batch = self.create_running_batch(running_reqs)
|
|
adder = self.create_adder(running_batch)
|
|
|
|
self.assertEqual(adder.rem_total_token_offset, 225)
|
|
|
|
self.mock_token_allocator.full_available_size.return_value = (
|
|
225 # full occupation of GRam
|
|
)
|
|
self.mock_token_allocator.available_size.return_value = 225
|
|
|
|
new_req = self.create_mock_req("new1", priority=1, max_new_tokens=49)
|
|
|
|
success = adder.preempt_to_schedule(new_req, mock_server_args)
|
|
|
|
self.assertTrue(success)
|
|
self.assertIn(running_reqs[0], adder.preempt_list)
|
|
self.assertEqual(adder.rem_total_token_offset, 175) # 50 + 75 + 100 - 50 = 175
|
|
running_batch.release_req.assert_called_once()
|
|
|
|
def test_preempt_success_low_priority_values_first(self):
|
|
params = [
|
|
("run1", 0, 50),
|
|
("run2", 1, 75),
|
|
("run3", 2, 100),
|
|
]
|
|
running_reqs = [
|
|
self.create_mock_req(rid, priority, max_new_tokens)
|
|
for rid, priority, max_new_tokens in params
|
|
]
|
|
mock_server_args = self.create_server_args(
|
|
schedule_low_priority_values_first=True
|
|
)
|
|
running_batch = self.create_running_batch(running_reqs)
|
|
adder = self.create_adder(running_batch)
|
|
|
|
self.assertEqual(adder.rem_total_token_offset, 225)
|
|
|
|
self.mock_token_allocator.full_available_size.return_value = (
|
|
225 # full occupation of GRam
|
|
)
|
|
self.mock_token_allocator.available_size.return_value = 225
|
|
|
|
new_req = self.create_mock_req("new1", priority=1, max_new_tokens=49)
|
|
|
|
success = adder.preempt_to_schedule(new_req, mock_server_args)
|
|
|
|
self.assertTrue(success)
|
|
self.assertIn(running_reqs[2], adder.preempt_list)
|
|
self.assertEqual(adder.rem_total_token_offset, 125) # 50 + 75 + 100 - 100 = 125
|
|
running_batch.release_req.assert_called_once()
|
|
|
|
def test_preempt_fail_low_priority_values_first(self):
|
|
params = [
|
|
("run1", 0, 50),
|
|
("run2", 1, 75),
|
|
("run3", 2, 100),
|
|
]
|
|
running_reqs = [
|
|
self.create_mock_req(rid, priority, max_new_tokens)
|
|
for rid, priority, max_new_tokens in params
|
|
]
|
|
mock_server_args = self.create_server_args(
|
|
schedule_low_priority_values_first=True
|
|
)
|
|
running_batch = self.create_running_batch(running_reqs)
|
|
adder = self.create_adder(running_batch)
|
|
|
|
self.assertEqual(adder.rem_total_token_offset, 225)
|
|
|
|
self.mock_token_allocator.full_available_size.return_value = (
|
|
225 # full occupation of GRam
|
|
)
|
|
self.mock_token_allocator.available_size.return_value = 225
|
|
|
|
new_req_fail_by_priority_check = self.create_mock_req(
|
|
"new1", priority=2, max_new_tokens=49
|
|
)
|
|
|
|
success_by_priority_check = adder.preempt_to_schedule(
|
|
new_req_fail_by_priority_check, mock_server_args
|
|
)
|
|
self.assertFalse(success_by_priority_check)
|
|
|
|
new_req_fail_by_priority_check = self.create_mock_req(
|
|
"new2", priority=1, max_new_tokens=110
|
|
)
|
|
success_by_capacity_check = adder.preempt_to_schedule(
|
|
new_req_fail_by_priority_check, mock_server_args
|
|
)
|
|
self.assertFalse(success_by_capacity_check)
|
|
|
|
def test_preempt_fail_high_priority_values_first(self):
|
|
params = [
|
|
("run1", 0, 50),
|
|
("run2", 1, 75),
|
|
("run3", 2, 100),
|
|
]
|
|
running_reqs = [
|
|
self.create_mock_req(rid, priority, max_new_tokens)
|
|
for rid, priority, max_new_tokens in params
|
|
]
|
|
mock_server_args = self.create_server_args(
|
|
schedule_low_priority_values_first=False
|
|
)
|
|
running_batch = self.create_running_batch(running_reqs)
|
|
adder = self.create_adder(running_batch)
|
|
|
|
self.assertEqual(adder.rem_total_token_offset, 225)
|
|
|
|
self.mock_token_allocator.full_available_size.return_value = (
|
|
225 # full occupation of GRam
|
|
)
|
|
self.mock_token_allocator.available_size.return_value = 225
|
|
|
|
new_req_fail_by_priority_check = self.create_mock_req(
|
|
"new1", priority=0, max_new_tokens=49
|
|
)
|
|
|
|
success_by_priority_check = adder.preempt_to_schedule(
|
|
new_req_fail_by_priority_check, mock_server_args
|
|
)
|
|
self.assertFalse(success_by_priority_check)
|
|
|
|
new_req_fail_by_priority_check = self.create_mock_req(
|
|
"new2", priority=-1, max_new_tokens=110
|
|
)
|
|
success_by_capacity_check = adder.preempt_to_schedule(
|
|
new_req_fail_by_priority_check, mock_server_args
|
|
)
|
|
self.assertFalse(success_by_capacity_check)
|
|
|
|
def test_preempt_success_low_priority_values_first_exact_once(self):
|
|
params = [
|
|
("run1", 0, 50),
|
|
("run2", 1, 75),
|
|
("run3", 2, 100),
|
|
("run4", 2, 125),
|
|
("run4", 2, 125),
|
|
]
|
|
running_reqs = [
|
|
self.create_mock_req(rid, priority, max_new_tokens)
|
|
for rid, priority, max_new_tokens in params
|
|
]
|
|
mock_server_args = self.create_server_args(
|
|
schedule_low_priority_values_first=True
|
|
)
|
|
running_batch = self.create_running_batch(running_reqs)
|
|
adder = self.create_adder(running_batch)
|
|
|
|
self.assertEqual(adder.rem_total_token_offset, 475)
|
|
|
|
self.mock_token_allocator.full_available_size.return_value = (
|
|
475 # full occupation of GRam
|
|
)
|
|
self.mock_token_allocator.available_size.return_value = 475
|
|
|
|
new_req = self.create_mock_req("new1", priority=1, max_new_tokens=75)
|
|
|
|
success = adder.preempt_to_schedule(new_req, mock_server_args)
|
|
self.assertTrue(success)
|
|
self.assertIn(running_reqs[2], adder.preempt_list)
|
|
self.assertEqual(
|
|
adder.rem_total_token_offset, 375
|
|
) # 50 + 75 + 100 + 125 + 125 - 100 = 375
|
|
running_batch.release_req.assert_called_once()
|
|
|
|
def test_preempt_success_low_priority_values_first_exact_twice(self):
|
|
params = [
|
|
("run1", 0, 50),
|
|
("run2", 1, 75),
|
|
("run3", 2, 100),
|
|
("run4", 2, 125),
|
|
("run4", 2, 125),
|
|
]
|
|
running_reqs = [
|
|
self.create_mock_req(rid, priority, max_new_tokens)
|
|
for rid, priority, max_new_tokens in params
|
|
]
|
|
mock_server_args = self.create_server_args(
|
|
schedule_low_priority_values_first=True
|
|
)
|
|
running_batch = self.create_running_batch(running_reqs)
|
|
adder = self.create_adder(running_batch)
|
|
|
|
self.assertEqual(adder.rem_total_token_offset, 475)
|
|
|
|
self.mock_token_allocator.full_available_size.return_value = (
|
|
475 # full occupation of GRam
|
|
)
|
|
self.mock_token_allocator.available_size.return_value = 475
|
|
|
|
new_req = self.create_mock_req("new1", priority=1, max_new_tokens=200)
|
|
|
|
success = adder.preempt_to_schedule(new_req, mock_server_args)
|
|
self.assertTrue(success)
|
|
self.assertIn(running_reqs[2], adder.preempt_list)
|
|
self.assertIn(running_reqs[3], adder.preempt_list)
|
|
self.assertEqual(
|
|
adder.rem_total_token_offset, 250
|
|
) # 50 + 75 + 100 + 125 + 125 - 100 - 125 = 250
|
|
self.assertEqual(running_batch.release_req.call_count, 2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|