minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Test] Extract common PD server setup into base fixture (#22080)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-04-03 16:37:12 -07:00
committed by GitHub
parent 9593d434c4
commit 95cdbce34f
2 changed files with 64 additions and 255 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/test/server_fixtures/disaggregation_fixture.py
+54
View File
@@ -12,6 +12,7 @@ from sglang.test.test_utils import (
DEFAULT_URL_FOR_TEST,
CustomTestCase,
is_in_ci,
popen_launch_pd_server,
popen_with_error_check,
)
from sglang.utils import wait_for_http_ready
@@ -56,6 +57,59 @@ class PDDisaggregationServerBase(CustomTestCase):
msg = "No RDMA devices specified for disaggregation test, using default settings."
warnings.warn(msg)
# Subclasses can set these to customize server args
extra_prefill_args = []
extra_decode_args = []
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--disaggregation-bootstrap-port",
cls.bootstrap_port,
"--tp",
"1",
] + list(cls.extra_prefill_args)
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--disaggregation-bootstrap-port",
cls.bootstrap_port,
"--tp",
"1",
"--base-gpu-id",
"1",
] + list(cls.extra_decode_args)
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
@classmethod
def launch_all(cls):
"""Start prefill, decode, wait for health, and launch LB."""
cls.start_prefill()
cls.start_decode()
cls.wait_server_ready(cls.prefill_url + "/health", process=cls.process_prefill)
cls.wait_server_ready(cls.decode_url + "/health", process=cls.process_decode)
cls.launch_lb()
@classmethod
def launch_lb(cls):
lb_command = [
test/registered/disaggregation/test_disaggregation_basic.py
+10 -255
View File
@@ -18,8 +18,6 @@ from sglang.test.test_utils import (
DEFAULT_DRAFT_MODEL_EAGLE3,
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_TARGET_MODEL_EAGLE3,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
popen_launch_pd_server,
)
register_cuda_ci(est_time=400, suite="stage-b-test-2-gpu-large")
@@ -30,56 +28,7 @@ class TestDisaggregationAccuracy(PDDisaggregationServerBase):
def setUpClass(cls):
super().setUpClass()
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
# Non blocking start servers
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health", process=cls.process_prefill)
cls.wait_server_ready(cls.decode_url + "/health", process=cls.process_decode)
cls.launch_lb()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--disaggregation-bootstrap-port",
cls.bootstrap_port,
"--tp",
"1",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--disaggregation-bootstrap-port",
cls.bootstrap_port,
"--tp",
"1",
"--base-gpu-id",
"1",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
cls.launch_all()
def test_gsm8k(self):
args = SimpleNamespace(
@@ -201,64 +150,14 @@ class TestDisaggregationMooncakeFailure(PDDisaggregationServerBase):
super().setUpClass()
# set DISAGGREGATION_TEST_FAILURE_PROB to simulate failure
os.environ["DISAGGREGATION_TEST_FAILURE_PROB"] = "0.05"
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
# Non blocking start servers
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health", process=cls.process_prefill)
cls.wait_server_ready(cls.decode_url + "/health", process=cls.process_decode)
cls.launch_lb()
cls.launch_all()
@classmethod
def tearDownClass(cls):
os.environ.pop("DISAGGREGATION_TEST_FAILURE_PROB")
super().tearDownClass()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--disaggregation-bootstrap-port",
cls.bootstrap_port,
"--tp",
"1",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--disaggregation-bootstrap-port",
cls.bootstrap_port,
"--tp",
"1",
"--base-gpu-id",
"1",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_gsm8k(self):
args = SimpleNamespace(
base_url=f"http://{self.base_host}:{self.lb_port}",
@@ -287,17 +186,15 @@ class TestDisaggregationMooncakeFailure(PDDisaggregationServerBase):
class TestDisaggregationMooncakeSpec(PDDisaggregationServerBase):
@classmethod
def setUpClass(cls):
super().setUpClass()
cls.model = DEFAULT_TARGET_MODEL_EAGLE3
cls.draft_model = DEFAULT_DRAFT_MODEL_EAGLE3
cls.spec_args = [
spec_args = [
"--speculative-algorithm",
"EAGLE",
"--speculative-draft-model-path",
cls.draft_model,
DEFAULT_DRAFT_MODEL_EAGLE3,
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
@@ -308,57 +205,9 @@ class TestDisaggregationMooncakeSpec(PDDisaggregationServerBase):
"8",
"--dtype=float16",
]
print(f"{cls.base_host=} {cls.lb_port=} {cls.prefill_port=} {cls.decode_port=}")
# Non blocking start servers
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health", process=cls.process_prefill)
cls.wait_server_ready(cls.decode_url + "/health", process=cls.process_decode)
cls.launch_lb()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--disaggregation-bootstrap-port",
cls.bootstrap_port,
"--tp",
"1",
] + cls.spec_args
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--disaggregation-bootstrap-port",
cls.bootstrap_port,
"--tp",
"1",
"--base-gpu-id",
"1",
] + cls.spec_args
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
cls.extra_prefill_args = spec_args
cls.extra_decode_args = spec_args
cls.launch_all()
def test_gsm8k(self):
args = SimpleNamespace(
@@ -381,62 +230,13 @@ class TestDisaggregationSimulatedRetract(PDDisaggregationServerBase):
super().setUpClass()
os.environ["SGLANG_TEST_RETRACT"] = "true"
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
# Non blocking start servers
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health", process=cls.process_prefill)
cls.wait_server_ready(cls.decode_url + "/health", process=cls.process_decode)
cls.launch_lb()
cls.launch_all()
@classmethod
def tearDownClass(cls):
os.environ.pop("SGLANG_TEST_RETRACT")
super().tearDownClass()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--disaggregation-bootstrap-port",
cls.bootstrap_port,
"--tp",
"1",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--disaggregation-bootstrap-port",
cls.bootstrap_port,
"--tp",
"1",
"--base-gpu-id",
"1",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_gsm8k(self):
args = SimpleNamespace(
base_url=f"http://{self.base_host}:{self.lb_port}",
@@ -463,57 +263,12 @@ class TestDisaggregationPauseResumePrefillLeak(PDDisaggregationServerBase):
def setUpClass(cls):
super().setUpClass()
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
cls.start_prefill()
cls.start_decode()
cls.wait_server_ready(cls.prefill_url + "/health", process=cls.process_prefill)
cls.wait_server_ready(cls.decode_url + "/health", process=cls.process_decode)
cls.launch_lb()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--disaggregation-bootstrap-port",
cls.bootstrap_port,
"--tp",
"1",
cls.extra_prefill_args = [
"--max-running-requests",
str(cls.MAX_RUNNING),
"--enable-metrics",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--disaggregation-bootstrap-port",
cls.bootstrap_port,
"--tp",
"1",
"--base-gpu-id",
"1",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
cls.launch_all()
def test_retract_pause_no_leak_on_prefill(self):
"""Retract-mode pause on a disagg prefill node must not leak prefill