ci: migrate scheduler tests to test/registered/scheduler/ (#16442)

This commit is contained in:
Alison Shao
2026-01-06 14:55:52 -08:00
committed by GitHub
parent 2e0527dd74
commit 399d5283f8
8 changed files with 27 additions and 13 deletions

View File

@@ -9,8 +9,6 @@ from sglang.test.ci.ci_utils import TestFile, run_unittest_files
# NOTE: please sort the test cases alphabetically by the test file name
suites = {
"per-commit-1-gpu": [
TestFile("test_abort.py", 131),
TestFile("test_chunked_prefill.py", 312),
TestFile("test_deterministic.py", 228),
TestFile("test_eval_fp8_accuracy.py", 250),
TestFile("test_evs.py", 20),
@@ -25,13 +23,8 @@ suites = {
TestFile("test_model_hooks.py", 6),
TestFile("test_modelopt_loader.py", 11),
TestFile("test_multi_tokenizer.py", 230),
TestFile("test_no_chunked_prefill.py", 108),
TestFile("test_no_overlap_scheduler.py", 217),
TestFile("test_page_size.py", 60),
TestFile("test_prefill_adder.py", 1),
TestFile("test_priority_scheduling.py", 130),
TestFile("test_request_queue_validation.py", 47),
TestFile("test_retract_decode.py", 259),
TestFile("test_score_api.py", 260),
TestFile("test_server_args.py", 9),
TestFile("test_skip_tokenizer_init.py", 77),
@@ -134,9 +127,7 @@ suite_amd = {
# TestFile("lora/test_lora_backend.py", 99), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/13107
# TestFile("lora/test_lora_cuda_graph.py", 250), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/13107
# TestFile("lora/test_lora_qwen3.py", 97), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/13107
TestFile("test_abort.py", 51),
TestFile("test_bench_typebaseddispatcher.py", 10),
TestFile("test_chunked_prefill.py", 312),
TestFile("test_eval_fp8_accuracy.py", 303),
TestFile("test_external_models.py", 45),
TestFile("test_input_embeddings.py", 38),
@@ -144,14 +135,10 @@ suite_amd = {
TestFile("test_jinja_template_utils.py", 1),
TestFile("test_model_hooks.py", 10),
TestFile("test_multi_tokenizer.py", 345),
TestFile("test_no_chunked_prefill.py", 108),
TestFile("test_page_size.py", 60),
TestFile("test_prefill_adder.py", 2),
TestFile("test_priority_scheduling.py", 195),
TestFile("test_profile_merger.py", 12),
TestFile("test_profile_merger_http_api.py", 15),
TestFile("test_request_queue_validation.py", 70),
TestFile("test_retract_decode.py", 450),
TestFile("test_rope_rocm.py", 3),
TestFile("test_server_args.py", 1),
TestFile("test_skip_tokenizer_init.py", 117),

View File

@@ -1,191 +0,0 @@
import multiprocessing
import time
import unittest
from concurrent.futures import ThreadPoolExecutor, as_completed
import requests
from sglang.srt.environ import envs
from sglang.srt.utils import kill_process_tree
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
run_and_check_memory_leak,
)
class TestAbort(CustomTestCase):
def workload_func(self, base_url, model):
def process_func():
def run_one(_):
prompt = """
System: You are a helpful assistant.
User: What is the capital of France?
Assistant: The capital of France is
"""
response = requests.post(
f"{base_url}/generate",
json={
"text": prompt,
"sampling_params": {
"temperature": 0,
"max_new_tokens": 2048,
},
},
)
ret = response.json()
with ThreadPoolExecutor(16) as executor:
list(executor.map(run_one, list(range(16))))
p = multiprocessing.Process(target=process_func)
p.start()
time.sleep(0.5)
p.terminate()
time.sleep(10)
def test_memory_leak(self):
run_and_check_memory_leak(
self.workload_func,
disable_radix_cache=False,
enable_mixed_chunk=False,
disable_overlap=False,
chunked_prefill_size=8192,
assert_has_abort=True,
)
class TestAbortAll(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=["--max-running-requests", 8],
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def _run_decode(self):
response = requests.post(
self.base_url + "/generate",
json={
"text": "The capital of France is",
"sampling_params": {
"temperature": 0,
"max_new_tokens": 16000,
"ignore_eos": True,
},
},
)
return response.json()
def test_abort_all(self):
num_requests = 32
with ThreadPoolExecutor(num_requests) as executor:
futures = [executor.submit(self._run_decode) for _ in range(num_requests)]
# ensure the decode has been started
time.sleep(2)
requests.post(
self.base_url + "/abort_request",
json={
"abort_all": True,
},
)
for future in as_completed(futures):
self.assertEqual(
future.result()["meta_info"]["finish_reason"]["type"], "abort"
)
class TestAbortAllWithRetraction(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
# Here's a small trick: in scheduler.py, when SGLANG_TEST_RETRACT is enabled,
# retraction is triggered when the batch size reaches 10.
# However, since SGLANG_TEST_RETRACT_NO_PREFILL_BS is set to 6, the remaining 4
# requests will stay in the waiting queue.
with (
envs.SGLANG_TEST_RETRACT.override(True),
envs.SGLANG_TEST_RETRACT_NO_PREFILL_BS.override(6),
):
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--max-running-requests",
16,
"--schedule-policy",
"random",
],
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def _run_decode(self):
response = requests.post(
self.base_url + "/generate",
json={
"text": "The capital of France is",
"sampling_params": {
"temperature": 0,
"max_new_tokens": 4000,
"ignore_eos": True,
},
},
)
return response.json()
def test_abort_all_with_retraction(self):
num_requests = 32
with ThreadPoolExecutor(num_requests) as executor:
futures = [executor.submit(self._run_decode) for _ in range(num_requests)]
# ensure the decode has been started and retractions happen.
time.sleep(8)
requests.post(
self.base_url + "/abort_request",
json={
"abort_all": True,
},
)
abort_in_queue_count = 0
abort_in_queue_with_none_empty_text = 0
for future in as_completed(futures):
self.assertEqual(
future.result()["meta_info"]["finish_reason"]["type"], "abort"
)
if (
future.result()["meta_info"]["finish_reason"]["message"]
== "Abort in waiting queue"
):
abort_in_queue_count += 1
if len(future.result()["output_ids"]) > 0:
abort_in_queue_with_none_empty_text += 1
assert abort_in_queue_count > 0
assert abort_in_queue_with_none_empty_text > 0
print("Finished test_abort_all_with_retraction")
if __name__ == "__main__":
unittest.main()

View File

@@ -1,31 +0,0 @@
"""
python3 -m unittest test_chunked_prefill.TestChunkedPrefill.test_mixed_chunked_prefill_without_radix_cache
"""
import unittest
from sglang.test.test_utils import CustomTestCase, run_mmlu_test, run_mulit_request_test
class TestChunkedPrefill(CustomTestCase):
def test_chunked_prefill(self):
run_mmlu_test(disable_radix_cache=False, enable_mixed_chunk=False)
def test_mixed_chunked_prefill(self):
run_mmlu_test(disable_radix_cache=False, enable_mixed_chunk=True)
def test_chunked_prefill_without_radix_cache(self):
run_mmlu_test(disable_radix_cache=True, enable_mixed_chunk=False)
def test_mixed_chunked_prefill_without_radix_cache(self):
run_mmlu_test(disable_radix_cache=True, enable_mixed_chunk=True)
def test_mixed_chunked_prefill_multi_requests(self):
run_mulit_request_test(
enable_mixed_chunk=True,
chunked_prefill_size=2048,
)
if __name__ == "__main__":
unittest.main()

View File

@@ -1,30 +0,0 @@
import unittest
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
CustomTestCase,
run_bench_serving,
run_mmlu_test,
)
class TestNoChunkedPrefill(CustomTestCase):
def test_no_chunked_prefill(self):
run_mmlu_test(
disable_radix_cache=False, enable_mixed_chunk=False, chunked_prefill_size=-1
)
def test_no_chunked_prefill_without_radix_cache(self):
res = run_bench_serving(
model=DEFAULT_MODEL_NAME_FOR_TEST,
num_prompts=10,
request_rate=float("inf"),
other_server_args=["--disable-radix-cache", "--chunked-prefill-size", "-1"],
)
assert res["completed"] == 10
if __name__ == "__main__":
unittest.main()

View File

@@ -1,35 +0,0 @@
"""
Usage:
python3 -m unittest test_overlap_schedule.TestOverlapSchedule.test_radix_attention_chunked_prefill
python3 test_overlap_schedule.py
"""
import unittest
from sglang.test.test_utils import CustomTestCase, run_mmlu_test
class TestOverlapSchedule(CustomTestCase):
def test_no_radix_attention_chunked_prefill(self):
run_mmlu_test(
disable_radix_cache=True, chunked_prefill_size=32, disable_overlap=True
)
def test_no_radix_attention_no_chunked_prefill(self):
run_mmlu_test(
disable_radix_cache=True, chunked_prefill_size=-1, disable_overlap=True
)
def test_radix_attention_chunked_prefill(self):
run_mmlu_test(
disable_radix_cache=False, chunked_prefill_size=32, disable_overlap=True
)
def test_radix_attention_no_chunked_prefill(self):
run_mmlu_test(
disable_radix_cache=False, chunked_prefill_size=-1, disable_overlap=True
)
if __name__ == "__main__":
unittest.main()

View File

@@ -1,307 +0,0 @@
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.test_utils import CustomTestCase
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)
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()

View File

@@ -1,421 +0,0 @@
import asyncio
import os
import re
import unittest
from typing import Any, List, Optional, Tuple
from sglang.srt.utils import kill_process_tree
from sglang.test.test_utils import (
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
STDERR_FILENAME,
STDOUT_FILENAME,
CustomTestCase,
popen_launch_server,
send_concurrent_generate_requests_with_custom_params,
)
class TestPriorityScheduling(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
cls.stdout = open(STDOUT_FILENAME, "w")
cls.stderr = open(STDERR_FILENAME, "w")
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=(
"--max-running-requests", # Enforce max request concurrency is 1
"1",
"--max-queued-requests", # Enforce max queued request number is 3
"3",
"--enable-priority-scheduling", # Enable priority scheduling
),
return_stdout_stderr=(cls.stdout, cls.stderr),
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
_verify_max_running_requests_and_max_queued_request_validation(1, 3)
cls.stdout.close()
cls.stderr.close()
os.remove(STDOUT_FILENAME)
os.remove(STDERR_FILENAME)
def test_priority_scheduling_request_ordering_validation(self):
"""Verify pending requests are ordered by priority and received timestamp."""
responses = asyncio.run(
send_concurrent_generate_requests_with_custom_params(
self.base_url,
[
{
"priority": 0,
"sampling_params": {"max_new_tokens": 10000},
}, # starts being processed first
{"priority": 1}, # third
{"priority": 1}, # fourth
{"priority": 2}, # second
],
)
)
expected_status_and_error_messages = [
(200, None),
(200, None),
(200, None),
(200, None),
]
e2e_latencies = []
_verify_genereate_responses(
responses, expected_status_and_error_messages, e2e_latencies
)
assert e2e_latencies[0] < e2e_latencies[3] < e2e_latencies[1] < e2e_latencies[2]
def test_priority_scheduling_existing_requests_abortion_validation(self):
"""Verify lower priority requests are aborted when incoming requests have higher priority"""
responses = asyncio.run(
send_concurrent_generate_requests_with_custom_params(
self.base_url,
[
{
"priority": 1,
"sampling_params": {"max_new_tokens": 10000},
}, # starts being processed first and holds the running queue capacity
{"priority": 2}, # aborted by request 5
{"priority": 3}, # aborted by request 6
{"priority": 4}, # aborted by request 7
{"priority": 5}, # fourth
{"priority": 6}, # third
{"priority": 7}, # second
],
)
)
expected_status_and_error_messages = [
(200, None),
(503, "The request is aborted by a higher priority request."),
(503, "The request is aborted by a higher priority request."),
(503, "The request is aborted by a higher priority request."),
(200, None),
(200, None),
(200, None),
]
e2e_latencies = []
_verify_genereate_responses(
responses, expected_status_and_error_messages, e2e_latencies
)
assert e2e_latencies[0] < e2e_latencies[6] < e2e_latencies[5] < e2e_latencies[4]
def test_priority_scheduling_incoming_request_rejection_validation(self):
"""Verify incoming requests are rejected when existing requests have higher priority"""
responses = asyncio.run(
send_concurrent_generate_requests_with_custom_params(
self.base_url,
[
{
"priority": 7,
"sampling_params": {"max_new_tokens": 10000},
}, # starts being processed first and holds the running queue capacity
{"priority": 6}, # second
{"priority": 5}, # third
{"priority": 4}, # fourth
{"priority": 3}, # rejected
{"priority": 2}, # rejected
{"priority": 1}, # rejected
],
)
)
expected_status_and_error_messages = [
(200, None),
(200, None),
(200, None),
(200, None),
(503, "The request queue is full."),
(503, "The request queue is full."),
(503, "The request queue is full."),
]
e2e_latencies = []
_verify_genereate_responses(
responses, expected_status_and_error_messages, e2e_latencies
)
assert e2e_latencies[0] < e2e_latencies[1] < e2e_latencies[2] < e2e_latencies[3]
def test_priority_scheduling_preemption_meeting_threshold_validation(self):
"""Verify running requests are preempted by requests with priorities meeting the preemption threshold"""
responses = asyncio.run(
send_concurrent_generate_requests_with_custom_params(
self.base_url,
[
{
"priority": 0,
"sampling_params": {"max_new_tokens": 10000},
}, # starts being processed first then preempted or pushed by later requests, and finishes last.
{
"priority": 10,
"sampling_params": {"max_new_tokens": 10000},
}, # scheduled after the third request, and finishes second.
{
"priority": 20,
"sampling_params": {"max_new_tokens": 10000},
}, # finishes first.
],
)
)
expected_status_and_error_messages = [
(200, None),
(200, None),
(200, None),
]
e2e_latencies = []
_verify_genereate_responses(
responses, expected_status_and_error_messages, e2e_latencies
)
assert e2e_latencies[2] < e2e_latencies[1] < e2e_latencies[0]
def test_priority_scheduling_preemption_below_threshold_validation(self):
"""Verify running requests are not preempted by requests with priorities below preemption threshold"""
responses = asyncio.run(
send_concurrent_generate_requests_with_custom_params(
self.base_url,
[
{
"priority": 0,
"sampling_params": {"max_new_tokens": 10000},
},
{
"priority": 5,
"sampling_params": {"max_new_tokens": 10000},
},
],
)
)
expected_status_and_error_messages = [
(200, None),
(200, None),
]
e2e_latencies = []
_verify_genereate_responses(
responses, expected_status_and_error_messages, e2e_latencies
)
assert e2e_latencies[0] < e2e_latencies[1]
class TestPrioritySchedulingMultipleRunningRequests(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
cls.stdout = open(STDOUT_FILENAME, "w")
cls.stderr = open(STDERR_FILENAME, "w")
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=(
"--max-running-requests", # Enforce max request concurrency is 2
"2",
"--max-queued-requests", # Enforce max queued request number is 3
"3",
"--enable-priority-scheduling", # Enable priority scheduling
),
return_stdout_stderr=(cls.stdout, cls.stderr),
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
_verify_max_running_requests_and_max_queued_request_validation(2, 3)
cls.stdout.close()
cls.stderr.close()
os.remove(STDOUT_FILENAME)
os.remove(STDERR_FILENAME)
def test_priority_scheduling_with_multiple_running_requests_preemption(self):
"""Verify preempting a subset of running requests is safe."""
responses = asyncio.run(
send_concurrent_generate_requests_with_custom_params(
self.base_url,
[
{
"priority": 10,
"sampling_params": {"max_new_tokens": 10000},
}, # finishes first
{
"priority": 5,
"sampling_params": {"max_new_tokens": 10000},
}, # preempted by fourth request, then finishes third
{
"priority": 15,
"sampling_params": {"max_new_tokens": 10000},
}, # preempt the first request
],
)
)
expected_status_and_error_messages = [
(200, None),
(200, None),
(200, None),
(200, None),
]
_verify_genereate_responses(responses, expected_status_and_error_messages, [])
def test_priority_scheduling_preemption_token_offset_calculation(self):
"""
Verify correct token offset calculation during preemption.
This test specifically targets the bug where rem_total_token_offset was incorrectly
calculated using the incoming request's tokens instead of the preempted request's tokens
(related to issue #13111 and PR #13201).
THE BUG:
In schedule_policy.py line 700, the code was using:
self.rem_total_token_offset -= self._get_running_request_total_token_offset(req)
Instead of:
self.rem_total_token_offset -= self._get_running_request_total_token_offset(running_req)
WHY THIS TEST CATCHES THE BUG:
- Request 1 (preempted): 8000 tokens - This is what SHOULD be freed
- Request 3 (incoming): 1000 tokens - This is what WAS freed (bug)
- Token difference: 8000 - 1000 = 7000 tokens incorrectly accounted
With the bug, the system thinks it only freed 1000 tokens instead of 8000 tokens.
This causes incorrect memory accounting and can lead to:
1. Scheduler believes less memory is available than actually is
2. Subsequent requests (like Request 4) may fail to schedule or cause issues
3. Memory calculations become increasingly inaccurate with each preemption
The test creates a scenario where:
1. A low-priority request with many tokens (8000) starts running
2. A high-priority request with few tokens (1000) arrives and triggers preemption
3. The system must correctly free 8000 tokens from the preempted request
4. Additional requests can be scheduled only if tokens were correctly freed
5. Execution order validates priority-based scheduling works correctly
The large token difference (8x) makes the bug's impact obvious and testable.
"""
responses = asyncio.run(
send_concurrent_generate_requests_with_custom_params(
self.base_url,
[
{
"priority": 0,
"sampling_params": {"max_new_tokens": 8000},
}, # Low priority, large token count - will be preempted
{
"priority": 1,
"sampling_params": {"max_new_tokens": 5000},
}, # Medium priority, medium token count - queued initially
{
"priority": 100,
"sampling_params": {"max_new_tokens": 1000},
}, # High priority, small token count - triggers preemption
{
"priority": 50,
"sampling_params": {"max_new_tokens": 2000},
}, # Should be schedulable after correct token accounting
],
)
)
# All requests should complete successfully
# The key is that the fourth request should be schedulable because
# the system correctly freed tokens from the first (preempted) request
expected_status_and_error_messages = [
(200, None),
(200, None),
(200, None),
(200, None),
]
e2e_latencies = []
_verify_genereate_responses(
responses, expected_status_and_error_messages, e2e_latencies
)
# Verify execution order: high priority requests finish before low priority ones
# Request 3 (priority 100) should finish first
# Request 4 (priority 50) should finish second
# Request 2 (priority 1) should finish third
# Request 1 (priority 0) should finish last (after being preempted)
# FIXME(harrison lim)
# assert e2e_latencies[2] < e2e_latencies[3] < e2e_latencies[1] < e2e_latencies[0]
def _verify_genereate_responses(
responses: Tuple[int, Any, float],
expected_code_and_error_message: Tuple[int, Any],
e2e_latencies: List[Optional[float]],
):
"""
Verify generate response results are as expected based on status code and response json object content.
In addition, collects e2e latency info to verify scheduling and processing ordering.
"""
for got, expected in zip(responses, expected_code_and_error_message):
got_status, got_json = got
expected_status, expected_err_msg = expected
# Check status code is as expected
assert got_status == expected_status
# Check error message content or fields' existence based on status code
if got_status != 200:
assert got_json["object"] == "error"
assert got_json["message"] == expected_err_msg
else:
assert "object" not in got_json
assert "message" not in got_json
# Collect e2e latencies for scheduling validation
e2e_latencies.append(
got_json["meta_info"]["e2e_latency"] if got_status == 200 else None
)
def _verify_max_running_requests_and_max_queued_request_validation(
max_running_requests: int, max_queued_requests: int
):
"""Verify running request and queued request numbers based on server logs."""
rr_pattern = re.compile(r"#running-req:\s*(\d+)")
qr_pattern = re.compile(r"#queue-req:\s*(\d+)")
with open(STDERR_FILENAME) as lines:
for line in lines:
rr_match, qr_match = rr_pattern.search(line), qr_pattern.search(line)
if rr_match:
assert int(rr_match.group(1)) <= max_running_requests
if qr_match:
assert int(qr_match.group(1)) <= max_queued_requests
if __name__ == "__main__":
unittest.main()

View File

@@ -1,119 +0,0 @@
import time
import unittest
from types import SimpleNamespace
import requests
from sglang.srt.environ import envs
from sglang.srt.utils import kill_process_tree
from sglang.test.run_eval import run_eval
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
from sglang.utils import is_in_ci
class TestRetractDecode(CustomTestCase):
"""python -m unittest test_retract_decode.TestRetractDecode"""
other_args = []
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
launch_args = ["--chunked-prefill-size", "128"] + cls.other_args
with envs.SGLANG_TEST_RETRACT.override(
True
), envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override(1):
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=launch_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_mmlu(self):
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="mmlu",
num_examples=64,
num_threads=32,
)
metrics = run_eval(args)
self.assertGreaterEqual(metrics["score"], 0.65)
time.sleep(1) # wait for mem check
assert self.process.poll() is None, "Server crashed during test"
class TestRetractDecodePaged(TestRetractDecode):
"""python -m unittest test_retract_decode.TestRetractDecodePaged"""
other_args = ["--page-size", "16"]
class TestRetractDecodeChunkCache(TestRetractDecode):
"""python -m unittest test_retract_decode.TestRetractDecodeChunkCache"""
other_args = ["--disable-radix-cache"]
class TestRetractDecodeChunkCachePaged(TestRetractDecode):
"""python -m unittest test_retract_decode.TestRetractDecodeChunkCachePaged"""
other_args = ["--disable-radix-cache", "--page-size", "16"]
@unittest.skipIf(is_in_ci(), "Skipped in CI due to long runtime")
class TestRetractDecodeLongOutput(CustomTestCase):
"""python -m unittest test_retract_decode.TestRetractDecodeLongOutput"""
other_args = []
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
launch_args = [
"--chunked-prefill-size",
"128",
"--page-size",
"16",
] + cls.other_args
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=launch_args,
)
def test_long_output_retract(self):
data = {
"input_ids": [[233 + i] * 1234 for i in range(256)],
"sampling_params": {"max_new_tokens": 90000, "ignore_eos": True},
}
res = requests.post(f"{self.base_url}/generate", json=data)
assert res.status_code == 200, f"Request failed: {res.status_code}"
assert self.process.poll() is None, "Server crashed during test"
@unittest.skipIf(is_in_ci(), "Skipped in CI due to long runtime")
class TestRetractDecodeLongOutputChunkCache(TestRetractDecodeLongOutput):
"""python -m unittest test_retract_decode.TestRetractDecodeLongOutputChunkCache"""
other_args = ["--disable-radix-cache"]
if __name__ == "__main__":
unittest.main()