diff --git a/tests/worker/test_omni_connector_mixin.py b/tests/worker/test_omni_connector_mixin.py index 7a2a0053b..af472fd95 100644 --- a/tests/worker/test_omni_connector_mixin.py +++ b/tests/worker/test_omni_connector_mixin.py @@ -105,6 +105,34 @@ class _FakeTPGroup: # ------------------------------------------------------------------ # +@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.""" diff --git a/vllm_omni/worker/omni_connector_model_runner_mixin.py b/vllm_omni/worker/omni_connector_model_runner_mixin.py index ae7a0c0fd..07a0327a6 100644 --- a/vllm_omni/worker/omni_connector_model_runner_mixin.py +++ b/vllm_omni/worker/omni_connector_model_runner_mixin.py @@ -25,7 +25,10 @@ from vllm.logger import init_logger from vllm_omni.data_entry_keys import OmniPayload from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory -from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec +from vllm_omni.distributed.omni_connectors.utils.config import ( + ConnectorSpec, + get_stage_connector_role, +) from vllm_omni.outputs import OmniConnectorOutput from vllm_omni.worker.payload_span import ( get_tensor_span, @@ -46,6 +49,19 @@ if TYPE_CHECKING: logger = init_logger(__name__) +def _should_create_payload_connector(model_config: Any) -> bool: + """Whether this stage owns runner payload transport for its edge. + + Sender edges may instead be owned solely by KV transfer. Receivers still + need a connector even though they do not declare a downstream payload hook. + """ + if get_stage_connector_role(model_config) != "sender": + return True + + next_stage_func = getattr(model_config, "custom_process_next_stage_input_func", None) + return isinstance(next_stage_func, str) and bool(next_stage_func) + + def should_accumulate_full_payload_output(model_config, custom_process_func) -> bool: """Producer-side structural gate. @@ -96,7 +112,9 @@ class OmniConnectorModelRunnerMixin: model_config: Stage-level model config with connector settings. kv_transfer_manager: Existing KV transfer manager to delegate to. """ - self._omni_connector: OmniConnectorBase | None = self._create_connector(model_config) + self._omni_connector: OmniConnectorBase | None = ( + self._create_connector(model_config) if _should_create_payload_connector(model_config) else None + ) self._kv_transfer_manager = kv_transfer_manager self._async_chunk: bool = getattr(model_config, "async_chunk", False)