Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / plan / observation.py: 94%

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"""Observe only native-plan facts required for source emission.""" 

16 

17from __future__ import annotations 

18 

19from dataclasses import dataclass, field, fields 

20from typing import Any 

21 

22from hyper_parallel.codegen.plan.placements import ( 

23 named_placement_to_dict, 

24 value_to_data, 

25) 

26 

27 

28@dataclass 

29class EmissionPlanObservation: 

30 """Plain-data projection consumed by Codegen emitters.""" 

31 

32 param_plan: dict[str, dict[str, Any]] = field(default_factory=dict) 

33 injections: list[dict[str, Any]] = field(default_factory=list) 

34 boundary_classes: dict[str, str] = field(default_factory=dict) 

35 boundary_types: dict[str, str] = field(default_factory=dict) 

36 nested_components: list[dict[str, str]] = field(default_factory=list) 

37 mesh_dim_names: tuple[str, ...] = () 

38 

39 def to_dict(self) -> dict[str, Any]: 

40 """Return the mutable mapping consumed by the session boundary.""" 

41 return {item.name: getattr(self, item.name) for item in fields(self)} 

42 

43 

44def observe_emission_plan(plan: Any, model: Any) -> EmissionPlanObservation: 

45 """Project a native sharding plan without changing or applying it.""" 

46 if plan is None: 

47 return EmissionPlanObservation() 

48 return EmissionPlanObservation( 

49 param_plan=_observe_boundary_plan(plan), 

50 injections=_observe_injections(plan), 

51 boundary_classes=_observe_boundary_classes(plan, model), 

52 boundary_types=_observe_boundary_types(plan, model), 

53 nested_components=_observe_nested_components(plan, model), 

54 mesh_dim_names=tuple(getattr(plan, "mesh_dim_names", ()) or ()), 

55 ) 

56 

57 

58def _observe_boundary_plan(plan: Any) -> dict[str, dict[str, Any]]: 

59 """Collect boundary layout contracts required for source emission.""" 

60 result: dict[str, dict[str, Any]] = {} 

61 for fqn, spec in (getattr(plan, "modules", {}) or {}).items(): 

62 entry: dict[str, Any] = {} 

63 for name in ("in_src", "in_dst", "out_src", "out_dst"): 

64 value = getattr(spec, name, None) 

65 if value is not None: 

66 entry[name] = named_placement_to_dict(value) 

67 for name in ("out_names", "_deferred_bias_params"): 

68 value = getattr(spec, name, None) 

69 if value: 

70 key = "deferred_bias_params" if name.startswith("_") else name 

71 entry[key] = list(value) 

72 for name in ("is_boundary", "region_dispatch"): 

73 value = getattr(spec, name, None) 

74 if value is not None: 

75 entry[name] = value 

76 result[fqn] = entry 

77 return result 

78 

79 

80def _observe_injections(plan: Any) -> list[dict[str, Any]]: 

81 """Collect compute and wrapper injections in stable module-path order.""" 

82 result = [] 

83 for fqn, spec in sorted((getattr(plan, "modules", {}) or {}).items()): 

84 entry: dict[str, Any] = {} 

85 for name in ("inner_wrapper", "inner_out_src", "local_compute_fn"): 

86 value = getattr(spec, name, None) 

87 if value is not None: 

88 entry[name] = value_to_data(value) 

89 inner_target = getattr(spec, "inner_target", None) 

90 if inner_target is not None: 

91 entry["inner_target"] = inner_target 

92 if entry: 

93 result.append({"match": fqn, **entry}) 

94 return result 

95 

96 

97def _observe_boundary_classes(plan: Any, model: Any) -> dict[str, str]: 

98 """Map planned module paths to their source class names.""" 

99 if model is None: 

100 return {} 

101 return { 

102 fqn: type(model.get_submodule(fqn)).__name__ 

103 for fqn in (getattr(plan, "modules", {}) or {}) 

104 } 

105 

106 

107def _observe_boundary_types(plan: Any, model: Any) -> dict[str, str]: 

108 """Map planned module paths to fully qualified runtime type identities.""" 

109 if model is None: 

110 return {} 

111 result = {} 

112 for fqn in (getattr(plan, "modules", {}) or {}): 

113 module_type = type(model.get_submodule(fqn)) 

114 result[fqn] = f"{module_type.__module__}:{module_type.__qualname__}" 

115 return result 

116 

117 

118def _observe_nested_components(plan: Any, model: Any) -> list[dict[str, str]]: 

119 """Identify nested model source owners needed for source composition.""" 

120 if model is None: 

121 return [] 

122 root_module = type(model).__module__ 

123 components: dict[tuple[str, str], dict[str, str]] = {} 

124 for boundary_fqn in (getattr(plan, "modules", {}) or {}): 

125 owner_module = type(model.get_submodule(boundary_fqn)).__module__ 

126 if ( 

127 owner_module == root_module 

128 or owner_module.startswith("torch.") 

129 or owner_module.startswith("hyper_parallel.") 

130 ): 

131 continue 

132 candidates = [] 

133 parts = boundary_fqn.split(".") 

134 for index in range(1, len(parts) + 1): 

135 ancestor_fqn = ".".join(parts[:index]) 

136 ancestor_type = type(model.get_submodule(ancestor_fqn)) 

137 if ancestor_type.__module__ == owner_module: 

138 candidates.append((ancestor_fqn, ancestor_type)) 

139 if not candidates: 

140 continue 

141 component_fqn, component_type = candidates[0] 

142 identity = f"{component_type.__module__}:{component_type.__qualname__}" 

143 components[(component_fqn, identity)] = { 

144 "fqn": component_fqn, 

145 "type": identity, 

146 "owner_module": owner_module, 

147 } 

148 return [components[key] for key in sorted(components)] 

149 

150 

151__all__ = ["EmissionPlanObservation", "observe_emission_plan"]