[Bugfix] Handle DoS Overflow Cases (#6598)

Signed-off-by: Alex Brooks <albrooks@redhat.com>

tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py

@@ -157,8 +157,8 @@ def test_speech_missing_required_fields(omni_server: OmniServer, online_client:
                 "ref_audio": None,
                 "ref_text": None,
             },
-            ("CustomVoice", "does not support"),
-            id="customvoice_task_not_supported",
+            ("Invalid voice", "vivian"),
+            id="customvoice_invalid_voice",
         ),
         pytest.param(
             {"task_type": "InvalidEnum"}, ("task_type", "literal_error", "CustomVoice"), id="task_type_invalid"
@@ -181,13 +181,17 @@ def test_speech_missing_required_fields(omni_server: OmniServer, online_client:
             ("mutually exclusive", "speaker_embedding", "ref_audio"),
             id="speaker_embedding_with_ref_audio_mutually_exclusive",
         ),
-        pytest.param({"max_new_tokens": 0}, ("max_new_tokens", "least 1"), id="max_new_tokens_below_min"),
+        pytest.param({"max_new_tokens": 0}, ("max_new_tokens", "greater_than_equal"), id="max_new_tokens_below_min"),
         pytest.param(
             {"max_new_tokens": _SPEECH_API_MAX_NEW_TOKENS + 1},
             ("max_new_tokens", "exceed 4096"),
             id="max_new_tokens_above_max",
         ),
-        pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative"),
+        pytest.param(
+            {"seed": -(2**63) - 1},
+            ("seed", "greater_than_equal", "-9223372036854775808"),
+            id="seed_below_min",
+        ),
         pytest.param({"seed": 2**63}, ("seed", "less_than_equal", "9223372036854775807"), id="seed_above_max"),
         pytest.param(
             {"initial_codec_chunk_frames": -1},
@@ -275,7 +279,9 @@ def test_speech_invalid_field_values(
         pytest.param({"input": "", "response_format": "wav"}, ("input", "empty"), id="input_empty"),
         pytest.param({"input": "   ", "response_format": "wav"}, ("input", "empty"), id="input_whitespace_only"),
         pytest.param(
-            {"input": "Hello world.", "max_new_tokens": 0}, ("max_new_tokens", "least 1"), id="max_new_tokens_below_min"
+            {"input": "Hello world.", "max_new_tokens": 0},
+            ("max_new_tokens", "greater_than_equal"),
+            id="max_new_tokens_below_min",
         ),
         pytest.param(
             {
@@ -554,14 +560,14 @@ def test_speech_batch_missing_items(omni_server: OmniServer, online_client: Onli
         pytest.param(
             "batch",
             {"max_new_tokens": 0},
-            ("max_new_tokens", "least 1"),
+            ("max_new_tokens", "greater_than_equal"),
             id="batch_max_new_tokens_below_min",
             marks=_SKIP_ISSUE_3649,
         ),
         pytest.param(
             "item",
             {"max_new_tokens": 0},
-            ("max_new_tokens", "least 1"),
+            ("max_new_tokens", "greater_than_equal"),
             id="item_max_new_tokens_below_min",
             marks=_SKIP_ISSUE_3649,
         ),

tests/engine/test_orchestrator_error_handling.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Tests for error propagation paths within the Orchestrator.
 
 Covers:
@@ -16,7 +16,9 @@ from types import SimpleNamespace
 
 import janus
 import pytest
+from vllm.sampling_params import SamplingParams
 from vllm.v1.engine.exceptions import EngineDeadError
+from vllm.v1.serial_utils import MsgpackEncoder
 
 from vllm_omni.engine.messages import (
     AddCompanionRequestMessage,
@@ -46,8 +48,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 def _sampling_params(max_tokens: int = 4):
-    from vllm.sampling_params import SamplingParams
-
     return SamplingParams(max_tokens=max_tokens)
 
 
@@ -188,6 +188,45 @@ class _UnavailableDiffusionStageClient(FakeStageClient):
         raise StageUnavailableError(f"stage {self.stage_id} has no live replica")
 
 
+class _OverflowOnEncodeStageClient(FakeStageClient):
+    """Msgpack-encodes the request's sampling params on dispatch, as the real
+    StageEngineCoreClient does. We use this to make sure overflow sampling params
+    do not break the orchestrator."""
+
+    async def add_request_async(self, request, *args, **kwargs) -> None:
+        MsgpackEncoder().encode(request.sampling_params)
+        await super().add_request_async(request, *args, **kwargs)  # unreachable
+
+
+@pytest.mark.asyncio
+async def test_encode_overflow_should_fail_request_not_loop(orchestrator_factory) -> None:
+    """Regression test for over-range seed handling."""
+    stage0 = _OverflowOnEncodeStageClient(stage_type="llm", final_output=True)
+    orchestrator_fixture = orchestrator_factory([stage0])
+
+    # seed > 2**63-1 cannot be msgpack-serialized -> the malicious input under test.
+    overflow_params = SamplingParams(seed=2**70)
+
+    await _enqueue_add_request(
+        orchestrator_fixture,
+        request_id="req-overflow",
+        prompt=SimpleNamespace(
+            request_id="req-overflow",
+            prompt_token_ids=[1, 2],
+            sampling_params=overflow_params,
+        ),
+        original_prompt={"prompt": "hi"},
+        sampling_params_list=[overflow_params],
+        final_stage_id=0,
+    )
+
+    error_msg = await _wait_for_error_message(orchestrator_fixture, request_id="req-overflow")
+    assert error_msg.fatal is False
+    assert error_msg.status_code == 400
+    assert orchestrator_fixture.thread.is_alive()
+    assert "req-overflow" not in orchestrator_fixture.orchestrator.request_states
+
+
 @pytest.mark.asyncio
 async def test_engine_dead_error_fails_only_dead_replica_requests(orchestrator_factory) -> None:
     """Multi-replica stage: one replica dying must fail only the requests bound

tests/entrypoints/openai_api/test_serving_speech.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # tests/entrypoints/openai/test_serving_speech.py
 import asyncio
 import base64
@@ -4147,16 +4150,15 @@ class TestCosyVoice3Serving:
         error = cosyvoice3_server._validate_tts_request(request)
         assert error is None
 
-    def test_validate_cosyvoice3_max_new_tokens_range(self, cosyvoice3_server):
-        request = OpenAICreateSpeechRequest(
-            input="Hello",
-            ref_audio="data:audio/wav;base64,abc",
-            ref_text="hello",
-            max_new_tokens=0,
-        )
-        error = cosyvoice3_server._validate_tts_request(request)
-        assert error is not None
-        assert "max_new_tokens" in error
+    def test_validate_cosyvoice3_max_new_tokens_range(self):
+        """Ensure max_new_tokens below the minimum is rejected during request validation."""
+        with pytest.raises(ValidationError, match="max_new_tokens"):
+            OpenAICreateSpeechRequest(
+                input="Hello",
+                ref_audio="data:audio/wav;base64,abc",
+                ref_text="hello",
+                max_new_tokens=0,
+            )
 
     @pytest.mark.parametrize(
         ("max_new_tokens", "expected_min_tokens", "expected_max_tokens"),

vllm_omni/engine/orchestrator.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """
 Orchestrator for vLLM-Omni multi-stage runtime.
 
@@ -363,8 +366,8 @@ class Orchestrator:
     def __init__(
         self,
         request_async_queue: janus.AsyncQueue[EngineQueueMessage],
-        output_async_queue: janus.AsyncQueue[dict[str, Any]],
-        rpc_async_queue: janus.AsyncQueue[dict[str, Any]],
+        output_async_queue: janus.AsyncQueue[EngineQueueMessage],
+        rpc_async_queue: janus.AsyncQueue[EngineQueueMessage],
         stage_pools: list[StagePool],
         *,
         async_chunk: bool = False,
@@ -1173,6 +1176,35 @@ class Orchestrator:
             abort=True,
         )
 
+    async def _fail_request_client_error(
+        self,
+        req_id: str,
+        stage_id: int,
+        error: str,
+        *,
+        close_duplex_sessions: bool = False,
+    ) -> None:
+        """Fail one request with a non-fatal 400 (bad input, engine survives).
+
+        The non-fatal counterpart of `_fail_request_dead_stage`: emits a
+        client-error ErrorMessage (default `fatal=False`) so the engine keeps
+        serving, then releases the request's state across every stage pool.
+        """
+        await self.output_async_queue.put(
+            ErrorMessage(
+                error=error,
+                status_code=HTTPStatus.BAD_REQUEST.value,
+                error_type=DEFAULT_CLIENT_ERROR_TYPE,
+                request_id=req_id,
+                stage_id=stage_id,
+            )
+        )
+        await self._cleanup_request_ids(
+            [req_id, *self._cfg_tracker.cleanup_parent(req_id)],
+            abort=True,
+            close_duplex_sessions=close_duplex_sessions,
+        )
+
     async def _dispatch_or_fail_request(
         self,
         dispatch: Callable[[], Awaitable[Any]],
@@ -1194,7 +1226,9 @@ class Orchestrator:
         StageUnavailableError (no live replica / evicted slot) or
         EngineDeadError (replica died mid-submit before the poll loop evicted
         it). Both mean "this request cannot be placed", not "the server is
-        broken": fail the request and keep serving (#4285). Unrelated errors
+        broken": fail the request and keep serving (#4285). An OverflowError
+        from encoding an out-of-range request value is likewise failed as a
+        client error (400) rather than killing the loop. Other unrelated errors
         propagate. Returns False when the request was failed.
 
         ``req_id`` is the request the failure is attributed to; ``dispatch_req_id``
@@ -1237,6 +1271,24 @@ class Orchestrator:
             if dead_replica is not None and dead_replica in pool.live_replica_ids():
                 await self._handle_dead_replica(stage_id, dead_replica, e)
             return False
+        except OverflowError as e:
+            # Catch overflow errors in the request payload to ensure that msgpack
+            # encoding can't kill the engine; instead, we just fail this request.
+            # The frontend should generally validate this as well, but this is needed
+            # in case we get values like seed > 2**64-1 from sampling params.
+            logger.warning(
+                "[Orchestrator] %s dispatch for req=%s hit an unserializable value on stage-%s: %s",
+                operation,
+                req_id,
+                stage_id,
+                e,
+            )
+            await self._fail_request_client_error(
+                req_id,
+                stage_id,
+                f"Request contains a value that cannot be serialized: {e}",
+            )
+            return False
 
     # ---- Shared helpers ----
 
@@ -2292,21 +2344,11 @@ class Orchestrator:
                 "and none were provided; failing the request",
                 request_id,
             )
-            await self.output_async_queue.put(
-                ErrorMessage(
-                    error=(
-                        "async_chunk requires prompt_token_ids on the stage-0 request; "
-                        "an embeds-only prompt cannot prewarm downstream stages"
-                    ),
-                    status_code=HTTPStatus.BAD_REQUEST.value,
-                    error_type=DEFAULT_CLIENT_ERROR_TYPE,
-                    request_id=request_id,
-                    stage_id=0,
-                )
-            )
-            await self._cleanup_request_ids(
-                [request_id, *self._cfg_tracker.cleanup_parent(request_id)],
-                abort=True,
+            await self._fail_request_client_error(
+                request_id,
+                0,
+                "async_chunk requires prompt_token_ids on the stage-0 request; "
+                "an embeds-only prompt cannot prewarm downstream stages",
                 close_duplex_sessions=True,
             )
             return False

vllm_omni/entrypoints/openai/protocol/audio.py

@@ -1,9 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 import math
 from typing import Any, Literal
 
 import numpy as np
 from pydantic import AliasChoices, BaseModel, Field, field_validator, model_validator
 
+# Bound int request fields to avoid overflow issues.
+_INT64_MIN = -(2**63)
+_INT64_MAX = 2**63 - 1
+
 _MAX_EMBEDDING_DIM = 8192
 
 SUPPORTED_AUDIO_FORMATS: frozenset[str] = frozenset({"wav", "pcm", "flac", "mp3", "opus"})
@@ -142,18 +148,21 @@ class OpenAICreateSpeechRequest(BaseModel):
     )
     max_new_tokens: int | None = Field(
         default=None,
+        ge=1,
+        le=_INT64_MAX,
         description="Maximum tokens to generate",
     )
     seed: int | None = Field(
         default=None,
-        ge=0,
-        le=2**63 - 1,
+        ge=_INT64_MIN,
+        le=_INT64_MAX,
         description="Random seed for reproducible generation. When set, ensures "
         "deterministic output for the same input text and seed value.",
     )
     initial_codec_chunk_frames: int | None = Field(
         default=None,
         ge=0,
+        le=_INT64_MAX,
         description="Per-request initial chunk size override. If null, computed dynamically based on server load.",
     )
     non_streaming_mode: bool | None = Field(
@@ -365,10 +374,14 @@ class OpenAICreateAudioGenerateRequest(BaseModel):
     )
     num_inference_steps: int | None = Field(
         default=None,
+        ge=1,
+        le=_INT64_MAX,
         description="Number of inference steps",
     )
     seed: int | None = Field(
         default=None,
+        ge=_INT64_MIN,
+        le=_INT64_MAX,
         description="Random seed for reproducibility",
     )
 
@@ -413,8 +426,8 @@ class SpeechBatchItem(BaseModel):
     ref_audio: str | None = None
     ref_text: str | None = None
     x_vector_only_mode: bool | None = None
-    max_new_tokens: int | None = None
-    initial_codec_chunk_frames: int | None = Field(default=None, ge=0)
+    max_new_tokens: int | None = Field(default=None, ge=1, le=_INT64_MAX)
+    initial_codec_chunk_frames: int | None = Field(default=None, ge=0, le=_INT64_MAX)
     non_streaming_mode: bool | None = None
 
 
@@ -433,8 +446,8 @@ class BatchSpeechRequest(BaseModel):
     ref_audio: str | None = None
     ref_text: str | None = None
     x_vector_only_mode: bool | None = None
-    max_new_tokens: int | None = None
-    initial_codec_chunk_frames: int | None = Field(default=None, ge=0)
+    max_new_tokens: int | None = Field(default=None, ge=1, le=_INT64_MAX)
+    initial_codec_chunk_frames: int | None = Field(default=None, ge=0, le=_INT64_MAX)
     non_streaming_mode: bool | None = None
 
 
@@ -517,10 +530,11 @@ class StreamingSpeechSessionConfig(BaseModel):
     instructions: str | None = None
     response_format: Literal["wav", "pcm", "flac", "mp3", "opus"] = DEFAULT_AUDIO_FORMAT
     speed: float | None = Field(default=1.0, ge=0.25, le=4.0)
-    max_new_tokens: int | None = Field(default=None, ge=1)
+    max_new_tokens: int | None = Field(default=None, ge=1, le=_INT64_MAX)
     initial_codec_chunk_frames: int | None = Field(
         default=None,
         ge=0,
+        le=_INT64_MAX,
         description="Initial chunk size for reduced TTFA. Overrides the deploy config for this session.",
     )
     non_streaming_mode: bool | None = Field(

vllm_omni/entrypoints/openai/protocol/images.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """
 OpenAI-compatible protocol definitions for image generation.
 
@@ -21,6 +21,10 @@ from pydantic import BaseModel, Field, field_validator
 
 from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers
 
+# Bound int request fields to avoid overflow issues.
+_INT64_MIN = -(2**63)
+_INT64_MAX = 2**63 - 1
+
 _IMAGE_FILE_METADATA = {
     "jpg": ("jpg", "image/jpeg"),
     "jpeg": ("jpeg", "image/jpeg"),
@@ -144,7 +148,7 @@ class ImageGenerationRequest(BaseModel):
         default=None,
         description="Optional model-specific parameters passed directly to the model's extra_args.",
     )
-    seed: int | None = Field(default=None, description="Random seed for reproducibility")
+    seed: int | None = Field(default=None, ge=_INT64_MIN, le=_INT64_MAX, description="Random seed for reproducibility")
     generator_device: str | None = Field(
         default=None,
         description="Device for the seeded torch.Generator (e.g. 'cpu', 'cuda'). Defaults to the runner's device.",

vllm_omni/entrypoints/openai/protocol/videos.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """
 OpenAI-compatible protocol definitions for video generation.
 
@@ -20,6 +20,10 @@ from pydantic import BaseModel, ConfigDict, Field, StringConstraints, field_vali
 from vllm_omni.entrypoints.openai.image_api_utils import parse_size
 from vllm_omni.inputs.data import DIFFUSION_QUALITY_LEVELS
 
+# Bound int request fields to avoid overflow issues.
+_INT64_MIN = -(2**63)
+_INT64_MAX = 2**63 - 1
+
 
 class VideoGenerationStatus(str, Enum):
     QUEUED = "queued"
@@ -48,10 +52,10 @@ def file_extension(media_type: str):
 class VideoParams(BaseModel):
     """Optional block for video-specific parameters."""
 
-    width: int | None = Field(default=None, ge=1, description="Video width in pixels")
-    height: int | None = Field(default=None, ge=1, description="Video height in pixels")
-    num_frames: int | None = Field(default=None, ge=1, description="Number of frames")
-    fps: int | None = Field(default=None, ge=1, description="Frames per second for output video")
+    width: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Video width in pixels")
+    height: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Video height in pixels")
+    num_frames: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Number of frames")
+    fps: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Frames per second for output video")
 
     @property
     def size(self) -> str | None:
@@ -140,10 +144,10 @@ class VideoGenerationRequest(BaseModel):
     user: str | None = Field(default=None, description="User identifier for tracking")
 
     # Video-specific fields (top-level for OpenAI-style compatibility)
-    width: int | None = Field(default=None, ge=1, description="Video width in pixels")
-    height: int | None = Field(default=None, ge=1, description="Video height in pixels")
-    fps: int | None = Field(default=None, ge=1, description="Frames per second for output video")
-    num_frames: int | None = Field(default=None, ge=1, description="Number of frames to generate")
+    width: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Video width in pixels")
+    height: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Video height in pixels")
+    fps: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Frames per second for output video")
+    num_frames: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Number of frames to generate")
     aspect_ratio: str | None = Field(
         default=None,
         description=(
@@ -151,7 +155,12 @@ class VideoGenerationRequest(BaseModel):
             "FL2VA follows the input image; Ref2VA defaults to 16:9."
         ),
     )
-    short_edge: int | None = Field(default=None, ge=1, description="MiniMax H3 output short edge in pixels")
+    short_edge: int | None = Field(
+        default=None,
+        ge=1,
+        le=_INT64_MAX,
+        description="MiniMax H3 output short edge in pixels",
+    )
     num_outputs_per_prompt: int = Field(
         default=1,
         ge=1,
@@ -218,7 +227,7 @@ class VideoGenerationRequest(BaseModel):
         le=20.0,
         description="True CFG scale (model-specific parameter, may be ignored if not supported)",
     )
-    seed: int | None = Field(default=None, description="Random seed for reproducibility")
+    seed: int | None = Field(default=None, ge=_INT64_MIN, le=_INT64_MAX, description="Random seed for reproducibility")
     generate_sound: bool = Field(
         default=False,
         description="Request model-generated audio for video models that support sound generation.",
@@ -237,6 +246,7 @@ class VideoGenerationRequest(BaseModel):
     frame_interpolation_exp: int = Field(
         default=1,
         ge=1,
+        le=_INT64_MAX,
         description="Interpolation exponent: 1=2x temporal resolution, 2=4x, etc.",
     )
     frame_interpolation_scale: float = Field(