Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / family.py: 95%

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"""Construction-time, allowlisted module selection without device dependencies.""" 

16 

17from __future__ import annotations 

18 

19from dataclasses import dataclass, field, replace 

20from importlib import import_module 

21from inspect import Parameter, signature 

22from math import isfinite 

23from types import MappingProxyType 

24from typing import Any, Callable, Mapping, Optional 

25 

26 

27@dataclass(frozen=True) 

28class SelectorSpec: 

29 """Exact scalar type and baseline default for one selection parameter.""" 

30 

31 value_type: type 

32 default: Any 

33 aliases: tuple[str, ...] = () 

34 

35 

36@dataclass(frozen=True) 

37class VariantRule: 

38 """A complete named selector combination and its declared variant name.""" 

39 

40 when: Mapping[str, Any] 

41 variant: str 

42 

43 

44@dataclass(frozen=True) 

45class ModuleFamily: 

46 """One baseline and explicitly allowed construction-time alternatives.""" 

47 

48 name: str 

49 baseline: Callable 

50 variants: Mapping[str, str] 

51 selector_specs: Mapping[str, SelectorSpec] = field(default_factory=dict) 

52 variant_rules: tuple[VariantRule, ...] = () 

53 default_variant: str | None = None 

54 parallel_axis: str | None = None 

55 sharding: Callable | None = None 

56 validate: Callable | None = None 

57 

58 

59# Keep the baseline in the value so its identity cannot be reused while registered. 

60_FAMILIES: dict[int, ModuleFamily] = {} 

61_SELECTOR_TYPES = (bool, int, float, str) 

62_RESERVED_SELECTORS = frozenset({"_target_", "target_path", "module", "module_fqn", "context", "mesh"}) 

63 

64 

65def _validate_selector_value(family: ModuleFamily, name: str, value: Any) -> None: 

66 """Validate exact types; non-finite floats cannot define stable rules.""" 

67 expected = family.selector_specs[name].value_type 

68 actual_type = type(value) 

69 if actual_type is not expected: # bool must not also match int. 

70 raise TypeError( 

71 f"ModuleFamily {family.name}: selector {name!r} must be {expected.__name__}, " 

72 f"got {type(value).__name__}: {value!r}" 

73 ) 

74 if expected is float and not isfinite(value): 

75 raise ValueError(f"ModuleFamily {family.name}: selector {name!r} must be finite") 

76 

77 

78def _validate_selector_specs(family: ModuleFamily) -> dict[str, Any]: 

79 """Validate scalar declarations and return the baseline selector values.""" 

80 defaults = {} 

81 names = set(family.selector_specs) 

82 for name, spec in family.selector_specs.items(): 

83 if not isinstance(name, str) or not name.isidentifier() or name in _RESERVED_SELECTORS: 

84 raise ValueError(f"ModuleFamily {family.name}: invalid selector name {name!r}") 

85 if not isinstance(spec, SelectorSpec) or spec.value_type not in _SELECTOR_TYPES: 

86 raise TypeError(f"ModuleFamily {family.name}: {name!r} requires a scalar SelectorSpec") 

87 _validate_selector_value(family, name, spec.default) 

88 for alias in spec.aliases: 

89 if (not isinstance(alias, str) or not alias.isidentifier() 

90 or alias in names or alias in _RESERVED_SELECTORS): 

91 raise ValueError(f"ModuleFamily {family.name}: invalid selector alias {alias!r}") 

92 names.add(alias) 

93 defaults[name] = spec.default 

94 return defaults 

95 

96 

97def _validate_variant_paths(family: ModuleFamily) -> None: 

98 """Check lazy import declarations without loading device implementations.""" 

99 for name, path in family.variants.items(): 

100 if not isinstance(name, str) or not name.strip(): 

101 raise ValueError(f"ModuleFamily {family.name}: variant names must be non-empty strings") 

102 if not isinstance(path, str) or "." not in path: 

103 raise ValueError(f"ModuleFamily {family.name}: invalid variant import path {path!r}") 

104 if not all(part.isidentifier() for part in path.split(".")): 

