From 95cdbce34fa983f2761788bdde72520952db2391 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Fri, 3 Apr 2026 16:37:12 -0700 Subject: [PATCH] [Test] Extract common PD server setup into base fixture (#22080) --- .../server_fixtures/disaggregation_fixture.py | 54 ++++ .../test_disaggregation_basic.py | 265 +----------------- 2 files changed, 64 insertions(+), 255 deletions(-) diff --git a/python/sglang/test/server_fixtures/disaggregation_fixture.py b/python/sglang/test/server_fixtures/disaggregation_fixture.py index eda4004e9..dce859fae 100644 --- a/python/sglang/test/server_fixtures/disaggregation_fixture.py +++ b/python/sglang/test/server_fixtures/disaggregation_fixture.py @@ -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 = [ diff --git a/test/registered/disaggregation/test_disaggregation_basic.py b/test/registered/disaggregation/test_disaggregation_basic.py index a8056f9ad..a598ae438 100644 --- a/test/registered/disaggregation/test_disaggregation_basic.py +++ b/test/registered/disaggregation/test_disaggregation_basic.py @@ -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