[diffusion] CI: fix nightly CI (#25241)

This commit is contained in:
Mick
2026-05-16 16:55:30 +08:00
committed by GitHub
parent af26b71ae8
commit 7f37ffae9d
4 changed files with 112 additions and 27 deletions
@@ -553,6 +553,8 @@ class LTX2SnapshotResidencyStrategy(LTX2TwoStageResidencyStrategy):
if self._module_is_on_gpu(target_module):
self._record_component_ready("transformer")
elif not self._snapshot_strategy.is_ready("transformer"):
if self._snapshot_low_vram_mode:
self._release_stage2_for_low_vram()
self._snapshot_strategy.prefetch_component("transformer", target_module)
else:
self._record_component_ready("transformer")
@@ -623,6 +625,16 @@ class LTX2SnapshotResidencyStrategy(LTX2TwoStageResidencyStrategy):
if stage1_param is not None and stage1_param.device.type == "cuda":
self._release_module_to_cpu_snapshot("transformer")
def _release_stage2_for_low_vram(self) -> None:
stage2_module = self.pipeline.get_module("transformer_2")
stage2_param = (
next(stage2_module.parameters(), None)
if stage2_module is not None
else None
)
if stage2_param is not None and stage2_param.device.type == "cuda":
self._release_module_to_cpu_snapshot("transformer_2")
def _record_component_ready(self, module_name: str) -> None:
self._snapshot_strategy.record_ready(
module_name, self.pipeline.get_module(module_name)
@@ -677,19 +677,37 @@ class ServerArgs(DisaggArgsMixin):
)
if self.strict_ports:
requested_ports = []
if needs_http:
self._require_port(self.port, "HTTP")
self._require_port(self.scheduler_port, "Scheduler")
requested_ports.append((self.port, "HTTP"))
requested_ports.append((self.scheduler_port, "Scheduler"))
if self.master_port is not None:
self._require_port(self.master_port, "Master")
requested_ports.append((self.master_port, "Master"))
seen_ports: dict[int, str] = {}
for port, name in requested_ports:
if port in seen_ports:
raise RuntimeError(
f"{name} port {port} duplicates {seen_ports[port]} port and "
"--strict-ports is enabled."
)
seen_ports[port] = name
self._require_port(port, name)
else:
settled_ports: set[int] = set()
if needs_http:
self.port = self.settle_port(self.port)
settled_ports.add(self.port)
initial_scheduler_port = self.scheduler_port + (
random.randint(0, 100) if self.scheduler_port == 5555 else 0
)
self.scheduler_port = self.settle_port(initial_scheduler_port)
self.master_port = self.settle_port(self.master_port, 37)
self.scheduler_port = self.settle_port(
initial_scheduler_port, avoid=settled_ports
)
settled_ports.add(self.scheduler_port)
if self.master_port is not None:
self.master_port = self.settle_port(
self.master_port, 37, avoid=settled_ports
)
def _adjust_parallelism(self):
sp_unspecified = self.sp_degree is None
@@ -1467,16 +1485,21 @@ class ServerArgs(DisaggArgsMixin):
return f"tcp://{scheduler_host}:{self.scheduler_port}"
def settle_port(
self, port: int, port_inc: int = 42, max_attempts: int = 100
self,
port: int,
port_inc: int = 42,
max_attempts: int = 100,
avoid: set[int] | None = None,
) -> int:
"""
Find an available port with retry logic.
"""
attempts = 0
original_port = port
avoid = avoid or set()
while attempts < max_attempts:
if is_port_available(port):
if port not in avoid and is_port_available(port):
if attempts > 0:
logger.info(
f"Port {original_port} was unavailable, using port {port} instead"
@@ -148,7 +148,7 @@
"num_gpus": 2,
"frameworks": {
"sglang": {
"serve_args": "--enable-torch-compile --warmup --pipeline-class-name LTX2TwoStagePipeline",
"serve_args": "--enable-torch-compile --warmup --pipeline-class-name LTX2TwoStagePipeline --cfg-parallel-size 2",
"extra_env": {}
}
}
+69 -19
View File
@@ -23,6 +23,7 @@ import io
import json
import os
import signal
import socket
import subprocess
import sys
import tempfile
@@ -40,11 +41,17 @@ CONFIGS_PATH = Path(__file__).parent / "comparison_configs.json"
INSTALL_SCRIPT = Path(__file__).parents[1] / "install_comparison_frameworks.sh"
DEFAULT_HOST = "127.0.0.1"
DEFAULT_PORT = 30000
SGLANG_MASTER_PORT_OFFSET = 5
SGLANG_SCHEDULER_PORT_OFFSET = 55
HEALTH_TIMEOUT = (
2400 # seconds (40 min — FLUX.2-dev needs ~10 min download + torch.compile)
)
REQUEST_TIMEOUT = 1200 # seconds
GPU_CLEAR_WAIT = 15 # seconds between framework runs
SERVER_FATAL_ERROR_PATTERNS = (
"CUDA out of memory",
"torch.OutOfMemoryError",
)
# Frameworks that need separate installation (conflict with sglang's deps)
INSTALLABLE_FRAMEWORKS = {"vllm-omni", "lightx2v"}
@@ -69,6 +76,11 @@ def _build_sglang_cmd(case: dict, fw_cfg: dict, port: int) -> list[str]:
str(port),
"--host",
DEFAULT_HOST,
"--strict-ports",
"--master-port",
str(port + SGLANG_MASTER_PORT_OFFSET),
"--scheduler-port",
str(port + SGLANG_SCHEDULER_PORT_OFFSET),
]
if case["num_gpus"] > 1:
cmd += ["--num-gpus", str(case["num_gpus"])]
@@ -240,23 +252,23 @@ def wait_for_health(
KILLALL_SCRIPT = Path(__file__).parents[3] / "killall_sglang.sh"
def kill_server(proc: subprocess.Popen) -> None:
"""Kill server process tree and clean up GPU processes."""
if proc.poll() is not None:
return
try:
os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
except (ProcessLookupError, PermissionError):
pass
try:
proc.wait(timeout=30)
except subprocess.TimeoutExpired:
def _is_port_available(port: int) -> bool:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
try:
os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
except (ProcessLookupError, PermissionError):
pass
proc.wait(timeout=10)
# Use killall_sglang.sh for thorough cleanup (esp. multi-GPU workers)
sock.bind((DEFAULT_HOST, port))
except OSError:
return False
return True
def _require_ports_available(ports: list[int]) -> None:
unavailable = [port for port in ports if not _is_port_available(port)]
if unavailable:
raise RuntimeError(f"Required port(s) unavailable before launch: {unavailable}")
def _cleanup_sglang_processes() -> None:
if KILLALL_SCRIPT.exists():
subprocess.run(
["bash", str(KILLALL_SCRIPT)],
@@ -265,6 +277,25 @@ def kill_server(proc: subprocess.Popen) -> None:
)
def kill_server(proc: subprocess.Popen) -> None:
"""Kill server process tree and clean up GPU processes."""
if proc.poll() is None:
try:
os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
except (ProcessLookupError, PermissionError):
pass
try:
proc.wait(timeout=30)
except subprocess.TimeoutExpired:
try:
os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
except (ProcessLookupError, PermissionError):
pass
proc.wait(timeout=10)
# Use killall_sglang.sh for thorough cleanup (esp. multi-GPU workers)
_cleanup_sglang_processes()
# ---------------------------------------------------------------------------
# Reference image helpers
# ---------------------------------------------------------------------------
@@ -711,9 +742,20 @@ def run_single(
log_file = log_dir / f"{case['id']}_{framework}.log"
log_fh = open(log_file, "w", encoding="utf-8", buffering=1)
log_thread = None
server_error = {}
proc = None
try:
if framework == "sglang":
_cleanup_sglang_processes()
_require_ports_available(
[
port,
port + SGLANG_MASTER_PORT_OFFSET,
port + SGLANG_SCHEDULER_PORT_OFFSET,
]
)
proc = subprocess.Popen(
cmd,
stdout=subprocess.PIPE,
@@ -731,6 +773,14 @@ def run_single(
sys.stdout.write(f" [server] {line}")
sys.stdout.flush()
fh.write(line)
if not server_error and any(
pattern in line for pattern in SERVER_FATAL_ERROR_PATTERNS
):
server_error["message"] = line.strip()
try:
os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
except (ProcessLookupError, PermissionError):
pass
except ValueError:
pass # pipe closed
@@ -752,7 +802,7 @@ def run_single(
try:
send_request(base_url, warmup_case, framework, config)
except Exception as e:
print(f" Warmup request {wi} failed (non-fatal): {e}")
raise RuntimeError(f"Warmup request {wi} failed: {e}") from e
# Measured request — pass perf_dump_path for SGLang server-side timing
if perf_dump_path and os.path.exists(perf_dump_path):
@@ -764,8 +814,8 @@ def run_single(
result["latency_s"] = round(latency, 3)
except Exception as e:
result["error"] = str(e)
print(f" ERROR: {e}")
result["error"] = server_error.get("message", str(e))
print(f" ERROR: {result['error']}")
finally:
if proc:
kill_server(proc)