diff --git a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py index bb8bfb5a4..459234fa9 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index 295a52779..8479ad6ae 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -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" diff --git a/scripts/ci/utils/diffusion/comparison_configs.json b/scripts/ci/utils/diffusion/comparison_configs.json index a2eb4b408..aa4171aca 100644 --- a/scripts/ci/utils/diffusion/comparison_configs.json +++ b/scripts/ci/utils/diffusion/comparison_configs.json @@ -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": {} } } diff --git a/scripts/ci/utils/diffusion/run_comparison.py b/scripts/ci/utils/diffusion/run_comparison.py index 8edfa9a67..c25df75ce 100644 --- a/scripts/ci/utils/diffusion/run_comparison.py +++ b/scripts/ci/utils/diffusion/run_comparison.py @@ -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)