Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / source / composition.py: 90%
107 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"""Compose nested AutoModel sources into one generated modeling unit."""
17from __future__ import annotations
19import ast
20import importlib.util
21import keyword
22import os
23import re
24from typing import Any
27def compose_nested_model_sources(source_text: str, meta: Any) -> str:
28 """Inline source modules that own boundaries below a nested model root.
30 The native model remains unchanged. This is an artifact-only composition:
31 it replaces a nested ``AutoModel.from_config`` construction with the
32 observed concrete class and appends that class's original source module to
33 the same generated file before normal TP/EP lowering runs.
35 Args:
36 source_text: Primary model source.
37 meta: Frozen codegen metadata carrying observed nested components.
39 Returns:
40 One source string containing the primary and nested model definitions.
41 """
42 components = getattr(meta, "nested_components", None) or ()
43 if not components:
44 return source_text
45 rendered = source_text
46 included_modules: set[str] = set()
47 for component in components:
48 identity = component.get("type")
49 if not isinstance(identity, str) or ":" not in identity:
50 raise ValueError(f"invalid nested component identity: {identity!r}")
51 module_name, class_name = identity.split(":", 1)
52 class_name = class_name.rsplit(".", 1)[-1]
53 replaced = _replace_auto_model_construction(
54 rendered,
55 component_fqn=str(component.get("fqn") or ""),
56 class_name=class_name,
57 )
58 if replaced is None:
59 continue
60 rendered = replaced
61 if module_name not in included_modules:
62 nested_source = _read_module_source(module_name)
63 rendered = _append_source_module(rendered, nested_source, module_name)
64 included_modules.add(module_name)
65 return rendered
68def _read_module_source(module_name: str) -> str:
69 """Read an installed module without importing its runtime dependencies."""
70 spec = importlib.util.find_spec(module_name)
71 path = None if spec is None else spec.origin
72 if not path or not os.path.isfile(path):
73 raise FileNotFoundError(
74 f"codegen: nested model source {module_name!r} is not readable"
75 )
76 with open(path, "r", encoding="utf-8", errors="replace") as handle:
77 return handle.read()
80def _append_source_module(primary: str, nested: str, module_name: str) -> str:
81 """Append one nested module after removing module-only syntax."""
82 tree = ast.parse(nested, filename=module_name)
83 package = module_name.rpartition(".")[0]
84 tree = _AbsoluteImports(package).visit(tree)
85 ast.fix_missing_locations(tree)
86 primary_tree = ast.parse(primary)
87 primary_bindings = _top_level_definitions(primary_tree)
88 nested_bindings = _top_level_definitions(tree)
89 collisions = {
90 name for name in primary_bindings.keys() & nested_bindings.keys()
91 if not isinstance(nested_bindings[name], ast.ClassDef)
92 }
93 if collisions:
94 replacements = {name: f"_hp_nested_{name}" for name in collisions}
95 for node in tree.body:
96 if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
97 node.name = replacements.get(node.name, node.name)
98 tree = _RenameBindings(replacements).visit(tree)
99 ast.fix_missing_locations(tree)
100 nodes = [
101 node for node in tree.body
102 if not isinstance(node, ast.Expr) or not _is_docstring(node)
103 ]
104 nodes = [
105 node for node in nodes
106 if not isinstance(node, ast.ImportFrom) or node.module != "__future__"
107 ]
108 body = "\n\n".join(ast.unparse(node) for node in nodes)
109 marker = (
110 "# [HYPER NESTED MODEL] source merged from "
111 f"{module_name}; TP/EP lowering below targets these concrete classes."
112 )
113 return f"{primary.rstrip()}\n\n\n{marker}\n{body}\n"
116def _top_level_definitions(tree: ast.Module) -> dict[str, ast.AST]:
117 """Return named functions, classes, and simple assignments."""
118 result: dict[str, ast.AST] = {}
119 for node in tree.body:
120 if isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):
121 result[node.name] = node
122 elif isinstance(node, (ast.Assign, ast.AnnAssign)):
123 targets = node.targets if isinstance(node, ast.Assign) else (node.target,)
124 for target in targets:
125 if isinstance(target, ast.Name):
126 result[target.id] = node
127 return result
130def _is_docstring(node: ast.Expr) -> bool:
131 return isinstance(node.value, ast.Constant) and isinstance(node.value.value, str)
134class _RenameBindings(ast.NodeTransformer):
135 """Rename colliding nested-module globals and their references."""
137 def __init__(self, names: dict[str, str]) -> None:
138 """Create a renamer for the supplied module-global bindings."""
139 self._names = names
141 def visit_Name(self, node: ast.Name) -> ast.AST: # pylint: disable=invalid-name
142 """Rename a matching load or store while preserving its AST context."""
143 replacement = self._names.get(node.id)
144 if replacement is None:
145 return node
146 return ast.copy_location(ast.Name(id=replacement, ctx=node.ctx), node)
149class _AbsoluteImports(ast.NodeTransformer):
150 """Anchor nested-module relative imports on their original package."""
152 def __init__(self, package: str) -> None:
153 """Create a transformer for one installed source package."""
154 self._package = package
156 def visit_ImportFrom( # pylint: disable=invalid-name
157 self,
158 node: ast.ImportFrom,
159 ) -> ast.AST | list[ast.AST]:
160 """Convert one relative import into its installed absolute module."""
161 if node.level == 0:
162 return node
163 relative = "." * node.level + (node.module or "")
164 module = importlib.util.resolve_name(relative, self._package)
165 if not all(part.isidentifier() and not keyword.iskeyword(part) for part in module.split(".")):
166 return self._dynamic_imports(node, module)
167 return ast.copy_location(
168 ast.ImportFrom(module=module, names=node.names, level=0),
169 node,
170 )
172 @staticmethod
173 def _dynamic_imports(node: ast.ImportFrom, module: str) -> list[ast.AST]:
174 """Render imports whose dynamic package contains non-Python identifiers."""
175 assignments = []
176 for alias in node.names:
177 if alias.name == "*":
178 raise ValueError(
179 f"codegen cannot compose a wildcard import from dynamic module {module!r}"
180 )
181 imported = ast.Call(
182 func=ast.Name(id="__import__", ctx=ast.Load()),
183 args=[ast.Constant(module)],
184 keywords=[ast.keyword(arg="fromlist", value=ast.List(
185 elts=[ast.Constant(alias.name)],
186 ctx=ast.Load(),
187 ))],
188 )
189 assignment = ast.Assign(
190 targets=[ast.Name(id=alias.asname or alias.name, ctx=ast.Store())],
191 value=ast.Attribute(value=imported, attr=alias.name, ctx=ast.Load()),
192 )
193 assignments.append(ast.copy_location(assignment, node))
194 return assignments
197def _replace_auto_model_construction(
198 source_text: str,
199 *,
200 component_fqn: str,
201 class_name: str,
202) -> str | None:
203 """Replace an observed AutoModel construction, or reject the candidate."""
204 attribute = component_fqn.rsplit(".", 1)[-1]
205 if not attribute:
206 raise ValueError("nested component FQN must not be empty")
207 pattern = re.compile(
208 rf"(?m)^(?P<indent>\s*)self\.{re.escape(attribute)}\s*=\s*"
209 r"AutoModel\.from_config\((?P<config>[^\n]+)\)\s*$"
210 )
211 replacement = rf"\g<indent>self.{attribute} = {class_name}(\g<config>)"
212 rendered, count = pattern.subn(replacement, source_text, count=1)
213 return rendered if count == 1 else None
216__all__ = ["compose_nested_model_sources"]