Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / filesystem_storage.py: 87%
217 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-21 04:29 +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"""File system storage implementations for checkpoint save and load."""
16import os
17import pickle
18from collections import Counter
19from pathlib import Path
20from typing import Any, Optional, Union
22from safetensors import safe_open
24from hyper_parallel.core.distributed_checkpoint.metadata import Metadata, MetadataIndex
25from hyper_parallel.core.distributed_checkpoint.planner import (
26 LoadPlan,
27 LoadPlanner,
28 ReadItem,
29 SavePlan,
30 SavePlanner,
31 WriteItem,
32)
33from hyper_parallel.core.distributed_checkpoint.storage import (
34 StorageInfo,
35 StorageReader,
36 StorageWriter,
37 WriteResult,
38 METADATA_FILE_NAME,
39)
40from hyper_parallel.core.distributed_checkpoint.util import narrow_tensor_by_index
41from hyper_parallel.platform import get_platform
42from hyper_parallel.platform.platform import PlatformType
45class FileSystemWriter(StorageWriter):
46 """
47 File system storage writer implementation.
49 Saves checkpoint data to the local file system, organizing tensors
50 into safetensors files and bytes into separate files.
51 """
53 def __init__(self, checkpoint_dir: Union[Path, str]):
54 self.checkpoint_dir = Path(checkpoint_dir) if isinstance(checkpoint_dir, str) else checkpoint_dir
55 self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
56 self.rank: int = 0
57 self.is_coordinator: bool = False
58 self.use_collectives: bool = True
60 def initialize_writer(self, checkpoint_id: Optional[Union[Path, str]] = None) -> None:
61 """
62 Initialize storage writer with new checkpoint directory.
64 Args:
65 checkpoint_id (Optional[Union[Path, str]]): New checkpoint directory path. Default None.
66 """
67 if checkpoint_id:
68 self.checkpoint_dir = Path(checkpoint_id) if isinstance(checkpoint_id, str) else checkpoint_id
69 self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
71 def configure_writer(self, is_coordinator: bool, **kwargs) -> None:
72 """
73 Configure storage writer.
75 Args:
76 is_coordinator (bool): Whether this rank is the coordinator.
77 **kwargs: Additional keyword arguments (e.g., rank, use_collectives).
78 """
79 self.is_coordinator = is_coordinator
80 self.rank = kwargs.get("rank") if "rank" in kwargs else get_platform().get_rank()
81 self.use_collectives = kwargs.get("use_collectives", True)
83 def optimize_local_plan(self, plan: SavePlan) -> SavePlan:
84 """
85 Optimize local plan.
87 Args:
88 plan (SavePlan): Local save plan.
90 Returns:
91 SavePlan: Optimized local plan.
92 """
93 return plan
95 def optimize_global_plan(self, plans: list[SavePlan]) -> list[SavePlan]:
96 """
97 Optimize global plan.
99 Args:
100 plans (list[SavePlan]): List of local plans from all ranks.
102 Returns:
103 list[SavePlan]: Optimized global plans.
104 """
105 return plans
108 def _serialize_bytes_item(self, item: WriteItem, planner: SavePlanner) -> bytes:
109 """Serialize a BYTE_IO item payload while preserving current behavior."""
110 data = planner.get_data(item)
111 if isinstance(data, bytes):
112 return data
113 return pickle.dumps(data)
116 def _write_bytes_items(self, plan: SavePlan, planner: SavePlanner) -> list[WriteResult]:
117 """
118 Write all BYTE_IO items into one per-rank bytes file.
120 Args:
121 plan (SavePlan): Save plan containing WriteItems.
122 planner (SavePlanner): Save planner used to resolve runtime data.
124 Returns:
125 list[WriteResult]: Write results for BYTE_IO items.
126 """
127 byte_items = [item for item in plan.items if item.type.value == "byte_io"]
128 if not byte_items:
129 return []
131 file_name = f"_rank{self.rank}_.bytes"
132 file_path = self.checkpoint_dir / file_name
134 results: list[WriteResult] = []
136 with open(file_path, "wb") as f:
137 for item in byte_items:
138 payload = self._serialize_bytes_item(item, planner)
139 offset = f.tell()
140 f.write(payload)
141 length = len(payload)
142 storage_info = StorageInfo(
143 relative_path=file_name,
144 offset=offset,
145 length=length,
146 )
147 results.append(
148 WriteResult(
149 index=item.index,
150 storage_data=storage_info,
151 )
152 )
154 return results
156 def _collect_tensors(
157 self, plan: SavePlan, planner: SavePlanner
158 ) -> tuple[dict[str, Any], dict[MetadataIndex, str]]:
159 """
160 Collect tensor data from planner runtime lookup.
162 Args:
163 plan (SavePlan): Save plan containing WriteItems.
164 planner (SavePlanner): Save planner.
166 Returns:
167 tuple[dict[str, Any], dict[MetadataIndex, str]]: Tensor data keyed by
168 physical safetensors key, plus logical index-to-key mapping.
170 Raises:
171 RuntimeError: If tensor data cannot be resolved for an item.
172 """
173 tensor_items = [
174 item for item in plan.items
175 if item.type.value == "tensor" and item.tensor_data
176 ]
177 fqn_counts = Counter(item.index.fqn for item in tensor_items)
178 reserved_keys = set(fqn_counts)
179 used_keys: set[str] = set()
180 next_chunk_index: dict[str, int] = {}
181 tensor_dict: dict[str, Any] = {}
182 tensor_keys: dict[MetadataIndex, str] = {}
184 for item in tensor_items:
185 tensor = planner.get_data(item)
186 if tensor is None:
187 raise RuntimeError(
188 f"Tensor data could not be resolved for index {item.index}. "
189 f"FQN: {item.index.fqn}"
190 )
192 fqn = item.index.fqn
193 tensor_key = fqn
194 if fqn_counts[fqn] > 1:
195 chunk_index = next_chunk_index.get(fqn, 0)
196 next_chunk_index[fqn] = chunk_index + 1
197 tensor_key = f"{fqn}.__dcp_chunk_{chunk_index}"
198 while tensor_key in reserved_keys or tensor_key in used_keys:
199 tensor_key += "_"
201 used_keys.add(tensor_key)
202 tensor_dict[tensor_key] = tensor
203 tensor_keys[item.index] = tensor_key
204 return tensor_dict, tensor_keys
206 def _write_tensors(
207 self,
208 plan: SavePlan,
209 tensor_dict: dict[str, Any],
210 tensor_keys: dict[MetadataIndex, str],
211 ) -> list[WriteResult]:
212 """
213 Write all tensors to safetensors file and create WriteResults.
215 Args:
216 plan (SavePlan): Save plan containing WriteItems.
217 tensor_dict (dict[str, Any]): Dictionary mapping physical keys to tensor data.
218 tensor_keys (dict[MetadataIndex, str]): Logical index-to-key mapping.
220 Returns:
221 list[WriteResult]: List of write results for tensor items.
222 """
223 if not tensor_dict:
224 return []
226 platform = get_platform()
227 file_name = f"_rank{self.rank}_.safetensors"
228 file_path = self.checkpoint_dir / file_name
229 platform.save_checkpoint(tensor_dict, str(file_path))
231 # Record StorageInfo for each tensor
232 # Note: we don't know per-tensor byte offsets, so offset=0, length=-1
233 results: list[WriteResult] = []
234 for item in plan.items:
235 if item.type.value == "tensor" and item.tensor_data:
236 storage_info = StorageInfo(
237 relative_path=file_name,
238 offset=0,
239 length=-1,
240 tensor_key=tensor_keys[item.index],
241 )
242 results.append(
243 WriteResult(
244 index=item.index,
245 storage_data=storage_info,
246 )
247 )
248 return results
250 def execute_write(self, plan: SavePlan, planner: SavePlanner) -> list[WriteResult]:
251 """
252 Write data to storage and return per-item storage metadata.
254 Group tensors into safetensors files and bytes into separate files, recording StorageInfo for each item.
256 Args:
257 plan (SavePlan): Save plan containing WriteItems.
258 planner (SavePlanner): Save planner.
260 Returns:
261 list[WriteResult]: List of write results with storage metadata.
262 """
263 results: list[WriteResult] = []
265 # Write all BYTE_IO items into one file per rank
266 results.extend(self._write_bytes_items(plan, planner))
268 # Collect and write tensors
269 tensor_dict, tensor_keys = self._collect_tensors(plan, planner)
270 results.extend(self._write_tensors(plan, tensor_dict, tensor_keys))
272 return results
274 def finalize_checkpoint(self, metadata: Metadata, results: list[list[WriteResult]]) -> None:
275 """
276 Finish writing checkpoint and populate metadata.storage_data.
278 When use_collectives=True: only coordinator saves global metadata to .metadata.
279 When use_collectives=False: each rank saves its own metadata to .rank{rank}_metadata,
280 no cross-rank interaction.
282 Args:
283 metadata (Metadata): Checkpoint metadata to update.
284 results (list[list[WriteResult]]): Write results from all ranks (or single rank when use_collectives=False).
285 """
286 should_save = not self.use_collectives or (self.use_collectives and self.is_coordinator)
287 if not should_save:
288 return
290 # Build storage_data: map MetadataIndex -> StorageInfo
291 storage_md: dict[MetadataIndex, StorageInfo] = {}
292 for wr_list in results:
293 for wr in wr_list:
294 storage_md[wr.index] = wr.storage_data
295 metadata.storage_data = storage_md
297 # Save metadata file
298 if self.use_collectives:
299 metadata_file = self.checkpoint_dir / METADATA_FILE_NAME
300 else:
301 metadata_file = self.checkpoint_dir / f"{self.rank}{METADATA_FILE_NAME}"
302 with open(metadata_file, "wb") as f:
303 pickle.dump(metadata, f)
306def _copy_tensor_to_target(
307 req: ReadItem, tensor: Any, target_tensor: Any, planner: LoadPlanner
308) -> None:
309 """
310 Copy tensor data to target tensor and commit.
312 Args:
313 req (ReadItem): ReadItem request.
314 tensor (Any): Source tensor (tensor-like object).
315 target_tensor (Any): Target tensor (tensor-like object).
316 planner (LoadPlanner): Load planner for committing.
317 """
318 if hasattr(target_tensor, "copy_"):
319 target_tensor.copy_(tensor)
320 planner.apply_tensor(req, target_tensor)
321 else:
322 # mindspore or non-tensor: copy via commit path
323 planner.apply_tensor(req, tensor)
326def _load_bytes_file(
327 path: str,
328 reqs: list[ReadItem],
329 planner: LoadPlanner,
330 storage_data: dict[MetadataIndex, StorageInfo],
331) -> None:
332 """
333 Load bytes from a file.
335 Args:
336 path (str): Path to the bytes file.
337 reqs (list[ReadItem]): List of ReadItems for this file.
338 planner (LoadPlanner): Load planner for loading bytes.
339 """
340 with open(path, "rb") as f:
341 for req in reqs:
342 storage_info = storage_data.get(req.storage_index)
343 if storage_info is None:
344 raise KeyError(
345 f"StorageInfo not found for index {req.storage_index}"
346 )
347 f.seek(storage_info.offset)
348 value = f.read(storage_info.length)
349 planner.apply_bytes(req, value)
352def _get_tensor_size(tensor: Any) -> Optional[tuple]:
353 """
354 Get size/shape of a tensor.
356 Args:
357 tensor (Any): Tensor object (tensor-like with shape/size attribute).
359 Returns:
360 Optional[tuple]: Tuple of tensor size or None if not available.
361 """
362 if hasattr(tensor, "size") and callable(tensor.size):
363 return tuple(tensor.size())
364 return getattr(tensor, "shape", None)
367def _get_storage_info(
368 req: ReadItem,
369 storage_data: dict[MetadataIndex, StorageInfo],
370) -> StorageInfo:
371 """Return physical storage metadata for one read request."""
372 storage_info = storage_data.get(req.storage_index)
373 if storage_info is None:
374 raise KeyError(f"StorageInfo not found for index {req.storage_index}")
375 return storage_info
378def _validate_and_copy_tensor(
379 req: ReadItem,
380 tensor: Any,
381 planner: LoadPlanner,
382) -> None:
383 """Validate a loaded tensor slice and copy it to its planner destination."""
384 target_tensor = planner.acquire_tensor(req)
385 if hasattr(target_tensor, "detach"):
386 target_tensor = target_tensor.detach()
388 target_size = _get_tensor_size(target_tensor)
389 tensor_size = _get_tensor_size(tensor)
390 if target_size is not None and tensor_size is not None and target_size != tensor_size:
391 raise AssertionError(
392 f"req {req.storage_index} mismatch sizes "
393 f"{target_size} vs {tensor_size}"
394 )
395 _copy_tensor_to_target(req, tensor, target_tensor, planner)
398def _load_torch_tensor_file(
399 path: str,
400 reqs: list[ReadItem],
401 planner: LoadPlanner,
402 storage_data: dict[MetadataIndex, StorageInfo],
403) -> None:
404 """Load tensor slices from a Torch safetensors file."""
405 with safe_open(path, framework="pt", device="cpu") as tensor_file:
406 available_keys = set(tensor_file.keys())
407 for req in reqs:
408 storage_info = _get_storage_info(req, storage_data)
409 tensor_key = storage_info.tensor_key or req.storage_index.fqn
410 if tensor_key not in available_keys:
411 raise KeyError(f"Key {tensor_key} not found in checkpoint file {path}")
412 tensor_slices = tuple(
413 slice(int(off), int(off) + int(length))
414 for off, length in zip(req.storage_offsets, req.lengths)
415 )
416 if tensor_slices:
417 tensor = tensor_file.get_slice(tensor_key)[tensor_slices]
418 else:
419 tensor = narrow_tensor_by_index(
420 tensor_file.get_tensor(tensor_key),
421 req.storage_offsets,
422 req.lengths,
423 )
424 _validate_and_copy_tensor(req, tensor, planner)
427def _load_platform_tensor_file(
428 path: str,
429 reqs: list[ReadItem],
430 planner: LoadPlanner,
431 storage_data: dict[MetadataIndex, StorageInfo],
432 platform: Any,
433) -> None:
434 """Load tensor slices through the active non-Torch platform adapter."""
435 param_dict = platform.load_checkpoint(path)
436 for req in reqs:
437 storage_info = _get_storage_info(req, storage_data)
438 tensor_key = storage_info.tensor_key or req.storage_index.fqn
439 if tensor_key not in param_dict:
440 raise KeyError(f"Key {tensor_key} not found in checkpoint file {path}")
441 tensor = narrow_tensor_by_index(
442 param_dict[tensor_key],
443 req.storage_offsets,
444 req.lengths,
445 )
446 _validate_and_copy_tensor(req, tensor, planner)
449def _load_tensor_file(
450 path: str,
451 reqs: list[ReadItem],
452 planner: LoadPlanner,
453 storage_data: dict[MetadataIndex, StorageInfo],
454) -> None:
455 """
456 Load and process tensors from a safetensors file.
458 Args:
459 path (str): Path to the safetensors file.
460 reqs (list[ReadItem]): List of ReadItems for this file.
461 planner (LoadPlanner): Load planner for resolving and committing tensors.
462 storage_data (dict[MetadataIndex, StorageInfo]): Physical storage mapping.
463 """
464 platform = get_platform()
465 if platform.platform_type == PlatformType.PYTORCH:
466 _load_torch_tensor_file(path, reqs, planner, storage_data)
467 return
468 _load_platform_tensor_file(path, reqs, planner, storage_data, platform)
471class FileSystemReader(StorageReader):
472 """
473 File system storage reader implementation.
475 Reads checkpoint data from the local file system, loading tensors
476 from safetensors files and bytes from separate files.
477 """
479 def __init__(self, checkpoint_dir: Union[Path, str]):
480 self.checkpoint_dir = Path(checkpoint_dir) if isinstance(checkpoint_dir, str) else checkpoint_dir
481 # Cached storage layout: MetadataIndex -> StorageInfo (torch-aligned)
482 self.storage_data: Optional[dict[MetadataIndex, StorageInfo]] = None
483 self.rank: int = 0
484 self.is_coordinator: bool = False
486 def initialize_reader(self, checkpoint_id: Optional[Union[Path, str]] = None) -> None:
487 """
488 Initialize storage reader with new checkpoint directory.
490 Args:
491 checkpoint_id (Optional[Union[Path, str]]): New checkpoint directory path. Default None.
492 """
493 if checkpoint_id:
494 self.checkpoint_dir = Path(checkpoint_id) if isinstance(checkpoint_id, str) else checkpoint_id
496 def load_metadata(self, **kwargs) -> Metadata:
497 """
498 Load checkpoint metadata from file.
500 When rank is provided in kwargs: load rank-local metadata from .rank{rank}_metadata
501 (for checkpoints saved with use_collectives=False).
502 Otherwise: load global metadata from .metadata.
504 Args:
505 **kwargs: Optional arguments (e.g., rank for rank-local metadata).
507 Returns:
508 Metadata: Metadata object loaded from file.
509 """
510 rank = kwargs.get("rank")
511 if rank is not None:
512 metadata_file = self.checkpoint_dir / f"{rank}{METADATA_FILE_NAME}"
513 else:
514 metadata_file = self.checkpoint_dir / METADATA_FILE_NAME
516 if not metadata_file.exists():
517 raise FileNotFoundError(f"Metadata file not found: {metadata_file}")
518 with open(metadata_file, "rb") as f:
519 metadata = pickle.load(f)
520 return metadata
522 def configure_reader(self, metadata: Metadata, is_coordinator: bool, **kwargs) -> None:
523 """Configure storage reader."""
524 # Cache storage_data separately for quick lookup in execute_read.
525 # This mirrors torch.filesystem, where reader keeps a storage_data dict.
526 self.storage_data = getattr(metadata, "storage_data", None)
527 self.is_coordinator = is_coordinator
528 self.rank = kwargs.get("rank") if "rank" in kwargs else get_platform().get_rank()
530 def optimize_local_plan(self, plan: LoadPlan) -> LoadPlan:
531 """
532 Optimize local plan.
534 Args:
535 plan (LoadPlan): Local load plan.
537 Returns:
538 LoadPlan: Optimized local plan.
539 """
540 return plan
542 def optimize_global_plan(self, plans: list[LoadPlan]) -> list[LoadPlan]:
543 """
544 Optimize global plan.
546 Args:
547 plans (list[LoadPlan]): List of local plans from all ranks.
549 Returns:
550 list[LoadPlan]: Optimized global plans.
551 """
552 return plans
554 def _get_storage_path(self, read_item: ReadItem) -> str:
555 """
556 Get storage file path for a read item.
558 Args:
559 read_item (ReadItem): ReadItem to get path for.
561 Returns:
562 str: Absolute path to the storage file.
563 """
564 if self.storage_data is None:
565 raise KeyError("Checkpoint metadata.storage_data is required for filesystem read")
566 storage_info = self.storage_data.get(read_item.storage_index)
567 if storage_info is None:
568 raise KeyError(f"StorageInfo not found for index {read_item.storage_index}")
569 return str(self.checkpoint_dir / storage_info.relative_path)
571 def _group_items_by_file(self, plan: LoadPlan) -> dict[str, list]:
572 """
573 Group ReadItems by storage file path.
575 Args:
576 plan (LoadPlan): Load plan containing ReadItems.
578 Returns:
579 dict[str, list[ReadItem]]: Dictionary mapping file paths to lists of ReadItems.
580 """
581 per_file: dict[str, list] = {}
582 for read_item in plan.items:
583 path = self._get_storage_path(read_item)
584 per_file.setdefault(path, []).append(read_item)
585 return per_file
587 def execute_read(self, plan: LoadPlan, planner: LoadPlanner) -> None:
588 """
589 Read data from storage.
591 Aligned with torch filesystem read_data: groups ReadItems by file,
592 loads each file once, narrows tensors by storage_offsets/lengths for
593 resharding, then resolves/copies/commits data.
595 Args:
596 plan (LoadPlan): Load plan containing ReadItems.
597 planner (LoadPlanner): Load planner for resolving and committing tensors.
598 """
599 # Group ReadItems by storage file path (like torch per_file)
600 per_file = self._group_items_by_file(plan)
602 # Process each file
603 for path, reqs in per_file.items():
604 if not os.path.exists(path):
605 raise FileNotFoundError(f"Checkpoint file not found: {path}")
607 if path.endswith(".bytes"):
608 # BYTE_IO: one bytes file per rank with per-item offsets.
609 _load_bytes_file(path, reqs, planner, self.storage_data)
610 else:
611 # TENSOR: one safetensors file per rank
612 _load_tensor_file(path, reqs, planner, self.storage_data)