256 lines
8.3 KiB
Python
256 lines
8.3 KiB
Python
"""Manual single-host Elastic EP recovery test.
|
|
|
|
Run:
|
|
|
|
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -m pytest \
|
|
test/manual/ep/test_elastic_recover.py -v -s
|
|
"""
|
|
|
|
import os
|
|
import shlex
|
|
import subprocess
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
import requests
|
|
|
|
from sglang.srt.utils import kill_process_tree
|
|
from sglang.test.server_fixtures.disaggregation_fixture import get_rdma_devices_args
|
|
from sglang.test.test_utils import (
|
|
DEFAULT_MODEL_NAME_FOR_TEST_MLA,
|
|
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
CustomTestCase,
|
|
try_cached_model,
|
|
)
|
|
from sglang.utils import wait_for_http_ready
|
|
|
|
TEST_MODEL = os.environ.get(
|
|
"SGLANG_ELASTIC_RECOVER_TEST_MODEL",
|
|
try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST_MLA),
|
|
)
|
|
EP_SIZE = 8
|
|
LOCAL_EP_SIZE = 4
|
|
DIST_INIT_ADDR = os.environ.get("SGLANG_ELASTIC_RECOVER_DIST_INIT", "127.0.0.1:25555")
|
|
PRIMARY_PORT = int(os.environ.get("SGLANG_ELASTIC_RECOVER_PRIMARY_PORT", "21000"))
|
|
JOINER_PORT = int(os.environ.get("SGLANG_ELASTIC_RECOVER_JOINER_PORT", "22000"))
|
|
RECOVER_WAIT_SECONDS = float(os.environ.get("SGLANG_ELASTIC_RECOVER_WAIT_SECONDS", "5"))
|
|
RECOVER_TIMEOUT_SECONDS = float(
|
|
os.environ.get("SGLANG_ELASTIC_RECOVER_TIMEOUT_SECONDS", "300")
|
|
)
|
|
RANDOM_SEED = int(os.environ.get("SGLANG_ELASTIC_RECOVER_RANDOM_SEED", "42"))
|
|
ib_devices = get_rdma_devices_args()
|
|
|
|
|
|
def _visible_device_ids() -> list[str]:
|
|
visible = os.environ.get("CUDA_VISIBLE_DEVICES")
|
|
if visible:
|
|
return [device.strip() for device in visible.split(",") if device.strip()]
|
|
try:
|
|
import torch
|
|
|
|
return [str(index) for index in range(torch.cuda.device_count())]
|
|
except Exception:
|
|
return []
|
|
|
|
|
|
def _server_args(node_rank: int, port: int, recover: bool = False) -> list[str]:
|
|
args = [
|
|
"sglang",
|
|
"serve",
|
|
"--model-path",
|
|
TEST_MODEL,
|
|
"--host",
|
|
"127.0.0.1",
|
|
"--port",
|
|
str(port),
|
|
"--device",
|
|
"cuda",
|
|
"--trust-remote-code",
|
|
"--tp",
|
|
str(EP_SIZE),
|
|
"--dp",
|
|
str(EP_SIZE),
|
|
"--nnodes",
|
|
"2",
|
|
"--node-rank",
|
|
str(node_rank),
|
|
"--dist-init-addr",
|
|
DIST_INIT_ADDR,
|
|
"--random-seed",
|
|
str(RANDOM_SEED),
|
|
"--enable-dp-attention",
|
|
"--enable-dp-lm-head",
|
|
"--elastic-ep-backend",
|
|
"mooncake",
|
|
"--mooncake-ib-device",
|
|
ib_devices,
|
|
"--moe-a2a-backend",
|
|
"mooncake",
|
|
"--deepep-mode",
|
|
"low_latency",
|
|
"--moe-dense-tp-size",
|
|
"1",
|
|
"--disable-custom-all-reduce",
|
|
"--enable-eplb",
|
|
"--ep-num-redundant-experts",
|
|
"72",
|
|
"--chunked-prefill-size",
|
|
"512",
|
|
"--cuda-graph-max-bs-decode",
|
|
"16",
|
|
"--mem-fraction-static",
|
|
"0.5",
|
|
]
|
|
if recover:
|
|
args.extend(["--elastic-ep-join-mode", "recover"])
|
|
extra_args = os.environ.get("SGLANG_ELASTIC_RECOVER_EXTRA_SERVER_ARGS", "")
|
|
return args + shlex.split(extra_args)
|
|
|
|
|
|
@unittest.skipUnless(
|
|
len(_visible_device_ids()) >= EP_SIZE,
|
|
"Elastic EP recovery E2E needs 8 visible GPUs.",
|
|
)
|
|
class TestElasticRecover4To4(CustomTestCase):
|
|
"""Kill one four-rank node and recover it with a fresh process group."""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.base_url = f"http://127.0.0.1:{PRIMARY_PORT}"
|
|
cls.processes: list[subprocess.Popen] = []
|
|
cls.log_files = []
|
|
cls.log_paths: dict[str, Path] = {}
|
|
visible_devices = _visible_device_ids()
|
|
|
|
cls.primary = cls._launch(
|
|
node_rank=0,
|
|
port=PRIMARY_PORT,
|
|
visible_devices=visible_devices[:LOCAL_EP_SIZE],
|
|
name="primary",
|
|
)
|
|
cls.initial_joiner = cls._launch(
|
|
node_rank=1,
|
|
port=JOINER_PORT,
|
|
visible_devices=visible_devices[LOCAL_EP_SIZE:EP_SIZE],
|
|
name="initial_joiner",
|
|
)
|
|
wait_for_http_ready(
|
|
f"{cls.base_url}/health_generate",
|
|
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
process=cls.primary,
|
|
)
|
|
|
|
@classmethod
|
|
def _launch(
|
|
cls,
|
|
*,
|
|
node_rank: int,
|
|
port: int,
|
|
visible_devices: list[str],
|
|
name: str,
|
|
recover: bool = False,
|
|
) -> subprocess.Popen:
|
|
log_path = Path(f"/tmp/elastic_ep_recover_{name}_{int(time.time())}.log")
|
|
log_file = open(log_path, "w")
|
|
env = os.environ.copy()
|
|
env["CUDA_VISIBLE_DEVICES"] = ",".join(visible_devices)
|
|
process = subprocess.Popen(
|
|
_server_args(node_rank, port, recover),
|
|
env=env,
|
|
stdout=log_file,
|
|
stderr=subprocess.STDOUT,
|
|
)
|
|
cls.processes.append(process)
|
|
cls.log_files.append(log_file)
|
|
cls.log_paths[name] = log_path
|
|
print(f"Started {name}; log: {log_path}")
|
|
return process
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
for process in reversed(getattr(cls, "processes", [])):
|
|
if process.poll() is None:
|
|
kill_process_tree(process.pid, wait_timeout=60)
|
|
for log_file in getattr(cls, "log_files", []):
|
|
log_file.close()
|
|
|
|
def _generate(self, routed_dp_rank: int | None = None) -> requests.Response:
|
|
payload = {
|
|
"text": "The capital of France is",
|
|
"sampling_params": {"max_new_tokens": 4, "temperature": 0.0},
|
|
}
|
|
if routed_dp_rank is not None:
|
|
payload["routed_dp_rank"] = routed_dp_rank
|
|
return requests.post(f"{self.base_url}/generate", json=payload, timeout=90)
|
|
|
|
def _generate_ok(self, description: str, routed_dp_rank: int | None = None) -> None:
|
|
response = self._generate(routed_dp_rank)
|
|
self.assertEqual(response.status_code, 200, f"{description}: {response.text}")
|
|
payload = response.json()
|
|
generated_text = payload.get("text", "")
|
|
self.assertIn(
|
|
"paris",
|
|
generated_text.casefold(),
|
|
f"{description}: unexpected generation: {generated_text!r}",
|
|
)
|
|
|
|
def _wait_for_recover_capture(self) -> None:
|
|
deadline = time.monotonic() + RECOVER_TIMEOUT_SECONDS
|
|
log_path = self.log_paths["recover_joiner"]
|
|
marker = "Capture target decode CUDA graph end"
|
|
while time.monotonic() < deadline:
|
|
self.assertIsNone(
|
|
self.recover_joiner.poll(),
|
|
"Recover joiner exited during CUDA graph capture",
|
|
)
|
|
if log_path.exists() and log_path.read_text(errors="replace").count(
|
|
marker
|
|
) >= (EP_SIZE - LOCAL_EP_SIZE):
|
|
return
|
|
time.sleep(2)
|
|
self.fail(f"Timed out waiting for recover CUDA graph capture: {log_path}")
|
|
|
|
def _wait_for_recovered_ranks(self) -> None:
|
|
self._wait_for_recover_capture()
|
|
self._generate_ok("recovery trigger")
|
|
deadline = time.monotonic() + RECOVER_TIMEOUT_SECONDS
|
|
marker = f"recover ranks {list(range(LOCAL_EP_SIZE, EP_SIZE))} done"
|
|
primary_log = self.log_paths["primary"]
|
|
while time.monotonic() < deadline:
|
|
self.assertIsNone(
|
|
self.recover_joiner.poll(), "Recover joiner exited before rejoining"
|
|
)
|
|
if (
|
|
primary_log.exists()
|
|
and primary_log.read_text(errors="replace").count(marker)
|
|
>= LOCAL_EP_SIZE
|
|
):
|
|
for request_index in range(3):
|
|
self._generate_ok(f"post-recovery request {request_index + 1}")
|
|
return
|
|
time.sleep(2)
|
|
self.fail(f"Timed out waiting for recovery collective: {primary_log}")
|
|
|
|
def test_recover_four_ranks(self):
|
|
self._generate_ok("initial service")
|
|
|
|
kill_process_tree(self.initial_joiner.pid, wait_timeout=60)
|
|
# Give the terminated schedulers time to disappear before fault handling.
|
|
time.sleep(RECOVER_WAIT_SECONDS)
|
|
self._generate_ok("degraded service after node1 failure")
|
|
|
|
visible_devices = _visible_device_ids()
|
|
self.recover_joiner = self._launch(
|
|
node_rank=1,
|
|
port=JOINER_PORT,
|
|
visible_devices=visible_devices[LOCAL_EP_SIZE:EP_SIZE],
|
|
name="recover_joiner",
|
|
recover=True,
|
|
)
|
|
self._wait_for_recovered_ranks()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|