Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
89 changes: 89 additions & 0 deletions tests/test_omni_request.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Regression tests for OmniRequest substitutability with the base Request.

vllm-omni rebinds ``vllm.v1.request.Request`` to ``OmniRequest`` at import
time (see ``vllm_omni/patch.py``). vLLM core constructs ``Request`` positionally
in some paths (notably ``vllm/v1/worker/gpu/warmup.py::warmup_kernels`` for V2
model-runner architectures like Qwen3ForCausalLM). The omni-specific params
must therefore be keyword-only so positional construction stays Liskov-
substitutable with the base class — including base-style calls that pass
``prompt_embeds`` (itself a positional-capable base param) positionally.
"""

import inspect

import numpy as np
import pytest
import torch
from vllm.sampling_params import SamplingParams
from vllm.v1.request import Request

from vllm_omni.engine import PromptEmbedsPayload
from vllm_omni.request import OmniRequest

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


def test_omni_params_are_keyword_only():
"""The three omni params must be keyword-only after ``*args``.

Guards against re-introducing the import-time rebind bug where omni params
came first positionally and broke positional ``Request(...)`` construction.
"""
params = inspect.signature(OmniRequest.__init__).parameters
for name in ("prompt_embeds", "external_req_id", "additional_information"):
assert params[name].kind is inspect.Parameter.KEYWORD_ONLY, name


def test_positional_construction_matches_base_request():
"""Reproduces the warmup_kernels call: Request(id, tokens, sp, pp)."""
req = OmniRequest("req-0", [1, 2, 3], SamplingParams(), None)

assert isinstance(req, Request)
assert req.request_id == "req-0"
assert req.prompt_token_ids == [1, 2, 3]
# Omni params default cleanly when constructed positionally.
assert req.external_req_id is None
assert req.additional_information is None
assert req.prompt_embeds_payload is None


def test_positional_prompt_embeds_does_not_collide():
"""Base-style positional call that includes ``prompt_embeds``.

``prompt_embeds`` is a positional-capable base param (after
``arrival_time``). A base-style call must not trip
``got multiple values for argument 'prompt_embeds'`` — the subclass must
not also inject its keyword override when the value arrived positionally.
Positional order: request_id, prompt_token_ids, sampling_params,
pooling_params, client_index, arrival_time, prompt_embeds.
"""
embeds = torch.arange(6, dtype=torch.float32).reshape(2, 3)
req = OmniRequest("req-pe", [1, 2], SamplingParams(), None, 0, None, embeds)

assert isinstance(req, Request)
assert torch.equal(req.prompt_embeds, embeds)
# Positional tensor is not a serialized payload.
assert req.prompt_embeds_payload is None
assert req.external_req_id is None


def test_keyword_omni_params_round_trip():
"""Keyword omni params are preserved; serialized embeds are decoded."""
arr = np.arange(6, dtype=np.float32).reshape(2, 3)
payload = PromptEmbedsPayload(data=arr.tobytes(), shape=[2, 3], dtype="float32")

req = OmniRequest(
request_id="req-1",
prompt_token_ids=[7, 8],
sampling_params=SamplingParams(),
pooling_params=None,
prompt_embeds=payload,
external_req_id="ext-1",
)

assert req.external_req_id == "ext-1"
# The serialized payload is retained and decoded into a tensor on the base.
assert req.prompt_embeds_payload is payload
assert torch.equal(req.prompt_embeds, torch.from_numpy(arr))
7 changes: 4 additions & 3 deletions vllm_omni/request.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,15 +30,16 @@ class OmniRequest(Request):

def __init__(
self,
*args,
prompt_embeds: PromptEmbedsPayload | torch.Tensor | None = None,
# Optional external request ID for tracking
external_req_id: str | None = None,
additional_information: AdditionalInformationPayload | None = None,
*args,
**kwargs,
):
prompt_embeds_tensor = self._maybe_decode_prompt_embeds(prompt_embeds)
super().__init__(prompt_embeds=prompt_embeds_tensor, *args, **kwargs)
if prompt_embeds is not None:
kwargs["prompt_embeds"] = self._maybe_decode_prompt_embeds(prompt_embeds)
super().__init__(*args, **kwargs)
# Preserve serialized prompt embeddings payload (optional)
self.prompt_embeds_payload: PromptEmbedsPayload | None = (
prompt_embeds if isinstance(prompt_embeds, PromptEmbedsPayload) else None
Expand Down
Loading