[Bugfix] fix when diffusion model not set sleeping_stages (#3023)

Signed-off-by: rongfu.leng <lenronfu@gmail.com>

thanks

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()

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: