diff --git a/tests/diffusion/models/magi2/test_native_compile_cuda.py b/tests/diffusion/models/magi2/test_native_compile_cuda.py index 1db394750..0ca9eee46 100644 --- a/tests/diffusion/models/magi2/test_native_compile_cuda.py +++ b/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", + ), ] diff --git a/tests/e2e/online_serving/test_wan22_t2v.py b/tests/e2e/online_serving/test_wan22_t2v.py index f3f21ae7d..c4152e42b 100644 --- a/tests/e2e/online_serving/test_wan22_t2v.py +++ b/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, ), diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 0632db55d..2a17bf37c 100644 --- a/tests/helpers/runtime.py +++ b/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: diff --git a/tests/helpers/tests/test_runtime_server_ports.py b/tests/helpers/tests/test_runtime_server_ports.py index 1d8c0b2f1..d4634e960 100644 --- a/tests/helpers/tests/test_runtime_server_ports.py +++ b/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()