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