Fix ROCm diffusion CI platform guard and startup deadlines

Signed-off-by: Sy03 <1370724210@qq.com>

tests/diffusion/models/magi2/test_native_compile_cuda.py

@@ -14,12 +14,17 @@ from tests.diffusion.models.magi2.test_native_packing import (
 )
 from tests.helpers.mark import hardware_marks
 from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig
+from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [
     pytest.mark.diffusion,
     pytest.mark.core_model,
     *hardware_marks(res={"cuda": "L4"}, num_cards=1),
-    pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA"),
+    # ROCm also exposes torch.cuda, but this path uses vLLM CUDA FlashAttention.
+    pytest.mark.skipif(
+        not current_omni_platform.is_cuda() or not torch.cuda.is_available(),
+        reason="requires vLLM CUDA FlashAttention",
+    ),
 ]
 
 

tests/e2e/online_serving/test_wan22_t2v.py

@@ -19,6 +19,7 @@ import pytest
 
 from tests.helpers.mark import hardware_marks
 from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
+from vllm_omni.platforms import current_omni_platform
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
@@ -31,9 +32,11 @@ SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
 
 def _get_diffusion_feature_cases(model: str):
     """Return a single default ``OmniServerParams`` row (no extra ``server_args``)."""
+    # MI300 CI needs over 900 s to load both expert checkpoints from cold storage.
+    init_timeout = 1800 if current_omni_platform.is_rocm() else None
     return [
         pytest.param(
-            OmniServerParams(model=model),
+            OmniServerParams(model=model, init_timeout=init_timeout, stage_init_timeout=init_timeout),
             id="default",
             marks=SINGLE_CARD_FEATURE_MARKS,
         ),

tests/helpers/runtime.py

@@ -161,6 +161,7 @@ class OmniServer:
         port: int | None = None,
         env_dict: dict[str, str] | None = None,
         use_omni: bool = True,
+        startup_timeout: int = 1200,
     ) -> None:
         cleanup_test_environment()
         self.model = model
@@ -169,6 +170,7 @@ class OmniServer:
         self.log_stats = "--disable-log-stats" not in args and "--log-stats" in args
         self.env_dict = env_dict
         self.use_omni = use_omni
+        self.startup_timeout = startup_timeout
         self.proc: subprocess.Popen | None = None
         self.host = "127.0.0.1"
         self._auto_port = port is None
@@ -246,7 +248,7 @@ class OmniServer:
             cwd=_omni_subprocess_cwd(),
         )
 
-        max_wait = 1200
+        max_wait = self.startup_timeout
         start_time = time.time()
         while time.time() - start_time < max_wait:
             ret = self.proc.poll()
@@ -1046,6 +1048,8 @@ def iter_omni_server(
                     port=port,
                     env_dict=params.env_dict,
                     use_omni=params.use_omni,
+                    # Allow frontend imports/startup in addition to engine initialization.
+                    startup_timeout=max(1200, (params.init_timeout or 900) + 300),
                 )
                 if port
                 else OmniServer(
@@ -1053,6 +1057,8 @@ def iter_omni_server(
                     server_args,
                     env_dict=params.env_dict,
                     use_omni=params.use_omni,
+                    # Allow frontend imports/startup in addition to engine initialization.
+                    startup_timeout=max(1200, (params.init_timeout or 900) + 300),
                 )
             ) as server:
                 if model != original_model:

tests/helpers/tests/test_runtime_server_ports.py

@@ -6,6 +6,8 @@
 import socket
 import subprocess
 import sys
+import threading
+from types import SimpleNamespace
 
 import pytest
 from filelock import FileLock
@@ -107,3 +109,62 @@ def test_server_accepts_its_own_listener(monkeypatch, in_child):
         assert server.proc.poll() is None
         with socket.create_connection((server.host, server.port), timeout=1):
             pass
+
+
+@pytest.mark.parametrize("init_timeout, expected_wait", [(None, 1200), (900, 1200), (1800, 2100)])
+@pytest.mark.parametrize("port", [None, 18099])
+def test_server_fixture_wait_covers_engine_init(monkeypatch, init_timeout, expected_wait, port):
+    captured = {}
+
+    class FakeServer:
+        def __init__(self, model, serve_args, **kwargs):
+            captured.update(kwargs)
+            captured["serve_args"] = serve_args
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *exc):
+            return False
+
+    monkeypatch.setattr(runtime, "OmniServer", FakeServer)
+    monkeypatch.setattr(runtime, "release_audio_transcriber", lambda: None)
+    request = SimpleNamespace(
+        param=runtime.OmniServerParams(model="fake-model", port=port, init_timeout=init_timeout),
+        node=SimpleNamespace(get_closest_marker=lambda name: None),
+    )
+    generator = runtime.iter_omni_server(request, "core_model", threading.Lock())
+    try:
+        next(generator)
+        assert captured["startup_timeout"] == expected_wait
+        args = captured["serve_args"]
+        assert args[args.index("--init-timeout") + 1] == str(init_timeout or 900)
+    finally:
+        generator.close()
+
+
+def test_server_keeps_polling_after_default_deadline(monkeypatch):
+    monkeypatch.setattr(runtime, "cleanup_test_environment", lambda: None)
+    server = runtime.OmniServer("fake-model", [], startup_timeout=2100)
+    monkeypatch.setattr(server, "_reserve_port", lambda: None)
+    monkeypatch.setattr(server, "_owns_listening_port", lambda: True)
+    monkeypatch.setattr(runtime.subprocess, "Popen", lambda *a, **kw: SimpleNamespace(poll=lambda: None))
+    # The server becomes ready after the old 1200 s limit, still within 2100 s.
+    ticks = iter([0, 1500])
+    monkeypatch.setattr(runtime.time, "time", lambda: next(ticks))
+
+    class ReadySocket:
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *exc):
+            return False
+
+        def settimeout(self, _timeout):
+            pass
+
+        def connect_ex(self, _address):
+            return 0
+
+    monkeypatch.setattr(runtime.socket, "socket", lambda *a: ReadySocket())
+    server._start_server()