Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / inline / pipeline.py: 43%
143 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 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"""Two-stage inline modeling source pipeline."""
17from __future__ import annotations
19import ast
20from dataclasses import replace
21import importlib
22import textwrap
23from typing import Any
25from hyper_parallel.codegen.inline.ir import (
26 ImportPatch,
27 InlinePatchSet,
28 ModuleSnippetPatch,
29)
30from hyper_parallel.codegen.inline.meta_plan import normalize_inline_meta
31from hyper_parallel.codegen.inline.definition_edits import drop_shadowed_definitions
32from hyper_parallel.codegen.inline.patch_engine import apply_patch_set
33from hyper_parallel.codegen.inline.framework_spec import (
34 replacement_spec_for as replacement_spec,
35 strategy_spec_for as strategy_spec,
36)
37from hyper_parallel.codegen.inline.strategy_pass import build_strategy_patches
38from hyper_parallel.codegen.inline.observations import collect_source_observations
39from hyper_parallel.codegen.inline.source_plan import build_source_patch_plan
42def render_inline_modeling(
43 source_text: str,
44 meta: Any,
45 model_type: str | None = None,
46 *,
47 patch_plan: Any = None,
48) -> str:
49 """Render the inline-patched modeling file for a frozen plan.
51 This is the only artifact route. A YAML target with no adapter declaration
52 is a hard error rather than a silent fallback, so an inline coverage gap
53 surfaces at generation time instead of degrading the artifact. A plan with
54 no inline rules needs no patch and the source is returned unchanged.
56 Args:
57 source_text: Modeling source the inline patches are applied to.
58 meta: Frozen plan metadata recorded by the generator.
59 model_type: Model family key used to resolve framework declarations.
61 Returns:
62 The rendered modeling source.
63 """
65 rendered, _ = render_inline_modeling_with_plan(
66 source_text,
67 meta,
68 model_type,
69 patch_plan=patch_plan,
70 )
71 return rendered
74def render_inline_modeling_with_plan(
75 source_text: str,
76 meta: Any,
77 model_type: str | None = None,
78 *,
79 patch_plan: Any = None,
80) -> tuple[str, Any]:
81 """Render source and return the independently normalized emission plan."""
82 rules = collect_source_observations(meta)
83 if not rules:
84 return source_text, meta
85 patches = patch_plan
86 recorded_local_compute = bool(
87 getattr(patches, "local_computes", ()) if patches is not None else ()
88 )
89 strategy_rules = (
90 tuple(replace(rule, local_compute_target=None) for rule in rules)
91 if recorded_local_compute
92 else rules
93 )
94 _require_inline_coverage(strategy_rules, model_type)
95 patch_set = build_source_patch_plan(
96 rules,
97 strategy_rules,
98 model_type,
99 _observed_boundary_classes(meta, patch_plan),
100 getattr(meta, "external_state_classes", None),
101 )
102 _add_recorded_forward_host_imports(source_text, patch_set, patch_plan)
103 _materialize_imported_forward_hosts(source_text, patch_set)
104 rendered = apply_patch_set(source_text, patch_set)
105 # The copied modeling source brings its own decorated helpers that share
106 # names with the library functions the artifact imports, and a local
107 # definition wins over the import: the native path calls the library
108 # implementation (whose kernel binding is resolved in its own module), so a
109 # redefined copy makes the generated model run something else. Only names
110 # the base source defines with a decorator are candidates, so the classes and
111 # snippets the inline passes render (which share names with library imports
112 # on purpose) keep their definitions.
113 rendered = drop_shadowed_definitions(
114 rendered, _library_import_names(rendered) & _decorated_names(source_text)
115 )
116 normalized_meta = normalize_inline_meta(meta, strategy_rules, model_type)
117 return rendered, normalized_meta
120def _add_recorded_forward_host_imports(
121 source_text: str,
122 patch_set: InlinePatchSet,
123 patch_plan: Any,
124) -> None:
125 """Expose observed replacement hosts needed by a forward rewrite.
127 A replacement application can be observed after the root model has built,
128 while its constructor belongs to an imported decoder rather than the copied
129 top-level source. In that case the frozen replacement record is the source
130 of truth for the concrete host type. Add its real import so the generic
131 imported-host materializer can emit the CP dispatch class.
132 """
133 source_classes = {
134 node.name
135 for node in ast.parse(source_text).body
136 if isinstance(node, ast.ClassDef)
137 }
138 imported_names = {
139 name
140 for item in patch_set.imports
141 if not item.raw
142 for name in item.names
143 }
144 replacements = (
145 getattr(patch_plan, "module_replacements", ())
146 if patch_plan is not None
147 else ()
148 )
149 replacement_types = {
150 record.replacement_type.__name__: record.replacement_type
151 for record in replacements
152 }
153 for patch in patch_set.forward_extracts:
154 if patch.class_name in source_classes or patch.class_name in imported_names:
155 continue
156 replacement_type = replacement_types.get(patch.class_name)
157 if replacement_type is None:
158 continue
159 patch_set.imports.append(
160 ImportPatch(replacement_type.__module__, (replacement_type.__name__,))
161 )
162 imported_names.add(patch.class_name)
165def _observed_boundary_classes(meta: Any, patch_plan: Any) -> dict[str, str]:
166 """Return boundary owners after native inner-wrapper target resolution."""
167 classes = dict(getattr(meta, "boundary_classes", None) or {})
168 wrappers = (
169 getattr(patch_plan, "inner_wrappers", ())
170 if patch_plan is not None
171 else ()
172 )
173 for record in wrappers:
174 target_type = getattr(record, "target_type", None)
175 if target_type is None:
176 continue
177 target_fqn = (
178 record.module_fqn
179 if record.target_fqn == "self"
180 else f"{record.module_fqn}.{record.target_fqn}"
181 )
182 classes[target_fqn] = target_type.__name__
183 return classes
186def _materialize_imported_forward_hosts(
187 source_text: str, patch_set: InlinePatchSet
188) -> None:
189 """Render same-name subclasses when a strategy host is an imported component."""
190 source_classes = {
191 node.name
192 for node in ast.parse(source_text).body
193 if isinstance(node, ast.ClassDef)
194 }
195 remaining = []
196 rendered: set[str] = set()
197 for patch in patch_set.forward_extracts:
198 if patch.class_name in source_classes:
199 remaining.append(patch)
200 continue
201 import_index = next(
202 (
203 index
204 for index, item in enumerate(patch_set.imports)
205 if not item.raw and patch.class_name in item.names
206 ),
207 None,
208 )
209 if import_index is None:
210 remaining.append(patch)
211 continue
212 if patch.class_name in rendered:
213 continue
214 rendered.add(patch.class_name)
215 imported = patch_set.imports[import_index]
216 other_names = tuple(name for name in imported.names if name != patch.class_name)
217 alias = f"_{patch.class_name}Base"
218 replacement_imports = []
219 if other_names:
220 replacement_imports.append(ImportPatch(imported.module, other_names))
221 replacement_imports.append(ImportPatch(
222 module="",
223 names=(),
224 raw=f"from {imported.module} import {patch.class_name} as {alias}",
225 ))
226 patch_set.imports[import_index:import_index + 1] = replacement_imports
227 body = patch.body.replace(
228 "return self._forward_impl(__hp_forward_args__)",
229 "return super().forward(__hp_forward_args__)",
230 ).replace("__hp_forward_args__", "*args, **kwargs")
231 snippet = (
232 f"class {patch.class_name}({alias}):\n"
233 " \"\"\"Generated source-visible parallel form.\"\"\"\n\n"
234 " cp_enable = False\n\n"
235 " def forward(self, *args, **kwargs):\n"
236 f"{textwrap.indent(body.strip(), ' ')}\n"
237 )
238 patch_set.module_snippets.append(ModuleSnippetPatch(snippet))
239 patch_set.forward_extracts = remaining
242def render_inline_strategies_for_source_unit(
243 source_text: str,
244 meta: Any,
245 model_type: str | None = None,
246 *,
247 patch_plan: Any = None,
248) -> str:
249 """Render only strategy forwards owned by one sibling source unit."""
250 rules = collect_source_observations(meta)
251 if not rules:
252 return source_text
253 recorded_local_compute = bool(
254 getattr(patch_plan, "local_computes", ())
255 if patch_plan is not None
256 else ()
257 )
258 strategy_rules = (
259 tuple(replace(rule, local_compute_target=None) for rule in rules)
260 if recorded_local_compute
261 else rules
262 )
263 _require_inline_coverage(strategy_rules, model_type)
264 patches = build_strategy_patches(
265 strategy_rules,
266 model_type,
267 _observed_boundary_classes(meta, patch_plan),
268 getattr(meta, "external_state_classes", None),
269 )
270 source_classes = {
271 node.name
272 for node in ast.parse(source_text).body
273 if isinstance(node, ast.ClassDef)
274 }
275 patches.forward_extracts = [
276 patch for patch in patches.forward_extracts
277 if patch.class_name in source_classes
278 ]
279 if not patches.forward_extracts:
280 return source_text
281 patches.module_snippets.clear()
282 patches.constructor_replaces.clear()
283 patches.class_removals.clear()
284 return apply_patch_set(source_text, patches)
287def _library_import_names(text: str) -> set[str]:
288 """Names the artifact imports from the library, i.e. the winning bindings.
290 Args:
291 text: Rendered modeling source.
293 Returns:
294 The set of names bound by its ``hyper_parallel`` imports.
295 """
297 names: set[str] = set()
298 try:
299 tree = ast.parse(text)
300 except SyntaxError:
301 return names
302 for node in ast.walk(tree):
303 if isinstance(node, ast.ImportFrom) and (node.module or "").startswith("hyper_parallel"):
304 names.update(alias.asname or alias.name for alias in node.names)
305 return names
308def _decorated_names(text: str) -> set[str]:
309 """Names of the module-level definitions the base source decorates.
311 A decorator is what makes a copied helper behave differently from the library
312 function of the same name (a per-module kernel binding, for instance), so
313 only those definitions may be dropped.
315 Args:
316 text: Base modeling source before the inline patches are applied.
318 Returns:
319 The set of decorated top-level definition names.
320 """
322 try:
323 tree = ast.parse(text)
324 except SyntaxError:
325 return set()
326 return {
327 node.name
328 for node in tree.body
329 if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef))
330 and node.decorator_list
331 }
334def _is_local_compute_factory(target: str) -> bool:
335 """Whether a local_compute target is a ``@local_compute`` factory.
337 Structural injection discipline: the decorator stamps the factory with
338 ``_injection_meta`` (kind ``LOCAL_COMPUTE``), so a decorated factory is
339 recognizable without a table. Such a target carries no inline strategy
340 body -- its region is bound at install time (the emitted region template
341 for an in-source class, the runtime local-region skeleton for an imported
342 class), which resolves the factory exactly as the native applier does.
343 """
344 from hyper_parallel.distributed.recipe_spec import LOCAL_COMPUTE # pylint: disable=C0415
346 module_name, _, attr = target.rpartition(".")
347 if not module_name or not attr:
348 return False
349 try:
350 module = importlib.import_module(module_name)
351 except ImportError:
352 return False
353 factory = getattr(module, attr, None)
354 if not callable(factory):
355 return False
356 meta = getattr(factory, "_injection_meta", None)
357 return getattr(meta, "kind", None) == LOCAL_COMPUTE
360def _require_inline_coverage(rules: tuple[Any, ...], model_type: str | None = None) -> None:
361 """Fail when a YAML rule has no structural resolution behind it.
363 Covers both rule kinds: a replacement target needs a ``ReplacementSpec``,
364 and every strategy target needs a ``StrategySpec``. One exception: a
365 ``local_compute_fn`` target that is itself a ``@local_compute`` factory
366 needs no inline strategy -- its region is runtime-bound (see
367 :func:`_is_local_compute_factory`), so the factory's injection discipline
368 is the coverage proof. Reporting all gaps at once keeps one fix cycle
369 from hiding the next gap.
370 """
372 uncovered: list[str] = []
373 for rule in rules:
374 if (
375 rule.replace_target is not None
376 and replacement_spec(rule.replace_target, model_type, module_type=rule.module_type) is None
377 ):
378 uncovered.append(rule.replace_target)
379 for target in (rule.local_compute_target, rule.inner_wrapper_target):
380 if target is None:
381 continue
382 if strategy_spec(
383 target, model_type, inner_wrapper=target == rule.inner_wrapper_target
384 ) is not None:
385 continue
386 # An inner_wrapper target has no runtime channel and stays
387 # fail-closed; a local_compute target passes through the
388 # @local_compute decoration check.
389 if target == rule.local_compute_target and _is_local_compute_factory(target):
390 continue
391 uncovered.append(target)
392 if uncovered:
393 raise RuntimeError(
394 f"codegen: no inline declaration for {sorted(set(uncovered))} "
395 f"(model_type={model_type!r}); the target names no known generic "
396 "component or framework strategy, and a local_compute target "
397 "names no @local_compute factory"
398 )