Skip to content

vllm_omni.config.composable_parallel.translator

Translate a StrategySpec stack into vLLM-Omni parallel sizing.

Read a stack of strategy specs (one per mesh axis) and work out the tensor/data/pipeline parallel sizes, the dense-EP flag, and the number of stage replicas, plus who owns each axis's request routing. The result is a plain, CPU-computable sizing struct (:class:OmniParallelConfig) keyed by the real OmniEngineArgs / EngineArgs field names, so the deploy layer can splat it onto a stage's engine args.

The distinction this module enforces — engine data parallelism (a true vLLM intra-engine world dimension) vs. omni stage replicas (independent engines fanned out by omni's coordinator) — is documented on the axis validators that depend on it (see :func:_validate_dp and :func:_stage_replica_lb_policy). Routing we do not support yet (key-stable / affinity routing) raises NotImplementedError; any other invalid spec raises :class:AxisTranslationError.

L1Owner module-attribute 

L1Owner = Literal['delegated', 'engine']

logger module-attribute 

logger = init_logger(__name__)

AxisTranslationError

Bases: ValueError

Error for an invalid or unsupported strategy spec.

A single error type (rather than a tree of subclasses) keeps the public surface small and consistent with the rest of the codebase; the specific cause is in the message and is logged before the raise so it is visible even when the type is unavailable to a caller (e.g. across a server boundary).

Strategies that are valid but not built yet — key-stable / affinity routing today — raise NotImplementedError instead, to distinguish "we haven't implemented this" from "your config is wrong".

OmniParallelConfig dataclass

Result of translating a spec stack into omni parallel sizing.

data_parallel_size class-attribute instance-attribute 

data_parallel_size: int = 1

delegated_axes property 

delegated_axes: tuple[MeshAxisKind, ...]

enable_expert_parallel class-attribute instance-attribute 

enable_expert_parallel: bool = False

l1_owners class-attribute instance-attribute 

l1_owners: Mapping[MeshAxisKind, L1Owner] = field(
    default_factory=dict
)

omni_lb_policy class-attribute instance-attribute 

omni_lb_policy: str | None = None

pipeline_parallel_size class-attribute instance-attribute 

pipeline_parallel_size: int = 1

stage_replica_size class-attribute instance-attribute 

stage_replica_size: int = 1

tensor_parallel_size class-attribute instance-attribute 

tensor_parallel_size: int = 1

world_size property 

world_size: int

as_engine_kwargs 

as_engine_kwargs() -> dict[str, object]

Return per-stage kwargs keyed by real OmniEngineArgs/EngineArgs field names.

Two derived values are intentionally not emitted here because they are not per-stage engine args:

  • stage_replica_size — a per-stage deploy num_replicas knob (StageDeployConfig); the deploy layer consumes it separately.
  • omni_lb_policy — a pipeline-wide load-balancer policy the engine reads once at construction (see StrategyApplyResult.omni_lb_policy); it is applied at the orchestrator level, not folded into per-stage args.

translate_strategy_stack 

translate_strategy_stack(
    specs: Sequence[StrategySpec],
) -> OmniParallelConfig

Translate a spec stack into an OmniParallelConfig.

Supported kinds: dp (engine data parallel), tp, pp, (dense) ep, and stage_replica (omni replicas). Raises NotImplementedError for deferred (affinity / key-stable) routing, and :class:AxisTranslationError for any other invalid spec (a kind not yet translatable, a repeated kind, an owner incompatible with the axis kind, or unsupported routing). The EP degree must equal tensor_parallel_size * data_parallel_size.