Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / emit / modeling.py: 59%
309 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"""Assemble the generated HF-style modeling file.
17The generated file is the original modeling source, verbatim, with two
18appended layers:
20* a header banner recording where it came from,
21* the adapter-driven inline patch — replaced modules and the TP / CP / EP
22 strategies lowered into the copied source at their semantic sites.
24Boundary forwards are then lowered at the boundary classes. The artifact is
25both an auditable source export and a directly executable model training entry.
26Training still derives its live plan from the supplied YAML through the native
27Trainer; preflight validates syntax and the recorded generation contract.
28"""
30from __future__ import annotations
32import ast
33import io
34import keyword
35import os
36import tokenize
37from typing import Any
39from hyper_parallel.codegen.source.composition import compose_nested_model_sources
40from hyper_parallel.codegen.source.compat import sanitize_source_compat
41from hyper_parallel.codegen.emit.parallel import lower_forward_boundaries
42from hyper_parallel.codegen.emit.recorded import materialize_recorded_forwards
43from hyper_parallel.codegen.inline.pipeline import render_inline_modeling_with_plan
45# ---------------------------------------------------------------------------
46# source copy
47# ---------------------------------------------------------------------------
50def copy_original_modeling(source: Any) -> str:
51 """Read the original modeling source text from a ``ResolvedSource``-like.
53 ``source`` is ``meta.source`` — the ``ResolvedSource.to_dict()`` form with
54 a ``file_path`` key. Accepts either a dict or a ``ResolvedSource`` (the
55 dict form is what the manager records, but the helper tolerates the live
56 dataclass so it stays usable outside the meta path).
57 """
58 if isinstance(source, dict):
59 path = source.get("file_path")
60 else:
61 path = getattr(source, "file_path", None)
62 if not path or not os.path.isfile(path):
63 raise FileNotFoundError(
64 f"codegen emit: original modeling source is not a readable file "
65 f"({path!r}); generation must run where the model's modeling file "
66 "exists"
67 )
68 with open(path, "r", encoding="utf-8", errors="replace") as handle:
69 return handle.read()
72def rewrite_relative_imports(text: str, module_name: str) -> str:
73 """Rewrite ``from .xxx`` / ``from ...xxx`` imports to absolute imports.
75 HF modeling files commonly import siblings with a relative spec, e.g.
76 ``from .configuration_qwen3 import Qwen3Config`` or ``from ...activations
77 import ACT2FN``. The generated bundle is loaded by ``spec_from_file_location``
78 under a synthetic package (``hyper_parallel_generated.<digest>.<stem>``), so a
79 relative spec resolves against that synthetic package instead of the real
80 transformers package and fails with ``attempted relative import beyond
81 top-level package``.
83 This anchors the relative spec on ``module_name`` (the original module's
84 real dotted path, e.g. ``transformers.models.qwen3.modeling_qwen3``): N
85 leading dots climb N-1 package levels above the module's parent package, then
86 append the rest. ``from ...activations`` → ``from transformers.activations``
87 and ``from .configuration_qwen3`` →
88 ``from transformers.models.qwen3.configuration_qwen3``.
90 ``module_name`` is ``"<remote>"`` for trusted remote code, which has no real
91 package to anchor on — the rewrite is a no-op there (import handling for
92 remote code keeps its relative imports). A non-dotted or unanchorable ``module_name``
93 is likewise left untouched.
95 Only the module-spec part of each relative import is rewritten; the symbol
96 list, comments, and whitespace are preserved. It uses ``tokenize`` rather
97 than a regex so relative-looking text inside string literals or comments is
98 never touched.
99 """
100 unless_anchor = _relative_anchor(module_name)
101 if unless_anchor is None:
102 return text
103 text = _rewrite_dynamic_relative_imports(text, unless_anchor)
105 try:
106 tokens = list(tokenize.generate_tokens(io.StringIO(text).readline))
107 except tokenize.TokenError:
108 return text
109 if not tokens:
110 return text
112 edits: list[tuple[int, int, str]] = []
113 i = 0
114 while i < len(tokens):
115 edit, i = _relative_import_edit(tokens, i, unless_anchor)
116 if edit is not None:
117 edits.append(edit)
119 if not edits:
120 return text
121 return _splice(text, edits)
124def _relative_import_edit(
125 tokens: list[tokenize.TokenInfo], index: int, anchor: list[str]
126) -> tuple[tuple[tuple[int, int], tuple[int, int], str] | None, int]:
127 """Parse one token position into an absolute relative-import edit."""
128 token = tokens[index]
129 if token.type != tokenize.NAME or token.string != "from":
130 return None, index + 1
131 cursor = _skip_import_trivia(tokens, index + 1)
132 if cursor >= len(tokens) or not _is_dot_token(tokens[cursor]):
133 return None, index + 1
134 dots_token = tokens[cursor]
135 dots = 0
136 while cursor < len(tokens) and _is_dot_token(tokens[cursor]):
137 dots += len(tokens[cursor].string)
138 cursor += 1
139 path, last_token, cursor = _relative_module_path(tokens, cursor)
140 import_index = _skip_import_trivia(tokens, cursor)
141 if import_index >= len(tokens) or tokens[import_index].string != "import":
142 return None, index + 1
143 absolute = _absolute_spec(anchor, dots, path)
144 if absolute is None:
145 return None, index + 1
146 end = last_token.end if last_token is not None else dots_token.end
147 return (dots_token.start, end, absolute), cursor + 1
150def _skip_import_trivia(tokens: list[tokenize.TokenInfo], index: int) -> int:
151 """Skip comments and physical/logical newlines between import tokens."""
152 trivia = {tokenize.NL, tokenize.NEWLINE, tokenize.COMMENT}
153 while index < len(tokens) and tokens[index].type in trivia:
154 index += 1
155 return index
158def _is_dot_token(token: tokenize.TokenInfo) -> bool:
159 """Return whether a token contains only relative-import dots."""
160 return token.type == tokenize.OP and bool(token.string) and set(token.string) == {"."}
163def _relative_module_path(
164 tokens: list[tokenize.TokenInfo], index: int
165) -> tuple[list[str], tokenize.TokenInfo | None, int]:
166 """Collect the NAME (dot NAME)* portion of a relative import."""
167 path: list[str] = []
168 last_token = None
169 while index < len(tokens):
170 token = tokens[index]
171 if token.type != tokenize.NAME or token.string == "import":
172 break
173 path.append(token.string)
174 last_token = token
175 index += 1
176 if index >= len(tokens) or tokens[index].string != ".":
177 break
178 index += 1
179 return path, last_token, index
182def _rewrite_dynamic_relative_imports(text: str, anchor: list[str]) -> str:
183 """Rewrite imports whose resolved package is not valid Python syntax."""
184 try:
185 tree = ast.parse(text)
186 except SyntaxError:
187 return text
188 edits = []
189 for node in ast.walk(tree):
190 if not isinstance(node, ast.ImportFrom) or node.level == 0:
191 continue
192 absolute = _absolute_spec(anchor, node.level, (node.module or "").split("."))
193 if absolute is None or _is_import_syntax_safe(absolute):
194 continue
195 assignments = []
196 for alias in node.names:
197 if alias.name == "*":
198 raise ValueError(
199 f"codegen cannot rewrite a wildcard import from dynamic module {absolute!r}"
200 )
201 target = alias.asname or alias.name
202 assignments.append(
203 f"{target} = __import__({absolute!r}, fromlist=[{alias.name!r}]).{alias.name}"
204 )
205 indent = " " * node.col_offset
206 replacement = ("\n" + indent).join(assignments)
207 edits.append(
208 (
209 (node.lineno, node.col_offset),
210 (node.end_lineno, node.end_col_offset),
211 replacement,
212 )
213 )
214 return _splice(text, edits) if edits else text
217def _is_import_syntax_safe(module: str) -> bool:
218 """Return whether every dotted component can appear in an import statement."""
219 return all(
220 part.isidentifier() and not keyword.iskeyword(part)
221 for part in module.split(".")
222 )
225def _relative_anchor(module_name: str) -> list[str] | None:
226 """Return the dotted package segments a relative spec anchors on, or None.
228 A real installed-module path (``transformers.models.qwen3.modeling_qwen3``)
229 anchors. The ``"<remote>"`` sentinel (trusted remote code) and any
230 non-dotted / placeholder-laden name cannot anchor — the caller leaves such
231 imports untouched.
232 """
233 if not isinstance(module_name, str) or not module_name:
234 return None
235 if "<" in module_name or ">" in module_name:
236 return None
237 segs = module_name.split(".")
238 if len(segs) < 2 or any(not seg for seg in segs):
239 return None
240 return segs
243def _absolute_spec(anchor: list[str], dots: int, path: list[str]) -> str | None:
244 """Compute the absolute import spec for one relative import, or None.
246 ``anchor`` is the module's package segments (its path minus the final
247 module segment). ``dots`` leading dots climb ``dots - 1`` package levels
248 above the module's parent package, then append ``path``. Returns None when
249 the relative import climbs above the package root (invalid source — leave
250 it untouched rather than corrupt it).
251 """
252 parent = anchor[:-1]
253 climb = dots - 1
254 if climb > len(parent):
255 return None
256 base_parts = parent if climb <= 0 else parent[: len(parent) - climb]
257 parts = base_parts + path
258 return ".".join(part for part in parts if part)
261def _splice(
262 text: str, edits: list[tuple[tuple[int, int], tuple[int, int], str]]
263) -> str:
264 """Splice absolute specs back into ``text`` by token (row, col) positions.
266 Edits are applied last-first so the earlier offsets stay valid. A single
267 contiguous span is replaced in every case: the relative spec starts at the
268 dots token and ends at the last module-path token (or the dots token when
269 there is no path, i.e. ``from . import``).
270 """
271 lines = text.splitlines(keepends=True)
272 if not lines:
273 return text
274 line_offsets = []
275 acc = 0
276 for line in lines:
277 line_offsets.append(acc)
278 acc += len(line)
280 def offset(row: int, col: int) -> int:
281 if 1 <= row <= len(line_offsets):
282 return line_offsets[row - 1] + col
283 return acc
285 positioned = sorted(
286 ((offset(*start), offset(*end), new) for start, end, new in edits),
287 key=lambda item: item[0],
288 reverse=True,
289 )
290 out = text
291 for start, end, new in positioned:
292 if start < end <= len(out):
293 out = out[:start] + new + out[end:]
294 return out
297# ---------------------------------------------------------------------------
298# file assembly
299# ---------------------------------------------------------------------------
302def emit_modeling_file(meta: Any, *, patch_plan: Any = None) -> str:
303 """Assemble the full generated modeling file text.
305 Returns the text (not written to disk) so the caller can compute a diff
306 first. The header banner carries the generation provenance; the source is
307 copied verbatim, rewritten by the adapter-driven inline pipeline, and then
308 lowered at the boundary classes.
309 """
310 text = copy_original_modeling(meta.source)
311 # Stale-source sanitize: a trust-remote-code modeling file ships with the
312 # checkpoint and reflects the author's transformers, not the installed one.
313 # Removing ``is_torch_fx_available`` (a transformers>=5.0 casualty) here
314 # keeps the artifact importable without touching the source file. No-op
315 # when the source does not import the symbol.
316 text = sanitize_source_compat(text)
317 module_name = (
318 meta.source.get("module_name")
319 if isinstance(meta.source, dict)
320 else getattr(meta.source, "module_name", "")
321 )
322 text = rewrite_relative_imports(text, module_name)
323 text = compose_nested_model_sources(text, meta)
324 text, emission_meta = _apply_inline_modeling(text, meta, patch_plan=patch_plan)
325 source_path = (
326 meta.source.get("file_path")
327 if isinstance(meta.source, dict)
328 else getattr(meta.source, "file_path", None)
329 )
330 text = materialize_recorded_forwards(
331 text,
332 emission_meta,
333 patch_plan=patch_plan,
334 source_path=(
335 None if getattr(emission_meta, "nested_components", None) else source_path
336 ),
337 )
338 text = _lower_forward_boundaries(text, emission_meta, module_name)
339 _verify_recorded_patch_materialization(
340 text, emission_meta, patch_plan=patch_plan
341 )
342 return _banner(meta) + text
345def _verify_recorded_patch_materialization(
346 text: str,
347 meta: Any,
348 *,
349 patch_plan: Any = None,
350) -> None:
351 """Require every trainer-installed source-visible patch to be materialized."""
352 patches = patch_plan
353 local_records = (
354 getattr(patches, "local_computes", ()) if patches is not None else ()
355 )
356 replacement_records = (
357 getattr(patches, "module_replacements", ()) if patches is not None else ()
358 )
359 wrapper_records = (
360 getattr(patches, "inner_wrappers", ()) if patches is not None else ()
361 )
362 if not local_records and not replacement_records and not wrapper_records:
363 return
364 tree = ast.parse(text)
365 classes = _class_index(tree)
366 materialized = getattr(meta, "materialized_patches", None)
367 if isinstance(materialized, list):
368 materialized.clear()
369 missing = _verify_local_records(tree, classes, local_records, materialized)
370 missing += _verify_replacement_records(
371 tree, classes, replacement_records, materialized
372 )
373 missing += _verify_wrapper_records(
374 classes, wrapper_records, replacement_records, materialized
375 )
376 if missing:
377 raise RuntimeError(
378 "codegen: trainer-installed source-visible patches were not "
379 "materialized in the generated model source: " + ", ".join(missing)
380 )
383def _class_index(tree: ast.AST) -> dict[str, ast.ClassDef]:
384 """Index top-level classes by name."""
385 return {
386 node.name: node
387 for node in getattr(tree, "body", ())
388 if isinstance(node, ast.ClassDef)
389 }
392def _class_methods(host: ast.ClassDef) -> set[str]:
393 """Return method names declared directly by a class."""
394 return {
395 node.name
396 for node in host.body
397 if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
398 }
401def _class_state_names(host: ast.ClassDef) -> set[str]:
402 """Return simple class attributes declared by assignment."""
403 names = set()
404 for child in host.body:
405 targets = (
406 child.targets
407 if isinstance(child, ast.Assign)
408 else (child.target,)
409 if isinstance(child, ast.AnnAssign)
410 else ()
411 )
412 names.update(target.id for target in targets if isinstance(target, ast.Name))
413 return names
416def _record_materialized(
417 materialized: Any, fqn: str, host: str, kind: str, source_form: str
418) -> None:
419 """Append one materialization record when metadata is mutable."""
420 if isinstance(materialized, list):
421 materialized.append(
422 {"fqn": fqn, "host_class": host, "kind": kind, "source_form": source_form}
423 )
426def _verify_local_records(
427 tree: ast.AST, classes: dict[str, ast.ClassDef], records: Any, materialized: Any
428) -> list[str]:
429 """Verify local-compute records in one source unit."""
430 missing = []
431 for record in records:
432 class_name = record.module_type.__name__
433 if _bound_function_patch_materialized(tree, record):
434 _record_materialized(
435 materialized,
436 record.module_fqn,
437 class_name,
438 "local_compute",
439 "bound_method_toggle_forward",
440 )
441 continue
442 host = classes.get(class_name)
443 if host is None:
444 continue
445 valid = {"_forward_impl", "forward"} <= _class_methods(
446 host
447 ) and "ep_enable" in _class_state_names(host)
448 if valid:
449 _record_materialized(
450 materialized,
451 record.module_fqn,
452 class_name,
453 "local_compute",
454 "same_class_toggle_forward",
455 )
456 else:
457 missing.append(f"{record.module_fqn} ({class_name})")
458 return missing
461def _verify_replacement_records(
462 tree: ast.AST, classes: dict[str, ast.ClassDef], records: Any, materialized: Any
463) -> list[str]:
464 """Verify replacement constructors in one source unit."""
465 bound_names = {
466 node.id for node in ast.walk(tree) if isinstance(node, ast.Name)
467 } | set(classes)
468 missing = []
469 for record in records:
470 name = record.replacement_type.__name__
471 visible = name in bound_names and any(
472 isinstance(node, ast.Call) and _call_name(node.func) == name
473 for node in ast.walk(tree)
474 )
475 if visible:
476 _record_materialized(
477 materialized,
478 record.module_fqn,
479 name,
480 "module_replacement",
481 "constructor",
482 )
483 else:
484 missing.append(f"{record.module_fqn} (replacement {name})")
485 return missing
488def _wrapper_identity(record: Any) -> tuple[Any, str, str]:
489 """Return target type, host name, and target FQN for a wrapper record."""
490 target_type = getattr(record, "target_type", None)
491 host_name = (
492 target_type.__name__
493 if target_type is not None
494 else type(record.request.target).__name__
495 )
496 target_fqn = (
497 record.module_fqn
498 if record.target_fqn == "self"
499 else f"{record.module_fqn}.{record.target_fqn}"
500 )
501 return target_type, host_name, target_fqn
504def _wrapper_source_form(
505 classes: dict[str, ast.ClassDef], host: ast.ClassDef, record: Any, replacements: Any
506) -> str | None:
507 """Return the concrete source form that materializes a wrapper."""
508 if (
509 "forward" in _class_methods(host)
510 and "cp_enable" in _class_state_names(host)
511 and _class_dispatch_reads_state(host, "cp_enable")
512 ):
513 return "same_class_toggle_forward"
514 subclass = next(
515 (
516 node
517 for node in classes.values()
518 if any(
519 isinstance(base, ast.Name) and base.id == host.name
520 for base in node.bases
521 )
522 and "forward" in _class_methods(node)
523 ),
524 None,
525 )
526 if subclass is not None:
527 return "subclass_forward"
528 target_type = getattr(record, "target_type", None)
529 replacement_backed = any(
530 item.module_fqn == record.module_fqn and item.replacement_type is target_type
531 for item in replacements
532 )
533 if replacement_backed and {"_forward_impl", "forward"} <= _class_methods(host):
534 return "replacement_forward_runtime_binding"
535 return None
538def _verify_wrapper_records(
539 classes: dict[str, ast.ClassDef], records: Any, replacements: Any, materialized: Any
540) -> list[str]:
541 """Verify inner-wrapper records in one source unit."""
542 missing = []
543 for record in records:
544 _, host_name, target_fqn = _wrapper_identity(record)
545 host = classes.get(host_name)
546 if host is None:
547 continue
548 source_form = _wrapper_source_form(classes, host, record, replacements)
549 if source_form:
550 _record_materialized(
551 materialized, target_fqn, host_name, "inner_wrapper", source_form
552 )
553 else:
554 missing.append(f"{target_fqn} (inner wrapper for {host_name})")
555 return missing
558def _call_name(node: ast.AST) -> str | None:
559 """Return the terminal name of a constructor call expression."""
560 if isinstance(node, ast.Name):
561 return node.id
562 if isinstance(node, ast.Attribute):
563 return node.attr
564 return None
567def verify_recorded_patch_bundle(
568 files: dict[str, str],
569 *,
570 patch_plan: Any = None,
571) -> None:
572 """Verify recorded local-compute patches across all generated source units."""
573 records = (
574 getattr(patch_plan, "local_computes", ()) if patch_plan is not None else ()
575 )
576 wrappers = (
577 getattr(patch_plan, "inner_wrappers", ()) if patch_plan is not None else ()
578 )
579 trees = {
580 name: ast.parse(text) for name, text in files.items() if name.endswith(".py")
581 }
582 missing = [
583 f"{record.module_fqn} ({record.module_type.__name__})"
584 for record in records
585 if not any(_tree_has_local_record(tree, record) for tree in trees.values())
586 ]
587 if missing:
588 raise RuntimeError(
589 "codegen: recorded local-compute patches were not materialized "
590 "in their owning source units: " + ", ".join(missing)
591 )
592 replacements = getattr(patch_plan, "module_replacements", ())
593 missing_wrappers = [
594 f"{record.module_fqn} ({_wrapper_identity(record)[1]})"
595 for record in wrappers
596 if not any(
597 _tree_has_wrapper_record(tree, record, replacements)
598 for tree in trees.values()
599 )
600 ]
601 if missing_wrappers:
602 raise RuntimeError(
603 "codegen: recorded inner-wrapper patches were not materialized "
604 "in their owning source units: " + ", ".join(missing_wrappers)
605 )
606 if any("HPCodegen" in text for text in files.values()):
607 raise RuntimeError(
608 "codegen: generated bundle contains forbidden HPCodegen class names"
609 )
612def _tree_has_local_record(tree: ast.AST, record: Any) -> bool:
613 """Return whether one source tree materializes a local-compute record."""
614 if _bound_function_patch_materialized(tree, record):
615 return True
616 host = _class_index(tree).get(record.module_type.__name__)
617 return (
618 host is not None
619 and {"_forward_impl", "forward"} <= _class_methods(host)
620 and "ep_enable" in _class_state_names(host)
621 )
624def _tree_has_wrapper_record(tree: ast.AST, record: Any, replacements: Any) -> bool:
625 """Return whether one source tree materializes an inner-wrapper record."""
626 _, class_name, _ = _wrapper_identity(record)
627 host = _class_index(tree).get(class_name)
628 if host is None or "forward" not in _class_methods(host):
629 return False
630 if "cp_enable" in _class_state_names(host) and _class_dispatch_reads_state(
631 host, "cp_enable"
632 ):
633 return True
634 target_type = getattr(record, "target_type", None)
635 replacement_backed = any(
636 item.module_fqn == record.module_fqn and item.replacement_type is target_type
637 for item in replacements
638 )
639 return replacement_backed and {"_forward_impl", "forward"} <= _class_methods(host)
642def _class_dispatch_reads_state(host: ast.ClassDef, state_name: str) -> bool:
643 """Return whether a class dispatches on one state flag in either layer."""
644 methods = (
645 child
646 for child in host.body
647 if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef))
648 and child.name in {"forward", "_forward_impl"}
649 )
650 return any(
651 isinstance(node, ast.Attribute)
652 and isinstance(node.value, ast.Name)
653 and node.value.id == "self"
654 and node.attr == state_name
655 for method in methods
656 for node in ast.walk(method)
657 )
660def _bound_function_patch_materialized(tree: ast.AST, record: Any) -> bool:
661 """Return whether a ``MethodType`` forward has an explicit EP toggle."""
662 function_name = getattr(record.original_forward, "__name__", "")
663 if not function_name:
664 return False
665 impl_name = f"{function_name}_hp_impl"
666 functions = {
667 node.name: node
668 for node in getattr(tree, "body", ())
669 if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
670 }
671 wrapper = functions.get(function_name)
672 if wrapper is None or impl_name not in functions:
673 return False
674 return _function_reads_ep_toggle(wrapper) and _tree_binds_method_impl(
675 tree, impl_name
676 )
679def _function_reads_ep_toggle(wrapper: ast.AST) -> bool:
680 """Return whether a function reads the instance EP enable flag."""
681 return any(
682 isinstance(node, ast.Attribute)
683 and isinstance(node.value, ast.Name)
684 and node.value.id == "self"
685 and node.attr == "ep_enable"
686 for node in ast.walk(wrapper)
687 )
690def _tree_binds_method_impl(tree: ast.AST, impl_name: str) -> bool:
691 """Return whether a tree binds an implementation through MethodType."""
692 return any(
693 isinstance(node, ast.Assign)
694 and any(
695 isinstance(target, ast.Attribute) and target.attr == "_forward_impl"
696 for target in node.targets
697 )
698 and isinstance(node.value, ast.Call)
699 and _call_name(node.value.func) == "MethodType"
700 and node.value.args
701 and isinstance(node.value.args[0], ast.Name)
702 and node.value.args[0].id == impl_name
703 for node in ast.walk(tree)
704 )
707def _apply_inline_modeling(
708 text: str,
709 meta: Any,
710 *,
711 patch_plan: Any = None,
712) -> tuple[str, Any]:
713 """Render the artifact through the adapter-driven inline source pipeline."""
715 source = getattr(meta, "source", None)
716 architecture = (
717 source.get("architecture")
718 if isinstance(source, dict)
719 else getattr(source, "architecture", None)
720 )
721 return render_inline_modeling_with_plan(
722 text,
723 meta,
724 model_type=architecture,
725 patch_plan=patch_plan,
726 )
729def _lower_forward_boundaries(text: str, meta: Any, module_name: str) -> str:
730 """Sink boundary wrappers into the model's ``forward`` bodies.
732 Each boundary class's ``forward`` is rewritten to redistribute through
733 the instance-bound compiled plan and run the region's compute
734 explicitly. A boundary whose class cannot be resolved or
735 whose per-class contract differs across FQNs fails fast — a silently
736 dropped redistribution would change the model's numeric behavior.
738 Fresh on the unmodified source text, so offsets stay valid for the whole
739 call. The ``local_compute_fn`` / ``inner_wrapper`` members the rewritten
740 forward references are bound by the runtime parallelize step, so the
741 generated file stays runtime-importable without a mesh.
743 ``boundary_classes`` is the observed FQN -> class-name map. When it is
744 empty there are no source-level boundaries to lower, so the source remains
745 untouched.
746 """
747 boundary_classes = getattr(meta, "boundary_classes", None) or {}
748 if not boundary_classes:
749 return text
750 return lower_forward_boundaries(
751 text,
752 meta,
753 boundary_classes=boundary_classes,
754 module_name=module_name,
755 )
758def _banner(meta: Any) -> str:
759 """Render the provenance banner comment."""
760 source = meta.source if isinstance(meta.source, dict) else _source_dict(meta.source)
761 pd = meta.parallel_dims or {}
762 lines = [
763 "# Generated by codegen. DO NOT EDIT MANUALLY.",
764 "#",
765 f"# source : {source.get('module_name', '')}",
766 f"# architecture : {source.get('architecture', '')}",
767 f"# parallel_dims: tp={pd.get('tp_size', 1)} cp={pd.get('cp_size', 1)} "
768 f"ep={pd.get('ep_size', 1)} pp={pd.get('pp_size', 1)} "
769 f"sp={pd.get('sequence_parallel', False)} "
770 f"loss_parallel={pd.get('loss_parallel', False)}",
771 f"# injections : {len(getattr(meta, 'injections', []) or [])} rule(s)",
772 "# runtime : mesh / FSDP2 / PP / AC / compile still handled by runtime",
773 "#",
774 "# [HYPER PLAN] training derives a live plan from its YAML and this model;",
775 "# Regenerate this file from the training YAML after strategy changes.",
776 "#",
777 "",
778 ]
779 return "\n".join(lines)
782def _source_dict(source: Any) -> dict[str, Any]:
783 if source is None:
784 return {}
785 if isinstance(source, dict):
786 return source
787 return {
788 "module_name": getattr(source, "module_name", ""),
789 "file_path": getattr(source, "file_path", ""),
790 "architecture": getattr(source, "architecture", ""),
791 }
794__all__ = [
795 "copy_original_modeling",
796 "emit_modeling_file",
797 "rewrite_relative_imports",
798]