# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Apply declarative parallel strategies onto merged stage configs. This is the *override-after-merge* seam. ``merge_pipeline_deploy`` first fuses the pipeline topology + deploy YAML into a list of ``StageConfig`` objects; this module then overlays the engine sizing derived from a per-role ``StrategySpec`` stack onto those stages. Design rules (one writer per axis): * A strategy only writes the axes it actually declares. If a stack has no ``tp`` axis, this module never touches ``tensor_parallel_size`` — it does not force a default of 1. * **Conflict-on-explicit.** If a knob was already set explicitly by the deploy YAML and the strategy derives a different value, we raise rather than silently override. Equal values are a no-op; unset values are filled. * **Pre-spawn device check.** When a stage declares ``devices``, the device count must be consistent with the engine world size (``tp * dp * pp``) and the replica count, so misconfigurations fail at config time instead of at spawn. ``omni_lb_policy`` is *not* a per-stage config knob: omni reads it once at orchestrator construction (``AsyncOmniEngine``), not from stage configs. So a stage_replica axis's derived policy is surfaced on :class:`StrategyApplyResult` for the caller to pass to the orchestrator; it is never silently written into a stage where it would be ignored. """ from __future__ import annotations from collections.abc import Mapping, Sequence from dataclasses import dataclass, field from typing import Any, NoReturn from vllm.logger import init_logger from vllm_omni.config.composable_parallel.spec import StrategySpec from vllm_omni.config.composable_parallel.translator import ( OmniParallelConfig, translate_strategy_stack, ) logger = init_logger(__name__) # Role key -> the StageConfig fields each declared axis writes. _ENGINE_FIELD_BY_KIND: dict[str, str] = { "tp": "tensor_parallel_size", "dp": "data_parallel_size", "pp": "pipeline_parallel_size", } RoleKey = str class StrategyApplyError(ValueError): """Error for applying a strategy onto stage configs. A single error type (rather than per-cause subclasses) keeps the public surface small and consistent with the rest of the codebase; the specific cause — an unmatched role, a conflict with an explicit deploy value, or a device-count mismatch — is in the message and is logged before the raise so it stays visible even when the type is unavailable to a caller (e.g. across a server boundary). """ def _fail(msg: str) -> NoReturn: """Log and raise :class:`StrategyApplyError`.""" logger.error("[composable_parallel] %s", msg) raise StrategyApplyError(msg) @dataclass class StrategyApplyResult: """Outcome of applying strategies to a stage list. ``stages`` is the same list passed in (mutated in place and returned for convenience). ``omni_lb_policy`` is the pipeline-wide load-balancer policy derived from any stage_replica axes (``None`` if none declared one); the caller wires it into the orchestrator, since omni does not read it from stage configs. ``per_role_config`` is keyed by the role key the caller supplied — a ``model_stage`` name (str). ``per_stage_config`` is the same information keyed by the resolved integer ``stage_id`` — handy for re-validating a stage after later layers (e.g. CLI overrides) have run. """ stages: list[Any] omni_lb_policy: str | None = None per_role_config: dict[RoleKey, OmniParallelConfig] = field(default_factory=dict) per_stage_config: dict[int, OmniParallelConfig] = field(default_factory=dict) def _resolve_stage(stages: Sequence[Any], key: RoleKey) -> Any: """Find the single stage matching ``key`` by ``model_stage`` name. Role keys are ``model_stage`` names (str) — the same human-readable names that strategy.yaml authors use to refer to a stage (e.g. ``"thinker"``). """ matches = [s for s in stages if getattr(s, "model_stage", None) == key] descriptor = f"model_stage={key!r}" if not matches: available = ", ".join(f"{getattr(s, 'model_stage', '?')!r}(id={getattr(s, 'stage_id', '?')})" for s in stages) _fail(f"strategy role {descriptor} did not match any stage; available stages: {available}") if len(matches) > 1: _fail(f"strategy role {descriptor} is ambiguous; it matched {len(matches)} stages") return matches[0] def _set_explicit(engine_args: dict[str, Any], field_name: str, value: Any, *, role: RoleKey) -> None: """Write ``field_name`` with conflict-on-explicit semantics. ``engine_args`` is the raw mapping loaded straight from the deploy YAML (before the engine-args dataclass is constructed), so a key is *present* only when the YAML set it explicitly. We therefore key off membership rather than ``.get``: an explicit ``None`` in the YAML is a real, present value and must not be conflated with a key miss (which ``.get`` would report as ``None``). """ if field_name in engine_args and engine_args[field_name] != value: existing = engine_args[field_name] _fail( f"role {role!r}: strategy derives {field_name}={value} but the deploy config already " f"set {field_name}={existing}. Remove one of them (the strategy is the single writer " "for the axes it declares)." ) engine_args[field_name] = value def _set_num_replicas(runtime: dict[str, Any], value: int, *, role: RoleKey) -> None: """Write ``num_replicas`` with conflict-on-explicit (treating 1 as unset).""" existing = runtime.get("num_replicas", 1) if existing not in (1, value): _fail( f"role {role!r}: strategy derives num_replicas={value} but the deploy config already " f"set num_replicas={existing}. Remove one of them." ) runtime["num_replicas"] = value def _parse_device_count(devices: Any) -> int | None: """Return the number of device ids in a ``devices`` value, or None.""" if devices is None: return None if isinstance(devices, (list, tuple)): return len([d for d in devices]) text = str(devices).strip() if not text: return None return len([d for d in text.split(",") if d.strip()]) def check_device_layout( devices: Any, *, tensor_parallel_size: int, data_parallel_size: int, pipeline_parallel_size: int, num_replicas: int, role: RoleKey, ) -> None: """Validate a stage's device count against its world size (× replicas in pool mode). Raises :class:`StrategyApplyError` when the declared ``devices`` count is neither the per-replica world size nor the full ``num_replicas`` pool. Exposed publicly so the resolved (post-CLI) layout can be re-checked after later override layers, not just the strategy-derived snapshot. """ count = _parse_device_count(devices) if count is None: return world = int(tensor_parallel_size) * int(data_parallel_size) * int(pipeline_parallel_size) replicas = int(num_replicas or 1) # Two valid shapes: a single per-replica template (== world) or the full # pool across replicas (== replicas * world). This mirrors omni's own # per-replica device splitter. valid = {world, replicas * world} if count not in valid: _fail( f"role {role!r}: declared {count} device id(s) but the strategy world size is {world} " f"(tp={tensor_parallel_size} * dp={data_parallel_size} * pp={pipeline_parallel_size}). " f"Provide either {world} (per-replica template) or {replicas * world} " f"(num_replicas={replicas} pool)." ) def _check_devices(runtime: dict[str, Any], cfg: OmniParallelConfig, *, role: RoleKey) -> None: """Pre-spawn check on the strategy-derived snapshot (see :func:`check_device_layout`).""" check_device_layout( runtime.get("devices"), tensor_parallel_size=cfg.tensor_parallel_size, data_parallel_size=cfg.data_parallel_size, pipeline_parallel_size=cfg.pipeline_parallel_size, num_replicas=int(runtime.get("num_replicas", 1) or 1), role=role, ) def _apply_to_stage(stage: Any, cfg: OmniParallelConfig, *, role: RoleKey) -> None: engine_args = stage.yaml_engine_args runtime = stage.yaml_runtime declared = set(cfg.l1_owners.keys()) for kind, field_name in _ENGINE_FIELD_BY_KIND.items(): if kind in declared: _set_explicit(engine_args, field_name, getattr(cfg, field_name), role=role) if "ep" in declared and cfg.enable_expert_parallel: _set_explicit(engine_args, "enable_expert_parallel", True, role=role) if "stage_replica" in declared: _set_num_replicas(runtime, cfg.stage_replica_size, role=role) _check_devices(runtime, cfg, role=role) def apply_strategy_specs( stages: list[Any], strategy_specs: Mapping[RoleKey, Sequence[StrategySpec]], ) -> StrategyApplyResult: """Overlay per-role strategy specs onto a merged stage list. Args: stages: the ``list[StageConfig]`` returned by ``merge_pipeline_deploy``. strategy_specs: maps a role (a ``model_stage`` name, e.g. ``"thinker"``) to that stage's stack of ``StrategySpec`` (one per declared mesh axis). Role keys are ``model_stage`` names (str). Returns: A :class:`StrategyApplyResult` with the mutated stages and the derived pipeline-wide ``omni_lb_policy`` (if any stage_replica axis set one). Raises: StrategyApplyError: a role matched zero or multiple stages, a derived value conflicts with an explicit deploy value (including two roles deriving different ``omni_lb_policy`` values), or a stage's device count is inconsistent with its world size. AxisTranslationError: the spec stack is invalid/unsupported. NotImplementedError: the spec stack requests routing not built yet. """ result = StrategyApplyResult(stages=stages) lb_policy: str | None = None lb_owner: RoleKey | None = None for key, specs in strategy_specs.items(): stage = _resolve_stage(stages, key) cfg = translate_strategy_stack(specs) _apply_to_stage(stage, cfg, role=key) result.per_role_config[key] = cfg stage_id = getattr(stage, "stage_id", None) if stage_id is not None: result.per_stage_config[int(stage_id)] = cfg if cfg.stage_replica_size > 1 and cfg.omni_lb_policy is not None: if lb_policy is not None and lb_policy != cfg.omni_lb_policy: _fail( f"role {key!r} derives omni_lb_policy={cfg.omni_lb_policy!r} but role " f"{lb_owner!r} already derived {lb_policy!r}. The load-balancer policy is " "pipeline-wide; only one value is allowed." ) lb_policy = cfg.omni_lb_policy lb_owner = key result.omni_lb_policy = lb_policy return result