Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / inline / patch_engine.py: 12%
241 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"""Apply inline patch IR to copied modeling source."""
17from __future__ import annotations
19import ast
20from collections import defaultdict
21from typing import Iterable
23from hyper_parallel.codegen.astkit.edits import (
24 TextEdit,
25 apply_edits,
26 insert_after_imports,
27 replace_function_body,
28)
29from hyper_parallel.codegen.astkit.index import build_source_index
30from hyper_parallel.codegen.inline.ir import (
31 ClassRemovalPatch,
32 ConstructorReplacePatch,
33 ForwardExtractPatch,
34 ImportPatch,
35 InlinePatchSet,
36 ModuleSnippetPatch,
37)
40def apply_patch_set(source_text: str, patch_set: InlinePatchSet) -> str:
41 """Apply imports, constructor edits, class removals, markers, and forward bodies.
43 Args:
44 source_text: Modeling source the patches are applied to.
45 patch_set: Edits emitted by the inline passes.
47 Returns:
48 The patched source text.
49 """
51 first_pass: list[TextEdit] = []
52 index = build_source_index(source_text)
53 import_edit = _import_edit(source_text, index.import_end, index.docstring_end, patch_set.imports)
54 if import_edit is not None:
55 first_pass.append(import_edit)
56 snippet_edit = _snippet_edit(source_text, index.import_end, index.docstring_end, patch_set.module_snippets)
57 if snippet_edit is not None:
58 first_pass.append(snippet_edit)
59 first_pass.extend(_constructor_edits(source_text, patch_set.constructor_replaces))
60 text = apply_edits(source_text, first_pass) if first_pass else source_text
62 if patch_set.class_removals:
63 text = _apply_class_removals(text, patch_set.class_removals)
65 second_pass: list[TextEdit] = []
66 second_pass.extend(_forward_extract_edits(text, patch_set.forward_extracts))
67 return apply_edits(text, second_pass) if second_pass else text
70def _import_edit(
71 source_text: str,
72 import_end: int | None,
73 docstring_end: int | None,
74 imports: Iterable[ImportPatch],
75) -> TextEdit | None:
76 lines = _render_imports(source_text, imports)
77 if not lines:
78 return None
79 return insert_after_imports(
80 source_text,
81 import_end,
82 "\n" + "\n".join(lines) + "\n",
83 default_offset=docstring_end or 0,
84 )
87def _render_imports(source_text: str, imports: Iterable[ImportPatch]) -> list[str]:
88 """Render import patches as deduplicated, grouped import lines."""
89 grouped: dict[str, set[str]] = defaultdict(set)
90 lines: list[str] = []
91 for patch in imports:
92 if patch.raw:
93 if patch.raw not in source_text:
94 lines.append(patch.raw)
95 continue
96 grouped[patch.module].update(patch.names)
97 for module in sorted(grouped):
98 names = tuple(sorted(grouped[module]))
99 line = f"from {module} import {', '.join(names)}"
100 if line not in source_text:
101 lines.append(line)
102 return lines
105def _snippet_edit(
106 source_text: str,
107 import_end: int | None,
108 docstring_end: int | None,
109 snippets: Iterable[ModuleSnippetPatch],
110) -> TextEdit | None:
111 parts = [snippet.text.strip("\n") for snippet in snippets if snippet.text.strip()]
112 if not parts:
113 return None
114 return insert_after_imports(
115 source_text,
116 import_end,
117 "\n\n" + "\n\n".join(parts) + "\n",
118 default_offset=docstring_end or 0,
119 )
122def _constructor_edits(
123 source_text: str,
124 patches: Iterable[ConstructorReplacePatch],
125) -> list[TextEdit]:
126 """Build the edits that swap constructor calls per patch target."""
127 patch_list = list(patches)
128 patches_by_old: dict[str, list[ConstructorReplacePatch]] = defaultdict(list)
129 for patch in patch_list:
130 patches_by_old[patch.old_ctor].append(patch)
131 if not patches_by_old:
132 return []
133 scopes = _scopes_by_old_ctor(patch_list)
134 tree = ast.parse(source_text)
135 offsets = _line_offsets(source_text)
136 edits: list[TextEdit] = []
137 hits: dict[str, int] = defaultdict(int)
138 for node in ast.walk(tree):
139 if not isinstance(node, ast.Assign):
140 continue
141 attr = _assigned_attr(node)
142 if attr is None:
143 continue
144 call = node.value
145 if not isinstance(call, ast.Call):
146 continue
147 name = _call_name(call.func)
148 if name is None or name not in patches_by_old:
149 continue
150 scope = scopes[name]
151 if scope is not None and attr not in scope:
152 continue
153 patch = patches_by_old[name][0]
154 hits[name] += 1
155 if patch.mode == "name":
156 start, end = _node_span(call.func, offsets)
157 edits.append(TextEdit(start, end, patch.new_ctor))
158 elif patch.mode == "wrap_source":
159 start, end = _node_span(call, offsets)
160 call_source = source_text[start:end]
161 extra = "".join(f", {arg}" for arg in patch.keyword_args)
162 replacement = f"{patch.new_ctor}(module={call_source}{extra})"
163 edits.append(TextEdit(start, end, replacement))
164 else:
165 raise ValueError(f"unknown constructor replacement mode: {patch.mode!r}")
166 _require_scoped_rewrites(patch_list, hits)
167 return _drop_nested_edits(edits)
170def _scopes_by_old_ctor(
171 patches: Iterable[ConstructorReplacePatch],
172) -> dict[str, set[str] | None]:
173 """Merge each constructor's rewrite scope; ``None`` means unrestricted."""
175 scopes: dict[str, set[str] | None] = {}
176 for patch in patches:
177 if not patch.scope_attrs:
178 scopes[patch.old_ctor] = None
179 continue
180 merged = scopes.get(patch.old_ctor)
181 if merged is None and patch.old_ctor in scopes:
182 continue
183 if merged is None:
184 merged = scopes[patch.old_ctor] = set()
185 merged.update(patch.scope_attrs)
186 return scopes
189def _assigned_attr(node: ast.Assign) -> str | None:
190 """Return the stored attribute of an attribute store, else ``None``.
192 Module replacements are sunk at the assignment that installs the module, so
193 scoping by the assigned attribute is exactly the FQN granularity the native
194 plan used to decide the replacement. Families attach modules at more than
195 the ``self.<attr>`` store: a loop-local ``layer.<attr>`` upgrade and a
196 ModuleList element ``self.layers[i].<attr>`` install the same matched
197 modules, so any attribute store yields its stored attribute -- the rule's
198 scope set keeps the rewrite at the plan's matched-FQN granularity.
199 """
201 if len(node.targets) != 1:
202 return None
203 target = node.targets[0]
204 return target.attr if isinstance(target, ast.Attribute) else None
207def _require_scoped_rewrites(
208 patches: Iterable[ConstructorReplacePatch],
209 hits: dict[str, int],
210) -> None:
211 """Fail when a scoped rule matched FQNs but no source call site.
213 Silently dropping the rewrite would emit an artifact that no longer matches
214 the plan that decided the replacement, so a missing call site is a
215 generation error: the scope handling needs to learn the new shape rather
216 than the rewrite being widened back to the whole file.
217 """
219 missing = sorted(
220 {patch.old_ctor for patch in patches if patch.scope_attrs and not hits.get(patch.old_ctor)}
221 )
222 if missing:
223 raise RuntimeError(
224 "codegen: no `self.<attr> = <ctor>(...)` call site found for "
225 f"{missing}; the plan replaced modules of this type, so the "
226 "replacement cannot be sunk without diverging from the plan"
227 )
230def _drop_nested_edits(edits: list[TextEdit]) -> list[TextEdit]:
231 """Keep outer constructor rewrites when constructor calls are nested."""
232 kept: list[TextEdit] = []
233 covered_end = -1
234 for edit in sorted(edits, key=lambda item: (item.start, -item.end)):
235 if edit.start < covered_end:
236 continue
237 kept.append(edit)
238 covered_end = edit.end
239 return kept
242def drop_shadowed_definitions(text: str, imported_names: Iterable[str]) -> str:
243 """Remove definitions that shadow a library import of the same name.
245 The artifact imports the very helpers the native path calls, but the copied
246 modeling source defines its own module-level helpers with those names, and a
247 local definition wins over the import. For a *decorated* helper that means
248 a different implementation runs: ``use_kernel_func_from_hub("rotary_pos_emb")``
249 resolves its kernel binding per module, and the generated module is not the
250 module the binding was registered for. Both paths must therefore end up with
251 the imported implementation, so the shadowing copy is dropped.
253 Callers pass names that the *base source* already defines, so the classes and
254 snippets the inline passes render (which share names with library imports on
255 purpose) keep their definitions.
257 Args:
258 text: Source to drop the shadowing definitions from.
259 imported_names: Names the library imports bind in module scope.
261 Returns:
262 The source with the shadowing definitions removed.
263 """
265 names = {name for name in imported_names if name}
266 if not names:
267 return text
268 tree = ast.parse(text)
269 offsets = _line_offsets(text)
270 edits: list[TextEdit] = []
271 for node in tree.body:
272 if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
273 continue
274 if node.name not in names or not node.decorator_list:
275 continue
276 start, end = _node_span(node, offsets)
277 edits.append(TextEdit(start, _skip_trailing_blank_lines(text, end), ""))
278 if not edits:
279 return text
280 return apply_edits(text, edits)
283def _name_is_referenced(tree: ast.AST, name: str) -> bool:
284 """Return whether ``name`` still appears as a reference in ``tree``."""
286 for node in ast.walk(tree):
287 if isinstance(node, ast.Name) and node.id == name:
288 return True
289 if isinstance(node, ast.Attribute) and node.attr == name:
290 return True
291 return False
294def _apply_class_removals(
295 text: str,
296 patches: Iterable[ClassRemovalPatch],
297) -> str:
298 """Rewrite surviving bare references to removed classes, then delete them.
300 Runs as a **separate pass** after constructor rewrites so the AST reflects
301 the already-patched text: every ``old_name(...)`` call has become
302 ``new_name(...)`` and only *bare* references (dict values, ``isinstance``
303 targets, annotations) remain. Those are rewritten to ``new_name`` when
304 they fall outside a class-deletion range; references inside a class being
305 deleted are left alone because the entire block is removed.
307 Args:
308 text: Source text after the constructor rewrites.
309 patches: Class removals requested by the replacement specs.
311 Returns:
312 The text with unused classes removed.
313 """
314 patch_list = list(patches)
315 if not patch_list:
316 return text
317 tree = ast.parse(text)
318 offsets = _line_offsets(text)
319 names_to_remove, new_name_by_old = _removal_targets(tree, patch_list)
320 if not names_to_remove:
321 return text
322 deletion_ranges, classdef_nodes = _class_deletion_spans(tree, offsets, names_to_remove)
323 edits = _reference_edits(
324 tree, offsets, names_to_remove, new_name_by_old, deletion_ranges
325 )
326 edits.extend(_deletion_edits(text, offsets, patch_list, classdef_nodes))
327 return apply_edits(text, edits) if edits else text
330def _removal_targets(
331 tree: ast.AST,
332 patch_list: list[ClassRemovalPatch],
333) -> tuple[set[str], dict[str, str]]:
334 """Return the classes still unused after the rewrites, and their renames.
336 A scoped rewrite can leave the original class in use -- the recipe keeps
337 the HF Q/K norm while replacing the layer norms -- so a class is removed
338 only once the patched source no longer references it. Removing a class
339 that is still called would turn a partially replaced model into a
340 ``NameError`` instead of the model the plan described.
341 """
343 names = {patch.old_name for patch in patch_list}
344 unused = {name for name in names if not _name_is_referenced(tree, name)}
345 return unused, {patch.old_name: patch.new_name for patch in patch_list}
348def _class_deletion_spans(
349 tree: ast.AST,
350 offsets: tuple[int, ...],
351 names_to_remove: set[str],
352) -> tuple[list[tuple[int, int]], dict[str, ast.ClassDef]]:
353 """Return the source spans of the removable classes and their nodes."""
355 deletion_ranges: list[tuple[int, int]] = []
356 classdef_nodes: dict[str, ast.ClassDef] = {}
357 for node in tree.body:
358 if isinstance(node, ast.ClassDef) and node.name in names_to_remove:
359 classdef_nodes[node.name] = node
360 deletion_ranges.append(_node_span(node, offsets))
361 return deletion_ranges, classdef_nodes
364def _reference_edits(
365 tree: ast.AST,
366 offsets: tuple[int, ...],
367 names_to_remove: set[str],
368 new_name_by_old: dict[str, str],
369 deletion_ranges: list[tuple[int, int]],
370) -> list[TextEdit]:
371 """Return edits rewriting bare references that survive the deletion."""
373 edits: list[TextEdit] = []
374 for node in ast.walk(tree):
375 if not isinstance(node, ast.Name) or node.id not in names_to_remove:
376 continue
377 start, end = _node_span(node, offsets)
378 if any(span <= start < stop for span, stop in deletion_ranges):
379 continue
380 edits.append(TextEdit(start, end, new_name_by_old[node.id]))
381 return edits
384def _deletion_edits(
385 text: str,
386 offsets: tuple[int, ...],
387 patch_list: list[ClassRemovalPatch],
388 classdef_nodes: dict[str, ast.ClassDef],
389) -> list[TextEdit]:
390 """Return edits deleting the removable class blocks and trailing blanks."""
392 edits: list[TextEdit] = []
393 for patch in patch_list:
394 cls_node = classdef_nodes.get(patch.old_name)
395 if cls_node is None:
396 continue
397 start, end = _node_span(cls_node, offsets)
398 end = _skip_trailing_blank_lines(text, end)
399 edits.append(TextEdit(start, end, ""))
400 return edits
403def _skip_trailing_blank_lines(text: str, end: int) -> int:
404 """Advance ``end`` past spaces and blank lines following a deleted block."""
406 while end < len(text) and text[end] in " \t":
407 end += 1
408 while end < len(text) and text[end] == "\n":
409 end += 1
410 while end < len(text) and text[end] in " \t":
411 end += 1
412 return end
415def _forward_extract_edits(
416 source_text: str,
417 patches: Iterable[ForwardExtractPatch],
418) -> list[TextEdit]:
419 """Build the edits that extract patched forwards into impl functions."""
420 if not patches:
421 return []
422 from hyper_parallel.codegen.emit.parallel import ( # pylint: disable=import-outside-toplevel
423 _build_forward_impl_edit,
424 )
426 index = build_source_index(source_text)
427 edits: list[TextEdit] = []
428 for patch in patches:
429 cls = index.find_class(patch.class_name)
430 if cls is None:
431 continue
432 func = cls.methods.get(patch.method_name)
433 if func is None:
434 continue
435 impl = _build_forward_impl_edit(
436 source_text,
437 func,
438 )
439 if impl is not None:
440 edits.append(impl)
441 edits.extend(replace_function_body(source_text, func, patch.body, indent=" "))
442 return edits
445def _call_name(node: ast.AST) -> str | None:
446 if isinstance(node, ast.Name):
447 return node.id
448 if isinstance(node, ast.Attribute):
449 return node.attr
450 return None
453def _node_span(node: ast.AST, offsets: tuple[int, ...]) -> tuple[int, int]:
454 """Return the ``(start, end)`` source span of ``node``.
456 Decorators are part of the span: deleting only the ``def``/``class`` line
457 would leave a dangling decorator that then applies to whatever statement
458 follows the removed block.
459 """
461 if (
462 getattr(node, "lineno", None) is None
463 or getattr(node, "col_offset", None) is None
464 or getattr(node, "end_lineno", None) is None
465 or getattr(node, "end_col_offset", None) is None
466 ):
467 raise ValueError(f"AST node has no source span: {type(node).__name__}")
468 line = node.lineno
469 column = node.col_offset
470 decorators = getattr(node, "decorator_list", None)
471 if decorators:
472 first = min(decorators, key=lambda item: item.lineno)
473 if first.lineno < line:
474 # A decorator node starts at its expression, one column *after* the
475 # ``@``; the span must start at the ``@`` itself, otherwise deleting
476 # the block leaves a stray ``@`` behind.
477 line, column = first.lineno, max(0, first.col_offset - 1)
478 start = offsets[line - 1] + column
479 end = offsets[node.end_lineno - 1] + node.end_col_offset
480 return start, end
483def _line_offsets(source_text: str) -> tuple[int, ...]:
484 offsets = [0]
485 total = 0
486 for line in source_text.splitlines(keepends=True):
487 total += len(line)
488 offsets.append(total)
489 return tuple(offsets)