[Refactor][2/N Scheduler]Replace _FULL_PAYLOAD_INPUT_STAGES with a resolved stage transport capability. (#6149)

Signed-off-by: wy17003 <yaow17003@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/engine/test_arg_utils.py

@@ -8,6 +8,8 @@ import inspect
 import json
 import os
 import shutil
+import sys
+import types
 from types import SimpleNamespace
 from unittest.mock import Mock
 
@@ -20,6 +22,8 @@ from vllm.engine.arg_utils import EngineArgs
 from vllm_omni.config.model import OmniModelConfig
 from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+from vllm_omni.platforms import current_omni_platform
+from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -42,7 +46,102 @@ def test_default_stage_id_is_concrete_int():
     assert cfg.stage_id == 0
 
 
-def test_multimodal_kwarg_overrides(mocker):
+def test_full_payload_capability_reaches_omni_model_config(monkeypatch):
+    """Keep the pipeline capability intact through the engine adapter."""
+    captured: dict[str, object] = {}
+    baseline_config = Mock()
+    worker_module = types.ModuleType("test_connector_worker")
+
+    class ConnectorRunner(OmniConnectorModelRunnerMixin):
+        pass
+
+    class ConnectorWorker:
+        model_runner_cls = ConnectorRunner
+
+    worker_module.ConnectorWorker = ConnectorWorker
+    monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module)
+
+    monkeypatch.setattr(OmniEngineArgs, "_ensure_omni_models_registered", lambda _self: None)
+    monkeypatch.setattr(EngineArgs, "create_model_config", lambda _self: baseline_config)
+
+    def capture_omni_config(_cls, model_config, **omni_kwargs):
+        assert model_config is baseline_config
+        captured.update(omni_kwargs)
+        return baseline_config
+
+    monkeypatch.setattr(
+        OmniModelConfig,
+        "from_vllm_model_config",
+        classmethod(capture_omni_config),
+    )
+
+    OmniEngineArgs(
+        stage_id=1,
+        requires_full_payload_input=True,
+        worker_cls="test_connector_worker.ConnectorWorker",
+    ).create_model_config()
+
+    assert captured["requires_full_payload_input"] is True
+
+
+def test_stage_without_connector_configuration_accepts_plain_runner(monkeypatch):
+    worker_module = types.ModuleType("test_async_scheduler_worker")
+
+    class WorkerWithoutConnector:
+        model_runner_cls = object
+
+    worker_module.WorkerWithoutConnector = WorkerWithoutConnector
+    monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module)
+
+    args = OmniEngineArgs(
+        stage_id=1,
+        async_chunk=True,
+        worker_cls="test_async_scheduler_worker.WorkerWithoutConnector",
+    )
+
+    assert args.requires_full_payload_input is False
+
+
+def test_full_payload_capability_requires_selected_worker_connector(monkeypatch):
+    worker_module = types.ModuleType("test_worker_without_connector")
+
+    class WorkerWithoutConnector:
+        model_runner_cls = object
+
+    worker_module.WorkerWithoutConnector = WorkerWithoutConnector
+    monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module)
+
+    with pytest.raises(ValueError, match="does not provide an Omni connector model runner"):
+        OmniEngineArgs(
+            stage_id=1,
+            requires_full_payload_input=True,
+            worker_cls="test_worker_without_connector.WorkerWithoutConnector",
+        )
+
+
+def test_full_payload_capability_validates_platform_selected_worker(monkeypatch):
+    worker_module = types.ModuleType("test_platform_worker_without_connector")
+
+    class WorkerWithoutConnector:
+        model_runner_cls = object
+
+    worker_module.WorkerWithoutConnector = WorkerWithoutConnector
+    monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module)
+    monkeypatch.setattr(
+        current_omni_platform,
+        "get_omni_ar_worker_cls",
+        lambda: "test_platform_worker_without_connector.WorkerWithoutConnector",
+    )
+
+    with pytest.raises(ValueError, match="does not provide an Omni connector model runner"):
+        OmniEngineArgs(
+            stage_id=1,
+            requires_full_payload_input=True,
+            worker_type="ar",
+        )
+
+
+def test_multimodal_kwarg_overrides(monkeypatch):
     """Ensure that overrides in the multimodal config are preserved."""
     sig = inspect.signature(OmniEngineArgs)
     default_mm_cache = sig.parameters["mm_processor_cache_gb"].default
@@ -56,8 +155,8 @@ def test_multimodal_kwarg_overrides(mocker):
         assert self.mm_processor_cache_gb == override_val
         return fake_model_config
 
-    mocker.patch.object(EngineArgs, "create_model_config", _fake_parent_create_model_config)
-    mocker.patch.object(OmniModelConfig, "from_vllm_model_config", side_effect=lambda model_config, **_: model_config)
+    monkeypatch.setattr(EngineArgs, "create_model_config", _fake_parent_create_model_config)
+    monkeypatch.setattr(OmniModelConfig, "from_vllm_model_config", lambda model_config, **_: model_config)
 
     cfg = OmniEngineArgs(
         model="Qwen/Qwen2-VL-2B-Instruct",