Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / util.py: 80%

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"""Common utility functions.""" 

16import dataclasses 

17from collections import defaultdict 

18from collections.abc import Collection, Mapping 

19from pathlib import Path 

20from typing import Any, Union 

21 

22from hyper_parallel.core.distributed_checkpoint.metadata import ( 

23 ChunkStorageMetadata, 

24 MetadataIndex, 

25 CHUNK_INFO, 

26 ChunkInfo 

27) 

28from hyper_parallel.core.distributed_checkpoint.planner import SavePlan, WriteItem 

29from hyper_parallel.core.distributed_checkpoint.reshard import infer_slice_area_by_rank 

30from hyper_parallel.core.dtensor.dtensor import DTensor 

31from hyper_parallel.platform import get_platform 

32 

33 

34platform = get_platform() 

35Tensor = platform.Tensor 

36 

37 

38def check_path(path: Union[Path, str]) -> None: 

39 """ 

40 Check whether path is existing or not. 

41 

42 Args: 

43 path (Union[Path, str]): path to check. Can only a file name in current directory, a pure directory, or a file 

44 name with directory. When path contains a directory, the function will check whether the directory exists, if 

45 not, the directory will be created. 

46 """ 

47 path_obj = Path(path) if isinstance(path, str) else path 

48 

49 if path_obj.exists(): 

50 return 

51 

52 if path_obj.suffix: 

53 path_obj.parent.mkdir(parents=True, exist_ok=True) 

54 else: 

55 path_obj.mkdir(parents=True, exist_ok=True) 

56 

57 

58def has_valid_filename(path: Path) -> bool: 

59 """ 

60 Check whether path has valid filename. A filename should contain name and suffix, name and suffix must contain 

61 letters, and then can have numbers and underscores. 

62 

63 Args: 

64 path (Path): path to check. 

65 

66 Return: 

67 bool: whether path has a valid filename. 

68 """ 

69 conditions = ( 

70 path.name, 

71 path.suffix, 

72 len(path.suffix) > 1, 

73 path.stem, 

74 any(c.isalpha() for c in path.stem), 

75 any(c.isalpha() for c in path.suffix[1:]) 

76 ) 

77 return all(conditions) 

78 

79 

80def narrow_tensor_by_index(tensor: Any, offsets: tuple, lengths: tuple) -> Any: 

81 """ 

82 Narrow the tensor by (offsets, lengths) per dimension. 

83 

84 Used for resharding operations to extract a slice from a tensor. 

85 Compatible with both torch and mindspore (uses slice indexing). 

86 

87 Args: 

88 tensor (Any): The tensor to narrow (tensor-like object supporting indexing). 

89 offsets (tuple): Tuple of offsets per dimension. 

90 lengths (tuple): Tuple of lengths per dimension. 

91 

92 Returns: 

93 Any: The narrowed tensor slice (tensor-like object). 

94 """ 

95 if not offsets or not lengths: 

96 return tensor 

97 slices = tuple( 

98 slice(int(off), int(off) + int(ln)) 

99 for off, ln in zip(offsets, lengths) 

100 ) 

101 return tensor[slices] 

102 

103 

104def chunk_to_area(chunk: ChunkStorageMetadata) -> tuple[tuple[int, int], ...]: 

105 """ 

106 Convert ChunkStorageMetadata to (start, end) area per dimension. 

107 

108 Args: 

109 chunk (ChunkStorageMetadata): ChunkStorageMetadata instance with offsets and sizes. 

110 

111 Returns: 

112 tuple[tuple[int, int], ...]: Tuple of (start, end) tuples for each dimension. 

113 """ 

114 return tuple( 

115 (chunk.offsets[i], chunk.offsets[i] + chunk.sizes[i]) 

116 for i in range(len(chunk.offsets)) 

117 ) 

118 

119 

120def create_chunk_list_for_tensor(obj: Union[Tensor, DTensor]) -> list[ChunkStorageMetadata]: 

121 """ 

122 Create list of local chunks for the given object (DTensor or plain tensor). 

123 

124 Used to determine what this rank needs to load (resharding). 

125 

126 Args: 

127 obj (Union[Tensor, DTensor]): hyper DTensor or platform Tensor. 

128 

129 Returns: 

130 list[ChunkStorageMetadata]: List of ChunkStorageMetadata representing 

131 local chunks needed by this rank. 

132 """ 

133 if isinstance(obj, DTensor): 

134 layout = obj.layout 

135 if layout is None: 

