Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / util.py: 90%
147 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-25 04:27 +0800
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
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.dtensor.layout import infer_slice_area_by_layout
31from hyper_parallel.core.dtensor.dtensor import DTensor
32from hyper_parallel.platform import get_platform
35platform = get_platform()
36Tensor = platform.Tensor
39def check_path(path: Union[Path, str]) -> None:
40 """
41 Check whether path is existing or not.
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
50 if path_obj.exists():
51 return
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)
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.
64 Args:
65 path (Path): path to check.
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)
81def narrow_tensor_by_index(tensor: Any, offsets: tuple, lengths: tuple) -> Any:
82 """
83 Narrow the tensor by (offsets, lengths) per dimension.
85 Used for resharding operations to extract a slice from a tensor.
86 Compatible with both torch and mindspore (uses slice indexing).
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.
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]
105def chunk_to_area(chunk: ChunkStorageMetadata) -> tuple[tuple[int, int], ...]:
106 """
107 Convert ChunkStorageMetadata to (start, end) area per dimension.
109 Args:
110 chunk (ChunkStorageMetadata): ChunkStorageMetadata instance with offsets and sizes.
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 )
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).
125 Used to determine what this rank needs to load (resharding).
127 Args:
128 obj (Union[Tensor, DTensor]): hyper DTensor or platform Tensor.
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 ]
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)
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))]
153 current_rank = platform.get_rank()
154 if current_rank not in rank_list:
155 return []
157 inner_rank_id = rank_list.index(current_rank)
158 full_shape = obj.shape
159 slice_area = infer_slice_area_by_layout(
160 layout,
161 inner_rank_id,
162 full_shape,
163 )
164 offsets = tuple(s for s, _ in slice_area)
165 sizes = tuple(e - s for s, e in slice_area)
166 return [ChunkStorageMetadata(offsets=offsets, sizes=sizes)]
168 if isinstance(obj, Tensor):
169 # handle Tensor with shard information
170 if hasattr(obj, CHUNK_INFO):
171 if not isinstance(getattr(obj, CHUNK_INFO), ChunkInfo):
172 raise ValueError("The attr CHUNK_INFO should be a ChunkInfo instance")
173 chunk = getattr(obj, CHUNK_INFO).chunk
174 return [chunk]
175 # platform.Tensor has exactly one chunk in metadata (full tensor)
176 shape = tuple(obj.shape)
177 return [ChunkStorageMetadata(offsets=(0,) * len(shape), sizes=shape)]
179 raise ValueError(f"Not support type {type(obj)} for creating chunk list ")
182def remove_redundant_plans(
183 all_plans: list[SavePlan],
184 save_to_minimum_rank: bool = False,
185) -> list[SavePlan]:
186 """
187 Remove duplicate entries across SavePlans. For each duplicate, only one plan
188 keeps the entry. The selection prefers the smallest planned storage size
189 (or the minimum rank when save_to_minimum_rank is True).
191 Args:
192 all_plans (list[SavePlan]): List of save plans to deduplicate.
193 save_to_minimum_rank (bool): If True, assign duplicates to the minimum rank; else to plan with minimal storage.
194 Default False.
195 """
196 # Build mapping from item index to set of plan indices containing it
197 duplicate_map: dict[MetadataIndex, set[int]] = defaultdict(set)
198 # Registry to retrieve WriteItem by its index
199 item_registry: dict[MetadataIndex, WriteItem] = {}
200 # Track which items remain in each plan after deduplication
201 remaining_items: list[set[MetadataIndex]] = [
202 {entry.index for entry in plan.items} for plan in all_plans
203 ]
205 # Collect all items and their plan associations
206 for idx, plan in enumerate(all_plans):
207 for entry in plan.items:
208 duplicate_map[entry.index].add(idx)
209 item_registry[entry.index] = entry
211 storage_sizes = [0] * len(all_plans)
213 # Separate unique items (appear in only one plan) from duplicates
214 # Process unique items first to prevent them from affecting load balancing
215 single_plan_items: list[tuple[MetadataIndex, int]] = []
216 multi_plan_items: list[tuple[MetadataIndex, set[int]]] = []
218 for item_key, containing_plans in duplicate_map.items():
219 if len(containing_plans) == 1:
220 single_plan_items.append((item_key, next(iter(containing_plans))))
221 else:
222 multi_plan_items.append((item_key, containing_plans))
224 # First pass: handle items that appear in only one plan
225 for item_key, target_idx in single_plan_items:
226 entry = item_registry[item_key]
227 storage_sizes[target_idx] += entry.tensor_storage_size() or 1
229 # Second pass: assign duplicate items to the plan with minimal storage size
230 for item_key, containing_plans in multi_plan_items:
231 if save_to_minimum_rank:
232 target_plan = min(containing_plans)
233 else:
234 target_plan = min(
235 containing_plans, key=lambda p_idx: storage_sizes[p_idx]
236 )
238 entry = item_registry[item_key]
239 storage_sizes[target_plan] += entry.tensor_storage_size() or 1
240 # Remove this item from all other plans
241 for p_idx in containing_plans - {target_plan}:
242 remaining_items[p_idx].discard(item_key)
244 if len(all_plans) != len(remaining_items):
245 raise AssertionError("len(all_plans) != len(remaining_items)")
247 # Generate deduplicated plans with only remaining items
248 return [
249 dataclasses.replace(
250 plan, items=[entry for entry in plan.items if entry.index in item_set]
251 )
252 for plan, item_set in zip(all_plans, remaining_items)
253 ]
256def traverse_state_dict(
257 state_dict: Any,
258 visitor: Any,
259) -> None:
260 """
261 Invoke ``visitor`` for each value recursively in ``state_dict``.
262 Mapping will be traversed and ``visitor`` will be applied to the leaf elements.
263 ``visitor`` will only be applied to elements in a list or a tuple, if the
264 container contains tensors or mappings.
265 """
267 def _is_terminal(value: Any) -> bool:
268 """Leaf-like container: no nested mappings/lists/tuples/tensors to recurse into."""
269 values: Collection
270 if isinstance(value, Mapping):
271 return False
272 if isinstance(value, (list, tuple)):
273 values = value
274 else:
275 return True
277 for entry in values:
278 if isinstance(entry, (Mapping, list, tuple)) and not _is_terminal(entry):
279 return False
280 if isinstance(entry, Tensor):
281 return False
282 return True
284 def _traverse_obj(path: tuple[Any, ...], value: Any) -> None:
285 if isinstance(value, Mapping):
286 for k, v in value.items():
287 _traverse_obj(path + (str(k),), v)
288 elif _is_terminal(value):
289 visitor(path, value)
290 elif isinstance(value, (list, tuple)):
291 for i, v in enumerate(value):
292 _traverse_obj(path + (i,), v)
294 for key, value in state_dict.items():
295 _traverse_obj((str(key),), value)
298def flatten_state_dict(state_dict: Any) -> tuple[dict[str, Any], dict[str, tuple[Any, ...]]]:
299 """Flatten a nested state dict to dotted FQN keys; returns ``(flat_dict, fqn -> path)``."""
300 fqn_names: dict[str, Any] = {}
301 mappings: dict[str, tuple[Any, ...]] = {}
303 def flat_copy(path: tuple[Any, ...], value: Any) -> None:
304 new_fqn = ".".join(map(str, path))
305 if new_fqn in fqn_names:
306 raise ValueError(
307 f"Duplicate flattened FQN {new_fqn!r} when converting nested state_dict; "
308 "two different values map to the same dotted name."
309 )
310 fqn_names[new_fqn] = value
311 mappings[new_fqn] = path
313 traverse_state_dict(state_dict, flat_copy)
314 return fqn_names, mappings
317def set_element(root_dict: Any, path: tuple[Any, ...], value: Any) -> None:
318 """Set ``value`` in ``root_dict`` along the ``path`` object path."""
319 if not path:
320 raise ValueError("path must be non-empty")
321 cur_container: Any = root_dict
323 def extend_list(lst: list[Any], idx: int) -> None:
324 while len(lst) <= idx:
325 lst.append(None)
327 for i in range(1, len(path)):
328 prev_key = path[i - 1]
329 next_key = path[i]
330 def_val: Any = {} if isinstance(next_key, str) else []
332 if isinstance(cur_container, Mapping):
333 cur_container = cur_container.setdefault(prev_key, def_val)
334 else:
335 extend_list(cur_container, prev_key)
336 if cur_container[prev_key] is None:
337 cur_container[prev_key] = def_val
338 cur_container = cur_container[prev_key]
340 last_key = path[-1]
341 if isinstance(last_key, int):
342 extend_list(cur_container, last_key)
344 cur_container[last_key] = value