Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / spec / project.py: 19%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# Licensed under the Apache License, Version 2.0 

3# ============================================================================ 

4"""Project a resolved ``TrainerConfig`` into a ``CodegenSpec``. 

5 

6The projection is deliberately narrow: it only extracts the fields that can 

7change the generated modeling artifact, so unrelated YAML edits do not 

8invalidate a bundle signature. It is NOT a second user-facing config source — 

9everything here derives from the training YAML. 

10""" 

11from __future__ import annotations 

12 

13import os 

14from typing import Any, Optional 

15 

16from hyper_parallel.codegen.artifact import ( 

17 ArtifactLayout, 

18 resolve_artifact_layout, 

19) 

20from hyper_parallel.codegen.spec.types import ( 

21 CodegenSpec, 

22 ParallelDimsSpec, 

23 SourceSpec, 

24 TargetSpec, 

25 TransformSpec, 

26) 

27 

28# Parallel dims that the generated artifact depends on. Kept as one list so 

29# the spec projection and the signature cannot drift apart. 

30_PARALLEL_DIM_KEYS = ( 

31 "dp_size", 

32 "tp_size", 

33 "cp_size", 

34 "ep_size", 

35 "pp_size", 

36 "sequence_parallel", 

37 "loss_parallel", 

38) 

39 

40 

41def project_codegen_spec( 

42 config: Any, 

43 yaml_path: str | None = None, 

44 *, 

45 layout: ArtifactLayout | None = None, 

46) -> CodegenSpec: 

47 """Project a resolved TrainerConfig into the codegen input spec. 

48 

49 ``config`` may be a ``TrainerConfig`` (or any object exposing the same 

50 attributes) so the manager can call this without importing trainer. 

51 """ 

52 if layout is None: 

53 if yaml_path is None: 

54 raise ValueError("project_codegen_spec requires yaml_path or layout") 

55 layout = resolve_artifact_layout(yaml_path, model_name=_model_name(config)) 

56 source_spec = project_source_spec(config) 

57 return CodegenSpec( 

58 source=source_spec, 

59 parallel_dims=project_parallel_dims(config), 

60 overrides=project_overrides(config), 

61 transform_specs=project_transform_specs(config), 

62 target=TargetSpec( 

63 enabled=getattr(config, "codegen", False), 

64 artifact_dir=layout.artifact_dir, 

65 output_modeling_name=os.path.basename(layout.modeling_path), 

66 ), 

67 ) 

68 

69 

70def project_source_spec(config: Any) -> SourceSpec: 

71 """Project the model identity from the model target / config. 

72 

73 The modeling source file is resolved later by ``source.resolver``; here 

74 we only record the identity the source resolver will key on. The model 

75 Target holds its id under the key the caller actually used — most YAMLs and 

76 ``HyperAutoModel.from_pretrained`` use ``pretrained_model_name_or_path`` — 

77 so read that first and fall back to ``model_name_or_path``. 

78 """ 

79 model = getattr(config, "model", None) 

80 model_name_or_path = getattr(model, "pretrained_model_name_or_path", None) 

81 if model_name_or_path is None: 

82 model_name_or_path = getattr(model, "model_name_or_path", None) 

83 if model_name_or_path is None: 

84 # Target configured without a positional model id — the HF path will 

85 # derive it from from_pretrained's first arg at load time. 

86 model_name_or_path = "" 

87 return SourceSpec( 

88 model_name_or_path=model_name_or_path, 

89 architecture=None, 

90 model_type=None, 

91 sha256=None, 

92 transformers_version=None, 

93 ) 

94 

95 

96def project_parallel_dims(config: Any) -> ParallelDimsSpec: 

97 """Project ``accelerator.tp_size/cp_size/ep_size/pp_size`` etc.""" 

98 accelerator = getattr(config, "accelerator", None) 

99 if accelerator is None: 

100 return ParallelDimsSpec() 

101 kwargs: dict[str, Any] = {} 

102 for key in _PARALLEL_DIM_KEYS: 

103 if key == "dp_size": 

104 kwargs[key] = _project_dp_size(accelerator) 

105 else: 

106 kwargs[key] = getattr(accelerator, key, ParallelDimsSpec().__dataclass_fields__[key].default) 

107 return ParallelDimsSpec(**kwargs) 

