134 lines
4.2 KiB
Python
134 lines
4.2 KiB
Python
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
from vllm.v1.core.sched.output import NewRequestData
|
|
|
|
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
|
|
from vllm_omni.core.sched.output import OmniNewRequestData
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def test_omni_new_request_data_copies_payloads():
|
|
prompt_embeds = torch.randn(2, 3)
|
|
additional_information = {
|
|
"speaker": ["test"],
|
|
"codes": torch.tensor([1, 2], dtype=torch.int64),
|
|
}
|
|
request = SimpleNamespace(
|
|
request_id="req-1",
|
|
external_req_id="ext-1",
|
|
prompt_token_ids=[101, 102],
|
|
mm_features=None,
|
|
sampling_params=None,
|
|
pooling_params=None,
|
|
num_computed_tokens=0,
|
|
lora_request=None,
|
|
prompt_embeds=prompt_embeds,
|
|
additional_information=additional_information,
|
|
)
|
|
|
|
data = OmniNewRequestData.from_request(request, ([0, 1],), prefill_token_ids=[101, 102])
|
|
|
|
assert data.prompt_embeds is prompt_embeds
|
|
assert data.additional_information is additional_information
|
|
assert data.prefill_token_ids == [101, 102]
|
|
|
|
|
|
def test_omni_new_request_data_allows_missing_payloads():
|
|
request = SimpleNamespace(
|
|
request_id="req-2",
|
|
external_req_id="ext-2",
|
|
prompt_token_ids=[201, 202],
|
|
mm_features=None,
|
|
sampling_params=None,
|
|
pooling_params=None,
|
|
num_computed_tokens=0,
|
|
lora_request=None,
|
|
prompt_embeds=None,
|
|
additional_information=None,
|
|
)
|
|
|
|
data = OmniNewRequestData.from_request(request, ([0],), prefill_token_ids=None)
|
|
|
|
assert data.prompt_embeds is None
|
|
assert data.additional_information is None
|
|
|
|
|
|
def test_from_base_preserves_every_upstream_field_and_adds_omni_payloads():
|
|
prompt_embeds = torch.randn(2, 3)
|
|
base = NewRequestData(
|
|
req_id="req-base",
|
|
prompt_token_ids=[1, 2],
|
|
mm_features=[],
|
|
sampling_params=None,
|
|
pooling_params=None,
|
|
block_ids=([3, 4],),
|
|
num_computed_tokens=1,
|
|
lora_request=None,
|
|
prompt_embeds=prompt_embeds,
|
|
prompt_is_token_ids=False,
|
|
prefill_token_ids=[1],
|
|
)
|
|
request = SimpleNamespace(
|
|
external_req_id="external-base",
|
|
additional_information={"payload": "value"},
|
|
model_intermediate_buffer={"hidden": "state"},
|
|
)
|
|
|
|
data = OmniNewRequestData.from_base(base, request)
|
|
|
|
for field_name in NewRequestData.__dataclass_fields__:
|
|
assert getattr(data, field_name) is getattr(base, field_name)
|
|
assert data.external_req_id == "external-base"
|
|
assert data.additional_information == {"payload": "value"}
|
|
assert data.model_intermediate_buffer == {"hidden": "state"}
|
|
|
|
|
|
def test_rewrap_converts_fallback_entries_and_keeps_fast_path_entries():
|
|
base = NewRequestData(
|
|
req_id="req-fallback",
|
|
prompt_token_ids=[1],
|
|
mm_features=[],
|
|
sampling_params=None,
|
|
pooling_params=None,
|
|
block_ids=([2],),
|
|
num_computed_tokens=0,
|
|
lora_request=None,
|
|
prompt_embeds=None,
|
|
prompt_is_token_ids=True,
|
|
prefill_token_ids=[1],
|
|
)
|
|
fast_path = OmniNewRequestData(
|
|
req_id="req-fast",
|
|
prompt_token_ids=[3],
|
|
mm_features=[],
|
|
sampling_params=None,
|
|
pooling_params=None,
|
|
block_ids=([4],),
|
|
num_computed_tokens=0,
|
|
lora_request=None,
|
|
prompt_embeds=None,
|
|
prompt_is_token_ids=True,
|
|
prefill_token_ids=[3],
|
|
additional_information={"fast": True},
|
|
)
|
|
scheduler = OmniSchedulerMixin()
|
|
scheduler.requests = {
|
|
"req-fallback": SimpleNamespace(
|
|
external_req_id="external-fallback",
|
|
additional_information={"fallback": True},
|
|
model_intermediate_buffer=None,
|
|
)
|
|
}
|
|
output = SimpleNamespace(scheduled_new_reqs=[base, fast_path])
|
|
|
|
scheduler._rewrap_scheduled_new_reqs(output)
|
|
|
|
fallback_data, fast_path_data = output.scheduled_new_reqs
|
|
assert isinstance(fallback_data, OmniNewRequestData)
|
|
assert fallback_data.prefill_token_ids == [1]
|
|
assert fallback_data.additional_information == {"fallback": True}
|
|
assert fast_path_data is fast_path
|