[Elastic EP] Fix recovery lifecycle and add manual coverage (#31744)
Co-authored-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
Xun Sun
@@ -0,0 +1,255 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user