diff --git a/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py b/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py new file mode 100644 index 000000000..cf7d6e234 --- /dev/null +++ b/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py @@ -0,0 +1,289 @@ +# 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_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() diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 29788d958..69a6dd603 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -586,6 +586,7 @@ async def omni_init_app_state( state.engine_client = engine_client state.log_stats = not args.disable_log_stats state.args = args + state.sleeping_stages = set() # For omni models state.stage_configs = engine_client.stage_configs if hasattr(engine_client, "stage_configs") else None @@ -949,7 +950,6 @@ async def omni_init_app_state( state.enable_server_load_tracking = args.enable_server_load_tracking state.server_load_metrics = 0 - state.sleeping_stages = set() def Omnivideo(request: Request) -> OmniOpenAIServingVideo | None: