[Test] Extract common PD server setup into base fixture (#22080)
This commit is contained in:
@@ -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 = [
|
||||
|
||||
Reference in New Issue
Block a user