105 raise ValueError(f"ModuleFamily {family.name}: invalid variant import path {path!r}") 

106 

107 

108def _validate_family_rules(family: ModuleFamily) -> None: 

109 """Check complete, unique whitelist entries before registration.""" 

110 defaults = _validate_selector_specs(family) 

111 _validate_variant_paths(family) 

112 if family.parallel_axis not in (None, "cp", "ep"): 

113 raise ValueError("parallel_axis must be cp, ep or None") 

114 if family.default_variant is not None: 

115 if family.default_variant not in family.variants: 

116 raise ValueError("default_variant must name a declared concrete variant") 

117 if "parallel" in family.selector_specs: 

118 raise ValueError("parallel is derived from context, not a public selector") 

119 defaults = {**defaults, "parallel": False} 

120 conditions = [defaults] # The default combination is an explicit baseline allowance. 

121 for rule in family.variant_rules: 

122 if not isinstance(rule, VariantRule) or not isinstance(rule.when, Mapping): 

123 raise TypeError(f"ModuleFamily {family.name}: expected VariantRule with a mapping condition") 

124 if set(rule.when) != set(defaults): 

125 raise ValueError(f"ModuleFamily {family.name}: each rule must contain exactly {sorted(defaults)}") 

126 for name, value in rule.when.items(): 

127 if name == "parallel" and family.default_variant is not None: 

128 actual_type = type(value) 

129 if actual_type is not bool: 

130 raise TypeError("derived parallel conditions must be bool") 

131 else: 

132 _validate_selector_value(family, name, value) 

133 if rule.variant not in family.variants: 

134 raise ValueError(f"ModuleFamily {family.name}: undeclared variant {rule.variant!r}") 

135 condition = dict(rule.when) 

136 if condition in conditions: 

137 raise ValueError(f"ModuleFamily {family.name}: duplicate or baseline-default rule {condition!r}") 

138 conditions.append(condition) 

139 

140 

141def register_module_family(family: ModuleFamily) -> None: 

142 """Register a checked snapshot for an exact baseline callable. 

143 

144 Args: 

145 family: Scalar selector schema, complete rules and lazy variant paths. 

146 """ 

147 if not isinstance(family.name, str) or not family.name.strip(): 

148 raise ValueError("ModuleFamily name must be a non-empty string") 

149 if not callable(family.baseline): 

150 raise TypeError(f"ModuleFamily {family.name}: baseline must be callable") 

151 if id(family.baseline) in _FAMILIES: 

152 raise ValueError(f"Module family already registered: {family.name}") 

153 _validate_family_rules(family) 

154 snapshot = replace( 

155 family, 

156 variants=MappingProxyType(dict(family.variants)), 

157 selector_specs=MappingProxyType(dict(family.selector_specs)), 

158 variant_rules=tuple( 

159 VariantRule(MappingProxyType(dict(rule.when)), rule.variant) for rule in family.variant_rules 

160 ), 

161 ) 

162 _FAMILIES[id(family.baseline)] = snapshot 

163 

164 

165def get_module_family(target: Callable) -> Optional[ModuleFamily]: 

166 """Return the declaration for an exact baseline object. 

167 

168 Args: 

169 target: Requested callable; class subclasses require their own registration. 

170 

171 Returns: 

172 Registered declaration, or None for an ordinary callable. 

173 """ 

174 return _FAMILIES.get(id(target)) 

175 

176 

177def get_family_selector_specs(target: Callable) -> Mapping[str, SelectorSpec]: 

178 """Return the selector schema for the requested baseline. 

179 

180 Args: 

181 target: Requested replacement callable. 

182 

183 Returns: 

184 Read-only declarations including aliases, or an empty mapping. 

185 """ 

186 family = get_module_family(target) 

187 if family is None: 

188 return {} 

189 return MappingProxyType({ 

190 spelling: spec for name, spec in family.selector_specs.items() 

191 for spelling in (name, *spec.aliases) 

192 }) 

193 

194 

195def normalize_family_selectors(target: Callable, arguments: Mapping) -> dict: 

