Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / inline / recognition.py: 70%
249 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"""Framework structural recognition for inline component selection.
17Which generic component replaces a matched source module is decided by the
18*structure of the family's own replacement factory* -- the generic component it
19constructs, the module that class comes from, the constructor arguments it
20passes and the kernel entry it hands over -- never by the model's name and
21never by a component-name table. The factory is not codegen-specific: the YAML
22already names it for the native path, so reading it adds no family codegen
23declaration.
25This is what makes "native HF model + one YAML" sufficient: a family whose
26attention is fused-QKV, whose experts are batched, or whose norm is an RMSNorm
27is covered the moment its runtime factory exists.
28"""
30from __future__ import annotations
32import ast
33from dataclasses import dataclass
34from functools import lru_cache
35import inspect
36import os
37import sys
38import textwrap
39from typing import Any
42@dataclass(frozen=True)
43class ReplacementArgument:
44 """One constructor keyword a replacement factory passes.
46 ``value`` is the argument's source text; ``None`` marks an argument bound to
47 the factory's own parameter (a runtime value such as ``module_fqn``), whose
48 generated value is the component's constructor default.
49 """
51 name: str
52 value: str | None = None
55@dataclass(frozen=True)
56class ReplacementTarget:
57 """The component a replacement factory constructs, read from its source.
59 Every fact the generated replacement needs is here: ``component`` (the class
60 the factory builds) and ``module`` (where that class is imported from) give
61 the emitted import; ``wraps_source`` says whether the matched module is
62 handed over as the component's ``module`` argument (a wrap) or whether the
63 component drops in for the source class; ``arguments`` carries the remaining
64 constructor keywords; ``interface`` names the family symbol the factory
65 hands over -- a component that receives one is rendered rather than merely
66 imported, so its class and that interface stay visible in the artifact.
67 """
69 component: str
70 module: str
71 wraps_source: bool
72 arguments: tuple[ReplacementArgument, ...]
73 interface: str | None = None
76@dataclass(frozen=True)
77class InnerWrapperTarget:
78 """Source-visible helper and companion mesh attribute of an inner wrapper."""
80 helper_module: str | None = None
81 helper_name: str | None = None
82 state_attr: str | None = None
83 inline_body: str | None = None
86def class_name_of(qualified_type: str | None) -> str | None:
87 """Return the bare class name of a dotted ``module_type`` path."""
88 if not qualified_type:
89 return None
90 return qualified_type.rsplit(".", 1)[-1] or None
93def family_from_target(target: str | None) -> str | None:
94 """Infer the adapter family directory from a framework adapter path.
96 ``hyper_parallel.models.<family>.adapter...`` -> ``<family>``. Returns
97 ``None`` for targets that are not model-adapter paths (for example a
98 framework recipe), which callers treat as "no family provider".
99 """
100 parts = (target or "").split(".")
101 if len(parts) >= 3 and parts[:2] == ["hyper_parallel", "models"]:
102 return parts[2]
103 return None
106def adapter_path_parts(target: str | None) -> bool:
107 """Whether a target lives under the ``models/<family>/adapter`` convention."""
108 parts = (target or "").split(".")
109 return len(parts) >= 4 and parts[:2] == ["hyper_parallel", "models"] and parts[3] == "adapter"
112def replacement_target(factory: Any) -> ReplacementTarget | None:
113 """Read the component one ``@module_replacement`` factory constructs.
115 Only a factory that constructs an imported class directly (``Ctor(...)``)
116 yields a target: that is the call codegen can sink into the generated
117 ``__init__``. A factory that builds through a class method
118 (``SomeExperts.from_module(...)``) has no constructor call to sink, so the
119 rule is left untouched -- the same conservative behaviour a component-name
120 table gave, without the table.
121 """
122 try:
123 source = textwrap.dedent(inspect.getsource(factory))
124 tree = ast.parse(source)
125 except (OSError, TypeError, SyntaxError):
126 return None
127 imported = _factory_imported_modules(tree)
128 parameters = _parameter_names(factory)
129 for node in ast.walk(tree):
130 if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Name):
131 continue
132 component = node.func.id
133 module = imported.get(component) or _global_class_module(factory, component)
134 if module is None:
135 continue
136 return ReplacementTarget(
137 component=component,
138 module=module,
139 wraps_source=any(
140 keyword.arg == "module"
141 and isinstance(keyword.value, ast.Name)
142 and keyword.value.id in parameters
143 for keyword in node.keywords
144 ),
145 arguments=_factory_arguments(node, parameters),
146 interface=next(
147 (
148 keyword.value.id
149 for keyword in node.keywords
150 if isinstance(keyword.value, ast.Name)
151 and keyword.value.id in getattr(factory, "__globals__", {})
152 and callable(factory.__globals__[keyword.value.id])
153 ),
154 None,
155 ),
156 )
157 return None
160def _factory_imported_modules(tree: ast.AST) -> dict[str, str]:
161 """Module path per name the factory source imports (``from X import Y``)."""
162 imported: dict[str, str] = {}
163 for node in ast.walk(tree):
164 if isinstance(node, ast.ImportFrom) and node.module:
165 for alias in node.names:
166 imported.setdefault(alias.asname or alias.name, node.module)
167 elif isinstance(node, ast.Import):
168 for alias in node.names:
169 imported.setdefault(alias.asname or alias.name.split(".")[0], alias.name)
170 return imported
173def _global_class_module(factory: Any, name: str) -> str | None:
174 """Module path of a class the factory's own module imported at module level."""
175 value = getattr(factory, "__globals__", {}).get(name)
176 module = getattr(value, "__module__", None)
177 return module if isinstance(value, type) and module else None
180def _parameter_names(factory: Any) -> set[str]:
181 """Parameter names of the factory itself (runtime values, not literals)."""
182 try:
183 return set(inspect.signature(factory).parameters)
184 except (TypeError, ValueError):
185 return set()
188def _factory_arguments(
189 call: ast.Call, parameters: set[str]
190) -> tuple[ReplacementArgument, ...]:
191 """The constructor keywords a factory passes, minus the wrapped module.
193 A value that names one of the factory's own parameters is recorded without
194 its text (the component's constructor default applies); a value naming a
195 callable in the factory's globals -- a kernel entry handed over -- keeps its
196 symbol so the artifact stays readable; anything else keeps its own text.
197 """
198 arguments: list[ReplacementArgument] = []
199 for keyword in call.keywords:
200 if keyword.arg is None:
201 continue
202 if keyword.arg == "module" and isinstance(keyword.value, ast.Name):
203 continue
204 value: str | None = (
205 None
206 if isinstance(keyword.value, ast.Name) and keyword.value.id in parameters
207 else ast.unparse(keyword.value)
208 )
209 arguments.append(ReplacementArgument(keyword.arg, value))
210 return tuple(arguments)
213def inner_wrapper_target(factory_path: str | None) -> InnerWrapperTarget | None:
214 """Recover an inner wrapper's emitted helper contract from its source."""
215 if not factory_path:
216 return None
217 module_name, _, factory_name = factory_path.rpartition(".")
218 if not module_name or not factory_name:
219 return None
220 tree = _module_source_tree(module_name)
221 if tree is None:
222 return None
223 factory_node = _module_function(tree, factory_name)
224 if factory_node is None:
225 return None
226 request_owner = _request_owner(factory_node, tree)
227 replacement_name = _replacement_callable(request_owner)
228 helper_name = (
229 _forward_helper_name(request_owner, replacement_name)
230 if request_owner is not None and replacement_name
231 else None
232 )
233 module_functions = {
234 node.name
235 for node in tree.body
236 if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
237 }
238 if not helper_name or helper_name not in module_functions:
239 inline_body = _nested_forward_body(request_owner, replacement_name)
240 if inline_body is None:
241 return None
242 return InnerWrapperTarget(inline_body=inline_body)
243 return InnerWrapperTarget(
244 helper_module=module_name,
245 helper_name=helper_name,
246 state_attr=_cp_mesh_companion_attr(tree),
247 )
250def _module_function(tree: ast.Module, name: str) -> ast.AST | None:
251 """Return a named top-level function from a parsed module."""
252 return next(
253 (
254 node
255 for node in tree.body
256 if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
257 and node.name == name
258 ),
259 None,
260 )
263def _replacement_callable(request_owner: ast.AST | None) -> str | None:
264 """Resolve the callable installed by an inner-wrapper request."""
265 if request_owner is None:
266 return None
267 return _rewrite_request_callable(request_owner) or _returned_nested_callable(request_owner)
270def _request_owner(factory_node: ast.AST, tree: ast.Module) -> ast.AST | None:
271 """Return the factory or its module-level delegate that builds the request."""
272 if (
273 _rewrite_request_callable(factory_node) is not None
274 or _returned_nested_callable(factory_node) is not None
275 ):
276 return factory_node
277 functions = {
278 node.name: node
279 for node in tree.body
280 if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
281 }
282 for node in ast.walk(factory_node):
283 if not isinstance(node, ast.Return) or not isinstance(node.value, ast.Call):
284 continue
285 if not isinstance(node.value.func, ast.Name):
286 continue
287 delegate = functions.get(node.value.func.id)
288 if delegate is not None and _rewrite_request_callable(delegate) is not None:
289 return delegate
290 return None
293def _returned_nested_callable(tree: ast.AST) -> str | None:
294 """Return a nested callable returned directly by an inner wrapper."""
295 nested = {
296 node.name
297 for node in ast.iter_child_nodes(tree)
298 if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
299 }
300 for node in ast.walk(tree):
301 if isinstance(node, ast.Return) and isinstance(node.value, ast.Name):
302 if node.value.id in nested:
303 return node.value.id
304 return None
307@lru_cache(maxsize=None)
308def _module_source(module_name: str) -> str | None:
309 """Load module source once without importing optional accelerator deps."""
310 relative = os.path.join(*module_name.split(".")) + ".py"
311 for root in sys.path:
312 path = os.path.join(root or os.curdir, relative)
313 if not os.path.isfile(path):
314 continue
315 try:
316 with open(path, "r", encoding="utf-8") as handle:
317 return handle.read()
318 except (OSError, UnicodeError):
319 return None
320 return None
323def _module_source_tree(module_name: str) -> ast.Module | None:
324 """Parse cached module source into an isolated syntax tree."""
325 source = _module_source(module_name)
326 if source is None:
327 return None
328 try:
329 return ast.parse(source)
330 except SyntaxError:
331 return None
334def _rewrite_request_callable(tree: ast.AST) -> str | None:
335 """Return the replacement callable passed to ``_ForwardRewriteRequest``."""
336 for node in ast.walk(tree):
337 if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Name):
338 continue
339 if node.func.id == "_ForwardRewriteRequest" and len(node.args) >= 2:
340 if isinstance(node.args[1], ast.Name):
341 return node.args[1].id
342 return None
345def _forward_helper_name(tree: ast.AST, replacement_name: str) -> str | None:
346 """Return the module-level helper called by the replacement forward."""
347 for node in ast.walk(tree):
348 if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) or node.name != replacement_name:
349 continue
350 for child in ast.walk(node):
351 if isinstance(child, ast.Return) and isinstance(child.value, ast.Call):
352 if isinstance(child.value.func, ast.Name):
353 return child.value.func.id
354 return None
357def _cp_mesh_companion_attr(tree: ast.AST) -> str | None:
358 """Return the companion attribute whose value is the wrapper's CP mesh."""
359 return (
360 _cp_mesh_dict_attr(tree)
361 or _cp_mesh_setattr_attr(tree)
362 or _cp_mesh_assignment_attr(tree)
363 )
366def _cp_mesh_dict_attr(tree: ast.AST) -> str | None:
367 """Return the CP mesh attribute declared in companion_attrs."""
368 for node in ast.walk(tree):
369 if not isinstance(node, ast.keyword) or node.arg != "companion_attrs":
370 continue
371 if not isinstance(node.value, ast.Dict):
372 continue
373 for key, value in zip(node.value.keys, node.value.values):
374 if (
375 isinstance(key, ast.Constant)
376 and isinstance(key.value, str)
377 and isinstance(value, ast.Name)
378 and value.id == "cp_mesh"
379 ):
380 return key.value
381 return None
384def _cp_mesh_setattr_attr(tree: ast.AST) -> str | None:
385 """Return the CP mesh attribute installed through setattr."""
386 for node in ast.walk(tree):
387 if (
388 isinstance(node, ast.Call)
389 and isinstance(node.func, ast.Name)
390 and node.func.id == "setattr"
391 and len(node.args) == 3
392 and isinstance(node.args[1], ast.Constant)
393 and isinstance(node.args[1].value, str)
394 and isinstance(node.args[2], ast.Name)
395 and node.args[2].id == "cp_mesh"
396 ):
397 return node.args[1].value
398 return None
401def _cp_mesh_assignment_attr(tree: ast.AST) -> str | None:
402 """Return the CP mesh attribute installed through assignment."""
403 for node in ast.walk(tree):
404 if not isinstance(node, (ast.Assign, ast.AnnAssign)):
405 continue
406 value = node.value
407 targets = node.targets if isinstance(node, ast.Assign) else (node.target,)
408 if not isinstance(value, ast.Name) or value.id != "cp_mesh":
409 continue
410 for target in targets:
411 if isinstance(target, ast.Attribute) and isinstance(target.value, ast.Name):
412 return target.attr
413 return None
416def _nested_forward_body(tree: ast.AST, replacement_name: str | None) -> str | None:
417 """Lower a closure-backed rewrite into an imported-host dispatch body.
419 A custom wrapper may build its replacement forward as a nested closure
420 instead of delegating to a module-level helper. The closure remains
421 source-visible when its captured values are also published through the
422 request's ``companion_attrs`` mapping. In that form, replace captured
423 state reads with instance attributes and the saved original forward with
424 ``super().forward``; the imported-host materializer supplies the subclass.
425 """
426 if not replacement_name:
427 return None
428 forward = next(
429 (
430 node
431 for node in ast.walk(tree)
432 if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
433 and node.name == replacement_name
434 ),
435 None,
436 )
437 if forward is None:
438 return None
439 companion_attrs = _companion_attr_bindings(tree)
440 original_forward = _original_forward_binding(tree)
441 if original_forward is None:
442 return None
444 class ClosureBindings(ast.NodeTransformer):
445 """Rewrite captured closure bindings to generated-instance state."""
447 def visit_Name(self, node: ast.Name) -> ast.AST: # pylint: disable=invalid-name
448 """Rewrite a referenced symbol into its resolved source form."""
449 attr = companion_attrs.get(node.id)
450 if attr is None:
451 return node
452 return ast.copy_location(
453 ast.Attribute(
454 value=ast.Name(id="self", ctx=ast.Load()),
455 attr=attr,
456 ctx=node.ctx,
457 ),
458 node,
459 )
461 def visit_Call(self, node: ast.Call) -> ast.AST: # pylint: disable=invalid-name
462 """Rewrite a call after recursively resolving its arguments."""
463 node = self.generic_visit(node)
464 if isinstance(node.func, ast.Name) and node.func.id == original_forward:
465 node.func = ast.Attribute(
466 value=ast.Call(
467 func=ast.Name(id="super", ctx=ast.Load()),
468 args=[],
469 keywords=[],
470 ),
471 attr="forward",
472 ctx=ast.Load(),
473 )
474 return node
476 body = list(forward.body)
477 if (
478 body
479 and isinstance(body[0], ast.Expr)
480 and isinstance(body[0].value, ast.Constant)
481 and isinstance(body[0].value.value, str)
482 ):
483 body = body[1:]
484 lowered = [ClosureBindings().visit(ast.fix_missing_locations(node)) for node in body]
485 if not lowered:
486 return None
487 active = "\n".join(ast.unparse(node) for node in lowered)
488 return (
489 "if self.cp_enable:\n"
490 f"{textwrap.indent(active, ' ')}\n\n"
491 "return self._forward_impl(__hp_forward_args__)"
492 )
495def _companion_attr_bindings(tree: ast.AST) -> dict[str, str]:
496 """Map closure variable names to their published companion attributes."""
497 result: dict[str, str] = {}
498 for node in ast.walk(tree):
499 if not isinstance(node, ast.keyword) or node.arg != "companion_attrs":
500 continue
501 if not isinstance(node.value, ast.Dict):
502 continue
503 for key, value in zip(node.value.keys, node.value.values):
504 if (
505 isinstance(key, ast.Constant)
506 and isinstance(key.value, str)
507 and isinstance(value, ast.Name)
508 ):
509 result[value.id] = key.value
510 return result
513def _original_forward_binding(tree: ast.AST) -> str | None:
514 """Return the local name that saves ``target.forward`` before rewriting."""
515 for node in ast.walk(tree):
516 if not isinstance(node, ast.Assign) or len(node.targets) != 1:
517 continue
518 target = node.targets[0]
519 value = node.value
520 if (
521 isinstance(target, ast.Name)
522 and isinstance(value, ast.Attribute)
523 and value.attr == "forward"
524 ):
525 return target.id
526 return None
529__all__ = [
530 "InnerWrapperTarget",
531 "ReplacementArgument",
532 "ReplacementTarget",
533 "adapter_path_parts",
534 "class_name_of",
535 "family_from_target",
536 "inner_wrapper_target",
537 "replacement_target",
538]