Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / emit / recorded.py: 68%
183 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"""Materialize builder-observed forward patches into generated model source."""
17from __future__ import annotations
19import ast
20import inspect
21import os
22import textwrap
23from collections import defaultdict
24from types import SimpleNamespace
25from typing import Any
27from hyper_parallel.codegen.astkit.edits import (
28 TextEdit,
29 apply_edits,
30 replace_function_body,
31)
32from hyper_parallel.codegen.astkit.index import FunctionInfo, build_source_index
33from hyper_parallel.codegen.emit.parallel import (
34 _build_forward_impl_edit,
35 _build_forward_body,
36 _render_forward_call,
37 _render_input_redistribute,
38 _render_local_compute,
39 _render_output_redistribute,
40)
41from hyper_parallel.codegen.plan.boundary_forms import (
42 FORM_GENERIC,
43 FORM_REGION,
44 classify_boundary_form,
45)
48def materialize_recorded_forwards(
49 source_text: str,
50 meta: Any,
51 *,
52 patch_plan: Any = None,
53 source_path: str | None = None,
54) -> str:
55 """Lower observed local-compute decisions into their source-owned classes."""
56 patches = patch_plan
57 records = getattr(patches, "local_computes", ()) if patches is not None else ()
58 if not records:
59 return _materialize_declared_local_computes(source_text, meta)
61 _restore_recorded_boundary_ownership(meta, records)
63 groups: dict[tuple[type, str], list[Any]] = defaultdict(list)
64 for record in records:
65 if source_path is not None and not _record_belongs_to_source(
66 record, source_path
67 ):
68 continue
69 original = _unbound(record.original_forward)
70 groups[(record.module_type, _source_function_name(original))].append(record)
72 rendered = source_text
73 for (host_type, source_function), group in sorted(
74 groups.items(), key=lambda item: (item[0][0].__name__, item[0][1])
75 ):
76 _require_one_contract(group, meta)
77 if source_function and _has_module_function(rendered, source_function):
78 rendered = _materialize_bound_module_function(
79 rendered, source_function, group[0], meta
80 )
81 continue
82 if not _has_source_class(rendered, host_type.__name__):
83 continue
84 rendered = _materialize_source_class(
85 rendered, host_type.__name__, group[0], meta
86 )
87 _mark_class_source_lowering_skipped(meta, host_type.__name__)
88 return rendered
91def _materialize_declared_local_computes(source_text: str, meta: Any) -> str:
92 """Lower explicit frozen local-compute injections when no live record exists."""
93 boundary_classes = getattr(meta, "boundary_classes", None) or {}
94 injections = getattr(meta, "injections", None) or ()
95 source_classes = {
96 node.name
97 for node in ast.parse(source_text).body
98 if isinstance(node, ast.ClassDef)
99 }
100 grouped: dict[str, list[Any]] = defaultdict(list)
101 for injection in injections:
102 local = injection.get("local_compute_fn")
103 module_fqn = injection.get("match")
104 class_name = boundary_classes.get(module_fqn)
105 if local is None or not module_fqn or class_name not in source_classes:
106 continue
107 grouped[class_name].append(SimpleNamespace(
108 module_fqn=module_fqn,
109 factory=_callable_identity(local),
110 ))
111 rendered = source_text
112 materialized = getattr(meta, "materialized_patches", None)
113 for class_name, records in sorted(grouped.items()):
114 _require_one_contract(records, meta)
115 rendered = _materialize_source_class(
116 rendered, class_name, records[0], meta
117 )
118 _mark_class_source_lowering_skipped(meta, class_name)
119 for record in records:
120 if isinstance(materialized, list):
121 materialized.append({
122 "fqn": record.module_fqn,
123 "host_class": class_name,
124 "kind": "local_compute",
125 "source_form": "same_class_toggle_forward",
126 })
127 return rendered
130def _mark_class_source_lowering_skipped(meta: Any, class_name: str) -> None:
131 """Keep generic boundary lowering off a source class with an EP toggle.
133 A source-visible local-compute patch rewrites the class method, not just
134 one module instance. Every boundary instance of that class must therefore
135 use the runtime per-instance boundary wrapper. Lowering another instance's
136 boundary into the shared method would extract a second ``_forward_impl``
137 and wrap the EP toggle a second time.
138 """
139 param_plan = getattr(meta, "param_plan", None) or {}
140 boundary_classes = getattr(meta, "boundary_classes", None) or {}
141 for module_fqn, owner_class in boundary_classes.items():
142 if owner_class != class_name:
143 continue
144 entry = param_plan.get(module_fqn)
145 if isinstance(entry, dict) and entry.get("is_boundary"):
146 param_plan[module_fqn] = {**entry, "skip_source_lowering": True}
149def _callable_identity(value: Any) -> str:
150 """Return a stable identity for a frozen callable declaration."""
151 if isinstance(value, dict):
152 return str(value.get("_target_") or value.get("target") or value.get("name"))
153 return str(value)
156def _record_belongs_to_source(record: Any, source_path: str) -> bool:
157 """Return whether a recorded host is defined by this source unit."""
158 candidates = (
159 inspect.getsourcefile(record.module_type),
160 inspect.getsourcefile(_unbound(record.original_forward)),
161 )
162 expected = os.path.normcase(os.path.abspath(source_path))
163 return any(
164 candidate is not None
165 and os.path.normcase(os.path.abspath(candidate)) == expected
166 for candidate in candidates
167 )
170def _restore_recorded_boundary_ownership(meta: Any, records: Any) -> None:
171 """Make recorded local-region forwards runtime-owned boundaries again.
173 The generic inline pass initially classifies known EP strategy targets as
174 external-state classes. When builder recording supplies the actual forward,
175 this emitter materializes a compiled-boundary local region instead of that
176 generic inline body. Restore the matching frozen entries so runtime binding
177 installs ``_hp_boundary`` and the recorded compute method they reference.
179 Args:
180 meta: Mutable generation metadata.
181 records: Builder-recorded local-compute patches.
182 """
183 param_plan = getattr(meta, "param_plan", None)
184 boundary_classes = getattr(meta, "boundary_classes", None)
185 external_classes = set(getattr(meta, "external_state_classes", None) or ())
186 recorded_classes: set[str] = set()
187 for record in records:
188 class_name = record.module_type.__name__
189 recorded_classes.add(class_name)
190 if isinstance(param_plan, dict):
191 entry = param_plan.get(record.module_fqn)
192 if isinstance(entry, dict):
193 param_plan[record.module_fqn] = {**entry, "is_boundary": True}
194 if isinstance(boundary_classes, dict):
195 boundary_classes[record.module_fqn] = class_name
196 if hasattr(meta, "external_state_classes"):
197 meta.external_state_classes = sorted(external_classes - recorded_classes)
200def _unbound(value: Any) -> Any:
201 return getattr(value, "__func__", value)
204def _source_function_name(value: Any) -> str:
205 """Return a top-level source function name, excluding class methods."""
206 name = getattr(value, "__qualname__", "")
207 return name if name and "." not in name else ""
210def _has_module_function(source_text: str, name: str) -> bool:
211 return any(
212 isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == name
213 for node in ast.parse(source_text).body
214 )
217def _has_source_class(source_text: str, name: str) -> bool:
218 return any(
219 isinstance(node, ast.ClassDef) and node.name == name
220 for node in ast.parse(source_text).body
221 )
224def _require_one_contract(records: list[Any], meta: Any) -> None:
225 """A shared source forward can represent only one boundary contract."""
226 contracts = {
227 repr({
228 key: value
229 for key, value in _entry_for_record(meta, record).items()
230 if key != "params"
231 })
232 for record in records
233 }
234 factories = {record.factory for record in records}
235 if len(contracts) != 1 or len(factories) != 1:
236 fqns = ", ".join(sorted(record.module_fqn for record in records))
237 raise RuntimeError(
238 "codegen: recorded modules sharing one forward have incompatible "
239 f"patch contracts: {fqns}"
240 )
243def _entry_for_record(meta: Any, record: Any) -> dict[str, Any]:
244 entry = (getattr(meta, "param_plan", None) or {}).get(record.module_fqn)
245 if entry is None:
246 raise RuntimeError(
247 f"codegen: recorded patch {record.module_fqn!r} has no frozen plan entry"
248 )
249 return entry
252def _injection_for_record(meta: Any, record: Any) -> dict[str, Any] | None:
253 for injection in getattr(meta, "injections", None) or ():
254 if injection.get("match") == record.module_fqn:
255 return injection
256 return None
259def _function_info(node: ast.FunctionDef | ast.AsyncFunctionDef, source_text: str) -> FunctionInfo:
260 """Build source-index information for one recorded function."""
261 lines = source_text.splitlines(keepends=True)
262 offsets = [0]
263 for line in lines:
264 offsets.append(offsets[-1] + len(line))
266 def start(line_number: int) -> int:
267 return offsets[line_number - 1]
269 positional = [arg.arg for arg in (*node.args.posonlyargs, *node.args.args)]
270 if positional and positional[0] in ("self", "module"):
271 positional.pop(0)
272 first_line = min(
273 [node.lineno, *(item.lineno for item in node.decorator_list)]
274 )
275 return FunctionInfo(
276 name=node.name,
277 body_start=start(node.body[0].lineno),
278 body_end=offsets[node.end_lineno],
279 def_offset=start(first_line),
280 param_names=positional,
281 kwonly_names=[arg.arg for arg in node.args.kwonlyargs],
282 vararg_name=node.args.vararg.arg if node.args.vararg else None,
283 kwarg_name=node.args.kwarg.arg if node.args.kwarg else None,
284 )
287def _recorded_body(source_text: str, func: FunctionInfo, record: Any, meta: Any) -> str:
288 """Render the body of one observed local-compute patch."""
289 entry = _entry_for_record(meta, record)
290 injection = _injection_for_record(meta, record)
291 form = classify_boundary_form(
292 entry,
293 tuple(getattr(meta, "mesh_dim_names", None) or ()),
294 injection,
295 )
296 emitted = FORM_REGION if record.factory is not None else FORM_GENERIC
297 local = injection.get("local_compute_fn") if injection else None
298 if local is not None:
299 body = _render_local_compute(
300 local,
301 _render_input_redistribute(func),
302 _render_output_redistribute(),
303 form,
304 )
305 else:
306 body = _build_forward_body(source_text, injection, func, form, emitted)
307 fallback = _render_forward_call(func).replace("outputs =", "return", 1)
308 return "\n".join([
309 " if self.ep_enable:",
310 textwrap.indent(body, " "),
311 "",
312 fallback,
313 ])
316def _materialize_bound_module_function(
317 source_text: str,
318 function_name: str,
319 record: Any,
320 meta: Any,
321) -> str:
322 """Turn an instance-bound module function into impl + explicit wrapper."""
323 tree = ast.parse(source_text)
324 node = next(
325 item for item in tree.body
326 if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef))
327 and item.name == function_name
328 )
329 info = _function_info(node, source_text)
330 impl_name = f"{function_name}_hp_impl"
331 original = source_text[info.def_offset:info.body_end]
332 original = original.replace(f"def {function_name}(", f"def {impl_name}(", 1)
333 if not original.endswith("\n"):
334 original += "\n"
335 body = textwrap.dedent(_recorded_body(source_text, info, record, meta))
336 body_edit = TextEdit(info.body_start, info.body_end, textwrap.indent(body, " "))
337 rendered = apply_edits(
338 source_text,
339 [TextEdit(info.def_offset, info.def_offset, original + "\n"), body_edit],
340 )
341 return _bind_module_impl(rendered, function_name, impl_name)
344def _materialize_source_class(
345 source_text: str,
346 class_name: str,
347 record: Any,
348 meta: Any,
349) -> str:
350 """Rewrite a source-defined host even when old meta normalization skips it."""
351 func = build_source_index(source_text).find_forward(class_name)
352 if func is None:
353 raise RuntimeError(
354 f"codegen: recorded source host {class_name!r} defines no forward"
355 )
356 body = _recorded_body(source_text, func, record, meta)
357 impl = _build_forward_impl_edit(
358 source_text,
359 func,
360 prefix=" ep_enable = False\n\n",
361 )
362 edits = list(replace_function_body(source_text, func, body, indent=""))
363 if impl is not None:
364 edits.append(impl)
365 return apply_edits(source_text, edits)
368def _bind_module_impl(source_text: str, wrapper_name: str, impl_name: str) -> str:
369 """Add the instance ``_forward_impl`` binding beside MethodType(forward)."""
370 tree = ast.parse(source_text)
371 lines = source_text.splitlines(keepends=True)
372 offsets = [0]
373 for line in lines:
374 offsets.append(offsets[-1] + len(line))
375 edits: list[TextEdit] = []
376 for node in ast.walk(tree):
377 if not isinstance(node, ast.Assign) or len(node.targets) != 1:
378 continue
379 target = node.targets[0]
380 value = node.value
381 if not (
382 isinstance(target, ast.Attribute)
383 and target.attr == "forward"
384 and isinstance(value, ast.Call)
385 and isinstance(value.func, ast.Name)
386 and value.func.id == "MethodType"
387 and value.args
388 and isinstance(value.args[0], ast.Name)
389 and value.args[0].id == wrapper_name
390 ):
391 continue
392 owner = ast.unparse(target.value)
393 indent = " " * node.col_offset
394 binding = f"{indent}{owner}._forward_impl = MethodType({impl_name}, {owner})\n"
395 edits.append(TextEdit(offsets[node.lineno - 1], offsets[node.lineno - 1], binding))
396 if not edits:
397 raise RuntimeError(
398 f"codegen: source function {wrapper_name!r} is not installed with MethodType"
399 )
400 return apply_edits(source_text, edits)
403__all__ = ["materialize_recorded_forwards"]