136 shape = obj.shape if hasattr(obj, "shape") else obj.to_local().shape 

137 return [ChunkStorageMetadata(offsets=(0,) * len(shape), sizes=tuple(shape))] 

138 

139 mesh_shape = getattr(layout, "mesh_shape", None) or getattr(layout, "_mesh", None) 

140 tensor_map = getattr(layout, "tensor_map", None) or getattr(layout, "_tensor_map", None) 

141 rank_list = getattr(layout, "rank_list", None) or getattr(layout, "_rank_list", None) 

142 

143 if mesh_shape is None or tensor_map is None or rank_list is None: 

144 shape = obj.shape if hasattr(obj, "shape") else obj.to_local().shape 

145 return [ChunkStorageMetadata(offsets=(0,) * len(shape), sizes=tuple(shape))] 

146 

147 current_rank = platform.get_rank() 

148 if current_rank not in rank_list: 

149 return [] 

150 

151 inner_rank_id = rank_list.index(current_rank) 

152 full_shape = obj.shape 

153 slice_area = infer_slice_area_by_rank( 

154 mesh_shape=mesh_shape, 

155 tensor_map=tensor_map, 

156 rank_id=inner_rank_id, 

157 full_shape=full_shape, 

158 ) 

159 offsets = tuple(s for s, _ in slice_area) 

160 sizes = tuple(e - s for s, e in slice_area) 

161 return [ChunkStorageMetadata(offsets=offsets, sizes=sizes)] 

162 

163 if isinstance(obj, Tensor): 

164 # handle Tensor with shard information 

165 if hasattr(obj, CHUNK_INFO): 

166 if not isinstance(getattr(obj, CHUNK_INFO), ChunkInfo): 

167 raise ValueError("The attr CHUNK_INFO should be a ChunkInfo instance") 

168 chunk = getattr(obj, CHUNK_INFO).chunk 

169 return [chunk] 

170 # platform.Tensor has exactly one chunk in metadata (full tensor) 

171 shape = tuple(obj.shape) 

172 return [ChunkStorageMetadata(offsets=(0,) * len(shape), sizes=shape)] 

173 

174 raise ValueError(f"Not support type {type(obj)} for creating chunk list ") 

175 

176 

177def remove_redundant_plans( 

178 all_plans: list[SavePlan], 

179 save_to_minimum_rank: bool = False, 

180) -> list[SavePlan]: 

181 """ 

182 Remove duplicate entries across SavePlans. For each duplicate, only one plan 

183 keeps the entry. The selection prefers the smallest planned storage size 

184 (or the minimum rank when save_to_minimum_rank is True). 

185 

186 Args: 

187 all_plans (list[SavePlan]): List of save plans to deduplicate. 

188 save_to_minimum_rank (bool): If True, assign duplicates to the minimum rank; else to plan with minimal storage. 

189 Default False. 

190 """ 

191 # Build mapping from item index to set of plan indices containing it 

192 duplicate_map: dict[MetadataIndex, set[int]] = defaultdict(set) 

193 # Registry to retrieve WriteItem by its index 

194 item_registry: dict[MetadataIndex, WriteItem] = {} 

195 # Track which items remain in each plan after deduplication 

196 remaining_items: list[set[MetadataIndex]] = [ 

197 {entry.index for entry in plan.items} for plan in all_plans 

198 ] 

199 

200 # Collect all items and their plan associations 

201 for idx, plan in enumerate(all_plans): 

202 for entry in plan.items: 

203 duplicate_map[entry.index].add(idx) 

204 item_registry[entry.index] = entry 

205 

206 storage_sizes = [0] * len(all_plans) 

207 

208 # Separate unique items (appear in only one plan) from duplicates 

209 # Process unique items first to prevent them from affecting load balancing 

210 single_plan_items: list[tuple[MetadataIndex, int]] = [] 

211 multi_plan_items: list[tuple[MetadataIndex, set[int]]] = [] 

212 

213 for item_key, containing_plans in duplicate_map.items(): 

214 if len(containing_plans) == 1: 

215 single_plan_items.append((item_key, next(iter(containing_plans)))) 

216 else: 

217 multi_plan_items.append((item_key, containing_plans)) 

218 

219 # First pass: handle items that appear in only one plan 

220 for item_key, target_idx in single_plan_items: 

221 entry = item_registry[item_key] 

222 storage_sizes[target_idx] += entry.tensor_storage_size() or 1 

223 

224 # Second pass: assign duplicate items to the plan with minimal storage size 

