Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / emit / snapshot_patches.py: 60%
155 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"""Convert native-build snapshot observations into emitter patch records."""
17from __future__ import annotations
19import ast
20import importlib
21import os
22from typing import Any, Optional
24from hyper_parallel.codegen.inline.targets import import_target_path
25from hyper_parallel.codegen.source.resolver import resolve_remote_siblings
28def emission_patches_from_snapshot_payload(
29 payload: dict[str, Any],
30) -> dict[str, Any]:
31 """Translate resolved native observations into emitter patch records."""
32 applications = payload["replacements"]["applications"]
33 replacements = [
34 {
35 "module_fqn": item["primary_fqn"],
36 "source_type": item["source_type"],
37 "replacement_type": item["replacement_type"],
38 }
39 for item in applications
40 ]
41 replacement_types = {
42 module_fqn: item["replacement_type"]
43 for item in applications
44 for module_fqn in item.get("module_fqns", (item["primary_fqn"],))
45 }
46 source = payload["source"]
47 plan = payload.get("plan") or {}
48 injections = plan.get("injections") or ()
49 boundary_classes = plan.get("boundary_classes") or {}
50 boundary_types = plan.get("boundary_types") or {}
51 param_plan = plan.get("param_plan") or {}
52 local_computes, inner_wrappers = [], []
53 for injection in injections:
54 local, wrapper = _injection_patches(
55 injection,
56 source,
57 boundary_classes,
58 boundary_types,
59 replacement_types,
60 param_plan,
61 )
62 if local is not None:
63 local_computes.append(local)
64 if wrapper is not None:
65 inner_wrappers.append(wrapper)
66 return {
67 "module_replacements": replacements,
68 "local_computes": local_computes,
69 "inner_wrappers": inner_wrappers,
70 }
73def _injection_patches(
74 injection: dict[str, Any],
75 source: dict[str, Any],
76 boundary_classes: dict[str, Any],
77 boundary_types: dict[str, Any],
78 replacement_types: dict[str, str],
79 param_plan: dict[str, Any],
80) -> tuple[Optional[dict[str, Any]], Optional[dict[str, Any]]]:
81 """Translate one resolved injection into source-emission patches."""
82 module_fqn = injection.get("match")
83 class_name = boundary_classes.get(module_fqn)
84 module_name = source.get("module_name")
85 declared_type = (
86 f"{module_name}:{class_name}" if module_name and class_name else None
87 )
88 source_type = _source_type_identity(source, class_name)
89 observed_type = boundary_types.get(module_fqn)
90 if source_type is None and _is_remote_boundary_type(observed_type):
91 source_type = observed_type
92 bound_forward = _bound_forward_identity(source, module_fqn)
93 local = _local_compute_patch(
94 injection,
95 source,
96 module_fqn,
97 class_name,
98 declared_type,
99 source_type,
100 bound_forward,
101 param_plan,
102 )
103 wrapper = _inner_wrapper_patch(
104 injection,
105 module_fqn,
106 declared_type,
107 source_type,
108 replacement_types,
109 )
110 return local, wrapper
113def _is_remote_boundary_type(value: Any) -> bool:
114 """Return whether an observed type identifies non-framework source."""
115 return isinstance(value, str) and not value.startswith(
116 ("hyper_parallel.", "torch.")
117 )
120def _local_compute_patch(
121 injection: dict[str, Any],
122 source: dict[str, Any],
123 module_fqn: str,
124 class_name: str,
125 declared_type: Optional[str],
126 source_type: Optional[str],
127 bound_forward: Optional[str],
128 param_plan: dict[str, Any],
129) -> Optional[dict[str, Any]]:
130 """Build one local-compute emission record."""
131 local_compute = injection.get("local_compute_fn")
132 if local_compute is None:
133 return None
134 if declared_type is None:
135 raise ValueError(
136 "snapshot local_compute_fn requires a source module and "
137 f"boundary class for {module_fqn!r}"
138 )
139 if source_type is None and bound_forward is None:
140 return None
141 module_type = (
142 (_imported_type_identity(source, class_name) if bound_forward else source_type)
143 or source_type
144 or declared_type
145 )
146 return {
147 "module_fqn": module_fqn,
148 "module_type": module_type,
149 "factory": import_target_path(local_compute),
150 "region_dispatch": (param_plan.get(module_fqn) or {}).get("region_dispatch"),
151 "original_forward": bound_forward or f"{source_type}.forward",
152 }
155def _inner_wrapper_patch(
156 injection: dict[str, Any],
157 module_fqn: str,
158 declared_type: Optional[str],
159 source_type: Optional[str],
160 replacement_types: dict[str, str],
161) -> Optional[dict[str, Any]]:
162 """Build one inner-wrapper emission record."""
163 wrapper = injection.get("inner_wrapper")
164 if wrapper is None:
165 return None
166 target_type = replacement_types.get(module_fqn) or source_type or declared_type
167 if target_type is None:
168 raise ValueError(
169 "snapshot inner_wrapper requires an observed replacement or "
170 f"boundary class for {module_fqn!r}"
171 )
172 return {
173 "module_fqn": module_fqn,
174 "target_fqn": injection.get("inner_target") or "self",
175 "wrapper": import_target_path(wrapper),
176 "target_type": target_type,
177 }
180def _source_class_names(source: dict[str, Any]) -> set[str]:
181 """Return classes defined by the snapshot's modeling source file."""
182 file_path = source.get("file_path")
183 if not file_path:
184 return set()
185 try:
186 with open(file_path, "r", encoding="utf-8", errors="replace") as handle:
187 tree = ast.parse(handle.read(), filename=file_path)
188 except (OSError, SyntaxError):
189 return set()
190 return {node.name for node in tree.body if isinstance(node, ast.ClassDef)}
193def _source_type_identity(
194 source: dict[str, Any], class_name: Optional[str]
195) -> Optional[str]:
196 """Resolve a class to the remote source unit that actually defines it."""
197 if not class_name:
198 return None
199 module_name = source.get("module_name")
200 if not isinstance(module_name, str) or not module_name:
201 return None
202 if class_name in _source_class_names(source):
203 return f"{module_name}:{class_name}"
205 package, separator, _ = module_name.rpartition(".")
206 candidates = []
207 for path in resolve_remote_siblings(source).values():
208 sibling_source = {"file_path": path}
209 if class_name not in _source_class_names(sibling_source):
210 continue
211 stem = os.path.splitext(os.path.basename(path))[0]
212 sibling_module = f"{package}.{stem}" if separator else stem
213 candidates.append(f"{sibling_module}:{class_name}")
214 if len(candidates) > 1:
215 raise ValueError(
216 f"snapshot class {class_name!r} is defined by multiple bundled source units: "
217 f"{sorted(candidates)}"
218 )
219 return candidates[0] if candidates else None
222def _bound_forward_identity(
223 source: dict[str, Any], module_fqn: Optional[str]
224) -> Optional[str]:
225 """Resolve a model-owned function installed through ``MethodType``."""
226 file_path = source.get("file_path")
227 module_name = source.get("module_name")
228 if not file_path or not module_name or not module_fqn:
229 return None
230 try:
231 with open(file_path, "r", encoding="utf-8", errors="replace") as handle:
232 tree = ast.parse(handle.read(), filename=file_path)
233 except (OSError, SyntaxError):
234 return None
236 boundary_name = module_fqn.rsplit(".", 1)[-1]
237 candidates = set()
238 for node in ast.walk(tree):
239 candidate = _method_type_forward_name(node, boundary_name)
240 if candidate:
241 candidates.add(candidate)
242 if len(candidates) > 1:
243 raise ValueError(
244 f"snapshot boundary {module_fqn!r} has multiple MethodType forwards: "
245 f"{sorted(candidates)}"
246 )
247 return f"{module_name}:{next(iter(candidates))}" if candidates else None
250def _method_type_forward_name(node: ast.AST, boundary_name: str) -> Optional[str]:
251 """Return the function assigned as one boundary's MethodType forward."""
252 if not isinstance(node, ast.Assign) or len(node.targets) != 1:
253 return None
254 target, call = node.targets[0], node.value
255 valid_target = (
256 isinstance(target, ast.Attribute)
257 and target.attr == "forward"
258 and isinstance(target.value, ast.Attribute)
259 and target.value.attr == boundary_name
260 )
261 if not valid_target or not isinstance(call, ast.Call):
262 return None
263 if _ast_name(call.func) != "MethodType" or not call.args:
264 return None
265 return call.args[0].id if isinstance(call.args[0], ast.Name) else None
268def _imported_type_identity(
269 source: dict[str, Any], class_name: Optional[str]
270) -> Optional[str]:
271 """Resolve an imported boundary class from modules referenced by the source."""
272 file_path = source.get("file_path")
273 if not file_path or not class_name:
274 return None
275 try:
276 with open(file_path, "r", encoding="utf-8", errors="replace") as handle:
277 tree = ast.parse(handle.read(), filename=file_path)
278 except (OSError, SyntaxError):
279 return None
281 explicit_modules = _explicit_import_modules(tree, class_name)
282 if len(explicit_modules) > 1:
283 raise ValueError(
284 f"snapshot class {class_name!r} is imported from multiple modules: "
285 f"{sorted(explicit_modules)}"
286 )
287 if explicit_modules:
288 return f"{next(iter(explicit_modules))}:{class_name}"
289 candidates = _exporting_import_modules(tree, class_name)
290 if len(candidates) > 1:
291 raise ValueError(
292 f"snapshot class {class_name!r} is exported by multiple imported modules: "
293 f"{candidates}"
294 )
295 return candidates[0] if candidates else None
298def _explicit_import_modules(tree: ast.AST, class_name: str) -> set[str]:
299 """Return modules that explicitly import one class name."""
300 return {
301 node.module
302 for node in tree.body
303 if isinstance(node, ast.ImportFrom)
304 and node.level == 0
305 and node.module
306 and any(
307 alias.name == class_name and (alias.asname or alias.name) == class_name
308 for alias in node.names
309 )
310 }
313def _exporting_import_modules(tree: ast.AST, class_name: str) -> list[str]:
314 """Return imported modules that export one class name."""
315 modules = {
316 node.module
317 for node in tree.body
318 if isinstance(node, ast.ImportFrom) and node.level == 0 and node.module
319 }
320 candidates = []
321 for module_name in sorted(modules):
322 try:
323 module = importlib.import_module(module_name)
324 except (ImportError, AttributeError):
325 continue
326 if hasattr(module, class_name):
327 candidates.append(f"{module_name}:{class_name}")
328 return candidates
331def _ast_name(node: ast.AST) -> Optional[str]:
332 """Return the terminal name of a callable expression."""
333 if isinstance(node, ast.Name):
334 return node.id
335 if isinstance(node, ast.Attribute):
336 return node.attr
337 return None
340__all__ = ["emission_patches_from_snapshot_payload"]