Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / storage.py: 100%
43 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"""Storage interfaces for checkpoint save and load."""
16import abc
17from dataclasses import dataclass
18from pathlib import Path
19from typing import Optional, Union
21from hyper_parallel.core.distributed_checkpoint.metadata import Metadata, MetadataIndex
22from hyper_parallel.core.distributed_checkpoint.planner import LoadPlan, LoadPlanner, SavePlan, SavePlanner
24METADATA_FILE_NAME = ".metadata"
27@dataclass
28class StorageInfo:
29 """
30 Storage information for a single logical item.
32 Torch-aligned: matches torch.distributed.checkpoint.filesystem._StorageInfo.
34 Attributes:
35 relative_path: Path relative to checkpoint root.
36 offset: Byte offset within the file.
37 length: Byte length of the data (best-effort, may be -1 for tensors).
38 tensor_key: Optional physical key used inside a tensor container file.
39 """
40 relative_path: str
41 offset: int
42 length: int
43 tensor_key: Optional[str] = None
46@dataclass
47class WriteResult:
48 """
49 Result of writing a single logical item.
51 Torch-aligned: contains the metadata index and storage information.
52 """
53 index: MetadataIndex
54 storage_data: StorageInfo
57class StorageWriter(abc.ABC):
58 """
59 Abstract base class for storage writers.
61 Defines the interface for writing checkpoint data to storage backends.
62 """
64 @abc.abstractmethod
65 def initialize_writer(self, checkpoint_id: Optional[Union[Path, str]] = None) -> None:
66 """
67 Initialize storage writer with optional new checkpoint directory.
69 Args:
70 checkpoint_id (Optional[Union[Path, str]]): The ID/path of the checkpoint directory.
71 If None, uses the previously configured checkpoint directory. Default None.
72 """
74 @abc.abstractmethod
75 def configure_writer(self, is_coordinator: bool, **kwargs) -> None:
76 """
77 Configure storage writer with coordinator and rank information.
79 Args:
80 is_coordinator (bool): Whether this rank is the coordinator rank.
81 **kwargs: Additional keyword arguments (e.g., rank, use_collectives).
82 """
84 @abc.abstractmethod
85 def optimize_local_plan(self, plan: SavePlan) -> SavePlan:
86 """
87 Optimize local plan for storage-specific optimizations.
89 Args:
90 plan (SavePlan): The local save plan to optimize.
92 Returns:
93 SavePlan: The optimized local save plan.
94 """
96 @abc.abstractmethod
97 def optimize_global_plan(self, plans: list[SavePlan]) -> list[SavePlan]:
98 """
99 Optimize global plan from all local plans.
101 Args:
102 plans (list[SavePlan]): List of local save plans from all ranks.
104 Returns:
105 list[SavePlan]: List of optimized global save plans.
106 """
108 @abc.abstractmethod
109 def execute_write(self, plan: SavePlan, planner: SavePlanner) -> list[WriteResult]:
110 """
111 Execute write operation to storage and return write results.
113 Args:
114 plan (SavePlan): The save plan to execute.
115 planner (SavePlanner): The save planner instance for accessing tensor data.
117 Returns:
118 list[WriteResult]: List of write results containing storage information for each written item.
119 """
121 @abc.abstractmethod
122 def finalize_checkpoint(self, metadata: Metadata, results: list[list[WriteResult]]) -> None:
123 """
124 Finalize checkpoint writing and complete metadata.
126 Args:
127 metadata (Metadata): The checkpoint metadata to finalize.
128 results (list[list[WriteResult]]): List of write results from all ranks,
129 where each inner list contains WriteResults from one rank.
130 """
133class StorageReader(abc.ABC):
134 """
135 Abstract base class for storage readers.
137 Defines the interface for reading checkpoint data from storage backends.
138 """
140 @abc.abstractmethod
141 def initialize_reader(self, checkpoint_id: Optional[Union[Path, str]] = None) -> None:
142 """
143 Initialize storage reader with optional new checkpoint directory.
145 Args:
146 checkpoint_id (Optional[Union[Path, str]]): The ID/path of the checkpoint directory.
147 If None, uses the previously configured checkpoint directory. Default None.
148 """
150 @abc.abstractmethod
151 def load_metadata(self, **kwargs) -> Metadata:
152 """
153 Load checkpoint metadata from storage.
155 Args:
156 **kwargs: Additional keyword arguments (e.g., rank for rank-local metadata).
158 Returns:
159 Metadata: The loaded checkpoint metadata.
160 """
162 @abc.abstractmethod
163 def configure_reader(self, metadata: Metadata, is_coordinator: bool, **kwargs) -> None:
164 """
165 Configure storage reader with metadata and coordinator information.
167 Args:
168 metadata (Metadata): The checkpoint metadata.
169 is_coordinator (bool): Whether this rank is the coordinator rank.
170 **kwargs: Additional keyword arguments (e.g., rank, use_collectives).
171 """
173 @abc.abstractmethod
174 def optimize_local_plan(self, plan: LoadPlan) -> LoadPlan:
175 """
176 Optimize local plan for storage-specific optimizations.
178 Args:
179 plan (LoadPlan): The local load plan to optimize.
181 Returns:
182 LoadPlan: The optimized local load plan.
183 """
185 @abc.abstractmethod
186 def optimize_global_plan(self, plans: list[LoadPlan]) -> list[LoadPlan]:
187 """
188 Optimize global plan from all local plans.
190 Args:
191 plans (list[LoadPlan]): List of local load plans from all ranks.
193 Returns:
194 list[LoadPlan]: List of optimized global load plans.
195 """
197 @abc.abstractmethod
198 def execute_read(self, plan: LoadPlan, planner: LoadPlanner) -> None:
199 """
200 Execute read operation from storage according to the load plan.
202 Args:
203 plan (LoadPlan): The load plan to execute.
204 planner (LoadPlanner): The load planner instance for applying loaded data.
205 """