[1/N] elastic-ep: Add runtime EP scale-up (#30164)
This commit is contained in:
Yoray Zack
@@ -0,0 +1,479 @@
|
||||
"""Manual tests for elastic EP scale-up.
|
||||
|
||||
Test classes:
|
||||
TestElasticScaleUp4To6 primary + joiner scale-up (6 GPUs)
|
||||
TestElasticScaleUp4To5To6 two consecutive single-rank scale-ups
|
||||
TestElasticScaleUp4To8 full primary + joiner scale-up (8 GPUs)
|
||||
|
||||
Run (8-GPU full scale-up):
|
||||
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -m pytest \\
|
||||
test/manual/ep/test_elastic_scale.py::TestElasticScaleUp4To8 \\
|
||||
-v -s
|
||||
"""
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.run_eval import run_eval
|
||||
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,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
TEST_MODEL = os.environ.get("NIXL_EP_TEST_MODEL", DEFAULT_MODEL_NAME_FOR_TEST_MLA)
|
||||
os.environ.setdefault("SGLANG_NIXL_EP_NUM_MAX_DISPATCH_TOKENS_PER_RANK", "1024")
|
||||
|
||||
ib_devices = get_rdma_devices_args()
|
||||
|
||||
|
||||
def _extra_server_args() -> list[str]:
|
||||
"""Extra `--flag [value]` tokens appended to every spawned server.
|
||||
|
||||
Set via ``SGLANG_ELASTIC_EXTRA_SERVER_ARGS`` as a single space-separated
|
||||
string, e.g. ``--disable-overlap-schedule``.
|
||||
"""
|
||||
raw = os.environ.get("SGLANG_ELASTIC_EXTRA_SERVER_ARGS", "").strip()
|
||||
return raw.split() if raw else []
|
||||
|
||||
|
||||
DISABLED_CUDA_GRAPH_ARGS = [
|
||||
"--cuda-graph-backend-decode",
|
||||
"disabled",
|
||||
"--cuda-graph-backend-prefill",
|
||||
"disabled",
|
||||
]
|
||||
|
||||
|
||||
def _assert_generate_logprob_ok(testcase: unittest.TestCase, base_url: str) -> None:
|
||||
response = requests.post(
|
||||
f"{base_url}/generate",
|
||||
json={
|
||||
"text": "The answer is",
|
||||
"sampling_params": {"max_new_tokens": 1, "temperature": 0.0},
|
||||
"return_logprob": True,
|
||||
"top_logprobs_num": 1,
|
||||
"logprob_start_len": 0,
|
||||
},
|
||||
timeout=60,
|
||||
)
|
||||
testcase.assertEqual(response.status_code, 200, response.text)
|
||||
input_logprobs = response.json()["meta_info"]["input_token_logprobs"]
|
||||
testcase.assertGreater(len(input_logprobs), 0)
|
||||
|
||||
|
||||
def _count_visible_gpus() -> int:
|
||||
env = os.environ.get("CUDA_VISIBLE_DEVICES")
|
||||
if env:
|
||||
return len([x for x in env.split(",") if x.strip()])
|
||||
try:
|
||||
import torch
|
||||
|
||||
return torch.cuda.device_count() if torch.cuda.is_available() else 0
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
|
||||
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()]
|
||||
return [str(index) for index in range(_count_visible_gpus())]
|
||||
|
||||
|
||||
LAUNCH_EP_SIZE = 4
|
||||
MAX_EP_SIZE = 8
|
||||
|
||||
DIST_INIT_ADDR = os.environ.get("SGLANG_ELASTIC_SCALE_DIST_INIT", "127.0.0.1:24555")
|
||||
PORT_A = int(os.environ.get("SGLANG_ELASTIC_SCALE_PORT_A", "21000"))
|
||||
PORT_B = int(os.environ.get("SGLANG_ELASTIC_SCALE_PORT_B", "10000"))
|
||||
PORT_C = int(os.environ.get("SGLANG_ELASTIC_SCALE_PORT_C", "11000"))
|
||||
HOST_A = os.environ.get("SGLANG_ELASTIC_SCALE_HOST_A", "127.0.0.1")
|
||||
BASE_URL_A = f"http://{HOST_A}:{PORT_A}"
|
||||
PRE_SCALE_JOINER_DELAY_SEC = float(
|
||||
os.environ.get("SGLANG_ELASTIC_PRE_SCALE_JOINER_DELAY_SEC", "0")
|
||||
)
|
||||
|
||||
|
||||
def _scale_up_common_args(
|
||||
dist_init_addr: str,
|
||||
tp_size: int,
|
||||
nnodes: int,
|
||||
node_rank: int,
|
||||
cuda_graph_args: list[str],
|
||||
moe_dense_tp_size: int | None,
|
||||
) -> list[str]:
|
||||
args = [
|
||||
"--trust-remote-code",
|
||||
"--moe-a2a-backend",
|
||||
"nixl",
|
||||
"--deepep-mode",
|
||||
"low_latency",
|
||||
"--tp",
|
||||
str(tp_size),
|
||||
"--dp",
|
||||
str(tp_size),
|
||||
"--enable-dp-attention",
|
||||
"--enable-dp-lm-head",
|
||||
"--elastic-ep-backend",
|
||||
"mooncake",
|
||||
"--mooncake-ib-device",
|
||||
ib_devices,
|
||||
"--enable-eplb",
|
||||
"--ep-num-redundant-experts",
|
||||
"24",
|
||||
"--elastic-ep-initial-size",
|
||||
str(LAUNCH_EP_SIZE),
|
||||
"--max-ep-size",
|
||||
str(MAX_EP_SIZE),
|
||||
"--mem-fraction-static",
|
||||
"0.5",
|
||||
"--chunked-prefill-size",
|
||||
"1024",
|
||||
"--nnodes",
|
||||
str(nnodes),
|
||||
"--node-rank",
|
||||
str(node_rank),
|
||||
"--dist-init-addr",
|
||||
dist_init_addr,
|
||||
]
|
||||
if moe_dense_tp_size is not None:
|
||||
args.extend(["--moe-dense-tp-size", str(moe_dense_tp_size)])
|
||||
return args + cuda_graph_args + _extra_server_args()
|
||||
|
||||
|
||||
class _ElasticScaleUpEndToEndBase(CustomTestCase):
|
||||
"""Shared scale-up E2E plumbing. Subclasses set JOIN_TP/JOIN_NNODES/JOIN_NODE_RANK."""
|
||||
|
||||
JOIN_TP: int
|
||||
JOIN_NNODES: int
|
||||
JOIN_NODE_RANK: int
|
||||
TARGET_EP_SIZE: int
|
||||
CUDA_GRAPH_ARGS: list[str]
|
||||
MOE_DENSE_TP_SIZE: int | None = 1
|
||||
|
||||
def setUp(self):
|
||||
if (
|
||||
not hasattr(type(self), "JOIN_TP")
|
||||
or type(self) is _ElasticScaleUpEndToEndBase
|
||||
):
|
||||
self.skipTest("Abstract base — run a concrete subclass instead")
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
if cls is _ElasticScaleUpEndToEndBase:
|
||||
raise unittest.SkipTest("Abstract base")
|
||||
cls.model = TEST_MODEL
|
||||
cls.base_url = BASE_URL_A
|
||||
cls._joining_procs = []
|
||||
cls._joining_log_fhs = []
|
||||
|
||||
primary_args = _scale_up_common_args(
|
||||
DIST_INIT_ADDR,
|
||||
tp_size=LAUNCH_EP_SIZE,
|
||||
nnodes=1,
|
||||
node_rank=0,
|
||||
cuda_graph_args=cls.CUDA_GRAPH_ARGS,
|
||||
moe_dense_tp_size=cls.MOE_DENSE_TP_SIZE,
|
||||
)
|
||||
primary_env = os.environ.copy()
|
||||
visible_devices = _visible_device_ids()
|
||||
if len(visible_devices) < LAUNCH_EP_SIZE:
|
||||
raise RuntimeError(
|
||||
f"Scale-up requires {LAUNCH_EP_SIZE} visible GPUs, got "
|
||||
f"{len(visible_devices)}"
|
||||
)
|
||||
primary_env["CUDA_VISIBLE_DEVICES"] = ",".join(visible_devices[:LAUNCH_EP_SIZE])
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=primary_args,
|
||||
env=primary_env,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _launch_joining_group(
|
||||
cls,
|
||||
*,
|
||||
rank_offset: int,
|
||||
join_tp: int,
|
||||
port: int,
|
||||
) -> subprocess.Popen:
|
||||
cmd = [
|
||||
"sglang",
|
||||
"serve",
|
||||
"--model-path",
|
||||
cls.model,
|
||||
*_scale_up_common_args(
|
||||
DIST_INIT_ADDR,
|
||||
tp_size=join_tp,
|
||||
nnodes=cls.JOIN_NNODES,
|
||||
node_rank=cls.JOIN_NODE_RANK,
|
||||
cuda_graph_args=cls.CUDA_GRAPH_ARGS,
|
||||
moe_dense_tp_size=cls.MOE_DENSE_TP_SIZE,
|
||||
),
|
||||
"--elastic-ep-join-mode",
|
||||
"scale",
|
||||
"--elastic-ep-join-rank-offset",
|
||||
str(rank_offset),
|
||||
"--host",
|
||||
"127.0.0.1",
|
||||
"--port",
|
||||
str(port),
|
||||
"--device",
|
||||
"cuda",
|
||||
]
|
||||
env = os.environ.copy()
|
||||
visible_devices = _visible_device_ids()
|
||||
join_end = rank_offset + join_tp
|
||||
if join_end > len(visible_devices):
|
||||
raise RuntimeError(
|
||||
f"Scale-up requires {join_end} visible GPUs, got "
|
||||
f"{len(visible_devices)}"
|
||||
)
|
||||
env["CUDA_VISIBLE_DEVICES"] = ",".join(visible_devices[rank_offset:join_end])
|
||||
base_joining_log = os.environ.get(
|
||||
"SGLANG_ELASTIC_SCALE_JOINING_LOG",
|
||||
f"/tmp/elastic_scale_joining_nnodes{cls.JOIN_NNODES}_{int(time.time())}.log",
|
||||
)
|
||||
if cls._joining_procs:
|
||||
root, ext = os.path.splitext(base_joining_log)
|
||||
joining_log = f"{root}_step{len(cls._joining_procs) + 1}{ext}"
|
||||
else:
|
||||
joining_log = base_joining_log
|
||||
joining_log_fh = open(joining_log, "w")
|
||||
joining_proc = subprocess.Popen(
|
||||
cmd,
|
||||
env=env,
|
||||
stdout=joining_log_fh,
|
||||
stderr=subprocess.STDOUT,
|
||||
)
|
||||
cls._joining_procs.append(joining_proc)
|
||||
cls._joining_log_fhs.append(joining_log_fh)
|
||||
return joining_proc
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
processes = [
|
||||
*reversed(cls._joining_procs),
|
||||
getattr(cls, "process", None),
|
||||
]
|
||||
for proc in processes:
|
||||
if proc is None:
|
||||
continue
|
||||
try:
|
||||
kill_process_tree(proc.pid)
|
||||
except Exception:
|
||||
pass
|
||||
for proc in processes:
|
||||
if proc is None:
|
||||
continue
|
||||
try:
|
||||
proc.wait(timeout=15)
|
||||
except Exception:
|
||||
pass
|
||||
for fh in cls._joining_log_fhs:
|
||||
try:
|
||||
fh.close()
|
||||
except Exception:
|
||||
pass
|
||||
time.sleep(2)
|
||||
|
||||
def _post(self, path: str, **kwargs) -> requests.Response:
|
||||
return requests.post(f"{self.base_url}{path}", timeout=60, **kwargs)
|
||||
|
||||
def _generate_ok(self, msg_suffix: str, routed_dp_rank: int | None = None) -> None:
|
||||
payload = {
|
||||
"text": "Hello",
|
||||
"sampling_params": {"max_new_tokens": 4, "temperature": 0.0},
|
||||
}
|
||||
if routed_dp_rank is not None:
|
||||
payload["routed_dp_rank"] = routed_dp_rank
|
||||
resp = self._post(
|
||||
"/generate",
|
||||
json=payload,
|
||||
)
|
||||
self.assertEqual(
|
||||
resp.status_code,
|
||||
200,
|
||||
f"/generate {msg_suffix} failed: {resp.text}",
|
||||
)
|
||||
|
||||
def _generate_logprob_ok(self, msg_suffix: str) -> None:
|
||||
try:
|
||||
_assert_generate_logprob_ok(self, self.base_url)
|
||||
except AssertionError as exc:
|
||||
raise AssertionError(
|
||||
f"/generate logprob {msg_suffix} failed: {exc}"
|
||||
) from exc
|
||||
|
||||
def _scale_once(
|
||||
self,
|
||||
*,
|
||||
old_ep_size: int,
|
||||
target_ep_size: int,
|
||||
join_tp: int,
|
||||
port: int,
|
||||
) -> None:
|
||||
joining_proc = self._launch_joining_group(
|
||||
rank_offset=old_ep_size,
|
||||
join_tp=join_tp,
|
||||
port=port,
|
||||
)
|
||||
self.assertIsNone(
|
||||
joining_proc.poll(),
|
||||
"Joining group exited before scale request; see joining log",
|
||||
)
|
||||
if PRE_SCALE_JOINER_DELAY_SEC > 0:
|
||||
time.sleep(PRE_SCALE_JOINER_DELAY_SEC)
|
||||
self.assertIsNone(
|
||||
joining_proc.poll(),
|
||||
"Joining group exited before scale request; see joining log",
|
||||
)
|
||||
|
||||
resp = self._post("/scale_elastic_ep", json={"new_ep_size": target_ep_size})
|
||||
self.assertEqual(resp.status_code, 200, resp.text)
|
||||
body = resp.json()
|
||||
self.assertEqual(body["old_ep_size"], old_ep_size)
|
||||
self.assertEqual(body["new_ep_size"], target_ep_size)
|
||||
|
||||
deadline = time.time() + 300
|
||||
while time.time() < deadline:
|
||||
resp = requests.get(f"{self.base_url}/is_scaling_elastic_ep", timeout=60)
|
||||
state = resp.json() if resp.ok else None
|
||||
if state is not None and not state.get("is_scaling_elastic_ep", True):
|
||||
self.assertEqual(state.get("effective_ep_size"), target_ep_size)
|
||||
self.assertEqual(state.get("scale_phase"), "serving_expanded")
|
||||
self.assertIsNone(state.get("last_error"))
|
||||
self._generate_ok(
|
||||
"on newest joiner",
|
||||
routed_dp_rank=target_ep_size - 1,
|
||||
)
|
||||
return
|
||||
try:
|
||||
self._post(
|
||||
"/generate",
|
||||
json={
|
||||
"text": "ping",
|
||||
"sampling_params": {"max_new_tokens": 1, "temperature": 0.0},
|
||||
},
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
time.sleep(2)
|
||||
self.fail("Timed out waiting for scaling to complete (300s)")
|
||||
|
||||
def _run_post_scale_gsm8k(self) -> None:
|
||||
metrics = run_eval(
|
||||
SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="gsm8k",
|
||||
api="completion",
|
||||
max_tokens=512,
|
||||
num_examples=256,
|
||||
num_threads=50,
|
||||
)
|
||||
)
|
||||
print(f"[TEST] Post-scale GSM8K accuracy: {metrics['score']:.2%}")
|
||||
self.assertGreater(
|
||||
metrics["score"],
|
||||
0.50,
|
||||
f"Post-scale GSM8K accuracy too low: {metrics['score']:.2%}",
|
||||
)
|
||||
|
||||
def test_scale_up_on_demand(self):
|
||||
"""Scale the primary group to the configured target."""
|
||||
self._generate_ok("pre-scale")
|
||||
|
||||
self._scale_once(
|
||||
old_ep_size=LAUNCH_EP_SIZE,
|
||||
target_ep_size=self.TARGET_EP_SIZE,
|
||||
join_tp=self.JOIN_TP,
|
||||
port=PORT_B,
|
||||
)
|
||||
|
||||
self._generate_ok("post-scale")
|
||||
self._generate_logprob_ok("post-scale")
|
||||
|
||||
self._run_post_scale_gsm8k()
|
||||
|
||||
|
||||
@unittest.skipUnless(
|
||||
_count_visible_gpus() >= 6,
|
||||
"4-to-6 scale-up E2E needs 6 GPUs.",
|
||||
)
|
||||
class TestElasticScaleUp4To6(_ElasticScaleUpEndToEndBase):
|
||||
"""Scale from four to six ranks."""
|
||||
|
||||
JOIN_TP = 2
|
||||
JOIN_NNODES = 2
|
||||
JOIN_NODE_RANK = 1
|
||||
TARGET_EP_SIZE = 6
|
||||
CUDA_GRAPH_ARGS = DISABLED_CUDA_GRAPH_ARGS
|
||||
MOE_DENSE_TP_SIZE = None
|
||||
|
||||
|
||||
@unittest.skipUnless(
|
||||
_count_visible_gpus() >= 6,
|
||||
"4-to-5-to-6 scale-up E2E needs 6 GPUs.",
|
||||
)
|
||||
class TestElasticScaleUp4To5To6(_ElasticScaleUpEndToEndBase):
|
||||
"""Scale from four to five and then from five to six ranks."""
|
||||
|
||||
JOIN_TP = 1
|
||||
JOIN_NNODES = 2
|
||||
JOIN_NODE_RANK = 1
|
||||
TARGET_EP_SIZE = 6
|
||||
CUDA_GRAPH_ARGS = DISABLED_CUDA_GRAPH_ARGS
|
||||
|
||||
def test_scale_up_on_demand(self):
|
||||
self._generate_ok("pre-scale")
|
||||
|
||||
self._scale_once(
|
||||
old_ep_size=4,
|
||||
target_ep_size=5,
|
||||
join_tp=1,
|
||||
port=PORT_B,
|
||||
)
|
||||
self._generate_ok("after first scale")
|
||||
self._generate_logprob_ok("after first scale")
|
||||
|
||||
self._scale_once(
|
||||
old_ep_size=5,
|
||||
target_ep_size=6,
|
||||
join_tp=1,
|
||||
port=PORT_C,
|
||||
)
|
||||
self._generate_ok("after second scale")
|
||||
self._generate_logprob_ok("after second scale")
|
||||
self._run_post_scale_gsm8k()
|
||||
|
||||
|
||||
@unittest.skipUnless(
|
||||
_count_visible_gpus() >= MAX_EP_SIZE,
|
||||
f"Full scale-up E2E needs {MAX_EP_SIZE} GPUs.",
|
||||
)
|
||||
class TestElasticScaleUp4To8(_ElasticScaleUpEndToEndBase):
|
||||
"""Scale from four to eight ranks."""
|
||||
|
||||
JOIN_TP = LAUNCH_EP_SIZE
|
||||
JOIN_NNODES = 2
|
||||
JOIN_NODE_RANK = 1
|
||||
TARGET_EP_SIZE = MAX_EP_SIZE
|
||||
CUDA_GRAPH_ARGS = DISABLED_CUDA_GRAPH_ARGS
|
||||
MOE_DENSE_TP_SIZE = None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user