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-10-25 05:45 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd. All rights reserved.
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# ============================================================================
16"""Storage interfaces for checkpoint save and load."""
18import abc
19from dataclasses import dataclass
20from pathlib import Path
21from typing import Optional, Union
23from hyper_parallel.core.distributed_checkpoint.metadata import Metadata, MetadataIndex
24from hyper_parallel.core.distributed_checkpoint.planner import LoadPlan, LoadPlanner, SavePlan, SavePlanner
26METADATA_FILE_NAME = ".metadata"
29@dataclass
30class StorageInfo:
31 """
32 Storage information for a single logical item.
34 Torch-aligned: matches torch.distributed.checkpoint.filesystem._StorageInfo.
36 Attributes:
37 relative_path: Path relative to checkpoint root.
38 offset: Byte offset within the file.
39 length: Byte length of the data (best-effort, may be -1 for tensors).
40 tensor_key: Optional physical key used inside a tensor container file.
41 """
42 relative_path: str
43 offset: int
44 length: int
45 tensor_key: Optional[str] = None
48@dataclass
49class WriteResult:
50 """
51 Result of writing a single logical item.
53 Torch-aligned: contains the metadata index and storage information.
54 """
55 index: MetadataIndex
56 storage_data: StorageInfo
59class StorageWriter(abc.ABC):
60 """
61 Abstract base class for storage writers.
63 Defines the interface for writing checkpoint data to storage backends.
64 """
66 @abc.abstractmethod
67 def initialize_writer(self, checkpoint_id: Optional[Union[Path, str]] = None) -> None:
68 """
69 Initialize storage writer with optional new checkpoint directory.
71 Args:
72 checkpoint_id (Optional[Union[Path, str]]): The ID/path of the checkpoint directory.
73 If None, uses the previously configured checkpoint directory. Default None.
74 """
76 @abc.abstractmethod
77 def configure_writer(self, is_coordinator: bool, **kwargs) -> None:
78 """
79 Configure storage writer with coordinator and rank information.
81 Args:
82 is_coordinator (bool): Whether this rank is the coordinator rank.
83 **kwargs: Additional keyword arguments (e.g., rank, use_collectives).
84 """
86 @abc.abstractmethod
87 def optimize_local_plan(self, plan: SavePlan) -> SavePlan:
88 """
89 Optimize local plan for storage-specific optimizations.
91 Args:
92 plan (SavePlan): The local save plan to optimize.
94 Returns:
95 SavePlan: The optimized local save plan.
96 """
98 @abc.abstractmethod
99 def optimize_global_plan(self, plan: SavePlan) -> SavePlan:
100 """
101 Optimize this rank's save plan once the global plan has been built.
103 Args:
104 plan (SavePlan): This rank's save plan, with storage indices assigned.
106 Returns:
107 SavePlan: The optimized save plan.
108 """
110 @abc.abstractmethod
111 def execute_write(self, plan: SavePlan, planner: SavePlanner) -> list[WriteResult]:
112 """
113 Execute write operation to storage and return write results.
115 Args:
116 plan (SavePlan): The save plan to execute.
117 planner (SavePlanner): The save planner instance for accessing tensor data.
119 Returns:
120 list[WriteResult]: List of write results containing storage information for each written item.
121 """
123 @abc.abstractmethod
124 def finalize_checkpoint(self, metadata: Metadata, results: list[list[WriteResult]]) -> None:
125 """
126 Finalize checkpoint writing and complete metadata.
128 Args:
129 metadata (Metadata): The checkpoint metadata to finalize.
130 results (list[list[WriteResult]]): List of write results from all ranks,
131 where each inner list contains WriteResults from one rank.
132 """
135class StorageReader(abc.ABC):
136 """
137 Abstract base class for storage readers.
139 Defines the interface for reading checkpoint data from storage backends.
140 """
142 @abc.abstractmethod
143 def initialize_reader(self, checkpoint_id: Optional[Union[Path, str]] = None) -> None:
144 """
145 Initialize storage reader with optional new checkpoint directory.
147 Args:
148 checkpoint_id (Optional[Union[Path, str]]): The ID/path of the checkpoint directory.
149 If None, uses the previously configured checkpoint directory. Default None.
150 """
152 @abc.abstractmethod
153 def load_metadata(self, **kwargs) -> Metadata:
154 """
155 Load checkpoint metadata from storage.
157 Args:
158 **kwargs: Additional keyword arguments (e.g., rank for rank-local metadata).
160 Returns:
161 Metadata: The loaded checkpoint metadata.
162 """
164 @abc.abstractmethod
165 def configure_reader(self, metadata: Metadata, is_coordinator: bool, **kwargs) -> None:
166 """
167 Configure storage reader with metadata and coordinator information.
169 Args:
170 metadata (Metadata): The checkpoint metadata.
171 is_coordinator (bool): Whether this rank is the coordinator rank.
172 **kwargs: Additional keyword arguments (e.g., rank, use_collectives).
173 """
175 @abc.abstractmethod
176 def optimize_local_plan(self, plan: LoadPlan) -> LoadPlan:
177 """
178 Optimize local plan for storage-specific optimizations.
180 Args:
181 plan (LoadPlan): The local load plan to optimize.
183 Returns:
184 LoadPlan: The optimized local load plan.
185 """
187 @abc.abstractmethod
188 def optimize_global_plan(self, plan: LoadPlan) -> LoadPlan:
189 """
190 Optimize this rank's load plan once the global plan has been built.
192 Args:
193 plan (LoadPlan): This rank's load plan.
195 Returns:
196 LoadPlan: The optimized load plan.
197 """
199 @abc.abstractmethod
200 def execute_read(
201 self,
202 plan: LoadPlan,
203 planner: LoadPlanner,
204 broadcast_groups: Optional[dict] = None,
205 broadcast_batch_bytes: int = 0,
206 ) -> None:
207 """
208 Execute read operation from storage according to the load plan.
210 Args:
211 plan (LoadPlan): The load plan to execute.
212 planner (LoadPlanner): The load planner instance for applying loaded data.
213 broadcast_groups (Optional[dict]): Communication groups for broadcast.
214 Only consulted when the reader was configured with
215 the global plan marked as read by another rank; omit it for a plain read.
216 broadcast_batch_bytes (int): Shards smaller than this may be sent together
217 rather than one at a time. Zero sends every shard on its own.
218 """