From e823f0d909a163e9ea92253bf1204f39d1af3f18 Mon Sep 17 00:00:00 2001 From: Shaun Walsh <153730091+Shaun-Walsh@users.noreply.github.com> Date: Mon, 28 Sep 2026 16:30:30 +0100 Subject: [PATCH] [BugFix] Add field constraints to /v1/omni/sleep and /v1/omni/wakeup (#4740) Signed-off-by: Shaun Walsh Co-authored-by: Claude Opus 4.6 --- .../test_invalid_server_control.py | 9 ++------- .../openai_api/test_omni_sleep_wakeup.py | 18 ++++++------------ .../entrypoints/serve/omni_control/protocol.py | 8 ++++---- 3 files changed, 12 insertions(+), 23 deletions(-) diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py b/tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py index 4be838316..3303cd6ec 100644 --- a/tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py +++ b/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, ), ], ) diff --git a/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py b/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py index ba5925d8f..b1c5f1189 100644 --- a/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py +++ b/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() diff --git a/vllm_omni/entrypoints/serve/omni_control/protocol.py b/vllm_omni/entrypoints/serve/omni_control/protocol.py index f7ce72491..2e1a53188 100644 --- a/vllm_omni/entrypoints/serve/omni_control/protocol.py +++ b/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)