# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import dataclasses import pytest from fastapi import FastAPI from fastapi.testclient import TestClient pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @dataclasses.dataclass class FakeAck: stage_id: int result: str = "ok" def _make_app(engine_client): from vllm_omni.entrypoints.openai.api_server import router app = FastAPI() app.include_router(router) app.state.engine_client = engine_client app.state.sleeping_stages = set() return app @pytest.fixture def sleep_capable_engine(mocker): engine = mocker.MagicMock() engine.sleep = mocker.AsyncMock(return_value=[FakeAck(stage_id=0), FakeAck(stage_id=1)]) engine.wake_up = mocker.AsyncMock(return_value=[FakeAck(stage_id=0), FakeAck(stage_id=1)]) return engine @pytest.fixture def sleep_incapable_engine(mocker): return mocker.MagicMock(spec=[]) # no 'sleep' or 'wake_up' attributes # --------------------------------------------------------------------------- # /v1/omni/sleep # --------------------------------------------------------------------------- def test_sleep_success(sleep_capable_engine): app = _make_app(sleep_capable_engine) client = TestClient(app) response = client.post("/v1/omni/sleep", json={"stage_ids": [0, 1], "level": 2}) assert response.status_code == 200 data = response.json() assert data["status"] == "SUCCESS" assert [ack["stage_id"] for ack in data["acks"]] == [0, 1] assert app.state.sleeping_stages == {0, 1} sleep_capable_engine.sleep.assert_awaited_once_with(stage_ids=[0, 1], level=2) def test_sleep_default_level(sleep_capable_engine): app = _make_app(sleep_capable_engine) client = TestClient(app) response = client.post("/v1/omni/sleep", json={"stage_ids": [0]}) assert response.status_code == 200 sleep_capable_engine.sleep.assert_awaited_once_with(stage_ids=[0], level=2) def test_sleep_empty_stage_ids(sleep_capable_engine, mocker): """Empty stage_ids: engine is still called and returns SUCCESS with no acks.""" sleep_capable_engine.sleep = mocker.AsyncMock(return_value=[]) app = _make_app(sleep_capable_engine) client = TestClient(app) response = client.post("/v1/omni/sleep", json={"stage_ids": [], "level": 2}) assert response.status_code == 200 data = response.json() assert data["status"] == "SUCCESS" assert data["acks"] == [] assert app.state.sleeping_stages == set() sleep_capable_engine.sleep.assert_awaited_once_with(stage_ids=[], level=2) def test_sleep_updates_sleeping_set(sleep_capable_engine): app = _make_app(sleep_capable_engine) app.state.sleeping_stages = {5} # pre-existing entry client = TestClient(app) client.post("/v1/omni/sleep", json={"stage_ids": [0, 1], "level": 2}) assert app.state.sleeping_stages == {5, 0, 1} def test_sleep_engine_not_support(sleep_incapable_engine): app = _make_app(sleep_incapable_engine) client = TestClient(app) response = client.post("/v1/omni/sleep", json={"stage_ids": [0], "level": 1}) assert response.status_code == 501 assert "sleep" in response.json()["detail"] # --------------------------------------------------------------------------- # /v1/omni/wakeup # --------------------------------------------------------------------------- def test_wakeup_success(sleep_capable_engine): app = _make_app(sleep_capable_engine) app.state.sleeping_stages = {0, 1} client = TestClient(app) response = client.post("/v1/omni/wakeup", json={"stage_ids": [0, 1]}) assert response.status_code == 200 data = response.json() assert data["status"] == "SUCCESS" assert [ack["stage_id"] for ack in data["acks"]] == [0, 1] assert app.state.sleeping_stages == set() sleep_capable_engine.wake_up.assert_awaited_once_with(stage_ids=[0, 1]) def test_wakeup_skipped_when_not_sleeping(sleep_capable_engine): app = _make_app(sleep_capable_engine) app.state.sleeping_stages = set() # nothing sleeping client = TestClient(app) response = client.post("/v1/omni/wakeup", json={"stage_ids": [0, 1]}) assert response.status_code == 200 data = response.json() assert data["status"] == "SKIPPED" sleep_capable_engine.wake_up.assert_not_awaited() def test_wakeup_partial_sleeping(sleep_capable_engine): """Only stage 0 is sleeping; stage 1 is not. Should still proceed (partial match).""" app = _make_app(sleep_capable_engine) app.state.sleeping_stages = {0} client = TestClient(app) response = client.post("/v1/omni/wakeup", json={"stage_ids": [0, 1]}) assert response.status_code == 200 data = response.json() assert data["status"] == "SUCCESS" assert 0 not in app.state.sleeping_stages def test_wakeup_empty_stage_ids(sleep_capable_engine): """Empty stage_ids: any() on empty iterable is False → SKIPPED, engine not called.""" app = _make_app(sleep_capable_engine) app.state.sleeping_stages = {0, 1} client = TestClient(app) response = client.post("/v1/omni/wakeup", json={"stage_ids": []}) assert response.status_code == 200 assert response.json()["status"] == "SKIPPED" sleep_capable_engine.wake_up.assert_not_awaited() def test_wakeup_engine_not_support(sleep_incapable_engine): app = _make_app(sleep_incapable_engine) app.state.sleeping_stages = {0} client = TestClient(app) response = client.post("/v1/omni/wakeup", json={"stage_ids": [0]}) assert response.status_code == 501 assert "wake_up" in response.json()["detail"] def test_wakeup_route_propagates_not_implemented_error(sleep_capable_engine, mocker): """``/v1/omni/wakeup`` must not swallow ``NotImplementedError`` from the engine. Route exception-propagation only: the mock authors the exception on purpose. The real level-2 discarded-weight path and live HTTP 501 mapping are covered by ``test_wakeup_after_level2_sleep_fails``. """ sleep_capable_engine.wake_up = mocker.AsyncMock(side_effect=NotImplementedError("engine wake_up refused")) app = _make_app(sleep_capable_engine) app.state.sleeping_stages = {0} # Re-raise in-process; default TestClient would wrap this as a 500 response. client = TestClient(app, raise_server_exceptions=True) with pytest.raises(NotImplementedError, match="engine wake_up refused"): client.post("/v1/omni/wakeup", json={"stage_ids": [0]}) # Failed wake must not clear the sleeping set. assert app.state.sleeping_stages == {0} def test_wakeup_removes_only_requested_stages(sleep_capable_engine): app = _make_app(sleep_capable_engine) app.state.sleeping_stages = {0, 1, 2} client = TestClient(app) client.post("/v1/omni/wakeup", json={"stage_ids": [0]}) assert app.state.sleeping_stages == {1, 2} def test_sleep_then_wakeup_roundtrip(sleep_capable_engine): app = _make_app(sleep_capable_engine) client = TestClient(app) sleep_resp = client.post("/v1/omni/sleep", json={"stage_ids": [0, 1], "level": 2}) assert sleep_resp.status_code == 200 assert app.state.sleeping_stages == {0, 1} wakeup_resp = client.post("/v1/omni/wakeup", json={"stage_ids": [0, 1]}) assert wakeup_resp.status_code == 200 assert app.state.sleeping_stages == set() # --------------------------------------------------------------------------- # pure-diffusion branch: omni_init_app_state initialises sleeping_stages # --------------------------------------------------------------------------- @pytest.fixture def pure_diffusion_engine(mocker): engine = mocker.MagicMock() engine.stage_configs = [{"stage_type": "diffusion"}] engine.sleep = mocker.AsyncMock(return_value=[FakeAck(stage_id=0)]) engine.wake_up = mocker.AsyncMock(return_value=[FakeAck(stage_id=0)]) # Remove attributes that would make _get_vllm_config return a config del engine.get_vllm_config del engine.vllm_config return engine @pytest.fixture def pure_diffusion_app(pure_diffusion_engine, mocker): """App whose state was initialised via omni_init_app_state (pure diffusion path).""" import asyncio from argparse import Namespace from vllm_omni.entrypoints.openai.api_server import omni_init_app_state, router args = Namespace( served_model_name=None, model="fake-diffusion-model", enable_log_requests=False, disable_log_stats=True, max_log_len=0, enable_server_load_tracking=False, ) app = FastAPI() app.include_router(router) mocker.patch( "vllm_omni.entrypoints.openai.api_server.OmniOpenAIServingChat" ).for_diffusion.return_value = mocker.MagicMock() mocker.patch( "vllm_omni.entrypoints.openai.api_server.OmniOpenAIServingVideo" ).for_diffusion.return_value = mocker.MagicMock() mocker.patch( "vllm_omni.entrypoints.openai.api_server.OmniOpenAIServingSpeech" ).for_diffusion.return_value = mocker.MagicMock() mocker.patch("vllm_omni.entrypoints.openai.api_server._DiffusionServingModels") loop = asyncio.new_event_loop() try: loop.run_until_complete(omni_init_app_state(pure_diffusion_engine, app.state, args)) finally: loop.close() return app def test_pure_diffusion_sleep_success(pure_diffusion_app, pure_diffusion_engine): client = TestClient(pure_diffusion_app) response = client.post("/v1/omni/sleep", json={"stage_ids": [0], "level": 2}) assert response.status_code == 200 assert response.json()["status"] == "SUCCESS" assert pure_diffusion_app.state.sleeping_stages == {0} pure_diffusion_engine.sleep.assert_awaited_once_with(stage_ids=[0], level=2) def test_pure_diffusion_wakeup_skipped(pure_diffusion_app, pure_diffusion_engine): # sleeping_stages is empty after init — nothing is sleeping client = TestClient(pure_diffusion_app) response = client.post("/v1/omni/wakeup", json={"stage_ids": [0]}) assert response.status_code == 200 assert response.json()["status"] == "SKIPPED" pure_diffusion_engine.wake_up.assert_not_awaited() def test_pure_diffusion_sleep_then_wakeup(pure_diffusion_app, pure_diffusion_engine): client = TestClient(pure_diffusion_app) sleep_resp = client.post("/v1/omni/sleep", json={"stage_ids": [0], "level": 2}) assert sleep_resp.status_code == 200 assert pure_diffusion_app.state.sleeping_stages == {0} wakeup_resp = client.post("/v1/omni/wakeup", json={"stage_ids": [0]}) assert wakeup_resp.status_code == 200 assert pure_diffusion_app.state.sleeping_stages == set()