[Core] Split Omni connector model runner mixin (#6903)
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
This commit is contained in:
NATURE authored and GitHub committed
1 parent
e0d0879794
commit
5b002727b0
7 files changed
+2463
-2569
No files matched your search
@@ -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
|
||||
"""Unit tests for OmniConnectorModelRunnerMixin.
|
||||
|
||||
These tests use a mock connector (in-memory dict store) and do not require
|
||||
@@ -376,6 +376,19 @@ class TestFinishedLoadReqsDrain(unittest.TestCase):
|
||||
|
||||
|
||||
class TestLoadCustomFuncSelection(unittest.TestCase):
|
||||
def test_uses_validator_override_from_public_mixin(self):
|
||||
config = SimpleNamespace(
|
||||
async_chunk=True,
|
||||
custom_process_next_stage_input_func=f"{__name__}._make_request",
|
||||
)
|
||||
|
||||
with patch.object(MixinHost, "_is_connector_payload_builder", return_value=False) as validator:
|
||||
selected_path, func = MixinHost._load_custom_func(config)
|
||||
|
||||
validator.assert_called_once_with(_make_request)
|
||||
assert selected_path is None
|
||||
assert func is None
|
||||
|
||||
def test_skips_non_payload_stage_input_processors_for_full_payload_mode(self):
|
||||
incompatible_paths = [
|
||||
"vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav",
|
||||
@@ -861,7 +874,7 @@ class TestLocalPayloadCacheLifecycle(unittest.TestCase):
|
||||
host._omni_connector.get.return_value = connector_result
|
||||
tp_group = _FakeTPGroup(world_size=2, rank_in_group=0)
|
||||
|
||||
with patch("vllm_omni.worker.omni_connector_model_runner_mixin.get_tp_group", return_value=tp_group):
|
||||
with patch.object(host, "_get_local_tp_group", return_value=tp_group):
|
||||
made_progress = host._poll_single_request("r1")
|
||||
|
||||
self.assertTrue(made_progress)
|
||||
@@ -882,7 +895,7 @@ class TestLocalPayloadCacheLifecycle(unittest.TestCase):
|
||||
host._get_req_chunk["r1"] = 0
|
||||
tp_group = _FakeTPGroup(world_size=2, rank_in_group=1)
|
||||
|
||||
with patch("vllm_omni.worker.omni_connector_model_runner_mixin.get_tp_group", return_value=tp_group):
|
||||
with patch.object(host, "_get_local_tp_group", return_value=tp_group):
|
||||
made_progress = host._poll_single_request("r1")
|
||||
|
||||
self.assertFalse(made_progress)
|
||||
@@ -900,7 +913,7 @@ class TestLocalPayloadCacheLifecycle(unittest.TestCase):
|
||||
payload = {"tok": [10], "finished": torch.tensor(True)}
|
||||
tp_group = _FakeTPGroup(world_size=2, rank_in_group=1, follower_result={"r1": payload})
|
||||
|
||||
with patch("vllm_omni.worker.omni_connector_model_runner_mixin.get_tp_group", return_value=tp_group):
|
||||
with patch.object(host, "_get_local_tp_group", return_value=tp_group):
|
||||
results = host.recv_full_payload_inputs(scheduler_output=None)
|
||||
|
||||
self.assertEqual(results, {"r1": payload})
|
||||
@@ -936,7 +949,7 @@ class TestTPAsyncChunkFanout(unittest.TestCase):
|
||||
host._omni_connector.get.return_value = (payload, 123)
|
||||
tp_group = _FakeTPGroup(world_size=2, rank_in_group=0)
|
||||
|
||||
with patch("vllm_omni.worker.omni_connector_model_runner_mixin.get_tp_group", return_value=tp_group):
|
||||
with patch.object(host, "_get_local_tp_group", return_value=tp_group):
|
||||
made_progress = host._poll_single_request("r1")
|
||||
|
||||
self.assertTrue(made_progress)
|
||||
@@ -951,7 +964,7 @@ class TestTPAsyncChunkFanout(unittest.TestCase):
|
||||
host = self._make_host(rank=1)
|
||||
tp_group = _FakeTPGroup(world_size=2, rank_in_group=1)
|
||||
|
||||
with patch("vllm_omni.worker.omni_connector_model_runner_mixin.get_tp_group", return_value=tp_group):
|
||||
with patch.object(host, "_get_local_tp_group", return_value=tp_group):
|
||||
made_progress = host._poll_single_request("r1")
|
||||
|
||||
self.assertFalse(made_progress)
|
||||
@@ -976,7 +989,7 @@ class TestTPAsyncChunkFanout(unittest.TestCase):
|
||||
}
|
||||
tp_group = _FakeTPGroup(world_size=2, rank_in_group=1, follower_result=packet)
|
||||
|
||||
with patch("vllm_omni.worker.omni_connector_model_runner_mixin.get_tp_group", return_value=tp_group):
|
||||
with patch.object(host, "_get_local_tp_group", return_value=tp_group):
|
||||
output = host.get_omni_connector_output()
|
||||
|
||||
self.assertEqual(output.chunk_ready_req_ids, {"r1"})
|
||||
@@ -1049,6 +1062,45 @@ class TestKVTransferLifecycle(unittest.TestCase):
|
||||
class TestAsyncPayloadLifecycle(unittest.TestCase):
|
||||
"""Regression tests for async payload delivery lifecycle."""
|
||||
|
||||
def test_accumulate_payload_concatenates_chunks(self):
|
||||
host = MixinHost()
|
||||
host._send_side_request_payload = {}
|
||||
first = host._accumulate_payload(
|
||||
"r1",
|
||||
{
|
||||
"embed": {"decode": torch.tensor([[1.0, 2.0]])},
|
||||
"ids": {"output": [1]},
|
||||
"meta": {"finished": False},
|
||||
},
|
||||
)
|
||||
merged = host._accumulate_payload(
|
||||
"r1",
|
||||
{
|
||||
"embed": {"decode": torch.tensor([[3.0, 4.0], [5.0, 6.0]])},
|
||||
"ids": {"output": [2, 3]},
|
||||
"meta": {"finished": True},
|
||||
},
|
||||
)
|
||||
torch.testing.assert_close(merged["embed"]["decode"], torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]))
|
||||
self.assertEqual(merged["ids"]["output"], [1, 2, 3])
|
||||
self.assertIs(merged["meta"]["finished"], True)
|
||||
self.assertEqual(first["embed"]["decode"].shape, (1, 2))
|
||||
self.assertEqual(first["ids"]["output"], [1])
|
||||
self.assertIs(first["meta"]["finished"], False)
|
||||
|
||||
def test_accumulate_payload_replaces_override_keys(self):
|
||||
host = MixinHost()
|
||||
host._send_side_request_payload = {}
|
||||
host._accumulate_payload("r1", {"embed": {"decode": torch.ones(2, 2)}, "ids": {"output": [1, 2]}})
|
||||
payload = {
|
||||
"embed": {"decode": torch.zeros(1, 2)},
|
||||
"ids": {"output": [3]},
|
||||
"meta": {"override_keys": [["embed", "decode"], ["ids", "output"]]},
|
||||
}
|
||||
merged = host._accumulate_payload("r1", payload)
|
||||
torch.testing.assert_close(merged["embed"]["decode"], payload["embed"]["decode"])
|
||||
self.assertEqual(merged["ids"]["output"], [3])
|
||||
|
||||
def test_send_side_request_payload_not_cleared_before_payload_is_consumable(self):
|
||||
host = MixinHost()
|
||||
host.init_omni_connectors(
|
||||
|
||||
@@ -1,166 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Characterization tests for ``vllm_omni.worker.payload_span``.
|
||||
|
||||
These pin the CURRENT behaviour of the thinker-decode span helpers so the later
|
||||
worker refactors that touch span
|
||||
plumbing have a regression net. Pure CPU — no model, no CUDA.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.worker.payload_span import (
|
||||
get_tensor_span,
|
||||
get_tensor_span_row,
|
||||
merge_tensor_spans,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _rows(n: int, *, start_val: int = 0, hidden: int = 1) -> torch.Tensor:
|
||||
"""A ``[n, hidden]`` tensor whose first column counts up from ``start_val``.
|
||||
|
||||
``[n, hidden]`` mirrors the connector async-chunk shape (tokens x hidden).
|
||||
"""
|
||||
base = torch.arange(start_val, start_val + n, dtype=torch.float32).unsqueeze(1)
|
||||
return base.repeat(1, hidden)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# get_tensor_span #
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_get_tensor_span_valid_returns_triplet():
|
||||
tensor = _rows(3, hidden=4)
|
||||
payload = {"emb": tensor, "start": 5, "end": 8}
|
||||
|
||||
span = get_tensor_span(payload, tensor_key="emb", start_key="start", end_key="end")
|
||||
|
||||
assert span is not None
|
||||
got_tensor, start, end = span
|
||||
assert got_tensor is tensor
|
||||
assert (start, end) == (5, 8)
|
||||
|
||||
|
||||
def test_get_tensor_span_non_tensor_returns_none():
|
||||
payload = {"emb": None, "start": 0, "end": 0}
|
||||
assert get_tensor_span(payload, tensor_key="emb", start_key="start", end_key="end") is None
|
||||
|
||||
|
||||
def test_get_tensor_span_non_int_bounds_returns_none():
|
||||
tensor = _rows(2)
|
||||
# start is a float, not an int -> rejected.
|
||||
payload = {"emb": tensor, "start": 0.0, "end": 2}
|
||||
assert get_tensor_span(payload, tensor_key="emb", start_key="start", end_key="end") is None
|
||||
|
||||
|
||||
def test_get_tensor_span_negative_start_returns_none():
|
||||
tensor = _rows(1)
|
||||
payload = {"emb": tensor, "start": -1, "end": 0}
|
||||
assert get_tensor_span(payload, tensor_key="emb", start_key="start", end_key="end") is None
|
||||
|
||||
|
||||
def test_get_tensor_span_end_before_start_returns_none():
|
||||
tensor = _rows(2)
|
||||
payload = {"emb": tensor, "start": 5, "end": 3}
|
||||
assert get_tensor_span(payload, tensor_key="emb", start_key="start", end_key="end") is None
|
||||
|
||||
|
||||
def test_get_tensor_span_length_mismatch_returns_none():
|
||||
tensor = _rows(3) # tensor.shape[0] == 3
|
||||
payload = {"emb": tensor, "start": 0, "end": 2} # end - start == 2
|
||||
assert get_tensor_span(payload, tensor_key="emb", start_key="start", end_key="end") is None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# merge_tensor_spans #
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_merge_none_operand_returns_none():
|
||||
span = (_rows(2), 0, 2)
|
||||
assert merge_tensor_spans(None, span) is None
|
||||
assert merge_tensor_spans(span, None) is None
|
||||
|
||||
|
||||
def test_merge_adjacent_spans_concatenates():
|
||||
existing = (_rows(3, start_val=0), 0, 3)
|
||||
incoming = (_rows(2, start_val=3), 3, 5)
|
||||
|
||||
merged = merge_tensor_spans(existing, incoming)
|
||||
|
||||
assert merged is not None
|
||||
tensor, start, end = merged
|
||||
assert (start, end) == (0, 5)
|
||||
assert tensor.shape[0] == 5
|
||||
assert tensor[:, 0].tolist() == [0.0, 1.0, 2.0, 3.0, 4.0]
|
||||
|
||||
|
||||
def test_merge_overlapping_spans_trims_incoming():
|
||||
# existing covers [0, 3); incoming covers [2, 5) -> 1 row of overlap.
|
||||
existing = (_rows(3, start_val=0), 0, 3)
|
||||
incoming = (_rows(3, start_val=100), 2, 5)
|
||||
|
||||
merged = merge_tensor_spans(existing, incoming)
|
||||
|
||||
assert merged is not None
|
||||
tensor, start, end = merged
|
||||
assert (start, end) == (0, 5)
|
||||
# First 3 rows from existing, then incoming's tail (overlap=1 row trimmed).
|
||||
assert tensor[:, 0].tolist() == [0.0, 1.0, 2.0, 101.0, 102.0]
|
||||
|
||||
|
||||
def test_merge_fully_overlapping_incoming_keeps_existing():
|
||||
# incoming [1, 3) is fully covered by existing [0, 3): overlap >= incoming len.
|
||||
existing = (_rows(3, start_val=0), 0, 3)
|
||||
incoming = (_rows(2, start_val=100), 1, 3)
|
||||
|
||||
merged = merge_tensor_spans(existing, incoming)
|
||||
|
||||
assert merged is not None
|
||||
tensor, start, end = merged
|
||||
assert (start, end) == (0, 3)
|
||||
assert tensor is existing[0] # returned unchanged
|
||||
|
||||
|
||||
def test_merge_non_contiguous_spans_returns_none():
|
||||
# gap between existing end (3) and incoming start (5).
|
||||
existing = (_rows(3, start_val=0), 0, 3)
|
||||
incoming = (_rows(2, start_val=100), 5, 7)
|
||||
|
||||
assert merge_tensor_spans(existing, incoming) is None
|
||||
|
||||
|
||||
def test_merge_casts_incoming_to_existing_dtype_and_device():
|
||||
existing = (_rows(2, start_val=0).to(torch.float32), 0, 2)
|
||||
incoming = (torch.arange(2, 4, dtype=torch.int64).unsqueeze(1), 2, 4)
|
||||
|
||||
merged = merge_tensor_spans(existing, incoming)
|
||||
|
||||
assert merged is not None
|
||||
tensor, _, _ = merged
|
||||
assert tensor.dtype == torch.float32
|
||||
assert tensor.device == existing[0].device
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# get_tensor_span_row #
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_get_tensor_span_row_in_range():
|
||||
span = (_rows(3, start_val=10), 5, 8) # rows map to indices 5, 6, 7
|
||||
row = get_tensor_span_row(span, 6)
|
||||
assert row is not None
|
||||
assert row[0].item() == 11.0 # second row (start_val 10 + 1)
|
||||
|
||||
|
||||
def test_get_tensor_span_row_none_span_returns_none():
|
||||
assert get_tensor_span_row(None, 0) is None
|
||||
|
||||
|
||||
def test_get_tensor_span_row_below_start_returns_none():
|
||||
span = (_rows(3), 5, 8)
|
||||
assert get_tensor_span_row(span, 4) is None
|
||||
|
||||
|
||||
def test_get_tensor_span_row_at_or_above_end_returns_none():
|
||||
span = (_rows(3), 5, 8)
|
||||
assert get_tensor_span_row(span, 8) is None
|
||||
@@ -0,0 +1,3 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Internal model-runner components for Omni connector transport."""
|
||||
+1322
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
@@ -1,64 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Helpers for explicit thinker decode span metadata."""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
THINKER_DECODE_EMBEDDINGS_KEY = "thinker_decode_embeddings"
|
||||
THINKER_OUTPUT_TOKEN_IDS_KEY = "thinker_output_token_ids"
|
||||
THINKER_DECODE_TOKEN_START_KEY = "thinker_decode_embeddings_token_start"
|
||||
THINKER_DECODE_TOKEN_END_KEY = "thinker_decode_embeddings_token_end"
|
||||
|
||||
CACHED_THINKER_DECODE_EMBEDDINGS_KEY = "cached_thinker_decode_embeddings"
|
||||
CACHED_THINKER_DECODE_TOKEN_START_KEY = "cached_thinker_decode_embeddings_token_start"
|
||||
CACHED_THINKER_DECODE_TOKEN_END_KEY = "cached_thinker_decode_embeddings_token_end"
|
||||
|
||||
TensorSpan = tuple[torch.Tensor, int, int]
|
||||
|
||||
|
||||
def get_tensor_span(payload: Mapping[str, Any], *, tensor_key: str, start_key: str, end_key: str) -> TensorSpan | None:
|
||||
tensor = payload.get(tensor_key)
|
||||
start = payload.get(start_key)
|
||||
end = payload.get(end_key)
|
||||
if not isinstance(tensor, torch.Tensor):
|
||||
return None
|
||||
if not isinstance(start, int) or not isinstance(end, int):
|
||||
return None
|
||||
if start < 0 or end < start or (end - start) != int(tensor.shape[0]):
|
||||
return None
|
||||
return tensor, start, end
|
||||
|
||||
|
||||
def merge_tensor_spans(existing_span: TensorSpan | None, incoming_span: TensorSpan | None) -> TensorSpan | None:
|
||||
if existing_span is None or incoming_span is None:
|
||||
return None
|
||||
|
||||
existing_tensor, existing_start, existing_end = existing_span
|
||||
incoming_tensor, incoming_start, incoming_end = incoming_span
|
||||
if incoming_tensor.device != existing_tensor.device or incoming_tensor.dtype != existing_tensor.dtype:
|
||||
incoming_tensor = incoming_tensor.to(device=existing_tensor.device, dtype=existing_tensor.dtype)
|
||||
if incoming_start == existing_end:
|
||||
return torch.cat([existing_tensor, incoming_tensor], dim=0), existing_start, incoming_end
|
||||
if incoming_start < existing_end:
|
||||
overlap = existing_end - incoming_start
|
||||
if overlap >= int(incoming_tensor.shape[0]):
|
||||
return existing_tensor, existing_start, existing_end
|
||||
trimmed_tensor = incoming_tensor[overlap:]
|
||||
return (
|
||||
torch.cat([existing_tensor, trimmed_tensor], dim=0),
|
||||
existing_start,
|
||||
existing_end + int(trimmed_tensor.shape[0]),
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def get_tensor_span_row(span: TensorSpan | None, index: int) -> torch.Tensor | None:
|
||||
if span is None:
|
||||
return None
|
||||
tensor, start, end = span
|
||||
if index < start or index >= end:
|
||||
return None
|
||||
return tensor[index - start]
|
||||
Reference in new issue
Block a user