225 for item_key, containing_plans in multi_plan_items: 

226 if save_to_minimum_rank: 

227 target_plan = min(containing_plans) 

228 else: 

229 target_plan = min( 

230 containing_plans, key=lambda p_idx: storage_sizes[p_idx] 

231 ) 

232 

233 entry = item_registry[item_key] 

234 storage_sizes[target_plan] += entry.tensor_storage_size() or 1 

235 # Remove this item from all other plans 

236 for p_idx in containing_plans - {target_plan}: 

237 remaining_items[p_idx].discard(item_key) 

238 

239 if len(all_plans) != len(remaining_items): 

240 raise AssertionError("len(all_plans) != len(remaining_items)") 

241 

242 # Generate deduplicated plans with only remaining items 

243 return [ 

244 dataclasses.replace( 

245 plan, items=[entry for entry in plan.items if entry.index in item_set] 

246 ) 

247 for plan, item_set in zip(all_plans, remaining_items) 

248 ] 

249 

250 

251def traverse_state_dict( 

252 state_dict: Any, 

253 visitor: Any, 

254) -> None: 

255 """ 

256 Invoke ``visitor`` for each value recursively in ``state_dict``. 

257 Mapping will be traversed and ``visitor`` will be applied to the leaf elements. 

258 ``visitor`` will only be applied to elements in a list or a tuple, if the 

259 container contains tensors or mappings. 

260 """ 

261 

262 def _is_terminal(value: Any) -> bool: 

263 """Leaf-like container: no nested mappings/lists/tuples/tensors to recurse into.""" 

264 values: Collection 

265 if isinstance(value, Mapping): 

266 return False 

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

268 values = value 

269 else: 

270 return True 

271 

272 for entry in values: 

273 if isinstance(entry, (Mapping, list, tuple)) and not _is_terminal(entry): 

274 return False 

275 if isinstance(entry, Tensor): 

276 return False 

277 return True 

278 

279 def _traverse_obj(path: tuple[Any, ...], value: Any) -> None: 

280 if isinstance(value, Mapping): 

281 for k, v in value.items(): 

282 _traverse_obj(path + (str(k),), v) 

283 elif _is_terminal(value): 

284 visitor(path, value) 

285 elif isinstance(value, (list, tuple)): 

286 for i, v in enumerate(value): 

287 _traverse_obj(path + (i,), v) 

288 

289 for key, value in state_dict.items(): 

290 _traverse_obj((str(key),), value) 

291 

292 

293def flatten_state_dict(state_dict: Any) -> tuple[dict[str, Any], dict[str, tuple[Any, ...]]]: 

294 """Flatten a nested state dict to dotted FQN keys; returns ``(flat_dict, fqn -> path)``.""" 

295 fqn_names: dict[str, Any] = {} 

296 mappings: dict[str, tuple[Any, ...]] = {} 

297 

298 def flat_copy(path: tuple[Any, ...], value: Any) -> None: 

299 new_fqn = ".".join(map(str, path)) 

300 if new_fqn in fqn_names: 

301 raise ValueError( 

302 f"Duplicate flattened FQN {new_fqn!r} when converting nested state_dict; " 

303 "two different values map to the same dotted name." 

304 ) 

305 fqn_names[new_fqn] = value 

306 mappings[new_fqn] = path 

307 

308 traverse_state_dict(state_dict, flat_copy) 

309 return fqn_names, mappings 

310 

311 

312def set_element(root_dict: Any, path: tuple[Any, ...], value: Any) -> None: 

313 """Set ``value`` in ``root_dict`` along the ``path`` object path.""" 

314 if not path: 

315 raise ValueError("path must be non-empty") 

316 cur_container: Any = root_dict 

317 

318 def extend_list(lst: list[Any], idx: int) -> None: 

319 while len(lst) <= idx: 

320 lst.append(None) 

321 

322 for i in range(1, len(path)): 

323 prev_key = path[i - 1] 

324 next_key = path[i] 

325 def_val: Any = {} if isinstance(next_key, str) else [] 

326 

327 if isinstance(cur_container, Mapping): 

328 cur_container = cur_container.setdefault(prev_key, def_val) 

329 else: 

330 extend_list(cur_container, prev_key) 

331 if cur_container[prev_key] is None: 

332 cur_container[prev_key] = def_val 

333 cur_container = cur_container[prev_key] 

334 

335 last_key = path[-1] 

336 if isinstance(last_key, int): 

337 extend_list(cur_container, last_key) 

338 

339 cur_container[last_key] = value