196 """Replace declared aliases with canonical selector names. 

197 

198 Args: 

199 target: Requested replacement callable. 

200 arguments: Explicit selectors and constructor arguments. 

201 

202 Returns: 

203 New arguments with aliases consumed; conflicting spellings are rejected. 

204 """ 

205 args = dict(arguments) 

206 family = get_module_family(target) 

207 if family is None: 

208 return args 

209 for name, spec in family.selector_specs.items(): 

210 for alias in spec.aliases: 

211 if alias not in args: 

212 continue 

213 value = args.pop(alias) 

214 _validate_selector_value(family, name, value) 

215 if name in args and (type(args[name]) is not type(value) or args[name] != value): 

216 raise ValueError(f"ModuleFamily {family.name}: conflicting selectors {name!r} and {alias!r}") 

217 args[name] = value 

218 return args 

219 

220 

221def _select_variant( 

222 family: ModuleFamily, args: dict, *, parallel: bool = False, 

223) -> tuple[Callable, str | None, dict[str, Any]]: 

224 """Consume target selectors and resolve one concrete constructor for both entry paths.""" 

225 selectors = {name: args.pop(name, spec.default) for name, spec in family.selector_specs.items()} 

226 for name, value in selectors.items(): 

227 _validate_selector_value(family, name, value) 

228 defaults = {name: spec.default for name, spec in family.selector_specs.items()} 

229 if family.default_variant is not None: 

230 selectors["parallel"] = parallel 

231 defaults["parallel"] = False 

232 if selectors == defaults: 

233 variant = family.default_variant 

234 else: 

235 matches = [rule.variant for rule in family.variant_rules if selectors == dict(rule.when)] 

236 if len(matches) != 1: 

237 legal = [defaults, *(dict(rule.when) for rule in family.variant_rules)] 

238 raise ValueError( 

239 f"ModuleFamily {family.name}: unsupported selector combination {selectors}; legal combinations: {legal}" 

240 ) 

241 variant = matches[0] 

242 constructor = family.baseline 

243 if variant is not None: 

244 module_path, _, name = family.variants[variant].rpartition(".") 

245 constructor = getattr(import_module(module_path), name) 

246 return constructor, variant, selectors 

247 

248 

249def resolve_module_family_target(target: Callable, arguments: Mapping) -> tuple[Callable, dict]: 

250 """Select a registered implementation and validate its ordinary keywords. 

251 

252 Args: 

253 target: Registered baseline or an ordinary replacement callable. 

254 arguments: Explicit selector values and constructor arguments. 

255 

256 Returns: 

257 Selected callable and a new argument dictionary without selectors. 

258 Ordinary targets retain their existing argument-validation policy. 

259 """ 

260 args = normalize_family_selectors(target, arguments) 

261 family = get_module_family(target) 

262 if family is None: 

263 return target, args 

264 if family.default_variant is not None: 

265 _validate_deferred_arguments(family, args) 

266 return target, args 

267 target, _, _ = _select_variant(family, args) 

268 target_signature = signature(target) 

269 allowed = { 

270 name 

271 for name, parameter in target_signature.parameters.items() 

272 if parameter.kind in (Parameter.POSITIONAL_OR_KEYWORD, Parameter.KEYWORD_ONLY) 

273 } 

274 unknown = set(args) - allowed 

275 if unknown: 

276 raise TypeError(f"ModuleFamily {family.name}: unknown replacement arguments: {sorted(unknown)}") 

277 target_signature.bind_partial(**args) 

278 return target, args 

279 

280 

281def _validate_deferred_arguments(family: ModuleFamily, args: Mapping[str, Any]) -> None: 

282 """Validate a new-family request without choosing a topology-dependent member.""" 

283 allowed = set(signature(family.baseline).parameters) - {"module", "module_fqn", "context", "mesh"} 

284 unknown = set(args) - allowed 

285 if unknown: 

286 raise TypeError(f"ModuleFamily {family.name}: unknown replacement arguments: {sorted(unknown)}") 

287 for name in family.selector_specs: 

288 if name in args: 

289 _validate_selector_value(family, name, args[name])