Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / inline / recognition.py: 32%
84 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"""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
34import inspect
35import textwrap
36from typing import Any
39@dataclass(frozen=True)
40class ReplacementArgument:
41 """One constructor keyword a replacement factory passes.
43 ``value`` is the argument's source text; ``None`` marks an argument bound to
44 the factory's own parameter (a runtime value such as ``module_fqn``), whose
45 generated value is the component's constructor default.
46 """
48 name: str
49 value: str | None = None
52@dataclass(frozen=True)
53class ReplacementTarget:
54 """The component a replacement factory constructs, read from its source.
56 Every fact the generated replacement needs is here: ``component`` (the class
57 the factory builds) and ``module`` (where that class is imported from) give
58 the emitted import; ``wraps_source`` says whether the matched module is
59 handed over as the component's ``module`` argument (a wrap) or whether the
60 component drops in for the source class; ``arguments`` carries the remaining
61 constructor keywords; ``interface`` names the family symbol the factory
62 hands over -- a component that receives one is rendered rather than merely
63 imported, so its class and that interface stay visible in the artifact.
64 """
66 component: str
67 module: str
68 wraps_source: bool
69 arguments: tuple[ReplacementArgument, ...]
70 interface: str | None = None
73def class_name_of(qualified_type: str | None) -> str | None:
74 """Return the bare class name of a dotted ``module_type`` path."""
75 if not qualified_type:
76 return None
77 return qualified_type.rsplit(".", 1)[-1] or None
80def family_from_target(target: str | None) -> str | None:
81 """Infer the adapter family directory from a framework adapter path.
83 ``hyper_parallel.models.<family>.adapter...`` -> ``<family>``. Returns
84 ``None`` for targets that are not model-adapter paths (for example a
85 framework recipe), which callers treat as "no family provider".
86 """
87 parts = (target or "").split(".")
88 if len(parts) >= 3 and parts[:2] == ["hyper_parallel", "models"]:
89 return parts[2]
90 return None
93def adapter_path_parts(target: str | None) -> bool:
94 """Whether a target lives under the ``models/<family>/adapter`` convention."""
95 parts = (target or "").split(".")
96 return len(parts) >= 4 and parts[:2] == ["hyper_parallel", "models"] and parts[3] == "adapter"
99def replacement_target(factory: Any) -> ReplacementTarget | None:
100 """Read the component one ``@module_replacement`` factory constructs.
102 Only a factory that constructs an imported class directly (``Ctor(...)``)
103 yields a target: that is the call codegen can sink into the generated
104 ``__init__``. A factory that builds through a class method
105 (``SomeExperts.from_module(...)``) has no constructor call to sink, so the
106 rule is left untouched -- the same conservative behaviour a component-name
107 table gave, without the table.
108 """
109 try:
110 source = textwrap.dedent(inspect.getsource(factory))
111 tree = ast.parse(source)
112 except (OSError, TypeError, SyntaxError):
113 return None
114 imported = _factory_imported_modules(tree)
115 parameters = _parameter_names(factory)
116 for node in ast.walk(tree):
117 if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Name):
118 continue
119 component = node.func.id
120 module = imported.get(component) or _global_class_module(factory, component)
121 if module is None:
122 continue
123 return ReplacementTarget(
124 component=component,
125 module=module,
126 wraps_source=any(
127 keyword.arg == "module"
128 and isinstance(keyword.value, ast.Name)
129 and keyword.value.id in parameters
130 for keyword in node.keywords
131 ),
132 arguments=_factory_arguments(node, parameters, factory),
133 interface=next(
134 (
135 keyword.value.id
136 for keyword in node.keywords
137 if isinstance(keyword.value, ast.Name)
138 and keyword.value.id in getattr(factory, "__globals__", {})
139 and callable(factory.__globals__[keyword.value.id])
140 ),
141 None,
142 ),
143 )
144 return None
147def _factory_imported_modules(tree: ast.AST) -> dict[str, str]:
148 """Module path per name the factory source imports (``from X import Y``)."""
149 imported: dict[str, str] = {}
150 for node in ast.walk(tree):
151 if isinstance(node, ast.ImportFrom) and node.module:
152 for alias in node.names:
153 imported.setdefault(alias.asname or alias.name, node.module)
154 elif isinstance(node, ast.Import):
155 for alias in node.names:
156 imported.setdefault(alias.asname or alias.name.split(".")[0], alias.name)
157 return imported
160def _global_class_module(factory: Any, name: str) -> str | None:
161 """Module path of a class the factory's own module imported at module level."""
162 value = getattr(factory, "__globals__", {}).get(name)
163 module = getattr(value, "__module__", None)
164 return module if isinstance(value, type) and module else None
167def _parameter_names(factory: Any) -> set[str]:
168 """Parameter names of the factory itself (runtime values, not literals)."""
169 try:
170 return set(inspect.signature(factory).parameters)
171 except (TypeError, ValueError):
172 return set()
175def _factory_arguments(
176 call: ast.Call, parameters: set[str], factory: Any
177) -> tuple[ReplacementArgument, ...]:
178 """The constructor keywords a factory passes, minus the wrapped module.
180 A value that names one of the factory's own parameters is recorded without
181 its text (the component's constructor default applies); a value naming a
182 callable in the factory's globals -- a kernel entry handed over -- keeps its
183 symbol so the artifact stays readable; anything else keeps its own text.
184 """
185 arguments: list[ReplacementArgument] = []
186 for keyword in call.keywords:
187 if keyword.arg is None:
188 continue
189 if keyword.arg == "module" and isinstance(keyword.value, ast.Name):
190 continue
191 value: str | None = (
192 None
193 if isinstance(keyword.value, ast.Name) and keyword.value.id in parameters
194 else ast.unparse(keyword.value)
195 )
196 arguments.append(ReplacementArgument(keyword.arg, value))
197 return tuple(arguments)
200def archetype_for_factory(factory_path: str | None) -> Any | None:
201 """Reverse-lookup a MoE archetype by its compute-factory ``_target_``.
203 A YAML ``local_compute_fn._target_`` that names a known archetype factory
204 resolves to that archetype: the strategy kind / target class / boundary
205 sub-patterns / EP body keys are then structural, not declared per family.
206 """
207 if not factory_path:
208 return None
209 from hyper_parallel.distributed.expert_parallel.archetypes import ( # pylint: disable=C0415
210 moe_archetypes,
211 )
213 for archetype in moe_archetypes().values():
214 if archetype.compute_factory == factory_path:
215 return archetype
216 return None
219__all__ = [
220 "ReplacementArgument",
221 "ReplacementTarget",
222 "adapter_path_parts",
223 "archetype_for_factory",
224 "class_name_of",
225 "family_from_target",
226 "replacement_target",
227]