Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / filesystem_storage.py: 83%
198 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +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 pathlib import Path
19from typing import Any, Optional, Union
21from safetensors import safe_open
23from hyper_parallel.core.distributed_checkpoint.metadata import Metadata, MetadataIndex
24from hyper_parallel.core.distributed_checkpoint.planner import (
25 LoadPlan,
26 LoadPlanner,
27 ReadItem,
28 SavePlan,
29 SavePlanner,
30 WriteItem,
31)
32from hyper_parallel.core.distributed_checkpoint.storage import (
33 StorageInfo,
34 StorageReader,
35 StorageWriter,
36 WriteResult,
37 METADATA_FILE_NAME,
38)
39from hyper_parallel.core.distributed_checkpoint.util import narrow_tensor_by_index
40from hyper_parallel.platform import get_platform
41from hyper_parallel.platform.platform import PlatformType
44class FileSystemWriter(StorageWriter):
45 """
46 File system storage writer implementation.
48 Saves checkpoint data to the local file system, organizing tensors
49 into safetensors files and bytes into separate files.
50 """
52 def __init__(self, checkpoint_dir: Union[Path, str]):
53 self.checkpoint_dir = Path(checkpoint_dir) if isinstance(checkpoint_dir, str) else checkpoint_dir
54 self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
55 self.rank: int = 0
56 self.is_coordinator: bool = False
57 self.use_collectives: bool = True
59 def initialize_writer(self, checkpoint_id: Optional[Union[Path, str]] = None) -> None:
60 """
61 Initialize storage writer with new checkpoint directory.
63 Args:
64 checkpoint_id (Optional[Union[Path, str]]): New checkpoint directory path. Default None.
65 """
66 if checkpoint_id:
67 self.checkpoint_dir = Path(checkpoint_id) if isinstance(checkpoint_id, str) else checkpoint_id
68 self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
70 def configure_writer(self, is_coordinator: bool, **kwargs) -> None:
71 """
72 Configure storage writer.
74 Args:
75 is_coordinator (bool): Whether this rank is the coordinator.
76 **kwargs: Additional keyword arguments (e.g., rank, use_collectives).
77 """
78 self.is_coordinator = is_coordinator
79 self.rank = kwargs.get("rank") if "rank" in kwargs else get_platform().get_rank()
80 self.use_collectives = kwargs.get("use_collectives", True)
82 def optimize_local_plan(self, plan: SavePlan) -> SavePlan:
83 """
84 Optimize local plan.
86 Args:
87 plan (SavePlan): Local save plan.
89 Returns:
90 SavePlan: Optimized local plan.
91 """
92 return plan
94 def optimize_global_plan(self, plans: list[SavePlan]) -> list[SavePlan]:
95 """
96 Optimize global plan.
98 Args:
99 plans (list[SavePlan]): List of local plans from all ranks.
101 Returns:
102 list[SavePlan]: Optimized global plans.
103 """
104 return plans
107 def _serialize_bytes_item(self, item: WriteItem, planner: SavePlanner) -> bytes:
108 """Serialize a BYTE_IO item payload while preserving current behavior."""
109 data = planner.get_data(item)
110 if isinstance(data, bytes):
111 return data
112 return pickle.dumps(data)
115 def _write_bytes_items(self, plan: SavePlan, planner: SavePlanner) -> list[WriteResult]:
116 """
117 Write all BYTE_IO items into one per-rank bytes file.
119 Args:
120 plan (SavePlan): Save plan containing WriteItems.
121 planner (SavePlanner): Save planner used to resolve runtime data.
123 Returns:
124 list[WriteResult]: Write results for BYTE_IO items.
125 """
126 byte_items = [item for item in plan.items if item.type.value == "byte_io"]
127 if not byte_items:
128 return []
130 file_name = f"_rank{self.rank}_.bytes"
131 file_path = self.checkpoint_dir / file_name
133 results: list[WriteResult] = []
135 with open(file_path, "wb") as f:
136 for item in byte_items:
137 payload = self._serialize_bytes_item(item, planner)
138 offset = f.tell()
139 f.write(payload)
140 length = len(payload)
141 storage_info = StorageInfo(
142 relative_path=file_name,
143 offset=offset,
144 length=length,
145 )
146 results.append(
147 WriteResult(
148 index=item.index,
149 storage_data=storage_info,
150 )
151 )
153 return results
155 def _collect_tensors(self, plan: SavePlan, planner: SavePlanner) -> dict[str, Any]:
156 """
157 Collect tensor data from planner runtime lookup.
159 Args:
160 plan (SavePlan): Save plan containing WriteItems.
161 planner (SavePlanner): Save planner.
163 Returns:
164 dict[str, Any]: Dictionary mapping FQN to tensor data.
166 Raises:
167 RuntimeError: If tensor data cannot be resolved for an item.
168 """
169 tensor_dict: dict[str, Any] = {}
170 for item in plan.items:
171 if item.type.value == "tensor" and item.tensor_data:
172 tensor = planner.get_data(item)
173 if tensor is None:
174 raise RuntimeError(
175 f"Tensor data could not be resolved for index {item.index}. "
176 f"FQN: {item.index.fqn}"
177 )
178 fqn = item.index.fqn
179 tensor_dict[fqn] = tensor
180 return tensor_dict
182 def _write_tensors(self, plan: SavePlan, tensor_dict: dict[str, Any]) -> list[WriteResult]:
183 """
184 Write all tensors to safetensors file and create WriteResults.
186 Args:
187 plan (SavePlan): Save plan containing WriteItems.
188 tensor_dict (dict[str, Any]): Dictionary mapping FQN to tensor data.
190 Returns:
191 list[WriteResult]: List of write results for tensor items.
192 """
193 if not tensor_dict:
194 return []
196 platform = get_platform()
197 file_name = f"_rank{self.rank}_.safetensors"
198 file_path = self.checkpoint_dir / file_name
199 platform.save_checkpoint(tensor_dict, str(file_path))
201 # Record StorageInfo for each tensor
202 # Note: we don't know per-tensor byte offsets, so offset=0, length=-1
203 results: list[WriteResult] = []
204 for item in plan.items:
205 if item.type.value == "tensor" and item.tensor_data:
206 storage_info = StorageInfo(
207 relative_path=file_name,
208 offset=0,
209 length=-1,
210 )
211 results.append(
212 WriteResult(
213 index=item.index,
214 storage_data=storage_info,
215 )
216 )
217 return results
219 def execute_write(self, plan: SavePlan, planner: SavePlanner) -> list[WriteResult]:
220 """
221 Write data to storage and return per-item storage metadata.
223 Group tensors into safetensors files and bytes into separate files, recording StorageInfo for each item.
225 Args:
226 plan (SavePlan): Save plan containing WriteItems.
227 planner (SavePlanner): Save planner.
229 Returns:
230 list[WriteResult]: List of write results with storage metadata.
231 """
232 results: list[WriteResult] = []
234 # Write all BYTE_IO items into one file per rank
235 results.extend(self._write_bytes_items(plan, planner))
237 # Collect and write tensors
238 tensor_dict = self._collect_tensors(plan, planner)
239 results.extend(self._write_tensors(plan, tensor_dict))
241 return results
243 def finalize_checkpoint(self, metadata: Metadata, results: list[list[WriteResult]]) -> None:
244 """
245 Finish writing checkpoint and populate metadata.storage_data.
247 When use_collectives=True: only coordinator saves global metadata to .metadata.
248 When use_collectives=False: each rank saves its own metadata to .rank{rank}_metadata,
249 no cross-rank interaction.
251 Args:
252 metadata (Metadata): Checkpoint metadata to update.
253 results (list[list[WriteResult]]): Write results from all ranks (or single rank when use_collectives=False).
254 """
255 should_save = not self.use_collectives or (self.use_collectives and self.is_coordinator)
256 if not should_save:
257 return
259 # Build storage_data: map MetadataIndex -> StorageInfo
260 storage_md: dict[MetadataIndex, StorageInfo] = {}
261 for wr_list in results:
262 for wr in wr_list:
263 storage_md[wr.index] = wr.storage_data
264 metadata.storage_data = storage_md
266 # Save metadata file
267 if self.use_collectives:
268 metadata_file = self.checkpoint_dir / METADATA_FILE_NAME
269 else:
270 metadata_file = self.checkpoint_dir / f"{self.rank}{METADATA_FILE_NAME}"
271 with open(metadata_file, "wb") as f:
272 pickle.dump(metadata, f)
275def _copy_tensor_to_target(
276 req: ReadItem, tensor: Any, target_tensor: Any, planner: LoadPlanner
277) -> None:
278 """
279 Copy tensor data to target tensor and commit.
281 Args:
282 req (ReadItem): ReadItem request.
283 tensor (Any): Source tensor (tensor-like object).
284 target_tensor (Any): Target tensor (tensor-like object).
285 planner (LoadPlanner): Load planner for committing.
286 """
287 if hasattr(target_tensor, "copy_"):
288 target_tensor.copy_(tensor)
289 planner.apply_tensor(req, target_tensor)
290 else:
291 # mindspore or non-tensor: copy via commit path
292 planner.apply_tensor(req, tensor)
295def _load_bytes_file(
296 path: str,
297 reqs: list[ReadItem],
298 planner: LoadPlanner,
299 storage_data: dict[MetadataIndex, StorageInfo],
300) -> None:
301 """
302 Load bytes from a file.
304 Args:
305 path (str): Path to the bytes file.
306 reqs (list[ReadItem]): List of ReadItems for this file.
307 planner (LoadPlanner): Load planner for loading bytes.
308 """
309 with open(path, "rb") as f:
310 for req in reqs:
311 storage_info = storage_data.get(req.storage_index)
312 if storage_info is None:
313 raise KeyError(
314 f"StorageInfo not found for index {req.storage_index}"
315 )
316 f.seek(storage_info.offset)
317 value = f.read(storage_info.length)
318 planner.apply_bytes(req, value)
321def _get_tensor_size(tensor: Any) -> Optional[tuple]:
322 """
323 Get size/shape of a tensor.
325 Args:
326 tensor (Any): Tensor object (tensor-like with shape/size attribute).
328 Returns:
329 Optional[tuple]: Tuple of tensor size or None if not available.
330 """
331 if hasattr(tensor, "size") and callable(tensor.size):
332 return tuple(tensor.size())
333 return getattr(tensor, "shape", None)
336def _load_tensor_file(
337 path: str, reqs: list[ReadItem], planner: LoadPlanner
338) -> None:
339 """
340 Load and process tensors from a safetensors file.
342 Args:
343 path (str): Path to the safetensors file.
344 reqs (list[ReadItem]): List of ReadItems for this file.
345 planner (LoadPlanner): Load planner for resolving and committing tensors.
346 """
347 platform = get_platform()
349 if platform.platform_type == PlatformType.PYTORCH:
350 with safe_open(path, framework="pt", device="cpu") as tensor_file:
351 for req in reqs:
352 fqn = req.storage_index.fqn
353 if fqn not in tensor_file.keys():
354 raise KeyError(f"Key {fqn} not found in checkpoint file {path}")
355 tensor_slices = tuple(
356 slice(int(off), int(off) + int(length))
357 for off, length in zip(req.storage_offsets, req.lengths)
358 )
359 if tensor_slices:
360 tensor = tensor_file.get_slice(fqn)[tensor_slices]
361 else:
362 tensor = narrow_tensor_by_index(
363 tensor_file.get_tensor(fqn),
364 req.storage_offsets,
365 req.lengths,
366 )
368 target_tensor = planner.acquire_tensor(req)
369 if hasattr(target_tensor, "detach"):
370 target_tensor = target_tensor.detach()
372 # Size check (torch-aligned AssertionError)
373 target_size = _get_tensor_size(target_tensor)
374 tensor_size = _get_tensor_size(tensor)
375 if target_size is not None and tensor_size is not None:
376 if target_size != tensor_size:
377 raise AssertionError(
378 f"req {req.storage_index} mismatch sizes "
379 f"{target_size} vs {tensor_size}"
380 )
382 # Copy data to target
383 _copy_tensor_to_target(req, tensor, target_tensor, planner)
384 return
386 param_dict = platform.load_checkpoint(path)
387 for req in reqs:
388 fqn = req.storage_index.fqn
389 if fqn not in param_dict:
390 raise KeyError(f"Key {fqn} not found in checkpoint file {path}")
391 full_tensor = param_dict[fqn]
392 tensor = narrow_tensor_by_index(
393 full_tensor,
394 req.storage_offsets,
395 req.lengths,
396 )
398 target_tensor = planner.acquire_tensor(req)
399 if hasattr(target_tensor, "detach"):
400 target_tensor = target_tensor.detach()
402 # Size check (torch-aligned AssertionError)
403 target_size = _get_tensor_size(target_tensor)
404 tensor_size = _get_tensor_size(tensor)
405 if target_size is not None and tensor_size is not None:
406 if target_size != tensor_size:
407 raise AssertionError(
408 f"req {req.storage_index} mismatch sizes "
409 f"{target_size} vs {tensor_size}"
410 )
412 # Copy data to target
413 _copy_tensor_to_target(req, tensor, target_tensor, planner)
416class FileSystemReader(StorageReader):
417 """
418 File system storage reader implementation.
420 Reads checkpoint data from the local file system, loading tensors
421 from safetensors files and bytes from separate files.
422 """
424 def __init__(self, checkpoint_dir: Union[Path, str]):
425 self.checkpoint_dir = Path(checkpoint_dir) if isinstance(checkpoint_dir, str) else checkpoint_dir
426 # Cached storage layout: MetadataIndex -> StorageInfo (torch-aligned)
427 self.storage_data: Optional[dict[MetadataIndex, StorageInfo]] = None
428 self.rank: int = 0
429 self.is_coordinator: bool = False
431 def initialize_reader(self, checkpoint_id: Optional[Union[Path, str]] = None) -> None:
432 """
433 Initialize storage reader with new checkpoint directory.
435 Args:
436 checkpoint_id (Optional[Union[Path, str]]): New checkpoint directory path. Default None.
437 """
438 if checkpoint_id:
439 self.checkpoint_dir = Path(checkpoint_id) if isinstance(checkpoint_id, str) else checkpoint_id
441 def load_metadata(self, **kwargs) -> Metadata:
442 """
443 Load checkpoint metadata from file.
445 When rank is provided in kwargs: load rank-local metadata from .rank{rank}_metadata
446 (for checkpoints saved with use_collectives=False).
447 Otherwise: load global metadata from .metadata.
449 Args:
450 **kwargs: Optional arguments (e.g., rank for rank-local metadata).
452 Returns:
453 Metadata: Metadata object loaded from file.
454 """
455 rank = kwargs.get("rank")
456 if rank is not None:
457 metadata_file = self.checkpoint_dir / f"{rank}{METADATA_FILE_NAME}"
458 else:
459 metadata_file = self.checkpoint_dir / METADATA_FILE_NAME
461 if not metadata_file.exists():
462 raise FileNotFoundError(f"Metadata file not found: {metadata_file}")
463 with open(metadata_file, "rb") as f:
464 metadata = pickle.load(f)
465 return metadata
467 def configure_reader(self, metadata: Metadata, is_coordinator: bool, **kwargs) -> None:
468 """Configure storage reader."""
469 # Cache storage_data separately for quick lookup in execute_read.
470 # This mirrors torch.filesystem, where reader keeps a storage_data dict.
471 self.storage_data = getattr(metadata, "storage_data", None)
472 self.is_coordinator = is_coordinator
473 self.rank = kwargs.get("rank") if "rank" in kwargs else get_platform().get_rank()
475 def optimize_local_plan(self, plan: LoadPlan) -> LoadPlan:
476 """
477 Optimize local plan.
479 Args:
480 plan (LoadPlan): Local load plan.
482 Returns:
483 LoadPlan: Optimized local plan.
484 """
485 return plan
487 def optimize_global_plan(self, plans: list[LoadPlan]) -> list[LoadPlan]:
488 """
489 Optimize global plan.
491 Args:
492 plans (list[LoadPlan]): List of local plans from all ranks.
494 Returns:
495 list[LoadPlan]: Optimized global plans.
496 """
497 return plans
499 def _get_storage_path(self, read_item: ReadItem) -> str:
500 """
501 Get storage file path for a read item.
503 Args:
504 read_item (ReadItem): ReadItem to get path for.
506 Returns:
507 str: Absolute path to the storage file.
508 """
509 if self.storage_data is None:
510 raise KeyError("Checkpoint metadata.storage_data is required for filesystem read")
511 storage_info = self.storage_data.get(read_item.storage_index)
512 if storage_info is None:
513 raise KeyError(f"StorageInfo not found for index {read_item.storage_index}")
514 return str(self.checkpoint_dir / storage_info.relative_path)
516 def _group_items_by_file(self, plan: LoadPlan) -> dict[str, list]:
517 """
518 Group ReadItems by storage file path.
520 Args:
521 plan (LoadPlan): Load plan containing ReadItems.
523 Returns:
524 dict[str, list[ReadItem]]: Dictionary mapping file paths to lists of ReadItems.
525 """
526 per_file: dict[str, list] = {}
527 for read_item in plan.items:
528 path = self._get_storage_path(read_item)
529 per_file.setdefault(path, []).append(read_item)
530 return per_file
532 def execute_read(self, plan: LoadPlan, planner: LoadPlanner) -> None:
533 """
534 Read data from storage.
536 Aligned with torch filesystem read_data: groups ReadItems by file,
537 loads each file once, narrows tensors by storage_offsets/lengths for
538 resharding, then resolves/copies/commits data.
540 Args:
541 plan (LoadPlan): Load plan containing ReadItems.
542 planner (LoadPlanner): Load planner for resolving and committing tensors.
543 """
544 # Group ReadItems by storage file path (like torch per_file)
545 per_file = self._group_items_by_file(plan)
547 # Process each file
548 for path, reqs in per_file.items():
549 if not os.path.exists(path):
550 raise FileNotFoundError(f"Checkpoint file not found: {path}")
552 if path.endswith(".bytes"):
553 # BYTE_IO: one bytes file per rank with per-item offsets.
554 _load_bytes_file(path, reqs, planner, self.storage_data)
555 else:
556 # TENSOR: one safetensors file per rank
557 _load_tensor_file(path, reqs, planner)