Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / inline / expansion.py: 0%

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"""Expand real orchestration source, stopping at framework and kernel APIs.""" 

16 

17import ast 

18from dataclasses import dataclass 

19import inspect 

20import textwrap 

21from types import ModuleType 

22from typing import Any, Callable 

23 

24 

25_PRIMITIVE_PREFIXES = ( 

26 "hyper_parallel.core", 

27 "hyper_parallel.collectives", 

28 "hyper_parallel.components.functional", 

29 "hyper_parallel.components.modules", 

30 "hyper_parallel.distributed.expert_parallel.experts", 

31 "hyper_parallel.distributed.context_parallel.collectives", 

32) 

33 

34 

35@dataclass(frozen=True) 

36class ExpandedSource: 

37 """Source and imports required by one expanded dependency closure.""" 

38 

39 imports: tuple[str, ...] 

40 source: str 

41 

42 

43def expand_function(function: Callable[..., Any]) -> ExpandedSource: 

44 """Copy an orchestration function and recursively resolve its globals. 

45 

46 Constants use their source initializer, so generated cache state never 

47 captures live device tensors from the generation process. Unsupported 

48 globals fail at generation rather than becoming unresolved artifact names. 

49 """ 

50 expansion = _Expansion() 

51 expansion.add_function(function, function.__name__) 

52 return ExpandedSource(tuple(sorted(expansion.imports)), "\n\n".join(expansion.definitions)) 

53 

54 

55class _Expansion: 

56 """Resolve one closure with deterministic order and collision checking.""" 

57 

58 def __init__(self) -> None: 

59 self.imports: set[str] = set() 

60 self.definitions: list[str] = [] 

61 self.bindings: dict[str, Any] = {} 

62 

63 def _claim(self, name: str, value: Any) -> bool: 

64 if name in self.bindings: 

65 if self.bindings[name] is not value: 

66 raise ValueError(f"codegen source expansion: conflicting global {name!r}") 

67 return False 

68 self.bindings[name] = value 

69 return True 

70 

71 def add_function(self, function: Callable[..., Any], name: str) -> None: 

72 """Resolve dependencies before emitting a function definition.""" 

73 if not self._claim(name, function): 

74 return 

75 source = textwrap.dedent(inspect.getsource(function)) 

76 tree = ast.parse(source) 

77 node = tree.body[0] 

78 if not isinstance(node, ast.FunctionDef) or node.decorator_list: 

79 raise ValueError(f"codegen source expansion: unsupported decorated function {name!r}") 

80 global_names = {item.id for item in ast.walk(node) if isinstance(item, ast.Name) 

81 and isinstance(item.ctx, ast.Load)} 

82 local_names = set(function.__code__.co_varnames) 

83 for global_name in sorted(global_names - local_names): 

84 if global_name in function.__globals__: 

85 self.add_global(global_name, function.__globals__[global_name], inspect.getmodule(function)) 

86 node.name = name 

87 self.definitions.append(ast.unparse(node)) 

88 

89 def add_global(self, name: str, value: Any, owner: ModuleType) -> None: 

90 """Emit an import, dependency function, or literal source initializer.""" 

91 if inspect.isfunction(value) and _is_orchestration(value): 

92 self.add_function(value, name) 

93 return 

94 if not self._claim(name, value): 

95 return 

96 if isinstance(value, ModuleType): 

97 self.imports.add(f"import {value.__name__} as {name}") 

98 return 

99 module_name = getattr(value, "__module__", None) 

100 original_name = getattr(value, "__name__", None) 

101 if module_name and original_name: 

102 self.imports.add(f"from {module_name} import {original_name} as {name}") 

103 return 

104 initializer = _literal_initializer(owner, name) 

105 self.definitions.append(f"{name} = {initializer}") 

106 

107 

108def _is_orchestration(function: Callable[..., Any]) -> bool: 

109 module = function.__module__ 

110 return module.startswith("hyper_parallel.") and not any( 

111 module == prefix or module.startswith(prefix + ".") for prefix in _PRIMITIVE_PREFIXES 

112 ) 

113 

114 

115def _literal_initializer(owner: ModuleType, name: str) -> str: 

116 """Return the source text of the literal assigned to ``name`` in ``owner``.""" 

117 for node in ast.parse(inspect.getsource(owner)).body: 

118 if isinstance(node, ast.Assign): 

119 targets, value = node.targets, node.value 

120 elif isinstance(node, ast.AnnAssign): 

121 targets, value = [node.target], node.value 

122 else: 

123 continue 

124 if any(isinstance(target, ast.Name) and target.id == name for target in targets): 

125 try: 

126 ast.literal_eval(value) 

127 except (ValueError, TypeError) as exc: 

128 raise ValueError(f"codegen source expansion: nonliteral global {name!r}") from exc 

129 return ast.unparse(value) 

130 raise ValueError(f"codegen source expansion: no literal initializer for {name!r}")