1740 lines
65 KiB
Python
1740 lines
65 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# 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
|
||
GPU or vLLM runtime.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import time
|
||
import unittest
|
||
from copy import deepcopy
|
||
from types import SimpleNamespace
|
||
from typing import Any
|
||
from unittest.mock import MagicMock, patch
|
||
|
||
import pytest
|
||
import torch
|
||
|
||
from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
|
||
OmniKVTransferManager,
|
||
)
|
||
from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
|
||
from vllm_omni.distributed.omni_connectors.utils.initialization import (
|
||
resolve_connector_spec,
|
||
)
|
||
from vllm_omni.outputs import OmniConnectorOutput
|
||
from vllm_omni.worker.omni_connector_model_runner_mixin import (
|
||
OmniConnectorModelRunnerMixin,
|
||
)
|
||
|
||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Mock helpers
|
||
# ------------------------------------------------------------------ #
|
||
|
||
|
||
class MockConnector:
|
||
"""In-memory connector for testing (mimics OmniConnectorBase)."""
|
||
|
||
def __init__(self, stage_id: int = 0):
|
||
self.stage_id = stage_id
|
||
self._store: dict[str, Any] = {}
|
||
|
||
def put(self, from_stage, to_stage, put_key, data):
|
||
key = f"{from_stage}_{to_stage}_{put_key}"
|
||
self._store[key] = data
|
||
return True, len(str(data)), None
|
||
|
||
def get(self, from_stage, to_stage, get_key, metadata=None):
|
||
key = f"{from_stage}_{to_stage}_{get_key}"
|
||
data = self._store.pop(key, None)
|
||
if data is None:
|
||
return None
|
||
return data, len(str(data))
|
||
|
||
def close(self):
|
||
pass
|
||
|
||
|
||
def _make_model_config(
|
||
stage_id: int = 0,
|
||
async_chunk: bool = False,
|
||
worker_type: str = "ar",
|
||
custom_func: str | None = None,
|
||
) -> SimpleNamespace:
|
||
return SimpleNamespace(
|
||
stage_connector_config=None,
|
||
async_chunk=async_chunk,
|
||
worker_type=worker_type,
|
||
custom_process_next_stage_input_func=custom_func,
|
||
)
|
||
|
||
|
||
def _make_request(req_id: str, external_req_id: str | None = None):
|
||
r = SimpleNamespace(
|
||
request_id=req_id,
|
||
external_req_id=external_req_id or req_id,
|
||
additional_information=None,
|
||
prompt_token_ids=[],
|
||
num_computed_tokens=0,
|
||
)
|
||
return r
|
||
|
||
|
||
class MixinHost(OmniConnectorModelRunnerMixin):
|
||
"""Minimal class that mixes in the mixin for testing."""
|
||
|
||
pass
|
||
|
||
|
||
class _FakeTPGroup:
|
||
def __init__(self, *, world_size: int, rank_in_group: int, follower_result: Any = None):
|
||
self.world_size = world_size
|
||
self.rank_in_group = rank_in_group
|
||
self.follower_result = follower_result
|
||
self.broadcast_inputs: list[Any] = []
|
||
|
||
def broadcast_object(self, obj: Any | None = None, src: int = 0):
|
||
self.broadcast_inputs.append(obj)
|
||
if self.rank_in_group == src:
|
||
return obj
|
||
return self.follower_result
|
||
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Test cases
|
||
# ------------------------------------------------------------------ #
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("role", "custom_func", "expected"),
|
||
[
|
||
("sender", None, False),
|
||
("sender", "package.build_payload", True),
|
||
("receiver", None, True),
|
||
],
|
||
)
|
||
def test_init_payload_connector_ownership(role, custom_func, expected):
|
||
model_config = _make_model_config(custom_func=custom_func)
|
||
model_config.stage_connector_config = {
|
||
"name": "MooncakeTransferEngineConnector",
|
||
"extra": {"role": role},
|
||
}
|
||
|
||
host = MixinHost()
|
||
connector = MockConnector()
|
||
with (
|
||
patch.object(host, "_create_connector", return_value=connector) as create,
|
||
patch.object(host, "_load_custom_func", return_value=(None, None)),
|
||
):
|
||
host.init_omni_connectors(model_config)
|
||
|
||
assert (create.call_count == 1) is expected
|
||
assert (host._omni_connector is connector) is expected
|
||
host.shutdown_omni_connectors()
|
||
|
||
|
||
class TestMixinAsyncChunkSendRecv(unittest.TestCase):
|
||
"""Test 2: Async chunk send/recv + bg threads."""
|
||
|
||
def test_send_chunk_passes_is_finished_and_connector(self):
|
||
connector = MockConnector(stage_id=0)
|
||
|
||
sender = MixinHost()
|
||
sender.init_omni_connectors(
|
||
model_config=_make_model_config(stage_id=0, async_chunk=True),
|
||
)
|
||
sender._omni_connector = connector
|
||
sender._stage_id = 0
|
||
sender._async_chunk = True
|
||
|
||
seen = {}
|
||
|
||
def mock_process(transfer_manager, pooling_output, request, is_finished=False):
|
||
seen["connector"] = transfer_manager.connector
|
||
seen["is_finished"] = is_finished
|
||
return {"data": pooling_output, "finished": is_finished}
|
||
|
||
sender._custom_process_func = mock_process
|
||
|
||
request = _make_request("req-1", "ext-req-1")
|
||
request.is_finished = lambda: True
|
||
sender._send_single_request(
|
||
{
|
||
"stage_id": 0,
|
||
"next_stage_id": 1,
|
||
"request_id": "ext-req-1",
|
||
"request": request,
|
||
"pooling_output": {"value": 42},
|
||
}
|
||
)
|
||
self.assertIs(seen["connector"], connector)
|
||
self.assertTrue(seen["is_finished"])
|
||
|
||
sender.shutdown_omni_connectors()
|
||
|
||
def test_send_chunk_does_not_retry_real_type_error(self):
|
||
connector = MockConnector(stage_id=0)
|
||
|
||
sender = MixinHost()
|
||
sender.init_omni_connectors(
|
||
model_config=_make_model_config(stage_id=0, async_chunk=True),
|
||
)
|
||
sender._omni_connector = connector
|
||
sender._stage_id = 0
|
||
sender._async_chunk = True
|
||
|
||
seen = {"calls": 0}
|
||
|
||
def broken_process(transfer_manager, pooling_output, request, is_finished=""):
|
||
seen["calls"] += 1
|
||
return {"data": is_finished + "tail"}
|
||
|
||
sender._custom_process_func = broken_process
|
||
|
||
request = _make_request("req-1", "ext-req-1")
|
||
request.is_finished = lambda: True
|
||
ok = sender.send_chunk(request, pooling_output={"value": 42})
|
||
self.assertFalse(ok)
|
||
self.assertEqual(seen["calls"], 1)
|
||
|
||
sender.shutdown_omni_connectors()
|
||
|
||
|
||
class TestMixinKVCacheTransfer(unittest.TestCase):
|
||
"""Test 3: KV cache delegation to OmniKVTransferManager."""
|
||
|
||
def test_send_kv_delegates(self):
|
||
mock_kvm = MagicMock()
|
||
mock_kvm.handle_finished_requests_kv_transfer.return_value = ["req-1"]
|
||
|
||
host = MixinHost()
|
||
host.init_omni_connectors(
|
||
model_config=_make_model_config(),
|
||
kv_transfer_manager=mock_kvm,
|
||
)
|
||
|
||
result = host.send_kv_cache(
|
||
finished_reqs={"req-1": {"seq_len": 10, "block_ids": [0]}},
|
||
kv_caches=[],
|
||
block_size=16,
|
||
cache_dtype="float16",
|
||
)
|
||
self.assertEqual(result, ["req-1"])
|
||
mock_kvm.handle_finished_requests_kv_transfer.assert_called_once()
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_recv_kv_delegates(self):
|
||
mock_kvm = MagicMock()
|
||
mock_kvm.receive_kv_cache_for_request.return_value = ({"layer_blocks": {}}, 100)
|
||
|
||
host = MixinHost()
|
||
host.init_omni_connectors(
|
||
model_config=_make_model_config(),
|
||
kv_transfer_manager=mock_kvm,
|
||
)
|
||
|
||
data, size = host.recv_kv_cache("req-1")
|
||
self.assertIsNotNone(data)
|
||
self.assertEqual(size, 100)
|
||
mock_kvm.receive_kv_cache_for_request.assert_called_once()
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_receive_multi_kv_fetches_companions_via_mixin(self):
|
||
mock_kvm = MagicMock()
|
||
|
||
host = MixinHost()
|
||
host.init_omni_connectors(
|
||
model_config=_make_model_config(),
|
||
kv_transfer_manager=mock_kvm,
|
||
)
|
||
|
||
host.recv_kv_cache = MagicMock(
|
||
side_effect=[({"layer_blocks": {"k": [1]}}, 64), ({"layer_blocks": {"k": [2]}}, 32)]
|
||
)
|
||
seen = {}
|
||
|
||
def collect_cfg(request_id, cfg_role_payloads):
|
||
seen["request_id"] = request_id
|
||
seen["cfg_role_payloads"] = cfg_role_payloads
|
||
return {"cfg_text_kv_metadata": {"seq_len": 3}}
|
||
|
||
req = SimpleNamespace(
|
||
request_id="req-1",
|
||
sampling_params=SimpleNamespace(cfg_kv_request_ids={"cfg_text": "req-1__cfg_text"}),
|
||
)
|
||
ok = host.receive_multi_kv_cache(req, cfg_kv_collect_func=collect_cfg)
|
||
self.assertTrue(ok)
|
||
host.recv_kv_cache.assert_any_call("req-1", target_device=None)
|
||
host.recv_kv_cache.assert_any_call("req-1__cfg_text", target_device=None)
|
||
mock_kvm.apply_kv_cache_to_request.assert_called_once_with(req, {"layer_blocks": {"k": [1]}})
|
||
self.assertEqual(seen["request_id"], "req-1")
|
||
self.assertEqual(
|
||
seen["cfg_role_payloads"],
|
||
{"cfg_text": ({"layer_blocks": {"k": [2]}}, 32)},
|
||
)
|
||
self.assertEqual(req.sampling_params.cfg_text_kv_metadata, {"seq_len": 3})
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_receive_multi_kv_skips_inactive_request(self):
|
||
mock_kvm = MagicMock()
|
||
|
||
host = MixinHost()
|
||
host.init_omni_connectors(
|
||
model_config=_make_model_config(),
|
||
kv_transfer_manager=mock_kvm,
|
||
)
|
||
|
||
host.requests = {}
|
||
host.recv_kv_cache = MagicMock(return_value=({"layer_blocks": {"k": [1]}}, 64))
|
||
req = SimpleNamespace(request_id="req-1", sampling_params=None)
|
||
|
||
ok = host.receive_multi_kv_cache(req)
|
||
|
||
self.assertFalse(ok)
|
||
host.recv_kv_cache.assert_not_called()
|
||
mock_kvm.apply_kv_cache_to_request.assert_not_called()
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
|
||
class TestOmniConnectorOutput(unittest.TestCase):
|
||
"""Test 4: Output aggregation across transfer modes."""
|
||
|
||
def test_output_aggregation(self):
|
||
host = MixinHost()
|
||
host.init_omni_connectors(
|
||
model_config=_make_model_config(),
|
||
)
|
||
|
||
host._chunk_ready_req_ids.add("req-1")
|
||
host._chunk_finished_req_ids.add("req-2")
|
||
host._local_request_metadata["req-1"] = {"next_stage_prompt_len": 10}
|
||
host._stage_recv_req_ids.add("req-3")
|
||
|
||
output = host.get_omni_connector_output()
|
||
self.assertIsInstance(output, OmniConnectorOutput)
|
||
self.assertEqual(output.chunk_ready_req_ids, {"req-1"})
|
||
self.assertEqual(output.chunk_finished_req_ids, {"req-2"})
|
||
self.assertEqual(output.request_metadata, {"req-1": {"next_stage_prompt_len": 10}})
|
||
self.assertEqual(output.stage_recv_req_ids, {"req-3"})
|
||
|
||
output2 = host.get_omni_connector_output()
|
||
self.assertEqual(output2.chunk_ready_req_ids, set())
|
||
self.assertEqual(output2.request_metadata, {})
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
|
||
class TestMixinNoConnector(unittest.TestCase):
|
||
"""Edge case: mixin works gracefully without a connector."""
|
||
|
||
def test_no_connector(self):
|
||
host = MixinHost()
|
||
host.init_omni_connectors(
|
||
model_config=_make_model_config(),
|
||
)
|
||
self.assertIsNone(host._omni_connector)
|
||
|
||
results = host.recv_full_payload_inputs(scheduler_output=None)
|
||
self.assertIsNone(results)
|
||
|
||
sent = host.send_full_payload_outputs(None, {"req-1": {}})
|
||
self.assertEqual(sent, [])
|
||
|
||
ok = host.send_chunk(_make_request("req-1"), pooling_output={})
|
||
self.assertFalse(ok)
|
||
|
||
output = host.get_omni_connector_output()
|
||
self.assertIsInstance(output, OmniConnectorOutput)
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
|
||
class TestFinishedLoadReqsDrain(unittest.TestCase):
|
||
"""Test A1 fix: get_omni_connector_output drains _finished_load_reqs."""
|
||
|
||
def test_finished_load_reqs_flow_to_chunk_ready(self):
|
||
host = MixinHost()
|
||
host.init_omni_connectors(
|
||
model_config=_make_model_config(),
|
||
)
|
||
|
||
host._finished_load_reqs.add("req-1")
|
||
host._finished_load_reqs.add("req-2")
|
||
|
||
output = host.get_omni_connector_output()
|
||
self.assertIn("req-1", output.chunk_ready_req_ids)
|
||
self.assertIn("req-2", output.chunk_ready_req_ids)
|
||
|
||
self.assertEqual(len(host._finished_load_reqs), 0)
|
||
self.assertEqual(len(host._chunk_ready_req_ids), 0)
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
|
||
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",
|
||
"vllm_omni.model_executor.stage_input_processors.mammoth_moda2.ar2diffusion",
|
||
"vllm_omni.model_executor.stage_input_processors.cosyvoice3.text2flow",
|
||
"vllm_omni.model_executor.stage_input_processors.glm_image.ar2diffusion",
|
||
]
|
||
|
||
for func_path in incompatible_paths:
|
||
selected_path, func = MixinHost._load_custom_func(
|
||
SimpleNamespace(
|
||
async_chunk=False,
|
||
custom_process_input_func=func_path,
|
||
custom_process_next_stage_input_func=None,
|
||
)
|
||
)
|
||
assert selected_path != func_path
|
||
assert func is None or MixinHost._is_connector_payload_builder(func)
|
||
|
||
|
||
class TestFullPayloadSendWithCustomFunc(unittest.TestCase):
|
||
"""Test B4: send_full_payload_outputs with full_payload_mode custom process func."""
|
||
|
||
def test_full_payload_send_passes_is_finished_and_connector(self):
|
||
seen = {}
|
||
|
||
def full_payload_func(transfer_manager, pooling_output, request, is_finished=False):
|
||
seen["connector"] = transfer_manager.connector
|
||
seen["is_finished"] = is_finished
|
||
seen["data"] = pooling_output
|
||
seen["rid"] = request.request_id if request else None
|
||
return {"processed": True, "finished": is_finished}
|
||
|
||
host = MixinHost()
|
||
host.init_omni_connectors(
|
||
model_config=_make_model_config(),
|
||
)
|
||
host._omni_connector = MockConnector(stage_id=0)
|
||
host._stage_id = 0
|
||
host._custom_process_func = full_payload_func
|
||
|
||
req = _make_request("req-1")
|
||
req.is_finished = lambda: True
|
||
sent = host.send_full_payload_outputs(
|
||
scheduler_output=None,
|
||
outputs={"req-1": ({"raw": 100}, req)},
|
||
)
|
||
self.assertEqual(sent, ["req-1"])
|
||
self.assertEqual(
|
||
seen,
|
||
{
|
||
"connector": host._omni_connector,
|
||
"is_finished": True,
|
||
"data": {"raw": 100},
|
||
"rid": "req-1",
|
||
},
|
||
)
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_accumulate_and_flush(self):
|
||
call_log = []
|
||
|
||
def full_payload_func(transfer_manager, pooling_output, request):
|
||
call_log.append(request.request_id if request else None)
|
||
return {"processed": True}
|
||
|
||
host = MixinHost()
|
||
host.init_omni_connectors(
|
||
model_config=_make_model_config(),
|
||
)
|
||
host._omni_connector = MockConnector(stage_id=0)
|
||
host._stage_id = 0
|
||
host._custom_process_func = full_payload_func
|
||
|
||
req = _make_request("req-1")
|
||
host.accumulate_full_payload_output("req-1", {"raw": 42}, req)
|
||
self.assertEqual(len(host._pending_full_payload_send), 1)
|
||
|
||
host.flush_full_payload_outputs({"req-1"})
|
||
self.assertEqual(len(host._pending_full_payload_send), 0)
|
||
self.assertEqual(len(call_log), 1)
|
||
self.assertEqual(call_log[0], "req-1")
|
||
|
||
time.sleep(0.1)
|
||
host.shutdown_omni_connectors()
|
||
|
||
|
||
class TestKVSentReqIdsAccumulation(unittest.TestCase):
|
||
"""Test that kv_sent_req_ids accumulates results from send_kv_cache."""
|
||
|
||
def test_kv_sent_accumulation(self):
|
||
mock_kvm = MagicMock()
|
||
mock_kvm.handle_finished_requests_kv_transfer.return_value = ["req-1", "req-2"]
|
||
|
||
host = MixinHost()
|
||
host.init_omni_connectors(
|
||
model_config=_make_model_config(),
|
||
kv_transfer_manager=mock_kvm,
|
||
)
|
||
|
||
host.send_kv_cache(
|
||
finished_reqs={"req-1": {}, "req-2": {}},
|
||
kv_caches=[],
|
||
block_size=16,
|
||
cache_dtype="float16",
|
||
)
|
||
|
||
output = host.get_omni_connector_output()
|
||
self.assertIn("req-1", output.kv_sent_req_ids)
|
||
self.assertIn("req-2", output.kv_sent_req_ids)
|
||
|
||
output2 = host.get_omni_connector_output()
|
||
self.assertEqual(output2.kv_sent_req_ids, [])
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
|
||
class TestChunkStreamCompletedGuard(unittest.TestCase):
|
||
"""Test that register_chunk_recv is skipped after finish sentinel.
|
||
|
||
This validates the fix for the race condition where the scheduling
|
||
coordinator re-registers a request for chunk polling after its
|
||
upstream chunk stream has already finished (is_finished sentinel
|
||
received), causing the bg recv thread to poll for a non-existent
|
||
shared-memory segment (e.g. ``_0_7`` when only 7 chunks 0–6 exist).
|
||
"""
|
||
|
||
def _make_host(self, stage_id: int = 1) -> MixinHost:
|
||
host = MixinHost()
|
||
host.init_omni_connectors(
|
||
model_config=_make_model_config(stage_id=stage_id, async_chunk=True),
|
||
)
|
||
host._omni_connector = MockConnector(stage_id=stage_id)
|
||
host._stage_id = stage_id
|
||
host._async_chunk = True
|
||
return host
|
||
|
||
def test_register_blocked_after_finish_sentinel(self):
|
||
"""register_chunk_recv must be a no-op after the finish sentinel."""
|
||
host = self._make_host(stage_id=1)
|
||
|
||
req = _make_request("req-1", "ext-req-1")
|
||
|
||
# Simulate the bg thread having received the finish sentinel:
|
||
with host._lock:
|
||
host._chunk_stream_completed.add("req-1")
|
||
|
||
# Now try to re-register — this mimics the coordinator asking
|
||
# the model runner to poll for the next (non-existent) chunk.
|
||
host.register_chunk_recv(req)
|
||
|
||
# The request must NOT appear in _pending_load_reqs
|
||
self.assertNotIn(
|
||
"req-1",
|
||
host._pending_load_reqs,
|
||
"register_chunk_recv should skip requests whose chunk stream is already complete",
|
||
)
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_register_allowed_before_finish(self):
|
||
"""register_chunk_recv works normally before finish sentinel."""
|
||
host = self._make_host(stage_id=1)
|
||
req = _make_request("req-1", "ext-req-1")
|
||
|
||
host.register_chunk_recv(req)
|
||
self.assertIn(
|
||
"req-1",
|
||
host._pending_load_reqs,
|
||
"register_chunk_recv should add request to pending when stream is not yet complete",
|
||
)
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_poll_uses_request_scoped_payload_sender_endpoint(self):
|
||
host = self._make_host(stage_id=1)
|
||
host._omni_connector.get = MagicMock(return_value=None)
|
||
req = _make_request("req-1", "ext-req-1")
|
||
req.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50051}
|
||
|
||
host.register_chunk_recv(req)
|
||
host._poll_single_request("req-1")
|
||
|
||
host._omni_connector.get.assert_called_once_with(
|
||
"0",
|
||
"1",
|
||
"ext-req-1_0_0",
|
||
{"source_host": "10.0.0.1", "source_port": 50051},
|
||
)
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_concurrent_requests_keep_distinct_payload_sender_endpoints(self):
|
||
host = self._make_host(stage_id=1)
|
||
host._omni_connector.get = MagicMock(return_value=None)
|
||
first = _make_request("req-1", "ext-req-1")
|
||
first.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50051}
|
||
second = _make_request("req-2", "ext-req-2")
|
||
second.payload_sender_info = {"host": "10.0.0.2", "zmq_port": 51051}
|
||
|
||
host.register_chunk_recv(first)
|
||
host.register_chunk_recv(second)
|
||
host._poll_single_request("req-2")
|
||
host._poll_single_request("req-1")
|
||
|
||
self.assertEqual(
|
||
host._omni_connector.get.call_args_list,
|
||
[
|
||
unittest.mock.call(
|
||
"0",
|
||
"1",
|
||
"ext-req-2_0_0",
|
||
{"source_host": "10.0.0.2", "source_port": 51051},
|
||
),
|
||
unittest.mock.call(
|
||
"0",
|
||
"1",
|
||
"ext-req-1_0_0",
|
||
{"source_host": "10.0.0.1", "source_port": 50051},
|
||
),
|
||
],
|
||
)
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_finish_sentinel_populates_completed_set(self):
|
||
"""Receiving is_finished=True adds to _chunk_stream_completed."""
|
||
host = self._make_host(stage_id=1)
|
||
|
||
# Simulate _poll_single_request receiving is_finished=True
|
||
req_id = "req-1"
|
||
with host._lock:
|
||
host._chunk_finished_req_ids.add(req_id)
|
||
host._chunk_stream_completed.add(req_id)
|
||
host._local_stage_payload_cache[req_id] = {"finished": True}
|
||
host._local_request_metadata[req_id] = {}
|
||
host._finished_load_reqs.add(req_id)
|
||
host._pending_load_reqs.pop(req_id, None)
|
||
|
||
self.assertIn(req_id, host._chunk_stream_completed)
|
||
|
||
# Subsequent register_chunk_recv should be blocked
|
||
req = _make_request(req_id, f"ext-{req_id}")
|
||
host.register_chunk_recv(req)
|
||
self.assertNotIn(req_id, host._pending_load_reqs)
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_stage_0_always_skipped(self):
|
||
"""Stage-0 has no upstream, register_chunk_recv is always no-op."""
|
||
host = self._make_host(stage_id=0)
|
||
host._stage_id = 0
|
||
|
||
req = _make_request("req-1")
|
||
host.register_chunk_recv(req)
|
||
self.assertNotIn("req-1", host._pending_load_reqs)
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_full_payload_recv_guard_still_works(self):
|
||
"""Pre-existing guard: staged full-payload results prevent registration."""
|
||
host = self._make_host(stage_id=1)
|
||
|
||
with host._lock:
|
||
host._stage_recv_req_ids.add("req-1")
|
||
|
||
req = _make_request("req-1", "ext-req-1")
|
||
host.register_chunk_recv(req)
|
||
self.assertNotIn("req-1", host._pending_load_reqs)
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
|
||
class TestCleanupFinishedRequest(unittest.TestCase):
|
||
"""Test cleanup_finished_request frees per-request mixin state."""
|
||
|
||
def _make_host(self, stage_id: int = 1) -> MixinHost:
|
||
host = MixinHost()
|
||
host.init_omni_connectors(
|
||
model_config=_make_model_config(stage_id=stage_id, async_chunk=True),
|
||
)
|
||
host._omni_connector = MockConnector(stage_id=stage_id)
|
||
host._stage_id = stage_id
|
||
host._async_chunk = True
|
||
return host
|
||
|
||
def test_cleanup_removes_all_state(self):
|
||
"""cleanup_finished_request removes all tracking dicts/sets."""
|
||
host = self._make_host(stage_id=1)
|
||
req_id = "req-1"
|
||
ext_id = "ext-req-1"
|
||
|
||
# Simulate state accumulated during a request's lifetime
|
||
host._request_ids_mapping[req_id] = ext_id
|
||
host._put_req_chunk[ext_id] = 5
|
||
host._get_req_chunk[req_id] = 3
|
||
host._send_side_request_payload[ext_id] = {"some": "data"}
|
||
host._code_prompt_token_ids[ext_id] = [[1, 2, 3]]
|
||
host._cached_ic[ext_id] = 16
|
||
host._chunk_stream_completed.add(req_id)
|
||
host._stage_recv_req_ids.add(req_id)
|
||
host._local_stage_payload_cache[req_id] = {"engine_inputs": {}}
|
||
host._local_request_metadata[req_id] = {"prompt_len": 10}
|
||
|
||
# Cleanup
|
||
host.cleanup_finished_request(req_id)
|
||
|
||
# All state should be gone
|
||
self.assertNotIn(req_id, host._request_ids_mapping)
|
||
self.assertNotIn(ext_id, host._put_req_chunk)
|
||
self.assertNotIn(req_id, host._get_req_chunk)
|
||
self.assertNotIn(ext_id, host._send_side_request_payload)
|
||
self.assertNotIn(ext_id, host._code_prompt_token_ids)
|
||
self.assertNotIn(ext_id, host._cached_ic)
|
||
self.assertNotIn(req_id, host._chunk_stream_completed)
|
||
self.assertNotIn(req_id, host._stage_recv_req_ids)
|
||
self.assertNotIn(req_id, host._local_stage_payload_cache)
|
||
self.assertNotIn(req_id, host._local_request_metadata)
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_cleanup_removes_per_cycle_ready_state(self):
|
||
"""cleanup_finished_request clears ready/finished carry-over for req-id reuse."""
|
||
host = self._make_host(stage_id=1)
|
||
req_id = "req-1"
|
||
|
||
host._pending_load_reqs[req_id] = _make_request(req_id, "ext-req-1")
|
||
host._finished_load_reqs.add(req_id)
|
||
host._chunk_ready_req_ids.add(req_id)
|
||
host._chunk_finished_req_ids.add(req_id)
|
||
|
||
host.cleanup_finished_request(req_id)
|
||
|
||
self.assertNotIn(req_id, host._pending_load_reqs)
|
||
self.assertNotIn(req_id, host._finished_load_reqs)
|
||
self.assertNotIn(req_id, host._chunk_ready_req_ids)
|
||
self.assertNotIn(req_id, host._chunk_finished_req_ids)
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_cleanup_without_mapping(self):
|
||
"""cleanup works for Stage-0 where _request_ids_mapping isn't set."""
|
||
host = self._make_host(stage_id=0)
|
||
host._stage_id = 0
|
||
req_id = "req-1"
|
||
|
||
# Stage-0 uses req_id directly (no ext_id mapping)
|
||
host._put_req_chunk[req_id] = 3
|
||
host._get_req_chunk[req_id] = 0
|
||
host._cached_ic[req_id] = 4
|
||
|
||
host.cleanup_finished_request(req_id)
|
||
|
||
self.assertNotIn(req_id, host._put_req_chunk)
|
||
self.assertNotIn(req_id, host._get_req_chunk)
|
||
self.assertNotIn(req_id, host._cached_ic)
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_deferred_cleanup_removes_cached_ic(self):
|
||
host = self._make_host(stage_id=1)
|
||
req_id = "req-1"
|
||
ext_id = "ext-req-1"
|
||
|
||
host._request_ids_mapping[req_id] = ext_id
|
||
host._pending_save_counts[ext_id] = 1
|
||
host._cached_ic[ext_id] = 8
|
||
|
||
host.cleanup_finished_request(req_id)
|
||
|
||
self.assertIn(ext_id, host._deferred_send_cleanup)
|
||
self.assertIn(ext_id, host._cached_ic)
|
||
|
||
host._decrement_pending_save_count(ext_id)
|
||
|
||
self.assertNotIn(ext_id, host._deferred_send_cleanup)
|
||
self.assertNotIn(ext_id, host._cached_ic)
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_prune_inactive_requests_cleans_stale_state_but_keeps_active(self):
|
||
"""Inactive request IDs should be pruned without touching active ones."""
|
||
host = self._make_host(stage_id=1)
|
||
active_req_id = "req-active"
|
||
stale_req_id = "req-stale"
|
||
stale_ext_id = "ext-stale"
|
||
|
||
host._request_ids_mapping[active_req_id] = "ext-active"
|
||
host._request_ids_mapping[stale_req_id] = stale_ext_id
|
||
host._put_req_chunk[stale_ext_id] = 2
|
||
host._get_req_chunk[stale_req_id] = 1
|
||
host._finished_load_reqs.add(stale_req_id)
|
||
host._chunk_ready_req_ids.update({active_req_id, stale_req_id})
|
||
host._chunk_finished_req_ids.add(stale_req_id)
|
||
host._chunk_stream_completed.add(stale_req_id)
|
||
host._stage_recv_req_ids.add(active_req_id)
|
||
host._send_side_request_payload[stale_ext_id] = {"stale": True}
|
||
host._code_prompt_token_ids[stale_ext_id] = [[1, 2, 3]]
|
||
|
||
pruned = host.prune_inactive_requests({active_req_id})
|
||
|
||
self.assertEqual(pruned, {stale_req_id})
|
||
self.assertIn(active_req_id, host._request_ids_mapping)
|
||
self.assertIn(active_req_id, host._chunk_ready_req_ids)
|
||
self.assertIn(active_req_id, host._stage_recv_req_ids)
|
||
self.assertNotIn(stale_req_id, host._request_ids_mapping)
|
||
self.assertNotIn(stale_ext_id, host._put_req_chunk)
|
||
self.assertNotIn(stale_req_id, host._get_req_chunk)
|
||
self.assertNotIn(stale_req_id, host._pending_load_reqs)
|
||
self.assertNotIn(stale_req_id, host._finished_load_reqs)
|
||
self.assertNotIn(stale_req_id, host._chunk_ready_req_ids)
|
||
self.assertNotIn(stale_req_id, host._chunk_finished_req_ids)
|
||
self.assertNotIn(stale_req_id, host._chunk_stream_completed)
|
||
self.assertNotIn(stale_req_id, host._stage_recv_req_ids)
|
||
self.assertNotIn(stale_ext_id, host._send_side_request_payload)
|
||
self.assertNotIn(stale_ext_id, host._code_prompt_token_ids)
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_prune_inactive_requests_keeps_recently_received_full_payload_state(self):
|
||
"""Late bg-thread receives must survive until the scheduler catches up."""
|
||
host = self._make_host(stage_id=1)
|
||
req_id = "req-recv-race"
|
||
ext_id = "ext-recv-race"
|
||
|
||
host._request_ids_mapping[req_id] = ext_id
|
||
host._put_req_chunk[ext_id] = 1
|
||
host._local_stage_payload_cache[req_id] = {"engine_inputs": {"ids": [1, 2, 3]}}
|
||
host._local_request_metadata[req_id] = {"next_stage_prompt_len": 3}
|
||
host._stage_recv_req_ids.add(req_id)
|
||
|
||
pruned = host.prune_inactive_requests(set())
|
||
|
||
self.assertEqual(pruned, set())
|
||
self.assertIn(req_id, host._request_ids_mapping)
|
||
self.assertIn(req_id, host._local_stage_payload_cache)
|
||
self.assertIn(req_id, host._local_request_metadata)
|
||
self.assertIn(req_id, host._stage_recv_req_ids)
|
||
self.assertIn(ext_id, host._put_req_chunk)
|
||
|
||
# Once the scheduler has consumed the wake-up and the request really
|
||
# disappears from all protected sets, prune should clean it up.
|
||
host._stage_recv_req_ids.clear()
|
||
host._local_stage_payload_cache.clear()
|
||
host._local_request_metadata.clear()
|
||
|
||
pruned = host.prune_inactive_requests(set())
|
||
|
||
self.assertEqual(pruned, {req_id})
|
||
self.assertNotIn(req_id, host._request_ids_mapping)
|
||
self.assertNotIn(ext_id, host._put_req_chunk)
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
|
||
class TestSendChunkCachesMapping(unittest.TestCase):
|
||
"""Test that send_chunk caches internal→external req ID mapping."""
|
||
|
||
def test_send_chunk_populates_request_ids_mapping(self):
|
||
"""send_chunk should cache the internal→external mapping."""
|
||
host = MixinHost()
|
||
host.init_omni_connectors(
|
||
model_config=_make_model_config(stage_id=0, async_chunk=True),
|
||
)
|
||
host._omni_connector = MockConnector(stage_id=0)
|
||
host._stage_id = 0
|
||
host._async_chunk = True
|
||
|
||
def mock_process(transfer_manager, pooling_output, request):
|
||
return {"data": "test", "finished": False}
|
||
|
||
host._custom_process_func = mock_process
|
||
|
||
request = _make_request("internal-1", "external-1")
|
||
host.send_chunk(request, pooling_output={"v": 1})
|
||
|
||
# The mapping should be cached
|
||
self.assertEqual(
|
||
host._request_ids_mapping.get("internal-1"),
|
||
"external-1",
|
||
)
|
||
|
||
time.sleep(0.1)
|
||
host.shutdown_omni_connectors()
|
||
|
||
|
||
class TestLocalPayloadCacheLifecycle(unittest.TestCase):
|
||
"""Unit tests for the local payload cache API (RFC §2.4)."""
|
||
|
||
def _make_host(self) -> MixinHost:
|
||
host = MixinHost()
|
||
host.init_omni_connectors(
|
||
model_config=_make_model_config(stage_id=0),
|
||
)
|
||
host._omni_connector = MockConnector(stage_id=0)
|
||
host._stage_id = 0
|
||
return host
|
||
|
||
def test_put_get_pop(self):
|
||
host = self._make_host()
|
||
payload = {"engine_inputs": {"ids": [1, 2, 3]}}
|
||
host.put_local_stage_payload("r1", payload)
|
||
|
||
self.assertEqual(host.get_local_stage_payload("r1"), payload)
|
||
popped = host.pop_local_stage_payload("r1")
|
||
self.assertEqual(popped, payload)
|
||
self.assertIsNone(host.get_local_stage_payload("r1"))
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_recv_full_payload_inputs_populates_local_cache(self):
|
||
host = self._make_host()
|
||
host._omni_connector = MockConnector(stage_id=0)
|
||
host._stage_id = 0
|
||
|
||
# Simulate a full payload already staged by the bg recv path
|
||
with host._lock:
|
||
host._local_stage_payload_cache["r1"] = {"tok": [10]}
|
||
host._stage_recv_req_ids.add("r1")
|
||
|
||
host.recv_full_payload_inputs(scheduler_output=None)
|
||
self.assertEqual(host.get_local_stage_payload("r1"), {"tok": [10]})
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_rank0_only_polls_connector_for_tp_full_payload(self):
|
||
host = self._make_host()
|
||
host._omni_connector = MagicMock()
|
||
host._stage_id = 2
|
||
host._local_rank = 0
|
||
host._request_ids_mapping["r1"] = "ext-r1"
|
||
host._get_req_chunk["r1"] = 0
|
||
payload = {"tok": [10], "finished": torch.tensor(True)}
|
||
connector_result = (payload, 123)
|
||
host._omni_connector.get.return_value = connector_result
|
||
tp_group = _FakeTPGroup(world_size=2, rank_in_group=0)
|
||
|
||
with patch.object(host, "_get_local_tp_group", return_value=tp_group):
|
||
made_progress = host._poll_single_request("r1")
|
||
|
||
self.assertTrue(made_progress)
|
||
host._omni_connector.get.assert_called_once_with("1", "2", "ext-r1_1_0")
|
||
self.assertEqual(tp_group.broadcast_inputs, [])
|
||
self.assertEqual(host.get_local_stage_payload("r1"), payload)
|
||
self.assertIn("r1", host._full_payload_pending_broadcast_req_ids)
|
||
self.assertNotIn("r1", host._stage_recv_req_ids)
|
||
self.assertIsNone(host.get_local_request_metadata("r1"))
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_tp_follower_skips_connector_poll_for_full_payload(self):
|
||
host = self._make_host()
|
||
host._omni_connector = MagicMock()
|
||
host._stage_id = 2
|
||
host._local_rank = 1
|
||
host._request_ids_mapping["r1"] = "ext-r1"
|
||
host._get_req_chunk["r1"] = 0
|
||
tp_group = _FakeTPGroup(world_size=2, rank_in_group=1)
|
||
|
||
with patch.object(host, "_get_local_tp_group", return_value=tp_group):
|
||
made_progress = host._poll_single_request("r1")
|
||
|
||
self.assertFalse(made_progress)
|
||
host._omni_connector.get.assert_not_called()
|
||
self.assertEqual(tp_group.broadcast_inputs, [])
|
||
self.assertNotIn("r1", host._local_stage_payload_cache)
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_recv_full_payload_inputs_broadcasts_tp_leader_results_to_followers(self):
|
||
host = self._make_host()
|
||
host._omni_connector = MagicMock()
|
||
host._stage_id = 2
|
||
host._local_rank = 1
|
||
host._pending_load_reqs["r1"] = object()
|
||
payload = {"tok": [10], "finished": torch.tensor(True)}
|
||
tp_group = _FakeTPGroup(world_size=2, rank_in_group=1, follower_result={"r1": payload})
|
||
|
||
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})
|
||
self.assertEqual(host.get_local_stage_payload("r1"), payload)
|
||
self.assertEqual(host.get_local_request_metadata("r1"), {})
|
||
self.assertEqual(host._stage_recv_req_ids, {"r1"})
|
||
self.assertNotIn("r1", host._pending_load_reqs)
|
||
self.assertEqual(tp_group.broadcast_inputs, [None])
|
||
host.shutdown_omni_connectors()
|
||
|
||
|
||
class TestTPAsyncChunkFanout(unittest.TestCase):
|
||
def _make_host(self, rank: int) -> MixinHost:
|
||
host = MixinHost()
|
||
host.init_omni_connectors(
|
||
model_config=_make_model_config(stage_id=2, async_chunk=True, worker_type="gen"),
|
||
)
|
||
host._omni_connector = MagicMock()
|
||
host._stage_id = 2
|
||
host._async_chunk = True
|
||
host._model_mode = "gen"
|
||
host._local_rank = rank
|
||
host._request_ids_mapping["r1"] = "ext-r1"
|
||
host._get_req_chunk["r1"] = 0
|
||
return host
|
||
|
||
def test_rank0_only_polls_connector_for_tp_async_chunk(self):
|
||
host = self._make_host(rank=0)
|
||
payload = {
|
||
"codes": {"audio": [10, 11]},
|
||
"meta": {"left_context_size": 0, "finished": torch.tensor(False)},
|
||
}
|
||
host._omni_connector.get.return_value = (payload, 123)
|
||
tp_group = _FakeTPGroup(world_size=2, rank_in_group=0)
|
||
|
||
with patch.object(host, "_get_local_tp_group", return_value=tp_group):
|
||
made_progress = host._poll_single_request("r1")
|
||
|
||
self.assertTrue(made_progress)
|
||
host._omni_connector.get.assert_called_once_with("1", "2", "ext-r1_1_0")
|
||
self.assertEqual(host.get_local_stage_payload("r1"), payload)
|
||
self.assertIn("r1", host._finished_load_reqs)
|
||
self.assertIn("r1", host._async_chunk_updated_req_ids)
|
||
self.assertEqual(tp_group.broadcast_inputs, [])
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_tp_follower_skips_connector_poll_for_async_chunk(self):
|
||
host = self._make_host(rank=1)
|
||
tp_group = _FakeTPGroup(world_size=2, rank_in_group=1)
|
||
|
||
with patch.object(host, "_get_local_tp_group", return_value=tp_group):
|
||
made_progress = host._poll_single_request("r1")
|
||
|
||
self.assertFalse(made_progress)
|
||
host._omni_connector.get.assert_not_called()
|
||
self.assertIsNone(host.get_local_stage_payload("r1"))
|
||
self.assertEqual(tp_group.broadcast_inputs, [])
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_get_output_broadcasts_tp_async_chunk_payloads_to_followers(self):
|
||
host = self._make_host(rank=1)
|
||
host._pending_load_reqs["r1"] = object()
|
||
payload = {
|
||
"code_predictor_codes": [10, 11],
|
||
"left_context_size": 0,
|
||
"finished": torch.tensor(True),
|
||
}
|
||
packet = {
|
||
"staged_payloads": {"r1": payload},
|
||
"request_metadata": {"r1": {"code_predictor_codes": [10, 11], "left_context_size": 0}},
|
||
"newly_finished": {"r1"},
|
||
"chunk_finished": {"r1"},
|
||
}
|
||
tp_group = _FakeTPGroup(world_size=2, rank_in_group=1, follower_result=packet)
|
||
|
||
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"})
|
||
self.assertEqual(output.chunk_finished_req_ids, {"r1"})
|
||
self.assertEqual(
|
||
output.request_metadata,
|
||
{"r1": {"code_predictor_codes": [10, 11], "left_context_size": 0}},
|
||
)
|
||
self.assertEqual(host.get_local_stage_payload("r1"), payload)
|
||
self.assertNotIn("r1", host._pending_load_reqs)
|
||
self.assertIn("r1", host._chunk_stream_completed)
|
||
self.assertEqual(tp_group.broadcast_inputs, [None])
|
||
host.shutdown_omni_connectors()
|
||
|
||
|
||
class TestKVTransferLifecycle(unittest.TestCase):
|
||
"""Unit tests for KV transfer lifecycle methods."""
|
||
|
||
def _make_host(self) -> MixinHost:
|
||
host = MixinHost()
|
||
host.init_omni_connectors(
|
||
model_config=_make_model_config(stage_id=0),
|
||
)
|
||
return host
|
||
|
||
def test_mark_drain_ack_complete(self):
|
||
host = self._make_host()
|
||
self.assertFalse(host.has_pending_kv_work())
|
||
|
||
host.mark_kv_transfer("r1", seq_len=100, block_ids=[0, 1, 2])
|
||
self.assertTrue(host.has_pending_kv_work())
|
||
self.assertTrue(host.is_kv_transfer_triggered("r1"))
|
||
|
||
# Drain moves pending → active
|
||
pending = host.drain_pending_kv_transfers()
|
||
self.assertEqual(pending, {"r1": {"seq_len": 100, "block_ids": [0, 1, 2]}})
|
||
self.assertIn("r1", host._kv_active_transfers)
|
||
self.assertTrue(host.has_pending_kv_work())
|
||
|
||
# Ack moves active → completed
|
||
host.ack_kv_transfers(["r1"])
|
||
self.assertNotIn("r1", host._kv_active_transfers)
|
||
self.assertIn("r1", host._kv_completed_transfers)
|
||
|
||
# Drain completed
|
||
completed = host.drain_completed_kv_transfers()
|
||
self.assertEqual(completed, {"r1"})
|
||
self.assertFalse(host.has_pending_kv_work())
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_mark_dedup(self):
|
||
host = self._make_host()
|
||
host.mark_kv_transfer("r1", seq_len=100, block_ids=[0])
|
||
host.mark_kv_transfer("r1", seq_len=200, block_ids=[0, 1])
|
||
# Second mark is a no-op
|
||
self.assertEqual(host._kv_pending_transfers["r1"]["seq_len"], 100)
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_cleanup_removes_kv_state(self):
|
||
host = self._make_host()
|
||
host.mark_kv_transfer("r1", seq_len=50, block_ids=[0])
|
||
host.drain_pending_kv_transfers()
|
||
host.cleanup_finished_request("r1")
|
||
self.assertFalse(host.is_kv_transfer_triggered("r1"))
|
||
self.assertNotIn("r1", host._kv_active_transfers)
|
||
self.assertFalse(host.has_pending_kv_work())
|
||
host.shutdown_omni_connectors()
|
||
|
||
|
||
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(
|
||
model_config=_make_model_config(stage_id=1, async_chunk=True, worker_type="ar"),
|
||
)
|
||
host._request_ids_mapping["r1"] = "r1"
|
||
payload = {
|
||
"embed": {"decode": torch.ones(1, 2)},
|
||
"ids": {"output": [1]},
|
||
"meta": {
|
||
"finished": torch.tensor(False),
|
||
"override_keys": [["embed", "decode"], ["ids", "output"]],
|
||
},
|
||
}
|
||
|
||
host._accumulate_payload("r1", dict(payload))
|
||
with host._lock:
|
||
host._finished_load_reqs.add("r1")
|
||
|
||
host.get_omni_connector_output()
|
||
self.assertIn("r1", host._send_side_request_payload)
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_payload_consumable_ignores_token_horizon_only_updates(self):
|
||
host = MixinHost()
|
||
host.init_omni_connectors(
|
||
model_config=_make_model_config(stage_id=1, async_chunk=True, worker_type="ar"),
|
||
)
|
||
payload = {
|
||
"ids": {"output": [1, 2, 3]},
|
||
"embed": {"decode_token_start": 2, "decode_token_end": 3},
|
||
"meta": {
|
||
"finished": torch.tensor(False),
|
||
"override_keys": [
|
||
["ids", "output"],
|
||
["embed", "decode_token_start"],
|
||
["embed", "decode_token_end"],
|
||
],
|
||
},
|
||
}
|
||
self.assertFalse(host._payload_is_consumable(payload))
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_payload_consumable_accepts_decode_embeddings(self):
|
||
host = MixinHost()
|
||
host.init_omni_connectors(
|
||
model_config=_make_model_config(stage_id=1, async_chunk=True, worker_type="ar"),
|
||
)
|
||
payload = {
|
||
"ids": {"output": [1, 2, 3]},
|
||
"embed": {"decode": torch.ones(1, 2)},
|
||
"meta": {"finished": torch.tensor(False)},
|
||
}
|
||
self.assertTrue(host._payload_is_consumable(payload))
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_ar_metadata_only_followup_chunk_does_not_rewake_request(self):
|
||
host = MixinHost()
|
||
host.init_omni_connectors(
|
||
model_config=_make_model_config(stage_id=1, async_chunk=True, worker_type="ar"),
|
||
)
|
||
host._omni_connector = MagicMock()
|
||
host._stage_id = 1
|
||
host._async_chunk = True
|
||
host._model_mode = "ar"
|
||
host._request_ids_mapping["r1"] = "ext-r1"
|
||
host._get_req_chunk["r1"] = 0
|
||
|
||
host._omni_connector.get.side_effect = [
|
||
(
|
||
{
|
||
"embed": {"decode": torch.ones(1, 2)},
|
||
"meta": {"finished": torch.tensor(False)},
|
||
},
|
||
1,
|
||
),
|
||
(
|
||
{
|
||
"meta": {"next_stage_prompt_len": 7, "finished": torch.tensor(False)},
|
||
},
|
||
1,
|
||
),
|
||
]
|
||
|
||
host._poll_single_request("r1")
|
||
output1 = host.get_omni_connector_output()
|
||
self.assertEqual(output1.chunk_ready_req_ids, {"r1"})
|
||
|
||
host._poll_single_request("r1")
|
||
output2 = host.get_omni_connector_output()
|
||
self.assertEqual(output2.chunk_ready_req_ids, set())
|
||
self.assertEqual(output2.request_metadata, {"r1": {"next_stage_prompt_len": 7}})
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_non_ar_recv_does_not_overwrite_unconsumed_staged_chunk(self):
|
||
host = MixinHost()
|
||
host.init_omni_connectors(
|
||
model_config=_make_model_config(stage_id=2, async_chunk=True, worker_type="gen"),
|
||
)
|
||
host._omni_connector = MagicMock()
|
||
host._stage_id = 2
|
||
host._async_chunk = True
|
||
host._model_mode = "gen"
|
||
host._request_ids_mapping["r1"] = "ext-r1"
|
||
host._get_req_chunk["r1"] = 1
|
||
host._local_stage_payload_cache["r1"] = {
|
||
"code_predictor_codes": [1, 2, 3],
|
||
"left_context_size": 0,
|
||
"finished": torch.tensor(False),
|
||
}
|
||
|
||
made_progress = host._poll_single_request("r1")
|
||
|
||
self.assertFalse(made_progress)
|
||
host._omni_connector.get.assert_not_called()
|
||
self.assertEqual(host._get_req_chunk["r1"], 1)
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_non_ar_recv_waits_for_scheduler_handoff_before_fetching_next_chunk(self):
|
||
host = MixinHost()
|
||
host.init_omni_connectors(
|
||
model_config=_make_model_config(stage_id=2, async_chunk=True, worker_type="gen"),
|
||
)
|
||
host._omni_connector = MagicMock()
|
||
host._stage_id = 2
|
||
host._async_chunk = True
|
||
host._model_mode = "gen"
|
||
host._request_ids_mapping["r1"] = "ext-r1"
|
||
host._get_req_chunk["r1"] = 1
|
||
host._local_request_metadata["r1"] = {
|
||
"code_predictor_codes": [10, 11, 12],
|
||
"left_context_size": 0,
|
||
}
|
||
host._finished_load_reqs.add("r1")
|
||
|
||
made_progress = host._poll_single_request("r1")
|
||
|
||
self.assertFalse(made_progress)
|
||
host._omni_connector.get.assert_not_called()
|
||
self.assertEqual(host._get_req_chunk["r1"], 1)
|
||
|
||
output = host.get_omni_connector_output()
|
||
self.assertEqual(output.request_metadata["r1"]["code_predictor_codes"], [10, 11, 12])
|
||
self.assertEqual(output.chunk_ready_req_ids, {"r1"})
|
||
|
||
host._omni_connector.get.return_value = (
|
||
{
|
||
"codes": {"audio": [20, 21, 22]},
|
||
"meta": {"left_context_size": 0, "finished": torch.tensor(False)},
|
||
},
|
||
1,
|
||
)
|
||
made_progress = host._poll_single_request("r1")
|
||
|
||
self.assertTrue(made_progress)
|
||
host._omni_connector.get.assert_called_once()
|
||
self.assertEqual(host._get_req_chunk["r1"], 2)
|
||
|
||
host.shutdown_omni_connectors()
|
||
|
||
|
||
class TestRankAwareKVRouting(unittest.TestCase):
|
||
def _make_host(self, *, from_tp: int, to_tp: int, local_rank: int) -> MixinHost:
|
||
host = MixinHost()
|
||
host.init_omni_connectors(model_config=_make_model_config(stage_id=1))
|
||
host._from_tp = from_tp
|
||
host._to_tp = to_tp
|
||
host._local_rank = local_rank
|
||
return host
|
||
|
||
def test_recv_keys_use_remote_rank_as_from_rank(self):
|
||
host = self._make_host(from_tp=4, to_tp=2, local_rank=1)
|
||
self.assertEqual(
|
||
host.get_rank_aware_kv_keys("req", from_stage=0),
|
||
["req_0_0_2_1", "req_0_0_3_1"],
|
||
)
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_single_rank_hooks_use_default_manager_key(self):
|
||
host = self._make_host(from_tp=1, to_tp=1, local_rank=0)
|
||
expected = ["omni_0_to_1_kv_cache_req"]
|
||
|
||
self.assertEqual(host.get_rank_aware_kv_send_keys("req", from_stage=0, to_stage=1), expected)
|
||
self.assertEqual(host.get_rank_aware_kv_keys("req", from_stage=0, to_stage=1), expected)
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_init_uses_effective_manager_tp2_topology(self):
|
||
model_config = _make_model_config(stage_id=0)
|
||
model_config.omni_kv_config = {
|
||
"need_send_cache": True,
|
||
"omni_from_stage": "0",
|
||
"omni_to_stage": "1",
|
||
"rank_mapping": {"from_tp": 2, "to_tp": 2},
|
||
}
|
||
with patch(
|
||
"vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank",
|
||
return_value=1,
|
||
):
|
||
manager = OmniKVTransferManager.from_model_config(model_config)
|
||
|
||
host = MixinHost()
|
||
host.init_omni_connectors(model_config=model_config, kv_transfer_manager=manager)
|
||
|
||
expected = ["req_0_0_1_1"]
|
||
self.assertEqual(host.get_rank_aware_kv_send_keys("req", from_stage=0, to_stage=1), expected)
|
||
self.assertEqual(host.get_rank_aware_kv_keys("req", from_stage=0, to_stage=1), expected)
|
||
self.assertEqual(host.get_kv_rank_mapping()["from_tp"], 2)
|
||
self.assertEqual(host.get_kv_rank_mapping()["to_tp"], 2)
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_send_keys_route_from_rank_gt_to_rank(self):
|
||
host = self._make_host(from_tp=4, to_tp=2, local_rank=3)
|
||
self.assertEqual(host.get_rank_aware_kv_send_keys("req", from_stage=0), ["req_0_0_3_1"])
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_invalid_recv_rank_mapping_raises(self):
|
||
host = self._make_host(from_tp=3, to_tp=2, local_rank=1)
|
||
with self.assertRaises(ValueError):
|
||
host.get_rank_aware_kv_keys("req", from_stage=0)
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_invalid_send_rank_mapping_raises(self):
|
||
host = self._make_host(from_tp=3, to_tp=2, local_rank=1)
|
||
with self.assertRaises(ValueError):
|
||
host.get_rank_aware_kv_send_keys("req", from_stage=0)
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_merge_rank_sharded_payloads_concatenates_head_dimension(self):
|
||
host = self._make_host(from_tp=4, to_tp=2, local_rank=0)
|
||
payloads = [
|
||
{"layer_blocks": {"key_cache": [torch.ones(2, 1, 3)], "value_cache": [torch.ones(2, 1, 3)]}},
|
||
{"layer_blocks": {"key_cache": [torch.full((2, 1, 3), 2.0)], "value_cache": [torch.full((2, 1, 3), 2.0)]}},
|
||
]
|
||
merged = host._merge_rank_sharded_kv_payloads(payloads)
|
||
self.assertEqual(tuple(merged["layer_blocks"]["key_cache"][0].shape), (2, 2, 3))
|
||
self.assertTrue(torch.equal(merged["layer_blocks"]["key_cache"][0][:, 0], torch.ones(2, 3)))
|
||
self.assertTrue(torch.equal(merged["layer_blocks"]["key_cache"][0][:, 1], torch.full((2, 3), 2.0)))
|
||
host.shutdown_omni_connectors()
|
||
|
||
def test_slice_rank_sharded_payload_splits_head_dimension(self):
|
||
host = self._make_host(from_tp=2, to_tp=4, local_rank=1)
|
||
payload = {
|
||
"layer_blocks": {
|
||
"key_cache": [torch.arange(24, dtype=torch.float32).reshape(2, 4, 3)],
|
||
"value_cache": [torch.arange(24, dtype=torch.float32).reshape(2, 4, 3)],
|
||
},
|
||
"metadata": {},
|
||
}
|
||
sliced = host._slice_rank_sharded_kv_payload(payload)
|
||
self.assertEqual(tuple(sliced["layer_blocks"]["key_cache"][0].shape), (2, 2, 3))
|
||
expected = torch.arange(24, dtype=torch.float32).reshape(2, 4, 3)[:, 2:4, :]
|
||
self.assertTrue(torch.equal(sliced["layer_blocks"]["key_cache"][0], expected))
|
||
host.shutdown_omni_connectors()
|
||
|
||
|
||
class TestAttachOmniConnectorOutput(unittest.TestCase):
|
||
def test_wraps_empty_model_runner_output_when_signals_exist(self):
|
||
from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
|
||
|
||
host = MixinHost()
|
||
host.get_omni_connector_output = lambda: OmniConnectorOutput(chunk_ready_req_ids={"req-1"})
|
||
|
||
wrapped = host.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT)
|
||
|
||
self.assertIsNot(wrapped, EMPTY_MODEL_RUNNER_OUTPUT)
|
||
self.assertEqual(wrapped.omni_connector_output.chunk_ready_req_ids, {"req-1"})
|
||
|
||
|
||
class TestConnectorConfigValidation(unittest.TestCase):
|
||
def test_invalid_connector_name_raises(self):
|
||
host = MixinHost()
|
||
model_config = _make_model_config(stage_id=1)
|
||
model_config.stage_connector_config = {"name": " "}
|
||
|
||
with self.assertRaisesRegex(RuntimeError, "missing connector name"):
|
||
host.init_omni_connectors(model_config=model_config)
|
||
|
||
|
||
class TestRankAwareHandshakePort(unittest.TestCase):
|
||
"""A sender must bind the port the orchestrator advertises for its rank."""
|
||
|
||
@staticmethod
|
||
def _resolve(name, extra, local_rank=0):
|
||
spec = resolve_connector_spec(
|
||
ConnectorSpec(name=name, extra=extra),
|
||
stage_id=int(extra.get("stage_id", 0)),
|
||
role=extra.get("role"),
|
||
local_rank=local_rank,
|
||
)
|
||
return spec.extra
|
||
|
||
def test_sender_port_matches_the_orchestrator_formula(self):
|
||
for local_rank in (0, 1, 3):
|
||
with self.subTest(local_rank=local_rank):
|
||
extra = {"zmq_port": 50071, "stage_id": 2, "role": "sender"}
|
||
resolved = self._resolve("NixlConnector", extra, local_rank=local_rank)
|
||
self.assertEqual(resolved["zmq_port"], 50071 + 2 + local_rank * 16)
|
||
self.assertEqual(extra["zmq_port"], 50071, "the source config must not be mutated")
|
||
|
||
def test_receiver_gets_sender_endpoint_without_a_bind_port(self):
|
||
receiver = self._resolve(
|
||
"NixlConnector",
|
||
{"zmq_port": 50071, "stage_id": 3, "from_stage": 2, "role": "receiver"},
|
||
local_rank=1,
|
||
)
|
||
self.assertNotIn("zmq_port", receiver)
|
||
self.assertEqual(receiver["sender_zmq_port"], 50071 + 2 + 16)
|
||
|
||
def test_handshake_less_connector_keeps_its_config(self):
|
||
shm = self._resolve("SharedMemoryConnector", {"zmq_port": 50071, "role": "sender"}, local_rank=1)
|
||
self.assertEqual(shm["zmq_port"], 50071)
|
||
|
||
def test_unresolvable_port_is_rejected(self):
|
||
extra = {"zmq_port": "${MISSING_PORT_VAR}", "stage_id": 0, "role": "sender"}
|
||
with self.assertRaises(ValueError):
|
||
self._resolve("NixlConnector", extra)
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("extra", "expected"),
|
||
[
|
||
(
|
||
{"role": "sender", "from_stage": 2, "zmq_port": 50071},
|
||
{"zmq_port": 52153},
|
||
),
|
||
(
|
||
{"role": "receiver", "from_stage": 1, "zmq_port": 50071, "host": "producer"},
|
||
{"sender_zmq_port": 52152, "sender_host": "producer"},
|
||
),
|
||
(
|
||
{
|
||
"role": "receiver",
|
||
"from_stage": 1,
|
||
"zmq_port": 50071,
|
||
"sender_zmq_port": 49000,
|
||
"sender_host": "bound-producer",
|
||
"outgoing": {"from_stage": 2, "zmq_port": 51071, "host": "local-worker"},
|
||
},
|
||
{"sender_zmq_port": 49000, "sender_host": "bound-producer", "zmq_port": 53153, "host": "local-worker"},
|
||
),
|
||
],
|
||
)
|
||
def test_split_runtime_factory_resolves_nixl_endpoints(extra, expected):
|
||
"""The runner must pass the resolved spec, not the raw config, to the factory."""
|
||
from vllm_omni.distributed.omni_connectors.model_runner import omni_connector_runtime as runtime
|
||
|
||
original = deepcopy(extra)
|
||
config = SimpleNamespace(stage_id=2, stage_connector_config={"name": "NixlConnector", "extra": extra})
|
||
with (
|
||
patch.object(runtime, "get_local_tp_rank", return_value=2, create=True),
|
||
patch.object(runtime, "get_omni_replica_id", return_value=2, create=True),
|
||
patch.object(runtime.OmniConnectorFactory, "create_connector") as factory,
|
||
):
|
||
assert MixinHost._create_connector(config) is factory.return_value
|
||
|
||
spec = factory.call_args.args[0]
|
||
assert spec.name == "NixlConnector"
|
||
assert spec.extra["stage_id"] == 2
|
||
for key, value in expected.items():
|
||
assert spec.extra[key] == value
|
||
if extra["role"] == "receiver" and "outgoing" not in extra:
|
||
assert "zmq_port" not in spec.extra
|
||
assert extra == original
|
||
|
||
|
||
@pytest.mark.parametrize("async_chunk", [False, True])
|
||
def test_split_payload_endpoint_retry_cleanup_and_reuse(async_chunk):
|
||
host = MixinHost()
|
||
host.init_omni_connectors(_make_model_config(async_chunk=async_chunk))
|
||
host._stage_id = 1
|
||
host._omni_connector = MagicMock()
|
||
host._omni_connector.get.return_value = None
|
||
first = _make_request("req-1", "ext-req-1")
|
||
first.payload_sender_info = {"host": "producer-a", "zmq_port": "50051"}
|
||
second = _make_request("req-2", "ext-req-2")
|
||
second.payload_sender_info = {"host": "producer-b", "zmq_port": 51051}
|
||
|
||
try:
|
||
host.register_chunk_recv(first)
|
||
host.register_chunk_recv(second)
|
||
for req_id in ("req-1", "req-2", "req-1"):
|
||
assert not host._poll_single_request(req_id)
|
||
assert host._omni_connector.get.call_args_list == [
|
||
unittest.mock.call("0", "1", "ext-req-1_0_0", {"source_host": "producer-a", "source_port": 50051}),
|
||
unittest.mock.call("0", "1", "ext-req-2_0_0", {"source_host": "producer-b", "source_port": 51051}),
|
||
unittest.mock.call("0", "1", "ext-req-1_0_0", {"source_host": "producer-a", "source_port": 50051}),
|
||
]
|
||
host.cleanup_finished_request("req-1")
|
||
assert "req-1" not in host._pending_load_reqs
|
||
assert host._pending_load_reqs["req-2"] is second
|
||
host.register_chunk_recv(_make_request("req-1", "reused"))
|
||
assert not host._poll_single_request("req-1")
|
||
host._omni_connector.get.assert_called_with("0", "1", "reused_0_0")
|
||
finally:
|
||
host.shutdown_omni_connectors()
|
||
|
||
|
||
@pytest.mark.parametrize("method", ["_recv_full_payload_result", "_recv_async_chunk_result"])
|
||
@pytest.mark.parametrize("rank", [0, 1])
|
||
def test_split_payload_sender_metadata_respects_tp_leader(method, rank):
|
||
host = MixinHost()
|
||
connector = MagicMock()
|
||
metadata = {"source_host": "producer", "source_port": 50051}
|
||
with (
|
||
patch.object(host, "_get_local_tp_group", return_value=_FakeTPGroup(world_size=2, rank_in_group=rank)),
|
||
patch.object(host, "is_data_transfer_rank", return_value=rank == 0),
|
||
):
|
||
result = getattr(host, method)(connector, "0", "1", "req_0_0", metadata)
|
||
if rank == 0:
|
||
connector.get.assert_called_once_with("0", "1", "req_0_0", metadata)
|
||
assert result is connector.get.return_value
|
||
else:
|
||
connector.get.assert_not_called()
|
||
assert result is None
|
||
|
||
|
||
class _FailingConnector:
|
||
"""Connector whose put() fails a configurable number of times."""
|
||
|
||
def __init__(self, fail_count: int = 1, raise_on_fail: bool = False):
|
||
self._fail_count = fail_count
|
||
self._raise_on_fail = raise_on_fail
|
||
self.attempt = 0
|
||
|
||
def put(self, from_stage, to_stage, put_key, data):
|
||
self.attempt += 1
|
||
if self.attempt <= self._fail_count:
|
||
if self._raise_on_fail:
|
||
raise ConnectionError("transient connector error")
|
||
return False, 0, None
|
||
return True, len(str(data)), None
|
||
|
||
def get(self, *a, **kw):
|
||
return None
|
||
|
||
def close(self):
|
||
pass
|
||
|
||
|
||
class TestSendRetry(unittest.TestCase):
|
||
"""Tests for P1-2: failed connector sends must be retried."""
|
||
|
||
def _make_sender(self, connector):
|
||
sender = MixinHost()
|
||
sender.init_omni_connectors(
|
||
model_config=_make_model_config(stage_id=0, async_chunk=True),
|
||
)
|
||
sender._omni_connector = connector
|
||
sender._stage_id = 0
|
||
sender._async_chunk = True
|
||
return sender
|
||
|
||
def _make_task(self, req_id="r1"):
|
||
return {
|
||
"stage_id": 0,
|
||
"next_stage_id": 1,
|
||
"request_id": req_id,
|
||
"data": {"payload": "test"},
|
||
}
|
||
|
||
def test_send_single_request_returns_false_on_put_failure(self):
|
||
connector = _FailingConnector(fail_count=999)
|
||
sender = self._make_sender(connector)
|
||
|
||
result = sender._send_single_request(self._make_task())
|
||
self.assertFalse(result)
|
||
sender.shutdown_omni_connectors()
|
||
|
||
def test_send_single_request_does_not_decrement_on_failure(self):
|
||
connector = _FailingConnector(fail_count=999)
|
||
sender = self._make_sender(connector)
|
||
sender._pending_save_counts["r1"] = 1
|
||
|
||
sender._send_single_request(self._make_task())
|
||
self.assertEqual(sender._pending_save_counts.get("r1"), 1, "pending count must NOT be decremented on failure")
|
||
sender.shutdown_omni_connectors()
|
||
|
||
def test_send_single_request_decrements_on_success(self):
|
||
connector = MockConnector(stage_id=0)
|
||
sender = self._make_sender(connector)
|
||
sender._pending_save_counts["r1"] = 1
|
||
|
||
result = sender._send_single_request(self._make_task())
|
||
self.assertTrue(result)
|
||
self.assertNotIn("r1", sender._pending_save_counts, "pending count should be zero/removed on success")
|
||
sender.shutdown_omni_connectors()
|
||
|
||
def test_requeue_or_drop_requeues_on_first_failure(self):
|
||
sender = self._make_sender(MockConnector(stage_id=0))
|
||
task = self._make_task()
|
||
|
||
sender._requeue_or_drop_failed_send(task)
|
||
|
||
self.assertEqual(task.get("_retry_count"), 1)
|
||
with sender._lock:
|
||
dq = sender._pending_save_reqs.get("r1")
|
||
self.assertIsNotNone(dq)
|
||
self.assertEqual(len(dq), 1)
|
||
sender.shutdown_omni_connectors()
|
||
|
||
def test_requeue_or_drop_drops_after_max_retries(self):
|
||
sender = self._make_sender(MockConnector(stage_id=0))
|
||
sender._pending_save_counts["r1"] = 1
|
||
task = self._make_task()
|
||
task["_retry_count"] = sender._MAX_SEND_RETRIES # already at max
|
||
|
||
sender._requeue_or_drop_failed_send(task)
|
||
|
||
with sender._lock:
|
||
dq = sender._pending_save_reqs.get("r1")
|
||
self.assertTrue(dq is None or len(dq) == 0, "task should NOT be re-enqueued after max retries")
|
||
self.assertNotIn("r1", sender._pending_save_counts, "pending count should be cleaned up on final drop")
|
||
sender.shutdown_omni_connectors()
|
||
|
||
def test_save_loop_retries_on_exception(self):
|
||
"""Integration: _save_loop retries a task when put() raises."""
|
||
from collections import deque
|
||
|
||
connector = _FailingConnector(fail_count=1, raise_on_fail=True)
|
||
sender = self._make_sender(connector)
|
||
task = self._make_task()
|
||
|
||
with sender._lock:
|
||
sender._pending_save_reqs["r1"] = deque([task])
|
||
sender._pending_save_counts["r1"] = 1
|
||
|
||
sender._stop_event.clear()
|
||
|
||
def run_one_loop():
|
||
sender._save_loop()
|
||
|
||
sender._stop_event.set() # will exit after one iteration
|
||
# Run manually instead of threading
|
||
# Simulate: pop task, send fails, requeue
|
||
popped_task = None
|
||
with sender._lock:
|
||
dq = sender._pending_save_reqs.get("r1")
|
||
if dq:
|
||
popped_task = dq.popleft()
|
||
if not dq:
|
||
del sender._pending_save_reqs["r1"]
|
||
|
||
if popped_task is not None:
|
||
success = False
|
||
try:
|
||
success = sender._send_single_request(popped_task)
|
||
except Exception:
|
||
pass
|
||
if not success:
|
||
sender._requeue_or_drop_failed_send(popped_task)
|
||
|
||
# After first failure, task should be re-enqueued
|
||
with sender._lock:
|
||
dq = sender._pending_save_reqs.get("r1")
|
||
self.assertIsNotNone(dq)
|
||
self.assertEqual(len(dq), 1)
|
||
requeued = dq[0]
|
||
self.assertEqual(requeued.get("_retry_count"), 1)
|
||
|
||
# Second attempt should succeed (connector now returns True)
|
||
success = sender._send_single_request(requeued)
|
||
self.assertTrue(success)
|
||
sender.shutdown_omni_connectors()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"fallback,external,expected", [("r", None, "mapped"), (None, "ext", "ext"), (None, None, None)]
|
||
)
|
||
def test_resolve_request_id_requires_internal_or_external_id(fallback, external, expected):
|
||
host = SimpleNamespace(_request_ids_mapping={"r": "mapped"})
|
||
request = SimpleNamespace(external_req_id=external)
|
||
if expected is None:
|
||
with pytest.raises(ValueError, match="request ID"):
|
||
OmniConnectorModelRunnerMixin._resolve_external_req_id(host, request, fallback)
|
||
else:
|
||
assert OmniConnectorModelRunnerMixin._resolve_external_req_id(host, request, fallback) == expected
|