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

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.ragged_utils import compute_ragged_boxes 

30from hyper_parallel.core.distributed_checkpoint.reshard import infer_slice_area_by_rank 

31from hyper_parallel.core.dtensor.dtensor import DTensor 

32from hyper_parallel.platform import get_platform 

33 

34 

35platform = get_platform() 

36Tensor = platform.Tensor 

37 

38 

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

40 """ 

41 Check whether path is existing or not. 

42 

43 Args: 

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

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

46 not, the directory will be created. 

47 """ 

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

49 

50 if path_obj.exists(): 

51 return 

52 

53 if path_obj.suffix: 

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

55 else: 

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

57 

58 

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

60 """ 

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

62 letters, and then can have numbers and underscores. 

63 

64 Args: 

65 path (Path): path to check. 

66 

67 Return: 

68 bool: whether path has a valid filename. 

69 """ 

70 conditions = ( 

71 path.name, 

72 path.suffix, 

73 len(path.suffix) > 1, 

74 path.stem, 

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

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

77 ) 

78 return all(conditions) 

79 

80 

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

82 """ 

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

84 

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

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

87 

88 Args: 

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

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

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

92 

93 Returns: 

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

95 """ 

96 if not offsets or not lengths: 

97 return tensor 

98 slices = tuple( 

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

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

101 ) 

102 return tensor[slices] 

103 

104 

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

106 """ 

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

108 

109 Args: 

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

111 

112 Returns: 

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

114 """ 

115 return tuple( 

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

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

118 ) 

119 

120 

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

122 """ 

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

124 

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

126 

127 Args: 

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

129 

130 Returns: 

131 list[ChunkStorageMetadata]: List of ChunkStorageMetadata representing 

132 local chunks needed by this rank. 

133 """ 

134 if isinstance(obj, DTensor): 

135 layout = obj.layout 

136 if layout is None: 

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

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

139 if layout.ragged_shard is not None: 

140 return [ 

141 ChunkStorageMetadata(offsets=box.offsets, sizes=box.sizes) 

142 for box in compute_ragged_boxes(obj) 

143 ] 

144 

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

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

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

148 

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

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

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

152 

153 current_rank = platform.get_rank() 

154 if current_rank not in rank_list: 

155 return [] 

156 

157 inner_rank_id = rank_list.index(current_rank) 

158 full_shape = obj.shape 

159 slice_area = infer_slice_area_by_rank( 

160 mesh_shape=mesh_shape, 

161 tensor_map=tensor_map, 

162 rank_id=inner_rank_id, 

163 full_shape=full_shape, 

164 ) 

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

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

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

168 

169 if isinstance(obj, Tensor): 

170 # handle Tensor with shard information 

171 if hasattr(obj, CHUNK_INFO): 

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

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

174 chunk = getattr(obj, CHUNK_INFO).chunk 

175 return [chunk] 

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

177 shape = tuple(obj.shape) 

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

179 

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

181 

182 

183def remove_redundant_plans( 

184 all_plans: list[SavePlan], 

185 save_to_minimum_rank: bool = False, 

186) -> list[SavePlan]: 

187 """ 

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

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

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

191 

192 Args: 

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

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

195 Default False. 

196 """ 

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

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

199 # Registry to retrieve WriteItem by its index 

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

201 # Track which items remain in each plan after deduplication 

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

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

204 ] 

205 

206 # Collect all items and their plan associations 

207 for idx, plan in enumerate(all_plans): 

208 for entry in plan.items: 

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

210 item_registry[entry.index] = entry 

211 

212 storage_sizes = [0] * len(all_plans) 

213 

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

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

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

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

218 

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

220 if len(containing_plans) == 1: 

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

222 else: 

223 multi_plan_items.append((item_key, containing_plans)) 

224 

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

226 for item_key, target_idx in single_plan_items: 

227 entry = item_registry[item_key] 

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

229 

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

231 for item_key, containing_plans in multi_plan_items: 

232 if save_to_minimum_rank: 

233 target_plan = min(containing_plans) 

234 else: 

235 target_plan = min( 

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

237 ) 

238 

239 entry = item_registry[item_key] 

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

241 # Remove this item from all other plans 

242 for p_idx in containing_plans - {target_plan}: 

243 remaining_items[p_idx].discard(item_key) 

244 

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

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

247 

248 # Generate deduplicated plans with only remaining items 

249 return [ 

250 dataclasses.replace( 

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

252 ) 

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

254 ] 

255 

256 

257def traverse_state_dict( 

258 state_dict: Any, 

259 visitor: Any, 

260) -> None: 

261 """ 

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

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

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

265 container contains tensors or mappings. 

266 """ 

267 

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

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

270 values: Collection 

271 if isinstance(value, Mapping): 

272 return False 

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

274 values = value 

275 else: 

276 return True 

277 

278 for entry in values: 

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

280 return False 

281 if isinstance(entry, Tensor): 

282 return False 

283 return True 

284 

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

286 if isinstance(value, Mapping): 

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

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

289 elif _is_terminal(value): 

290 visitor(path, value) 

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

292 for i, v in enumerate(value): 

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

294 

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

296 _traverse_obj((str(key),), value) 

297 

298 

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

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

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

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

303 

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

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

306 if new_fqn in fqn_names: 

307 raise ValueError( 

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

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

310 ) 

311 fqn_names[new_fqn] = value 

312 mappings[new_fqn] = path 

313 

314 traverse_state_dict(state_dict, flat_copy) 

315 return fqn_names, mappings 

316 

317 

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

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

320 if not path: 

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

322 cur_container: Any = root_dict 

323 

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

325 while len(lst) <= idx: 

326 lst.append(None) 

327 

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

329 prev_key = path[i - 1] 

330 next_key = path[i] 

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

332 

333 if isinstance(cur_container, Mapping): 

334 cur_container = cur_container.setdefault(prev_key, def_val) 

335 else: 

336 extend_list(cur_container, prev_key) 

337 if cur_container[prev_key] is None: 

338 cur_container[prev_key] = def_val 

339 cur_container = cur_container[prev_key] 

340 

341 last_key = path[-1] 

342 if isinstance(last_key, int): 

343 extend_list(cur_container, last_key) 

344 

345 cur_container[last_key] = value