[Core] Composable parallel: Phase 1 strategy overlay (TP/DP/stage-replica) (RFC #4084) (#4281)

Signed-off-by: Kushan A. <kahmadian@nvidia.com>
Signed-off-by: kushanam <kushanam@users.noreply.github.com>
Signed-off-by: Kushan A <kahmadian@nvidia.com>
Signed-off-by: kahmadian <kahmadian@users.noreply.github.com>
Co-authored-by: kushanam <kushanam@users.noreply.github.com>
Co-authored-by: kahmadian <kahmadian@users.noreply.github.com>

tests/engine/test_async_omni_engine_stage_init.py

@@ -739,7 +739,7 @@ def test_resolve_stage_configs_injects_global_diffusion_attention_when_missing(m
     monkeypatch.setattr(
         engine_mod,
         "load_and_resolve_stage_configs",
-        lambda *args, **kwargs: ("dummy-config", [stage_cfg]),
+        lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
     )
 
     _config_path, stage_configs = engine._resolve_stage_configs(
@@ -772,7 +772,7 @@ def test_resolve_stage_configs_preserves_stage_diffusion_attention(monkeypatch):
     monkeypatch.setattr(
         engine_mod,
         "load_and_resolve_stage_configs",
-        lambda *args, **kwargs: ("dummy-config", [stage_cfg]),
+        lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
     )
 
     _config_path, stage_configs = engine._resolve_stage_configs(
@@ -798,7 +798,7 @@ def test_resolve_stage_configs_does_not_inject_diffusion_attention_into_llm_stag
     monkeypatch.setattr(
         engine_mod,
         "load_and_resolve_stage_configs",
-        lambda *args, **kwargs: ("dummy-config", [stage_cfg]),
+        lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
     )
 
     _config_path, stage_configs = engine._resolve_stage_configs(

tests/engine/test_strategy_lb_policy.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for strategy-derived omni_lb_policy threading into AsyncOmniEngine.
+
+``AsyncOmniEngine._apply_strategy_lb_policy`` only touches ``self._omni_lb_policy``
+and module-level logging, so we exercise it with a lightweight stub ``self`` to
+avoid standing up a full engine (which needs GPUs/model weights).
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _Stub:
+    def __init__(self, current: str = "random") -> None:
+        self._omni_lb_policy = current
+
+
+_apply = AsyncOmniEngine._apply_strategy_lb_policy
+
+
+def test_no_derived_policy_is_noop():
+    stub = _Stub("random")
+    _apply(stub, None, {})
+    assert stub._omni_lb_policy == "random"
+
+
+def test_derived_overrides_default_random():
+    # "random" is the engine default == "unset", so the strategy value wins.
+    stub = _Stub("random")
+    _apply(stub, "round-robin", {})
+    assert stub._omni_lb_policy == "round-robin"
+
+
+def test_derived_overrides_when_flag_absent():
+    stub = _Stub("random")
+    _apply(stub, "least-queue-length", {"omni_lb_policy": None})
+    assert stub._omni_lb_policy == "least-queue-length"
+
+
+def test_explicit_matching_flag_is_fine():
+    stub = _Stub("round-robin")
+    _apply(stub, "round-robin", {"omni_lb_policy": "round-robin"})
+    assert stub._omni_lb_policy == "round-robin"
+
+
+def test_explicit_conflicting_flag_raises():
+    stub = _Stub("round-robin")
+    with pytest.raises(ValueError, match="Conflicting load-balancer policy"):
+        _apply(stub, "least-queue-length", {"omni_lb_policy": "round-robin"})
+
+
+def test_explicit_random_is_treated_as_unset():
+    # An explicit --omni-lb-policy=random is indistinguishable from the default,
+    # so a strategy value still overrides it (documented Phase-1 semantics).
+    stub = _Stub("random")
+    _apply(stub, "round-robin", {"omni_lb_policy": "random"})
+    assert stub._omni_lb_policy == "round-robin"