Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / emit / modeling.py: 11%
140 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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
25publishes no plan globals: the runtime reads the frozen plan from
26``codegen_meta.json`` and shards from it (``runtime.parallelize_from_generated``
27falls back to ``_parallelize_inline_from_meta`` when the module defines no
28``hp_parallelize``). The preflight import checks that the assembled module
29remains importable in the training environment.
30"""
31from __future__ import annotations
33import io
34import os
35import tokenize
36from typing import Any
38# ---------------------------------------------------------------------------
39# source copy
40# ---------------------------------------------------------------------------
42def copy_original_modeling(source: Any) -> str:
43 """Read the original modeling source text from a ``ResolvedSource``-like.
45 ``source`` is ``meta.source`` — the ``ResolvedSource.to_dict()`` form with
46 a ``file_path`` key. Accepts either a dict or a ``ResolvedSource`` (the
47 dict form is what the manager records, but the helper tolerates the live
48 dataclass so it stays usable outside the meta path).
49 """
50 if isinstance(source, dict):
51 path = source.get("file_path")
52 else:
53 path = getattr(source, "file_path", None)
54 if not path or not os.path.isfile(path):
55 raise FileNotFoundError(
56 f"codegen emit: original modeling source is not a readable file "
57 f"({path!r}); generation must run where the model's modeling file "
58 "exists"
59 )
60 with open(path, "r", encoding="utf-8", errors="replace") as handle:
61 return handle.read()
64def rewrite_relative_imports(text: str, module_name: str) -> str:
65 """Rewrite ``from .xxx`` / ``from ...xxx`` imports to absolute imports.
67 HF modeling files commonly import siblings with a relative spec, e.g.
68 ``from .configuration_qwen3 import Qwen3Config`` or ``from ...activations
69 import ACT2FN``. The generated bundle is loaded by ``spec_from_file_location``
70 under a synthetic package (``hyper_parallel_generated.<digest>.<stem>``), so a
71 relative spec resolves against that synthetic package instead of the real
72 transformers package and fails with ``attempted relative import beyond
73 top-level package``.
75 This anchors the relative spec on ``module_name`` (the original module's
76 real dotted path, e.g. ``transformers.models.qwen3.modeling_qwen3``): N
77 leading dots climb N-1 package levels above the module's parent package, then
78 append the rest. ``from ...activations`` → ``from transformers.activations``
79 and ``from .configuration_qwen3`` →
80 ``from transformers.models.qwen3.configuration_qwen3``.
82 ``module_name`` is ``"<remote>"`` for trusted remote code, which has no real
83 package to anchor on — the rewrite is a no-op there (import handling for
84 remote code keeps its relative imports). A non-dotted or unanchorable ``module_name``
85 is likewise left untouched.
87 Only the module-spec part of each relative import is rewritten; the symbol
88 list, comments, and whitespace are preserved. It uses ``tokenize`` rather
89 than a regex so relative-looking text inside string literals or comments is
90 never touched.
91 """
92 unless_anchor = _relative_anchor(module_name)
93 if unless_anchor is None:
94 return text
96 try:
97 tokens = list(tokenize.generate_tokens(io.StringIO(text).readline))
98 except tokenize.TokenError:
99 return text
100 if not tokens:
101 return text
103 edits: list[tuple[int, int, str]] = []
104 i = 0
105 n = len(tokens)
106 while i < n:
107 tk = tokens[i]
108 if tk.type != tokenize.NAME or tk.string != "from":
109 i += 1
110 continue
111 # the spec must start with a relative (dots) operator
112 j = i + 1
113 while j < n and tokens[j].type in (tokenize.NL, tokenize.NEWLINE, tokenize.COMMENT):
114 j += 1
115 if j >= n or tokens[j].type != tokenize.OP or not set(tokens[j].string) <= {"."}:
116 i += 1
117 continue
118 dots_tok = tokens[j]
119 dots = len(dots_tok.string)
121 # collect the relative module path: NAME ('.' NAME)*
122 k = j + 1
123 path: list[str] = []
124 last_path_tok = None
125 while k < n:
126 t = tokens[k]
127 if t.type == tokenize.NAME and t.string != "import":
128 path.append(t.string)
129 last_path_tok = t
130 k += 1
131 if k < n and tokens[k].type == tokenize.OP and tokens[k].string == ".":
132 k += 1 # consume the separator; import follows next round
133 else:
134 break
135 else:
136 break
138 # the path must be followed by the ``import`` keyword (allow a comment
139 # line in between; a multiline symbol list comes *after* ``import``).
140 m = k
141 while m < n and tokens[m].type in (tokenize.NL, tokenize.NEWLINE, tokenize.COMMENT):
142 m += 1
143 if m >= n or tokens[m].type != tokenize.NAME or tokens[m].string != "import":
144 i += 1
145 continue
147 abs_spec = _absolute_spec(unless_anchor, dots, path)
148 if abs_spec is None:
149 i += 1
150 continue
152 start = dots_tok.start
153 end = last_path_tok.end if last_path_tok is not None else dots_tok.end
154 edits.append((start, end, abs_spec))
155 i = k + 1
157 if not edits:
158 return text
159 return _splice(text, edits)
162def _relative_anchor(module_name: str) -> list[str] | None:
163 """Return the dotted package segments a relative spec anchors on, or None.
165 A real installed-module path (``transformers.models.qwen3.modeling_qwen3``)
166 anchors. The ``"<remote>"`` sentinel (trusted remote code) and any
167 non-dotted / placeholder-laden name cannot anchor — the caller leaves such
168 imports untouched.
169 """
170 if not isinstance(module_name, str) or not module_name:
171 return None
172 if "<" in module_name or ">" in module_name:
173 return None
174 segs = module_name.split(".")
175 if len(segs) < 2 or any(not seg for seg in segs):
176 return None
177 return segs
180def _absolute_spec(anchor: list[str], dots: int, path: list[str]) -> str | None:
181 """Compute the absolute import spec for one relative import, or None.
183 ``anchor`` is the module's package segments (its path minus the final
184 module segment). ``dots`` leading dots climb ``dots - 1`` package levels
185 above the module's parent package, then append ``path``. Returns None when
186 the relative import climbs above the package root (invalid source — leave
187 it untouched rather than corrupt it).
188 """
189 parent = anchor[:-1]
190 climb = dots - 1
191 if climb > len(parent):
192 return None
193 base_parts = parent if climb <= 0 else parent[: len(parent) - climb]
194 parts = base_parts + path
195 return ".".join(part for part in parts if part)
198def _splice(text: str, edits: list[tuple[tuple[int, int], tuple[int, int], str]]) -> str:
199 """Splice absolute specs back into ``text`` by token (row, col) positions.
201 Edits are applied last-first so the earlier offsets stay valid. A single
202 contiguous span is replaced in every case: the relative spec starts at the
203 dots token and ends at the last module-path token (or the dots token when
204 there is no path, i.e. ``from . import``).
205 """
206 lines = text.splitlines(keepends=True)
207 if not lines:
208 return text
209 line_offsets = []
210 acc = 0
211 for line in lines:
212 line_offsets.append(acc)
213 acc += len(line)
215 def offset(row: int, col: int) -> int:
216 if 1 <= row <= len(line_offsets):
217 return line_offsets[row - 1] + col
218 return acc
220 positioned = sorted(
221 ((offset(*start), offset(*end), new) for start, end, new in edits),
222 key=lambda item: item[0],
223 reverse=True,
224 )
225 out = text
226 for start, end, new in positioned:
227 if start < end <= len(out):
228 out = out[:start] + new + out[end:]
229 return out
232# ---------------------------------------------------------------------------
233# file assembly
234# ---------------------------------------------------------------------------
236def emit_modeling_file(meta: Any) -> str:
237 """Assemble the full generated modeling file text.
239 Returns the text (not written to disk) so the caller can compute a diff
240 first. The header banner carries the generation provenance; the source is
241 copied verbatim, rewritten by the adapter-driven inline pipeline, and then
242 lowered at the boundary classes.
243 """
244 text = copy_original_modeling(meta.source)
245 # Stale-source sanitize: a trust-remote-code modeling file ships with the
246 # checkpoint and reflects the author's transformers, not the installed one.
247 # Removing ``is_torch_fx_available`` (a transformers>=5.0 casualty) here
248 # keeps the artifact importable without touching the source file. No-op
249 # when the source does not import the symbol.
250 from hyper_parallel.codegen.source.compat import sanitize_source_compat
252 text = sanitize_source_compat(text)
253 module_name = (
254 meta.source.get("module_name")
255 if isinstance(meta.source, dict)
256 else getattr(meta.source, "module_name", "")
257 )
258 text = rewrite_relative_imports(text, module_name)
259 text = _apply_inline_modeling(text, meta)
260 text = _lower_forward_boundaries(text, meta, module_name)
261 return _banner(meta) + text
264def _apply_inline_modeling(text: str, meta: Any) -> str:
265 """Render the artifact through the adapter-driven inline source pipeline."""
267 from hyper_parallel.codegen.inline import render_inline_modeling
269 source = getattr(meta, "source", None)
270 architecture = source.get("architecture") if isinstance(source, dict) else getattr(source, "architecture", None)
271 model_type = getattr(meta, "model_class", None) or architecture
272 return render_inline_modeling(text, meta, model_type=model_type)
275def _lower_forward_boundaries(text: str, meta: Any, module_name: str) -> str:
276 """Sink boundary wrappers into the model's ``forward`` bodies.
278 Each boundary class's ``forward`` is rewritten to redistribute through
279 the instance-bound compiled plan and run the region's compute
280 explicitly. A boundary whose class cannot be resolved or
281 whose per-class contract differs across FQNs fails fast — a silently
282 dropped redistribution would change the model's numeric behavior.
284 Fresh on the unmodified source text, so offsets stay valid for the whole
285 call. The ``local_compute_fn`` / ``inner_wrapper`` members the rewritten
286 forward references are bound by the runtime parallelize step, so the
287 generated file stays runtime-importable without a mesh.
289 Two generation modes exist; exactly one is chosen here per generation:
291 - the default :func:`lower_forward_boundaries` (the four parallel forms);
292 - the opt-in switch template :func:`lower_forward_boundaries_toggle` (the
293 "original forward + per-dimension ``{tp|cp|ep}_enable`` switches" form),
294 selected only when ``HP_CODEGEN_TOGGLE_FORWARD`` is set (non-empty). This
295 keeps the default artifact byte-identical while letting a generation opt
296 into the switch template without adding a second runtime path.
298 ``boundary_classes`` is the FQN -> class-name map frozen at freeze time. A
299 legacy or hand-built plan may have no such map; without
300 it the lowerer cannot resolve any boundary FQN to its owning class, so it
301 would fail fast on the first boundary instead of degrading to the earlier
302 behavior (params sharded by the native Phase A reuse, no forward rewrite).
303 When the map is absent, the source remains untouched and runtime dispatch
304 handles the boundaries.
305 """
306 from hyper_parallel.codegen.emit.parallel import (
307 lower_forward_boundaries,
308 lower_forward_boundaries_toggle,
309 )
311 boundary_classes = getattr(meta, "boundary_classes", None) or {}
312 if not boundary_classes:
313 return text
314 if os.environ.get("HP_CODEGEN_TOGGLE_FORWARD"):
315 return lower_forward_boundaries_toggle(
316 text,
317 meta,
318 boundary_classes=boundary_classes,
319 module_name=module_name,
320 )
321 return lower_forward_boundaries(
322 text,
323 meta,
324 boundary_classes=boundary_classes,
325 module_name=module_name,
326 )
329def _banner(meta: Any) -> str:
330 """Render the provenance banner comment."""
331 source = meta.source if isinstance(meta.source, dict) else _source_dict(meta.source)
332 pd = meta.parallel_dims or {}
333 lines = [
334 "# Generated by codegen. DO NOT EDIT MANUALLY.",
335 "#",
336 f"# config : {getattr(meta, 'yaml_path', '')}",
337 f"# signature : {getattr(meta, 'signature', '')}",
338 f"# source : {source.get('module_name', '')}",
339 f"# source sha256: {source.get('sha256', '')}",
340 f"# architecture : {source.get('architecture', '')}",
341 f"# parallel_dims: tp={pd.get('tp_size', 1)} cp={pd.get('cp_size', 1)} "
342 f"ep={pd.get('ep_size', 1)} pp={pd.get('pp_size', 1)} "
343 f"sp={pd.get('sequence_parallel', False)} "
344 f"loss_parallel={pd.get('loss_parallel', False)}",
345 f"# injections : {len(getattr(meta, 'injections', []) or [])} rule(s)",
346 "# runtime : mesh / FSDP2 / PP / AC / compile still handled by runtime",
347 "#",
348 "# [HYPER PLAN] the frozen sharding plan travels in codegen_meta.json;",
349 "# the runtime shards from it. Regenerate to change.",
350 "#",
351 "",
352 ]
353 return "\n".join(lines)
356def _source_dict(source: Any) -> dict[str, Any]:
357 if source is None:
358 return {}
359 if isinstance(source, dict):
360 return source
361 return {
362 "module_name": getattr(source, "module_name", ""),
363 "file_path": getattr(source, "file_path", ""),
364 "architecture": getattr(source, "architecture", ""),
365 "sha256": getattr(source, "sha256", ""),
366 }
369__all__ = [
370 "copy_original_modeling",
371 "emit_modeling_file",
372 "rewrite_relative_imports",
373]