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

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)

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")
+                requested_ports.append((self.port, "HTTP"))
-            self._require_port(self.scheduler_port, "Scheduler")
+            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.scheduler_port = self.settle_port(
-            self.master_port = self.settle_port(self.master_port, 37)
+                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"

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": {}
                 }
             }

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,10 +252,34 @@ def wait_for_health(
 KILLALL_SCRIPT = Path(__file__).parents[3] / "killall_sglang.sh"
 
 
+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:
+            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)],
+            timeout=30,
+            capture_output=True,
+        )
+
+
 def kill_server(proc: subprocess.Popen) -> None:
     """Kill server process tree and clean up GPU processes."""
-    if proc.poll() is not None:
+    if proc.poll() is None:
-        return
         try:
             os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
         except (ProcessLookupError, PermissionError):
@@ -257,12 +293,7 @@ def kill_server(proc: subprocess.Popen) -> None:
                 pass
             proc.wait(timeout=10)
     # Use killall_sglang.sh for thorough cleanup (esp. multi-GPU workers)
-    if KILLALL_SCRIPT.exists():
+    _cleanup_sglang_processes()
-        subprocess.run(
-            ["bash", str(KILLALL_SCRIPT)],
-            timeout=30,
-            capture_output=True,
-        )
 
 
 # ---------------------------------------------------------------------------
@@ -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)
+        result["error"] = server_error.get("message", str(e))
-        print(f"  ERROR: {e}")
+        print(f"  ERROR: {result['error']}")
     finally:
         if proc:
             kill_server(proc)