[BugFix] Add field constraints to /v1/omni/sleep and /v1/omni/wakeup (#4740)

Signed-off-by: Shaun Walsh <shaunwalsh24@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py

@@ -13,8 +13,6 @@ from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
 
 pytestmark = [pytest.mark.slow, pytest.mark.diffusion, pytest.mark.omni]
 
-_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649")
-
 # One module-scoped server for this file. ``--enable-sleep-mode`` is required by the
 # level-2 wake contract test; invalid-JSON cases only hit request validation and are
 # unaffected. Do not mix ``omni_server`` (module) with ``omni_server_function`` here:
@@ -44,9 +42,8 @@ _QWEN_IMAGE = [
         pytest.param(
             "sleep",
             {"stage_ids": [], "level": 2},
-            ("stage_ids", "Field required", "Missing"),
+            ("stage_ids", "too_short", "at least 1"),
             id="sleep_empty_stage_ids",
-            marks=_SKIP_ISSUE_3649,
         ),
         pytest.param(
             "sleep",
@@ -65,7 +62,6 @@ _QWEN_IMAGE = [
             {"stage_ids": [0], "level": -1},
             ("level", "greater"),
             id="sleep_negative_level",
-            marks=_SKIP_ISSUE_3649,
         ),
         pytest.param("wakeup", {}, ("stage_ids", "missing", "Field required"), id="wakeup_missing_stage_ids"),
         pytest.param(
@@ -77,9 +73,8 @@ _QWEN_IMAGE = [
         pytest.param(
             "wakeup",
             {"stage_ids": []},
-            ("stage_ids", "length"),
+            ("stage_ids", "too_short", "at least 1"),
             id="wakeup_empty_stage_ids",
-            marks=_SKIP_ISSUE_3649,
         ),
     ],
 )

tests/entrypoints/openai_api/test_omni_sleep_wakeup.py

@@ -69,20 +69,15 @@ def test_sleep_default_level(sleep_capable_engine):
     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=[])
+def test_sleep_empty_stage_ids(sleep_capable_engine):
+    """Empty stage_ids is rejected by min_length=1 validation."""
     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)
+    assert response.status_code == 422
+    sleep_capable_engine.sleep.assert_not_awaited()
 
 
 def test_sleep_updates_sleeping_set(sleep_capable_engine):
@@ -153,15 +148,14 @@ def test_wakeup_partial_sleeping(sleep_capable_engine):
 
 
 def test_wakeup_empty_stage_ids(sleep_capable_engine):
-    """Empty stage_ids: any() on empty iterable is False → SKIPPED, engine not called."""
+    """Empty stage_ids is rejected by min_length=1 validation."""
     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"
+    assert response.status_code == 422
     sleep_capable_engine.wake_up.assert_not_awaited()
 
 

vllm_omni/entrypoints/serve/omni_control/protocol.py

@@ -2,13 +2,13 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Request models for Omni server control routes."""
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 
 class OmniSleepRequest(BaseModel):
-    stage_ids: list[int]
-    level: int = 2
+    stage_ids: list[int] = Field(..., min_length=1)
+    level: int = Field(default=2, ge=0)
 
 
 class OmniWakeupRequest(BaseModel):
-    stage_ids: list[int]
+    stage_ids: list[int] = Field(..., min_length=1)