Skip to content

vllm_omni.config.composable_parallel.apply

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.

RoleKey module-attribute 

RoleKey = str

logger module-attribute 

logger = init_logger(__name__)

StrategyApplyError

Bases: 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).

StrategyApplyResult dataclass

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.

omni_lb_policy class-attribute instance-attribute 

omni_lb_policy: str | None = None

per_role_config class-attribute instance-attribute 

per_role_config: dict[RoleKey, OmniParallelConfig] = field(
    default_factory=dict
)

per_stage_config class-attribute instance-attribute 

per_stage_config: dict[int, OmniParallelConfig] = field(
    default_factory=dict
)

stages instance-attribute 

stages: list[Any]

apply_strategy_specs 

apply_strategy_specs(
    stages: list[Any],
    strategy_specs: Mapping[
        RoleKey, Sequence[StrategySpec]
    ],
) -> StrategyApplyResult

Overlay per-role strategy specs onto a merged stage list.

Parameters:

Name Type Description Default
stages list[Any]

the list[StageConfig] returned by merge_pipeline_deploy.

 required
strategy_specs Mapping[RoleKey, Sequence[StrategySpec]]

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).

 required

Returns:

Name Type Description
A StrategyApplyResult

class:StrategyApplyResult with the mutated stages and the derived
StrategyApplyResult

pipeline-wide omni_lb_policy (if any stage_replica axis set one).

Raises:

Type Description
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.

check_device_layout 

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.