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

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
@@ -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 = [