Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / recipe_spec.py: 95%
131 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2025-2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""recipe_spec: recipe input contract and placement DSL for the dual-mode DTensor strategy.
17Contains the user-visible plan/spec data model (05 §3.1/§3.2 canonical):
18- ``MeshAxisName``: mesh dimension name enum (canonical definition, imported and reused by later docs such as 06);
19- ``NamedPlacement``: alias for ``dict[MeshAxisName, Placement]``;
20- ``ModuleShardingSpec``: per-module I/O contract;
21- ``PlacementMismatchError``: error for placement declarations inconsistent with DTensor propagation;
22- ``resolve_placements`` / ``parse_placement`` / ``parse_named_placement`` / ``_normalize_out_fields``:
23 placement utilities;
24- the injection DSL (``local_compute`` / ``inner_wrapper`` decorators,
25 ``InjectionMeta``, ``require_injection_meta``): the recipe-side input
26 contract of the explicit injection mechanism.
28The model-level :class:`ShardingPlan` lives in ``distributed/plan.py``; the
29semantic-role templates live in ``distributed/_builder/default_templates.py``.
30"""
32import inspect
33import logging
34import re
35from dataclasses import dataclass, field
36from enum import Enum
37from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
39from hyper_parallel.core.dtensor.placement_types import (
40 Partial,
41 Placement,
42 Replicate,
43 Shard,
44)
46if TYPE_CHECKING:
47 from hyper_parallel.distributed.tensor_parallel.head_count import TpLocalAttrPlan
49logger = logging.getLogger(__name__)
52class MeshAxisName(str, Enum):
53 """Canonical enum of mesh dimension names.
55 A str enum: directly comparable to plain strings like "tp" and usable as
56 a dict key.
57 """
58 TP = "tp"
59 CP = "cp"
60 EP = "ep"
61 PP = "pp"
62 DP = "dp"
63 DP_REPLICATE = "dp_replicate"
64 DP_SHARD = "dp_shard"
65 DP_CP = "dp_cp"
66 EP_SHARD = "ep_shard"
69# Shorthand aliases for the {TP: ..., CP: ..., EP: ...} literals in templates and examples.
70# str-enum keys interoperate with plain string keys like "tp" (hash/eq are identical).
71TP = MeshAxisName.TP
72CP = MeshAxisName.CP
73EP = MeshAxisName.EP
74DP = MeshAxisName.DP
76# NamedPlacement = {MeshAxisName: Placement}.
77# The key is a mesh dimension name; the N in Shard(N) of a value is a tensor dimension index (05 §3.2.1).
78NamedPlacement = Dict[MeshAxisName, Placement]
81class PlacementMismatchError(ValueError):
82 """DTensor propagation result is inconsistent with the ModuleShardingSpec declaration (05 §5.3)."""
84 def __init__(self, module_name: str, expected: Any, actual: Any, stage: str) -> None:
85 """Initialize the mismatch error.
87 Args:
88 module_name: FQN of the module whose placement mismatched.
89 expected: Placement declared by the ShardingConfig.
90 actual: Placement produced by DTensor propagation.
91 stage: Which contract stage mismatched (e.g. "in_dst"/"out_src").
92 """
93 self.module_name = module_name
94 self.expected = expected
95 self.actual = actual
96 self.stage = stage
97 super().__init__(
98 f"[{module_name}] {stage} placement mismatch:\n"
99 f" Expected (from ShardingConfig.{stage}): {expected}\n"
100 f" Actual (from DTensor propagation): {actual}\n"
101 f" → Check the ShardingConfig for this module."
102 )
105@dataclass
106class ModuleShardingSpec:
107 """Complete DTensor contract for a single module (05 §3.2).
109 The four placement fields form the complete I/O contract — the runtime does no
110 inference and executes exactly as declared:
112 in_src: placement of the input when it arrives at the module boundary
113 (from the output of an upstream module or the dataloader)
114 in_dst: placement required by the module's internal computation
115 (a mismatch triggers communication)
116 out_src: placement naturally produced by the module's internal computation
117 (used by validate mode)
118 out_dst: placement expected by downstream modules (a mismatch triggers
119 communication)
121 plan_overrides input side only: the contract fields (params/in_src/in_dst/
122 out_src/out_dst/out_names) follow one rule — **"unset inherits, set wins"**
123 (2026-08-05): ``None`` (unset) or the sentinel ``"auto"`` means
124 inherit-from-template (merge only); an explicit value IS the final value,
125 including the empty dict ``{}`` (explicit "no sharding / no contract"
126 — the ViT ``params={}`` pattern); the sentinel ``"none"`` is a readable
127 alias for clearing (``{}`` for params/in_*, ``None`` for out_*). The
128 planner resolves sentinels and normalizes ``None`` at merge/insert time;
129 a spec inside a finished ShardingPlan always carries concrete dicts
130 (params/in_src/in_dst) or Optional dicts (out_*).
131 """
132 # ── Parameter sharding: submodule path → NamedPlacement ──
133 # Default None = "not declared" (merge: inherit); an explicit {} means
134 # "this boundary shards nothing" (I/O-stitch-only citizen).
135 params: Optional[Dict[str, NamedPlacement]] = None
137 # ── Input contract ──
138 in_src: Optional[Dict[str, NamedPlacement]] = None
139 in_dst: Optional[Dict[str, NamedPlacement]] = None
141 # ── Output contract ──
142 # out_src=None: no src validation; out_dst=None: output needs no redistribution.
143 # Single-output modules use {"output": NamedPlacement}; the scalar shorthand
144 # {TP: ...} is wrapped into {"output": ...} during the normalization phase
145 # (_normalize_out_fields).
146 out_src: Optional[Dict[str, NamedPlacement]] = None
147 out_dst: Optional[Dict[str, NamedPlacement]] = None
148 # out_names: output-name ordering for multi-output modules (returning a tuple),
149 # used to map the keys of out_src/out_dst to tuple positions
150 # (RedistOp.arg_index). Defaults to the key order of out_src.
151 out_names: Optional[List[str]] = None
153 # TP-local module-instance attributes explicitly declared by the user.
154 # None means "not declared" during override merge; [] explicitly clears
155 # an inherited glob declaration. The planner normalizes this field into
156 # _tp_local_attr_plan after all overrides have been merged.
157 tp_divide_attrs: Optional[List[str]] = None
159 # Internal planner output. init=False keeps it out of both the public
160 # programmatic constructor and the YAML transport interface.
161 # Annotated as a string: TpLocalAttrPlan lives in
162 # tensor_parallel/head_count.py, which imports resolve_placements from
163 # this module — a runtime import here would be circular.
164 _tp_local_attr_plan: Optional["TpLocalAttrPlan"] = field(
165 default=None, init=False, repr=False, compare=False,
166 )
168 # D-22 (rowwise bias defer): planner-computed tuple of bias param paths
169 # (relative to the boundary module, e.g. ("o_proj.bias",)) whose addition
170 # is deferred until AFTER the boundary exit TP reduction (Megatron
171 # RowParallelLinear semantics: the bias never enters F.linear inside the
172 # region, so the Partial reduce counts it exactly once). Internal — built
173 # by ShardingPlanner._finalize_deferred_biases after all overrides are
174 # merged (detection is anchored on the FINAL spec declarations, so
175 # derived / merge / insert / derive=False specs share one code path);
176 # user spec objects never carry this field.
177 _deferred_bias_params: Tuple[str, ...] = field(
178 default=(), init=False, repr=False, compare=False,
179 )
181 # ── Boundary flag ──
182 is_boundary: bool = True
184 # Codegen source-lowering opt-out: set by meta normalization (NOT the
185 # planner) after freeze, when the boundary's class is shared with a
186 # non-boundary instance — a class-level forward rewrite would leak the
187 # boundary body into instances that never receive ``_hp_boundary``.
188 # The runtime per-instance wrapper executes the same boundary contract
189 # instead. ``None`` (absent) is the only unmarked state; only an
190 # explicit ``True`` crosses the meta boundary (int_truthy freeze kind).
191 skip_source_lowering: Optional[bool] = None
193 # ── Structural flags (user-configurable) ──
194 # ┌─────────────────────────────────────────────────────────────────┐
195 # │ The user extension-point interfaces: region_dispatch / │
196 # │ local_compute_fn / inner_target / inner_wrapper │
197 # │ │
198 # │ One axiom governs every boundary: region computation is │
199 # │ dispatch-through by default (under validate: DTensors enter │
200 # │ directly, strategy propagation, real out_src validation); what │
201 # │ cannot dispatch explicitly declares region_dispatch=False │
202 # │ (skeleton black-box: to_local -> local execution -> │
203 # │ declarative re-wrap). │
204 # │ │
205 # │ [local-region family] (module level: skeleton unchanged, │
206 # │ content swapped) │
207 # │ region_dispatch / local_compute_fn │
208 # │ The skeleton = _wrap_local_region_forward: boundary │
209 # │ entry/exit stitching + local compute + validate dual-mode │
210 # │ fault tolerance (to_local/_temp_local_params/from_local │
211 # │ re-wrapping), shared by both modes. Whether a module runs │
212 # │ through the skeleton is **derived from a single resolution │
213 # │ chain** (not a stored bool) by _resolve_local_compute_fn: │
214 # │ chain link 1 local_compute_fn (user-defined computation: │
215 # │ callable, or a factory Target built at apply │
216 # │ time — e.g. a shipped EP archetype factory │
217 # │ (recipes.qwen2moe_ep_compute_fn)) │
218 # │ chain link 2 region_dispatch=False (the module's own │
219 # │ forward can't dispatch — it IS the compute) │
220 # │ none present → None (skeleton not used) │
221 # │ │
222 # │ ★ EXPLICIT INJECTION (since the rework): the shipped EP │
223 # │ compute is NEVER auto-injected. The planner only shards the │
224 # │ expert params ({EP: Shard(0)} + _ep_size/_ep_stack metadata); │
225 # │ the compute must be declared via plan_overrides / direct │
226 # │ spec assignment, WITH an explicit region_dispatch (no │
227 # │ default for injections — the tutorial teaches False for │
228 # │ CP/EP-style comm-carrying injections and why). An EP-sharded │
229 # │ boundary whose chain resolves to None fails fast at apply │
230 # │ time (_preflight_compute_injection). HF-native MoE layouts │
231 # │ (per-expert / batched) get the template's region_dispatch │
232 # │ CLEARED by the planner (their forward is not EP-aware); │
233 # │ custom-named pre-stacked modules (w1/w2/w3) keep False │
234 # │ (EP-aware by construction). │
235 # │ │
236 # │ [inner-wrap family] (submodule level: content and wrapping both │
237 # │ replaced) │
238 # │ inner_target / inner_wrapper │
239 # │ The mechanism is a generic "locate an inner submodule + │
240 # │ replace its forward": inner_target answers "replace whom", │
241 # │ inner_wrapper answers "replace with what". It is NOT CP-gated │
242 # │ (declaration == application); the shipped reference domain is │
243 # │ CP (K/V all-gather), and the four INNER_WRAPPER_REGISTRY │
244 # │ entries still require an active cp axis (fail-fast otherwise) │
245 # │ — custom wrappers receive cp_mesh=None and own their │
246 # │ semantics. │
247 # │ │
248 # │ ★ EXPLICIT INJECTION (since the rework): NO heuristic │
249 # │ dispatch. inner_wrapper must be declared explicitly: a │
250 # │ INNER_WRAPPER_REGISTRY name ("sdpa_qkv"/"sdpa_hf"/"flex_qkv"/ │
251 # │ "flex_hf"), a callable fn(target_module, cp_mesh), or a │
252 # │ Target pointing at a wrapper fn (e.g. │
253 # │ wrappers.sdpa_hf_cp_wrapper). _needs_cp_attn is now pure │
254 # │ METADATA (template recognition) — it no longer triggers any │
255 # │ injection; an attention boundary under an active cp mesh │
256 # │ without inner_wrapper fails fast at apply time. │
257 # │ │
258 # │ The two families are orthogonal and composable: the same module │
259 # │ may declare both inner_* (wrapping of an inner attention) and │
260 # │ local_* (module-level skeleton). │
261 # └─────────────────────────────────────────────────────────────────┘
262 #
263 # region_dispatch: **the validate execution-mode declaration for the
264 # region computation (whatever its source: the module's own forward or
265 # an injected function)** — one axiom across the whole framework: region
266 # computation is dispatch-through by default; what cannot dispatch
267 # explicitly declares False.
268 # - None (no injection): an ordinary boundary — under validate, DTensors
269 # go straight into the original forward, dispatch-through + real
270 # out_src validation (the axiom's default, not this field's "default
271 # value");
272 # - False (no injection): the module's own forward cannot dispatch
273 # (data-dependent logic inside the forward, e.g. the a2a of an
274 # in-house EP-aware MoE) → takes the local-region skeleton, compute =
275 # the module's own forward (derived and filled by the planner when the
276 # template recognizes such a module, visible in the logs);
277 # - With injection (local_compute_fn / inner_wrapper non-None): **must be
278 # declared explicitly** (no default) — True = the injected code is pure
279 # standard ops, dispatchable (validate dispatch-through + real
280 # out_src/inner_out_src validation); False = the injected code contains
281 # communication / custom kernels (validate black-box local execution +
282 # declarative re-wrap). Missing → fail-fast at apply time;
283 # - True (no injection): a redundant declaration → fail-fast (an
284 # ordinary boundary is inherently dispatch-through).
285 # production is not affected by this field (region is always local
286 # passthrough).
287 region_dispatch: Optional[bool] = None
289 # ── inner-wrap custom entry points (user-configurable, 05 §4.4.2/§8.6) ──
290 # inner_target: **pure location** — names the attribute of the inner
291 # submodule whose forward is to be replaced ("self" means the module
292 # itself). Automatic locating (_resolve_inner_target) fails fast, in which
293 # case the user must specify it via plan_overrides/injections. Declaring
294 # inner_target WITHOUT inner_wrapper is an error since the rework
295 # (location alone cannot pick a scheme — no heuristic dispatch).
296 # inner_wrapper: **pure behavior** — selects which scheme wraps the target.
297 # NOT CP-gated: declaration == application, whatever the parallel mode.
298 # - str: a name in the INNER_WRAPPER_REGISTRY registry ("sdpa_qkv"/"sdpa_hf"/
299 # "flex_qkv"/"flex_hf", or a user-registered name); explicitly pins a
300 # built-in scheme; an unknown name fails fast; the four SHIPPED names
301 # are CP schemes and require an active cp axis (fail-fast otherwise);
302 # - Callable: a fully custom wrapper, which MUST be decorated
303 # ``@inner_wrapper`` (injection discipline, see injection.py —
304 # undecorated callables fail fast). Context contract:
305 # fn(target_module, mesh, tp_mesh, cp_mesh, ep_mesh) — the anchor plus
306 # the mesh family are MANDATORY context params, ALL filled by the
307 # framework at apply time (None for inactive axes; the user just uses
308 # them); ``spec`` is the only optional context. The wrapper replaces
309 # target.forward in place (use collectives.flex_cp_allgather for K/V
310 # all-gather). The replaced forward must accept the original
311 # forward's params (validated at apply time). region_dispatch must be
312 # declared explicitly (see above): the injected code either targets
313 # local tensors only (False, the adapter manages the DTensor
314 # conversion) or is pure standard ops, dispatchable (True, validate
315 # dispatch-through with real validation);
316 # - Target (hyper_parallel.trainer.config.Target): a delayed wrapper
317 # reference resolved from YAML; the referenced fn must likewise be
318 # @inner_wrapper decorated (the shipped built-ins
319 # wrappers.sdpa_hf_cp_wrapper etc. satisfy this contract
320 # directly). A None return = in-place replacement done; a callable
321 # return (also decorated) = applied as a custom wrapper.
322 inner_target: Optional[str] = None
323 inner_wrapper: Optional[Union[str, Callable]] = None # or Target
324 # inner_out_src: the **explicit declaration** of the output placement for
325 # inner-wrap case B (inner_target points at a submodule rather than
326 # self) — the framework does zero derivation and zero guessing about the
327 # inner output layout:
328 # - sentinel "first_input": a layout-preserving declaration (the output
329 # layout == the runtime layout of the first DTensor input argument;
330 # used by attention-type wrappers; illegal for multiple outputs);
331 # - NamedPlacement ({axis: Placement}): an explicit single-output
332 # declaration;
333 # - {name: NamedPlacement}: a per-name declaration for multiple outputs
334 # (tuple positions follow the declaration key order).
335 # Case B without a declaration → fail-fast at apply time (case A
336 # target=self uses the boundary out_src and does not need this field).
337 # validate does no propagation check on the inner region: declaration
338 # errors are backstopped by post-rewrap global-shape consistency / the
339 # boundary out_src check / numerical comparison.
340 inner_out_src: Optional[Union[str, Dict]] = None
342 # ── local-region custom computation (user-configurable, 05 §4.4.3/§8.6) ──
343 # local_compute_fn: the custom compute FACTORY of the local region
344 # (chain link 1; single form since 2026-08-10 — the direct compute-fn
345 # form was retired, mesh family is always filled, use-it-or-not):
346 # - Callable: a ``@local_compute``-decorated factory
347 # fn(mesh, tp_mesh, cp_mesh, ep_mesh, [module], <config keys...>)
348 # -> compute_fn (programmatic direct pass; injection discipline:
349 # mesh family mandatory, no defaults, no *args/**kwargs);
350 # - Target: the same factory as a delayed reference resolved from
351 # YAML (config keys bound by name; typo fail-fast).
352 # The factory is built ONCE at apply time with the framework context
353 # filled by name (mesh / tp_mesh / cp_mesh / ep_mesh mandatory —
354 # ep_mesh is the same object the expert params were sharded on;
355 # module optional anchor) and must RETURN the region compute fn
356 # fn(module, *local_args) -> Tensor, whose params MUST match the
357 # module's forward params (validated at apply time) — executed on
358 # local tensors inside the _wrap_local_region_forward skeleton (the
359 # skeleton handles boundary entry/exit stitching + validate dual-mode
360 # to_local/_temp_local_params/from_local re-wrapping). Config keys
361 # carry DATA only (function-typed config is rejected by discipline —
362 # custom behavior = your own injected function, e.g. a MoE factory
363 # with its router written inline).
364 # Suitable for custom modules that want to reuse the skeleton with their
365 # own data-dependent logic: typically a custom MoE (router not in
366 # MOE_ROUTER_ADAPTERS / expert layout not using HF-standard naming /
367 # hooked to a DeepEP fused dispatcher).
368 # Priority: local_compute_fn > region_dispatch=False gate (the module's own
369 # forward) — a single resolution chain (_resolve_local_compute_fn);
370 # declaring it takes effect immediately: **there is no need — and it is
371 # wrong — to also set region_dispatch=False** unless the injected fn is
372 # itself non-dispatchable; the skeleton gate is derived from the
373 # resolution chain (non-None means the skeleton is used).
374 # Dual-mode convention: with region_dispatch=False the inputs are always
375 # local tensors and the return value is a local tensor; validate's
376 # DTensor unwrap/re-wrapping is done by the skeleton, so the compute_fn
377 # need not be aware of the mode. With region_dispatch=True the compute
378 # fn must be pure dispatchable ops (validate feeds DTensors straight
379 # through and truly validates out_src).
380 local_compute_fn: Optional[Callable] = None # or factory Target
382 # ── Internal flags (set automatically by ShardingPlanner / applier) ──
383 _is_terminal: bool = False # marked automatically during chained propagation
384 # _needs_cp_attn: template recognition METADATA (attention module: the
385 # inner attention needs a CP-aware forward under an active cp mesh). Since
386 # the explicit-injection rework it does NOT trigger any injection — it is
387 # used only by the apply-time preflight (an attention boundary under cp>1
388 # without inner_wrapper fails fast) and for introspection.
389 _needs_cp_attn: bool = False
390 # _resolved_inner_wrapper: written back by the applier after resolution (for
391 # introspection) — the inner wrapper actually injected:
392 # "sdpa_qkv"/"sdpa_hf"/"flex_qkv"/"flex_hf"/"custom"/<target_path>/None.
393 _resolved_inner_wrapper: Optional[str] = None
394 # _resolved_inner_target: written back by the applier after resolution —
395 # the inner-wrap target actually located: the child attribute name, or
396 # "self" when the boundary module itself is wrapped. Nothing is located
397 # silently: the resolved target is always visible here and in the INFO log.
398 _resolved_inner_target: Optional[str] = None
399 # D-09 (05 §6.4.7): EP pass-through for HF-native MoE. A non-empty _ep_stack
400 # means per-expert parameters must be pre-stacked into [E, ...] in Phase A.
401 # Since the explicit-injection rework the compute side is NOT auto-injected:
402 # declare local_compute_fn explicitly (e.g. Target → an EP archetype
403 # factory such as recipes.qwen2moe_ep_compute_fn).
404 _ep_stack: Dict[str, List[str]] = field(default_factory=dict)
405 # D-10 (05 §6.4.8): virtual/TP-extended EP parameter-sharding marker. When
406 # >0 this is the extended EP group size (= ep_size; the a2a communication
407 # domain includes TP/CP ranks) and a derived expert mesh (edp, ep) owns
408 # every parameter whose placement explicitly contains ``EP: Shard(...)``.
409 # Routed MoE experts and sparse lookup tables share this generic source-
410 # layout path. It drives ONLY parameter sharding/expert-mesh derivation —
411 # matching all-to-all compute must be injected explicitly.
412 _ep_size: int = 0
415def resolve_placements(
416 named: NamedPlacement,
417 mesh_dim_names: Tuple[str, ...],
418) -> List[Placement]:
419 """Arrange placements in mesh_dim_names order, fill missing axes with Replicate()."""
420 return [named.get(axis, Replicate()) for axis in mesh_dim_names]
423_PLACEMENT_DSL_DOC = (
424 '"replicate" / "partial" / "shard(N)" (N is the tensor dim index, e.g. "shard(0)")')
427def parse_placement(text: Any, *, path: str = "placement") -> Placement:
428 """Parse the YAML placement DSL string into a Placement object.
430 Grammar (closed set): ``replicate`` / ``partial`` / ``shard(N)`` — case
431 insensitive, whitespace tolerant. Anything else fails fast listing the
432 grammar (typo self-discovery, same philosophy as
433 ``_check_target_config_keys``).
434 """
435 if not isinstance(text, str):
436 raise ValueError(
437 f"{path}: placement must be a string ({_PLACEMENT_DSL_DOC}), "
438 f"got {type(text).__name__} {text!r}")
439 normalized = re.sub(r"\s+", "", text).lower()
440 if normalized == "replicate":
441 return Replicate()
442 if normalized == "partial":
443 return Partial()
444 match = re.fullmatch(r"shard\((\d+)\)", normalized)
445 if match:
446 return Shard(int(match.group(1)))
447 raise ValueError(
448 f"{path}: cannot parse placement {text!r} — valid grammar: {_PLACEMENT_DSL_DOC}")
451def parse_named_placement(raw: Any, *, path: str = "named_placement") -> NamedPlacement:
452 """Parse the YAML form ``{axis: placement_str}`` into a NamedPlacement.
454 Axis names stay plain strings (str-enum interop; custom mesh dims such as
455 "encoder_dp" are legal) — validity against the actual mesh is checked at
456 plan-merge time (ShardingPlanner._validate_override_axes), because the
457 mesh does not exist at config time.
458 """
459 if not isinstance(raw, dict) or not raw:
460 raise ValueError(
461 f"{path}: expected a non-empty mapping {{axis name: placement string}}, got {raw!r}")
462 named: NamedPlacement = {}
463 for axis, placement_text in raw.items():
464 if not isinstance(axis, str) or not axis:
465 raise ValueError(
466 f"{path}: axis name must be a non-empty string, got {axis!r}")
467 named[axis] = parse_placement(
468 placement_text, path=f"{path}.{axis}")
469 return named
472def _normalize_out_fields(spec: ModuleShardingSpec) -> ModuleShardingSpec:
473 """Normalize the scalar shorthand {TP: ...} into {'output': {TP: ...}} (05 §3.5).
475 Detection heuristic: if val is a non-None dict and any of its values is not a
476 dict, it is judged to be a scalar NamedPlacement shorthand. Idempotent — a
477 second call on an already-normalized dict contract changes nothing.
478 """
479 for attr in ("out_src", "out_dst"):
480 val = getattr(spec, attr, None)
481 if val and not all(isinstance(v, dict) for v in val.values()):
482 setattr(spec, attr, {"output": dict(val)})
483 return spec
486# ────────────────────────────────────────────────────────────────────────────
487# Injection DSL: template decorators for injected functions (the entry-point
488# discipline of the explicit injection mechanism).
489#
490# Design principle: **an injected function is called exactly once at apply
491# time, must declare the mesh family (filled by the framework by name; axes
492# that are not active are filled with None), and may use them or not**. Two
493# decorators cover all injection channels, each with a single canonical form:
494#
495# - ``@local_compute``: a regional-compute **factory**,
496# ``fn(mesh, tp_mesh, cp_mesh, ep_mesh, [module], <config keys...>)
497# -> compute_fn`` -- built once at apply time;
498# - ``@inner_wrapper``: an inner-forward wrapper, ``fn(target_module, mesh,
499# tp_mesh, cp_mesh, ep_mesh)`` -- returns the replacement forward (in-repo
500# discipline; the forward rewriter installs it) or replaces
501# ``target.forward`` in place and returns None (the external contract).
502#
503# Hard rules enforced by the decorators (fail-fast at import time): required
504# context must be declared in full; context parameters must not have default
505# values; injected functions forbid *args/**kwargs. The full design rationale
506# is documented in docs/design/ (05 §3.7) — the runtime-layer signature
507# validators live in ``distributed/_builder/forward_rewriter.py``.
508# ────────────────────────────────────────────────────────────────────────────
510# Injection kinds (meta.kind)
511LOCAL_COMPUTE = "local_compute" # regional-compute factory (the only form of local_compute_fn)
512INNER_WRAPPER = "inner_wrapper" # inner-forward wrapper
514# Reserved framework-context names per kind (declaring one means the
515# framework fills it by name; configuring a key with the same name errors)
516_MESH_FAMILY = frozenset({"mesh", "tp_mesh", "cp_mesh", "ep_mesh"})
517FACTORY_CONTEXT = frozenset({"module"}) | _MESH_FAMILY
518WRAPPER_CONTEXT = frozenset({"target_module"}) | _MESH_FAMILY
520# Required context per kind (the decorators enforce full declaration --
521# the mesh family is always passed by the framework and the user merely
522# uses it; the framework fills None for axes that are not active)
523FACTORY_REQUIRED = _MESH_FAMILY
524WRAPPER_REQUIRED = frozenset({"target_module"}) | _MESH_FAMILY
526_ALLOWED_CONTEXT = {
527 LOCAL_COMPUTE: FACTORY_CONTEXT, # optional module anchor + mesh family
528 INNER_WRAPPER: WRAPPER_CONTEXT,
529}
530_REQUIRED_CONTEXT = {
531 LOCAL_COMPUTE: FACTORY_REQUIRED, # mesh family is required
532 INNER_WRAPPER: WRAPPER_REQUIRED,
533}
535_DECORATOR_NAMES = {
536 LOCAL_COMPUTE: "@local_compute",
537 INNER_WRAPPER: "@inner_wrapper",
538}
541@dataclass(frozen=True)
542class InjectionMeta:
543 """Metadata written by the decorators onto the injected function object
544 (``fn._injection_meta``).
546 Attributes:
547 kind: The injection kind (LOCAL_COMPUTE / INNER_WRAPPER).
548 context: The declared framework-context keys.
549 """
550 kind: str # LOCAL_COMPUTE / INNER_WRAPPER
551 context: frozenset # declared framework-context keys
554def _make_decorator(kind: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
555 """Build the template decorator for one injection kind.
557 Args:
558 kind: The injection kind (LOCAL_COMPUTE or INNER_WRAPPER).
560 Returns:
561 A decorator that validates the injected function's signature at
562 import time and stamps it with :class:`InjectionMeta`.
563 """
564 allowed = _ALLOWED_CONTEXT[kind]
565 required = _REQUIRED_CONTEXT[kind]
567 def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
568 """Validate ``fn`` against the injection discipline and stamp its meta.
570 Args:
571 fn: The injected function to decorate.
573 Returns:
574 The same function, with ``_injection_meta`` attached.
576 Raises:
577 TypeError: If the signature cannot be introspected, uses
578 *args/**kwargs, gives a context parameter a default value,
579 or omits required context parameters.
580 """
581 fname = getattr(fn, "__name__", fn)
582 try:
583 sig = inspect.signature(fn)
584 except (TypeError, ValueError) as exc:
585 raise TypeError(
586 f"{_DECORATOR_NAMES[kind]}: cannot introspect the signature of "
587 f"{fn!r} (an injected function must be an introspectable "
588 "plain callable)") from exc
589 context = []
590 for name, p in sig.parameters.items():
591 if p.kind in (inspect.Parameter.VAR_POSITIONAL,
592 inspect.Parameter.VAR_KEYWORD):
593 raise TypeError(
594 f"{_DECORATOR_NAMES[kind]} injected function {fname}"
595 f" must not use *args/**kwargs (parameter {name!r}) -- "
596 "the signature of an injected function must be an explicit "
597 "parameter list: context is declared by name and filled by "
598 "the framework, config keys are bound by name, and "
599 "**kwargs would silently swallow typos (conflicting with "
600 "the _check_target_config_keys policy)")
601 if name in allowed:
602 if p.default is not inspect.Parameter.empty:
603 raise TypeError(
604 f"{_DECORATOR_NAMES[kind]} injected function {fname}: "
605 f"context parameter {name!r} must not have a default "
606 "value -- context names are reserved by the framework "
607 "and are always filled by name, so a default would "
608 "never take effect")
609 context.append(name)
610 missing = sorted(required - set(context))
611 if missing:
612 raise TypeError(
613 f"{_DECORATOR_NAMES[kind]} injected function {fname} is "
614 f"missing required context parameters {missing} -- the "
615 "injection discipline requires explicitly receiving "
616 f"{sorted(required)} (all filled by the framework by name at "
617 "apply time; the user merely uses them)")
618 fn._injection_meta = InjectionMeta(kind=kind, context=frozenset(context))
619 return fn
621 return decorator
624local_compute = _make_decorator(LOCAL_COMPUTE)
625inner_wrapper = _make_decorator(INNER_WRAPPER)
628def require_injection_meta(fn: Callable[..., Any], kind: str, *, source: str) -> InjectionMeta:
629 """Fetch an injected function's metadata; fail-fast (with an instructive
630 error) if it is undecorated or of the wrong kind.
632 Args:
633 fn: The function expected to carry injection metadata.
634 kind: The expected injection kind.
635 source: Call-site label prepended to error messages.
637 Returns:
638 The :class:`InjectionMeta` attached to ``fn``.
640 Raises:
641 TypeError: If ``fn`` lacks the template decorator or its decorator
642 kind does not match ``kind``.
643 """
644 meta = getattr(fn, "_injection_meta", None)
645 name = getattr(fn, "__name__", fn)
646 if meta is None:
647 raise TypeError(
648 f"{source}: injected function {name} is missing the "
649 f"{_DECORATOR_NAMES[kind]} decorator -- the explicit-injection "
650 "discipline requires every injected function to carry a template "
651 "decorator (the decorator declares the framework context it needs "
652 "and guarantees the signature is validatable): use @local_compute "
653 "for the runtime fn of local_compute_fn (a regional-compute "
654 "factory) and @inner_wrapper for inner_wrapper (exported by "
655 "hyper_parallel.distributed)")
656 if meta.kind != kind:
657 raise TypeError(
658 f"{source}: injected function {name} has the wrong decorator "
659 f"kind -- got {_DECORATOR_NAMES[meta.kind]}, expected "
660 f"{_DECORATOR_NAMES[kind]}")
661 return meta