Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / api.py: 72%
129 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"""Hyper Parallel Checkpoint API"""
16import multiprocessing as mp
17import queue
18import threading
19import traceback
20from concurrent.futures import Future
21from dataclasses import dataclass
22from enum import Enum, auto
23from pathlib import Path
24from typing import Any, Optional, Union
26from hyper_parallel.core.distributed_checkpoint.async_staging import build_staged_state_dict
27from hyper_parallel.core.distributed_checkpoint.standard_planner import StandardSavePlanner, StandardLoadPlanner
28from hyper_parallel.core.distributed_checkpoint.filesystem_storage import FileSystemReader, FileSystemWriter
29from hyper_parallel.core.distributed_checkpoint.metadata import Metadata
30from hyper_parallel.core.distributed_checkpoint.planner import SavePlanner, LoadPlanner
31from hyper_parallel.core.distributed_checkpoint.storage import StorageReader, StorageWriter
32from hyper_parallel.platform import get_platform
34platform = get_platform()
37class _AsyncPersistStatus(Enum):
38 """Queue payload status from :func:`_async_persist_worker` to the parent join thread."""
40 SUCCESS = auto()
41 FAILURE = auto()
44@dataclass
45class AsyncSaveResponse:
46 """Result of :func:`async_save`.
48 Host staging runs synchronously before :func:`async_save` returns; only checkpoint
49 **persistence** is asynchronous. ``persist_completion`` completes when the child
50 process finishes :func:`_save_impl` (plan, collectives, disk I/O) and supplies
51 :class:`Metadata`.
52 """
54 persist_completion: Future[Metadata]
57def _gather_from_all_ranks(
58 local_object: Any,
59 world_size: int,
60 use_collectives: bool,
61) -> list[Any]:
62 """
63 Gather objects from all ranks.
65 Args:
66 local_object (Any): Local object for current rank.
67 world_size (int): Total number of ranks.
68 use_collectives (bool): Whether to use collective communication.
70 Returns:
71 list[Any]: List of all objects from all ranks.
72 """
73 if use_collectives and world_size > 1:
74 all_objects = [None] * world_size
75 platform.all_gather_object(all_objects, local_object)
76 return all_objects
77 return [local_object]
80def _save_impl(
81 state_dict: dict[str, Any],
82 *,
83 checkpoint_id: Optional[Union[Path, str]] = None,
84 storage_writer: Optional[StorageWriter] = None,
85 planner: Optional[SavePlanner] = None,
86 no_dist: bool = False,
87 use_collectives: bool = True,
88) -> Metadata:
89 """Synchronous distributed checkpoint save (shared by :func:`save` and :func:`async_save`)."""
90 # Convert checkpoint_id to Path if it's a string
91 checkpoint_id = Path(checkpoint_id) if isinstance(checkpoint_id, str) else checkpoint_id
93 # Determine if we're in distributed mode
94 use_collectives = False if no_dist else use_collectives
96 # Set up storage writer
97 if storage_writer is None:
98 if checkpoint_id is None:
99 raise ValueError("Either storage_writer or checkpoint_id must be provided")
100 storage_writer = FileSystemWriter(checkpoint_id)
101 else:
102 if checkpoint_id:
103 storage_writer.initialize_writer(checkpoint_id)
105 # Set up planner
106 planner = StandardSavePlanner() if planner is None else planner
108 # Get rank and coordinator info
109 rank = platform.get_rank()
110 world_size = platform.get_world_size()
111 is_coordinator = rank == 0
113 # Configure planner
114 planner.configure_planner(
115 state_dict=state_dict,
116 is_coordinator=is_coordinator,
117 rank=rank,
118 use_collectives=use_collectives
119 )
121 # Configure storage writer (use_collectives for rank-local metadata when False)
122 storage_writer.configure_writer(
123 is_coordinator=is_coordinator,
124 rank=rank,
125 use_collectives=use_collectives
126 )
128 cached_res = planner.get_cached() if hasattr(planner, 'get_cached') else None
129 if cached_res:
130 # Get final plan and metadata from cache
131 final_plan, metadata = cached_res.final_plan, cached_res.metadata
133 else:
134 # Build local plan
135 local_plan = planner.build_local_plan()
136 local_plan = storage_writer.optimize_local_plan(local_plan)
138 # Gather all local plans and build global plan
139 all_local_plans = _gather_from_all_ranks(local_plan, world_size, use_collectives)
140 global_plans, metadata = planner.build_global_plan(all_local_plans)
141 global_plans = storage_writer.optimize_global_plan(global_plans)
143 # Select central plan for current rank
144 if use_collectives and world_size > 1 and global_plans:
145 central_plan = global_plans[rank]
146 elif global_plans:
147 central_plan = global_plans[0]
148 else:
149 central_plan = local_plan
151 # Finalize and cache plan
152 final_plan = planner.finalize_plan(central_plan)
153 # Add final plan and metadata to the cache
154 if hasattr(planner, 'cache_result'):
155 planner.cache_result(final_plan, metadata)
157 # Write data
158 write_results = storage_writer.execute_write(final_plan, planner)
160 # Finalize checkpoint
161 all_write_results = _gather_from_all_ranks(write_results, world_size, use_collectives)
162 storage_writer.finalize_checkpoint(metadata, all_write_results)
164 return metadata
167def _async_persist_worker(
168 result_queue: mp.Queue,
169 staged: dict[str, Any],
170 checkpoint_id: Optional[Union[Path, str]],
171 storage_writer: Optional[StorageWriter],
172 planner: Optional[SavePlanner],
173 no_dist: bool,
174 use_collectives: bool,
175) -> None:
176 """Child-process entry: run :func:`_save_impl` and report ``Metadata`` or an error string on ``result_queue``."""
177 try:
178 meta = _save_impl(
179 staged,
180 checkpoint_id=checkpoint_id,
181 storage_writer=storage_writer,
182 planner=planner,
183 no_dist=no_dist,
184 use_collectives=use_collectives,
185 )
186 result_queue.put((_AsyncPersistStatus.SUCCESS, meta))
187 except Exception: # pylint: disable=broad-except
188 result_queue.put((_AsyncPersistStatus.FAILURE, traceback.format_exc()))
191def _async_persist_wait_process(
192 proc: mp.Process,
193 result_queue: mp.Queue,
194 persist_future: Future[Metadata],
195) -> None:
196 """Join persist ``proc`` and complete ``persist_future`` (runs on a daemon thread)."""
197 proc.join()
198 if persist_future.done():
199 return
200 try:
201 status, payload = result_queue.get_nowait()
202 except queue.Empty:
203 persist_future.set_exception(
204 RuntimeError(
205 f"async_persist process exited with code {proc.exitcode} and no result on queue"
206 )
207 )
208 return
209 if status == _AsyncPersistStatus.SUCCESS:
210 persist_future.set_result(payload)
211 elif status == _AsyncPersistStatus.FAILURE:
212 persist_future.set_exception(RuntimeError(payload))
213 else:
214 persist_future.set_exception(
215 RuntimeError(f"async_persist queue returned unexpected status: {status!r}")
216 )
219def save(
220 state_dict: dict[str, Any],
221 *,
222 checkpoint_id: Optional[Union[Path, str]] = None,
223 storage_writer: Optional[StorageWriter] = None,
224 planner: Optional[SavePlanner] = None,
225 no_dist: bool = False,
226 use_collectives: bool = True,
227) -> Metadata:
228 """
229 Save a distributed checkpoint in SPMD style.
231 This function saves a state_dict containing DTensors, where each rank
232 only saves their local shards.
234 Args:
235 state_dict (dict[str, Any]): The state_dict to save.
236 checkpoint_id (Optional[Union[Path, str]]): The ID/path of this checkpoint instance (can be Path or str).
237 Default None.
238 storage_writer (Optional[StorageWriter]): Instance of StorageWriter. If None, FileSystemWriter
239 will be created based on checkpoint_id. Default None.
240 planner (Optional[SavePlanner]): Instance of SavePlanner. If None, StandardSavePlanner will be used.
241 Default None.
242 no_dist (bool): If True, save in single process mode. Default False.
243 use_collectives (bool): If True, use collective communication for coordination.
244 If False, each rank saves its own shard data and rank-local metadata (.metadata_rank{rank}),
245 with no cross-rank interaction. Default True.
247 Returns:
248 Metadata: Metadata object for the saved checkpoint.
249 """
250 metadata = _save_impl(
251 state_dict,
252 checkpoint_id=checkpoint_id,
253 storage_writer=storage_writer,
254 planner=planner,
255 no_dist=no_dist,
256 use_collectives=use_collectives,
257 )
258 platform.barrier()
259 return metadata
262def async_save(
263 state_dict: dict[str, Any],
264 *,
265 checkpoint_id: Optional[Union[Path, str]] = None,
266 storage_writer: Optional[StorageWriter] = None,
267 planner: Optional[SavePlanner] = None,
268 no_dist: bool = False,
269 use_collectives: bool = True,
270) -> AsyncSaveResponse:
271 """
272 Asynchronous version of :func:`save` using a **background child process** for persistence.
274 **Staging** (tensor / DTensor → host copy) runs **synchronously in the caller
275 process** via :func:`build_staged_state_dict`, so no process pool is used for
276 staging and the training stack sees a normal Python call path. When this
277 function returns successfully, host staging is done and the original
278 ``state_dict`` may be mutated.
280 **Persistence** (plan, collectives, disk I/O) runs in **one** background
281 :class:`multiprocessing.Process` that executes :func:`_save_impl` on the staged
282 dict. A small daemon **thread** only joins that process and fills
283 ``persist_completion``; it does not perform tensor work.
285 The staged dict and ``storage_writer`` / ``planner`` must be picklable for the
286 persist child process (same constraints as before for the worker path).
288 .. warning::
289 Experimental API. Always wait on ``persist_completion`` for a fully persisted checkpoint.
291 Args:
292 state_dict (dict[str, Any]): The state_dict to save.
293 checkpoint_id (Optional[Union[Path, str]]): Same as :func:`save`.
294 storage_writer (Optional[StorageWriter]): Same as :func:`save`.
295 planner (Optional[SavePlanner]): Same as :func:`save`.
296 no_dist (bool): Same as :func:`save`.
297 use_collectives (bool): Same as :func:`save`.
299 Returns:
300 AsyncSaveResponse: Contains ``persist_completion`` only; staging is synchronous.
301 """
302 persist_completion: Future[Metadata] = Future()
304 staged = build_staged_state_dict(state_dict)
306 result_queue: mp.Queue = mp.Queue(maxsize=1)
307 proc = mp.Process(
308 target=_async_persist_worker,
309 args=(
310 result_queue,
311 staged,
312 checkpoint_id,
313 storage_writer,
314 planner,
315 no_dist,
316 use_collectives,
317 ),
318 name="HPAsyncCheckpointPersist",
319 )
320 proc.start()
321 join_thread = threading.Thread(
322 target=_async_persist_wait_process,
323 args=(proc, result_queue, persist_completion),
324 daemon=True,
325 name="HPAsyncCheckpointPersistJoin",
326 )
327 join_thread.start()
328 return AsyncSaveResponse(persist_completion=persist_completion)
331def load(
332 state_dict: dict[str, Any],
333 *,
334 checkpoint_id: Optional[Union[Path, str]] = None,
335 storage_reader: Optional[StorageReader] = None,
336 planner: Optional[LoadPlanner] = None,
337 no_dist: bool = False,
338 use_collectives: bool = True,
339) -> None:
340 """
341 Load a distributed checkpoint into state_dict in SPMD style.
343 Each rank will try to read the least amount of data necessary
344 to fulfill the requested state_dict. When loading DTensor instances,
345 each rank only reads data for their local shards.
347 Args:
348 state_dict (dict[str, Any]): The state_dict to load the checkpoint into (modified in-place).
349 checkpoint_id (Optional[Union[Path, str]]): The ID/path of this checkpoint instance (can be Path or str).
350 Default None.
351 storage_reader (Optional[StorageReader]): Instance of StorageReader. If None, FileSystemReader
352 will be created based on checkpoint_id. Default None.
353 planner (Optional[LoadPlanner]): Instance of LoadPlanner. If None, StandardLoadPlanner will be used.
354 Default None.
355 no_dist (bool): If True, load without cross-rank synchronization. Default False.
356 use_collectives (bool): If False, load from rank-local metadata (.metadata_rank{rank}),
357 for checkpoints saved with save(use_collectives=False). No cross-rank interaction. Default True.
359 Returns:
360 None. The state_dict is modified in-place.
361 """
362 # Convert checkpoint_id to Path if it's a string
363 checkpoint_id = Path(checkpoint_id) if isinstance(checkpoint_id, str) else checkpoint_id
365 # Determine if we're in distributed mode
366 use_collectives = False if no_dist else use_collectives
368 # Set up storage reader
369 if storage_reader is None:
370 if checkpoint_id is None:
371 raise ValueError("Either storage_reader or checkpoint_id must be provided")
372 storage_reader = FileSystemReader(checkpoint_id)
373 else:
374 if checkpoint_id:
375 storage_reader.initialize_reader(checkpoint_id)
377 # Set up planner
378 planner = StandardLoadPlanner() if planner is None else planner
380 # Get rank and coordinator info
381 rank = platform.get_rank()
382 world_size = platform.get_world_size()
383 is_coordinator = rank == 0
385 # Load metadata
386 try:
387 metadata = storage_reader.load_metadata()
388 except FileNotFoundError:
389 # Fallback to rank-local metadata (e.g. checkpoint saved with use_collectives=False)
390 metadata = storage_reader.load_metadata(rank=rank)
391 use_collectives = False
393 # Configure planner
394 planner.configure_planner(
395 state_dict=state_dict,
396 metadata=metadata,
397 is_coordinator=is_coordinator,
398 rank=rank
399 )
401 # Configure storage reader
402 storage_reader.configure_reader(
403 metadata=metadata,
404 is_coordinator=is_coordinator,
405 rank=rank,
406 use_collectives=use_collectives
407 )
409 # Build local plan
410 local_plan = planner.build_local_plan()
411 local_plan = storage_reader.optimize_local_plan(local_plan)
413 # Gather all local plans and build global plan
414 all_local_plans = _gather_from_all_ranks(local_plan, world_size, use_collectives)
415 global_plans = planner.build_global_plan(all_local_plans)
416 global_plans = storage_reader.optimize_global_plan(global_plans)
418 # Select central plan for current rank
419 if use_collectives and world_size > 1 and global_plans:
420 central_plan = global_plans[rank]
421 elif global_plans:
422 central_plan = global_plans[0]
423 else:
424 central_plan = local_plan
426 # Finalize plan
427 final_plan = planner.finalize_plan(central_plan)
429 # Execute read
430 storage_reader.execute_read(final_plan, planner)