Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / emit / patches.py: 74%

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"""Emitter-only views reconstructed from frozen snapshot patch records.""" 

16 

17from __future__ import annotations 

18 

19import importlib 

20import keyword 

21from dataclasses import dataclass, field 

22from types import MappingProxyType, SimpleNamespace 

23from typing import Any, Mapping, Optional 

24 

25 

26@dataclass(frozen=True) 

27class ModuleReplacementPatch: 

28 """One already-resolved module replacement.""" 

29 

30 module_fqn: str 

31 source_type: type 

32 replacement_type: type 

33 factory: Any = None 

34 factory_kwargs: Mapping[str, Any] = field( 

35 default_factory=lambda: MappingProxyType({}) 

36 ) 

37 

38 

39@dataclass(frozen=True) 

40class LocalComputePatch: 

41 """One already-resolved local-compute source patch.""" 

42 

43 module_fqn: str 

44 module_type: type 

45 factory: Any 

46 factory_kwargs: Mapping[str, Any] 

47 resolved_compute: Any 

48 region_dispatch: Optional[bool] 

49 original_forward: Any 

50 installed_forward: Any 

51 boundary: Any = None 

52 excluded_subtrees: tuple[str, ...] = () 

53 

54 

55@dataclass(frozen=True) 

56class InnerWrapperPatch: 

57 """One already-resolved inner-wrapper source patch.""" 

58 

59 module_fqn: str 

60 target_fqn: str 

61 wrapper: Any 

62 request: Any 

63 target_type: Optional[type] = None 

64 secondary_requests: tuple[Any, ...] = () 

65 

66 

67@dataclass(frozen=True) 

68class EmissionPatchSet: 

69 """Immutable source-emission view of frozen native observations.""" 

70 

71 module_replacements: tuple[ModuleReplacementPatch, ...] = () 

72 local_computes: tuple[LocalComputePatch, ...] = () 

73 inner_wrappers: tuple[InnerWrapperPatch, ...] = () 

74 

75 

76def emission_patches_from_meta(meta: Any) -> Optional[EmissionPatchSet]: 

77 """Rehydrate emitter views from plain snapshot data.""" 

78 data = getattr(meta, "emission_patches", None) 

79 if not data: 

80 return None 

81 replacements = tuple( 

82 ModuleReplacementPatch( 

83 module_fqn=item["module_fqn"], 

84 source_type=_resolve_symbol(item["source_type"]), 

85 replacement_type=_resolve_symbol(item["replacement_type"]), 

86 ) 

87 for item in data.get("module_replacements", ()) 

88 ) 

89 local_computes = tuple( 

90 LocalComputePatch( 

91 module_fqn=item["module_fqn"], 

92 module_type=_resolve_symbol(item["module_type"]), 

93 factory=_resolve_symbol(item["factory"]), 

94 factory_kwargs=MappingProxyType({}), 

95 resolved_compute=None, 

96 region_dispatch=item.get("region_dispatch"), 

97 original_forward=_resolve_symbol(item["original_forward"]), 

98 installed_forward=None, 

99 ) 

100 for item in data.get("local_computes", ()) 

101 ) 

102 inner_wrappers = tuple( 

103 InnerWrapperPatch( 

104 module_fqn=item["module_fqn"], 

105 target_fqn=item["target_fqn"], 

106 wrapper=_resolve_symbol(item["wrapper"]), 

107 request=SimpleNamespace(target=None), 

108 target_type=_resolve_symbol(item["target_type"]), 

109 ) 

110 for item in data.get("inner_wrappers", ()) 

111 ) 

112 return EmissionPatchSet(replacements, local_computes, inner_wrappers) 

113 

114 

115def local_compute_binding_name(target: Any) -> str: 

116 """Return the explicit instance attribute used for one compute factory.""" 

117 if isinstance(target, Mapping): 

118 target = target.get("_target_") or target.get("target") or target.get("name") 

119 else: 

120 target = getattr(target, "_target_", target) 

121 if isinstance(target, str): 

122 name = target.rsplit(":", 1)[-1].rsplit(".", 1)[-1] 

123 else: 

124 name = getattr(target, "__name__", "") 

125 if not name or not name.isidentifier() or keyword.iskeyword(name): 

126 raise ValueError( 

127 f"codegen: local-compute target {target!r} has no valid binding name" 

128 ) 

129 return name 

130 

131 

132def _resolve_symbol(path: str) -> Any: 

133 module_name, separator, qualname = path.partition(":") 

134 if not separator: 

135 raise ValueError(f"codegen: invalid patch symbol path {path!r}") 

136 value = importlib.import_module(module_name) 

137 for segment in qualname.split("."): 

138 value = getattr(value, segment) 

139 return value 

140 

141 

142__all__ = ["EmissionPatchSet", "emission_patches_from_meta", "local_compute_binding_name"]