[diffusion] CI: fix nightly CI (#25241)
This commit is contained in:
@@ -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": {}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user