anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[FEAT] support multi-stage deployment (#2396)

Browse Source 
Signed-off-by: ZhengWG <zwg0606@gmail.com>
Signed-off-by: Zheng Wengang <zwg0606@gmail.com>
Signed-off-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com>
Signed-off-by: yinpe <11810305@mail.sustech.edu.cn>
Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
Co-authored-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com>
Co-authored-by: yinpe <11810305@mail.sustech.edu.cn>
Co-authored-by: yinpeiqi <yinpeiqi809@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Chenguang Zheng <645327136@qq.com>
This commit is contained in:
Zheng WengangPeiqi YinyinpeyinpeiqiHongsheng LiuGao HanChenguang Zheng
 authored and GitHub committed 2026-05-06 19:33:18 +08:00
1 parent e969d2e99f
commit 1e5f288a91
55 files changed
+4684 -3407
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/helpers/runtime.py
+64 -28
View File
@@ -385,7 +385,8 @@ class OmniServerStageCli(OmniServer):
self.stage_ids = stage_ids or self._load_stage_ids(resolved_cfg)
if 0 not in self.stage_ids:
raise ValueError(f"Stage CLI test requires stage_id=0 in config: {stage_config_path}")
self.stage_procs: dict[int, subprocess.Popen] = {}
self.stage_replica_counts = self._load_stage_replica_counts(resolved_cfg)
self.stage_procs: dict[tuple[int, int], subprocess.Popen] = {}
self.proc = None
@staticmethod
@@ -415,6 +416,19 @@ class OmniServerStageCli(OmniServer):
runtime_devices[int(stage_id)] = str(devices)
return runtime_devices
@staticmethod
def _load_stage_replica_counts(resolved_config: dict) -> dict[int, int]:
replica_counts: dict[int, int] = {}
for stage in OmniServerStageCli._stage_entries(resolved_config):
stage_id = stage.get("stage_id")
if stage_id is None:
continue
replica_counts[int(stage_id)] = max(
1,
int(stage.get("num_replicas") or stage.get("runtime", {}).get("num_replicas", 1)),
)
return replica_counts
@classmethod
def _parse_device_list(cls, devices: str | int) -> list[str]:
if isinstance(devices, int):
@@ -452,13 +466,29 @@ class OmniServerStageCli(OmniServer):
return ",".join(visible_device_list[idx] for idx in logical_ids)
def _set_stage_device_env(self, stage_id: int, env: dict[str, str], devices: str) -> None:
mapped_devices = self._map_stage_devices(stage_id, self.visible_device_list, devices)
def _devices_for_replica(self, stage_id: int, devices: str, replica_id: int) -> str:
replica_count = self.stage_replica_counts.get(stage_id, 1)
if replica_count == 1:
return devices
device_list = self._parse_device_list(devices)
if len(device_list) % replica_count != 0:
raise ValueError(
f"Stage {stage_id} has {len(device_list)} device(s) for {replica_count} replica(s); "
"device count must be divisible by replica count"
)
devices_per_replica = len(device_list) // replica_count
start = replica_id * devices_per_replica
return ",".join(device_list[start : start + devices_per_replica])
def _set_stage_device_env(self, stage_id: int, env: dict[str, str], devices: str, replica_id: int = 0) -> None:
replica_devices = self._devices_for_replica(stage_id, devices, replica_id)
mapped_devices = self._map_stage_devices(stage_id, self.visible_device_list, replica_devices)
env_var = getattr(current_omni_platform, "device_control_env_var", None)
if env_var:
env[env_var] = mapped_devices
def _build_stage_cmd(self, stage_id: int, *, headless: bool) -> list[str]:
def _build_stage_cmd(self, stage_id: int, *, headless: bool, replica_id: int = 0) -> list[str]:
cmd = [
sys.executable,
"-m",
@@ -474,6 +504,8 @@ class OmniServerStageCli(OmniServer):
self.host,
"--omni-master-port",
str(self.master_port),
"--replica-id",
str(replica_id),
]
if headless:
@@ -484,7 +516,7 @@ class OmniServerStageCli(OmniServer):
cmd += self.serve_args
return cmd
def _launch_stage(self, stage_id: int, *, headless: bool) -> None:
def _launch_stage(self, stage_id: int, *, headless: bool, replica_id: int = 0) -> None:
env = os.environ.copy()
env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
if self.env_dict is not None:
@@ -492,19 +524,20 @@ class OmniServerStageCli(OmniServer):
devices = self.stage_runtime_devices.get(stage_id)
if devices:
self._set_stage_device_env(stage_id, env, devices)
self._set_stage_device_env(stage_id, env, devices, replica_id=replica_id)
cmd = self._build_stage_cmd(stage_id, headless=headless)
print(f"Launching OmniServerStageCli stage {stage_id}: {' '.join(cmd)}")
cmd = self._build_stage_cmd(stage_id, headless=headless, replica_id=replica_id)
print(f"Launching OmniServerStageCli stage {stage_id} replica {replica_id}: {' '.join(cmd)}")
# Capture each subprocess's stdout+stderr to a per-stage log file so
# debugging "Stage N exited before API server ready" doesn't rely on
# guessing; the file is surfaced in the RuntimeError message.
log_path = Path(tempfile.gettempdir()) / f"omni_stage_{stage_id}_{self.master_port}.log"
log_path = Path(tempfile.gettempdir()) / f"omni_stage_{stage_id}_replica_{replica_id}_{self.master_port}.log"
self._stage_log_paths = getattr(self, "_stage_log_paths", {})
self._stage_log_paths[stage_id] = log_path
stage_key = (stage_id, replica_id)
self._stage_log_paths[stage_key] = log_path
log_fh = open(log_path, "w", buffering=1) # noqa: SIM115 - closed in __exit__
self._stage_log_files = getattr(self, "_stage_log_files", {})
self._stage_log_files[stage_id] = log_fh
self._stage_log_files[stage_key] = log_fh
proc = subprocess.Popen(
cmd,
env=env,
@@ -512,36 +545,39 @@ class OmniServerStageCli(OmniServer):
stdout=log_fh,
stderr=subprocess.STDOUT,
)
self.stage_procs[stage_id] = proc
if stage_id == 0:
self.stage_procs[stage_key] = proc
if stage_id == 0 and replica_id == 0:
self.proc = proc
def _ensure_stage_processes_alive(self) -> None:
for stage_id, proc in self.stage_procs.items():
for (stage_id, replica_id), proc in self.stage_procs.items():
ret = proc.poll()
if ret is not None:
log_path = getattr(self, "_stage_log_paths", {}).get(stage_id)
log_path = getattr(self, "_stage_log_paths", {}).get((stage_id, replica_id))
tail = ""
if log_path and log_path.exists():
try:
with open(log_path, encoding="utf-8", errors="replace") as f:
lines = f.readlines()
tail = "\n=== Last 60 lines of stage {} log ({}) ===\n{}".format(
stage_id, log_path, "".join(lines[-60:]) or "<empty>"
tail = "\n=== Last 60 lines of stage {} replica {} log ({}) ===\n{}".format(
stage_id, replica_id, log_path, "".join(lines[-60:]) or "<empty>"
)
except Exception as exc: # pragma: no cover - diagnostic only
tail = f"\n<failed to read stage log {log_path}: {exc}>"
raise RuntimeError(f"Stage {stage_id} exited with code {ret} before API server became ready.{tail}")
raise RuntimeError(
f"Stage {stage_id} replica {replica_id} exited with code {ret} before API server became ready.{tail}"
)
def _start_server(self) -> None:
ordered_stage_ids = [0, *[stage_id for stage_id in self.stage_ids if stage_id != 0]]
self._launch_stage(0, headless=False)
self._launch_stage(0, headless=False, replica_id=0)
time.sleep(2)
self._ensure_stage_processes_alive()
for stage_id in ordered_stage_ids[1:]:
self._launch_stage(stage_id, headless=True)
for replica_id in range(self.stage_replica_counts.get(stage_id, 1)):
self._launch_stage(stage_id, headless=True, replica_id=replica_id)
max_wait = 1200
start_time = time.time()
@@ -568,15 +604,15 @@ class OmniServerStageCli(OmniServer):
whatever state the stage was in when it was torn down.
"""
log_paths = getattr(self, "_stage_log_paths", {}) or {}
for stage_id in sorted(log_paths):
log_path = log_paths[stage_id]
for stage_id, replica_id in sorted(log_paths):
log_path = log_paths[(stage_id, replica_id)]
if not log_path or not log_path.exists():
continue
try:
with open(log_path, encoding="utf-8", errors="replace") as f:
lines = f.readlines()
except Exception as exc: # pragma: no cover - diagnostic only
print(f"[OmniServerStageCli] stage {stage_id} log read failed: {exc}", flush=True)
print(f"[OmniServerStageCli] stage {stage_id} replica {replica_id} log read failed: {exc}", flush=True)
continue
total = len(lines)
if total <= head_lines + tail_lines:
@@ -587,18 +623,18 @@ class OmniServerStageCli(OmniServer):
head_chunk = lines[:head_lines]
tail_chunk = lines[-tail_lines:]
elided = total - head_lines - tail_lines
print(f"\n=== stage {stage_id} log HEAD ({log_path}) ===", flush=True)
print(f"\n=== stage {stage_id} replica {replica_id} log HEAD ({log_path}) ===", flush=True)
print("".join(head_chunk).rstrip("\n"), flush=True)
if tail_chunk:
print(f"\n... [{elided} lines elided] ...", flush=True)
print(f"\n=== stage {stage_id} log TAIL ({log_path}) ===", flush=True)
print(f"\n=== stage {stage_id} replica {replica_id} log TAIL ({log_path}) ===", flush=True)
print("".join(tail_chunk).rstrip("\n"), flush=True)
print(f"=== end stage {stage_id} log ===\n", flush=True)
print(f"=== end stage {stage_id} replica {replica_id} log ===\n", flush=True)
def __exit__(self, exc_type, exc_val, exc_tb):
self._dump_stage_logs_for_debug()
for stage_id in sorted(self.stage_procs, reverse=True):
proc = self.stage_procs[stage_id]
for stage_key in sorted(self.stage_procs, reverse=True):
proc = self.stage_procs[stage_key]
if proc.poll() is None:
self._kill_process_tree(proc.pid)
run_forced_gpu_cleanup_round()