108 

109 

110def project_overrides(config: Any) -> list[dict[str, Any]]: 

111 """Project ``plan_overrides`` to stable canonical dicts. 

112 

113 The canonicalization lives here: keep the declarative fields, resolve 

114 ``Target`` values to their ``to_dict()`` form, and drop ``None`` fields so 

115 the canonical form is stable under YAML default omissions. 

116 """ 

117 entries = getattr(config, "plan_overrides", None) or [] 

118 return [_canonicalize_override(entry) for entry in entries] 

119 

120 

121def project_transform_specs(config: Any) -> TransformSpec: 

122 """Project codegen-affecting model-rewrite config. 

123 

124 TODO: Populate transform specs when operator, weight, and checkpoint 

125 transforms are supported. 

126 """ 

127 model = getattr(config, "model", None) 

128 config_overrides = getattr(model, "config_overrides", None) 

129 extra = {} 

130 if config_overrides: 

131 extra["config_overrides"] = _canonicalize_value(config_overrides) 

132 return TransformSpec(extra_codegen_options=extra) 

133 

134 

135def _model_name(config: Any) -> Optional[str]: 

136 """Best-effort model name for artifact file naming.""" 

137 model = getattr(config, "model", None) 

138 path = getattr(model, "_target_", None) 

139 if isinstance(path, str): 

140 name = path.rsplit(".", 1)[-1] 

141 if name.startswith(("AutoModel", "Pretrained", "from_")): 

142 return None 

143 return name 

144 return getattr(config, "model_name", None) 

145 

146 

147def _project_dp_size(accelerator: Any) -> int: 

148 """Derive runtime DP size from ``WORLD_SIZE`` and non-DP topology.""" 

149 explicit = getattr(accelerator, "dp_size", None) 

150 if explicit: 

151 return max(1, int(explicit)) 

152 world_size = int(os.environ.get("WORLD_SIZE") or os.environ.get("OMPI_COMM_WORLD_SIZE") or 1) 

153 tp_size = max(1, int(getattr(accelerator, "tp_size", 1))) 

154 cp_size = max(1, int(getattr(accelerator, "cp_size", 1))) 

155 pp_size = max(1, int(getattr(accelerator, "pp_size", 1))) 

156 non_dp_size = tp_size * cp_size * pp_size 

157 if world_size % non_dp_size != 0: 

158 raise ValueError( 

159 f"world_size {world_size} is not divisible by non-DP size {non_dp_size}" 

160 ) 

161 return max(1, world_size // non_dp_size) 

162 

163 

164def _canonicalize_override(entry: Any) -> dict[str, Any]: 

165 """Stable dict form of one PlanOverride entry. 

166 

167 Resolves Target values instead of stashing raw ``repr()``. 

168 ``replace_module`` is projected so changes invalidate the signature. 

169 """ 

170 result: dict[str, Any] = {"match": entry.match} 

171 for attr in ( 

172 "when", 

173 "module_type", 

174 "exact_type", 

175 "inner_target", 

176 "region_dispatch", 

177 "tp_divide_attrs", 

178 ): 

179 value = getattr(entry, attr, None) 

180 if value is not None: 

181 result[attr] = value 

182 for attr in ("inner_wrapper", "replace_module", "local_compute_fn"): 

183 value = getattr(entry, attr, None) 

184 if value is not None: 

185 result[attr] = _canonicalize_value(value) 

186 for attr in ("params", "in_src", "in_dst", "out_src", "out_dst"): 

187 value = getattr(entry, attr, None) 

188 if value is not None: 

189 result[attr] = _canonicalize_value(value) 

190 return result 

191 

192 

193def _canonicalize_value(value: Any) -> Any: 

194 """Resolve Target objects to dict form; recurse into containers.""" 

195 if hasattr(value, "to_dict"): 

196 return value.to_dict() 

197 if isinstance(value, dict): 

198 return {key: _canonicalize_value(item) for key, item in value.items()} 

199 if isinstance(value, (list, tuple)): 

200 return [_canonicalize_value(item) for item in value] 

201 return value 

202 

203 

204__all__ = [ 

205 "project_codegen_spec", 

206 "project_overrides", 

207 "project_parallel_dims", 

208 "project_source_spec", 

209 "project_transform_specs", 

210]