Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / filesystem_storage.py: 95%
266 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# ============================================================================
15"""File system storage implementations for checkpoint save and load."""
16import pickle
17from collections import Counter, deque
18from pathlib import Path
19from typing import Any, Optional, Union
21import torch.distributed as dist
22from safetensors import safe_open
24from hyper_parallel.core.distributed_checkpoint.metadata import (
25 Metadata,
26 MetadataIndex,
27)
28from hyper_parallel.core.distributed_checkpoint.planner import (
29 LoadItemType,
30 LoadPlan,
31 LoadPlanner,
32 ReadItem,
33 SavePlan,
34 SavePlanner,
35 WriteItem,
36)
37from hyper_parallel.core.distributed_checkpoint.storage import (
38 StorageInfo,
39 StorageReader,
40 StorageWriter,
41 WriteResult,
42 METADATA_FILE_NAME,
43)
44from hyper_parallel.core.distributed_checkpoint.broadcast import (
45 BroadcastBatcher,
46 wait_broadcasts,
47)
48from hyper_parallel.core.distributed_checkpoint.checkpoint_io import save_checkpoint_file
49from hyper_parallel.core.distributed_checkpoint.utils import (
50 dcp_timer_decorator,
51 logger,
52)
55class _MetadataUnpickler(pickle.Unpickler):
56 """Unpickler for the metadata file, remapping module names while loading.
58 ``_FOREIGN_MODULE_MAP`` maps a module name recorded in the pickle to the module to import
59 in its place, so that metadata written by a tree where these classes live under other
60 module names still loads. Add an entry per renamed module to support such a tree. It
61 is empty here: metadata written by this tree names bare ``hyper_parallel.*`` modules,
62 which are imported as they are.
63 """
65 _FOREIGN_MODULE_MAP = {}
67 def find_class(self, module: str, name: str) -> Any:
68 """
69 Resolve a pickled reference, sending the modules that moved to where they now live.
71 Args:
72 module (str): Module the pickle names.
73 name (str): Name to resolve within it.
75 Returns:
76 Any: The class or function the reference stands for.
77 """
78 if module in _MetadataUnpickler._FOREIGN_MODULE_MAP:
79 module = _MetadataUnpickler._FOREIGN_MODULE_MAP[module]
80 return super().find_class(module, name)
83class FileSystemWriter(StorageWriter):
84 """
85 File system storage writer implementation.
87 Saves checkpoint data to the local file system, organizing tensors
88 into safetensors files and bytes into separate files.
89 """
91 def __init__(self, checkpoint_dir: Union[Path, str]) -> None:
92 """
93 Args:
94 checkpoint_dir (Union[Path, str]): Directory the checkpoint is written into,
95 made if it is not there yet.
96 """
97 self.checkpoint_dir = Path(checkpoint_dir) if isinstance(checkpoint_dir, str) else checkpoint_dir
98 self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
99 self.rank: int = 0
100 self.is_coordinator: bool = False
101 self.use_collectives: bool = True
103 def initialize_writer(self, checkpoint_id: Optional[Union[Path, str]] = None) -> None:
104 """
105 Initialize storage writer with new checkpoint directory.
107 Args:
108 checkpoint_id (Optional[Union[Path, str]]): New checkpoint directory path. Default None.
109 """
110 if checkpoint_id:
111 self.checkpoint_dir = Path(checkpoint_id) if isinstance(checkpoint_id, str) else checkpoint_id
112 self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
114 def configure_writer(self, is_coordinator: bool, **kwargs: Any) -> None:
115 """
116 Configure storage writer.
118 Args:
119 is_coordinator (bool): Whether this rank is the coordinator.
120 **kwargs: Additional keyword arguments (e.g., rank, use_collectives).
121 """
122 self.is_coordinator = is_coordinator
123 self.rank = kwargs.get("rank") if "rank" in kwargs else dist.get_rank()
124 self.use_collectives = kwargs.get("use_collectives", True)
126 def optimize_local_plan(self, plan: SavePlan) -> SavePlan:
127 """
128 Optimize local plan.
130 Args:
131 plan (SavePlan): Local save plan.
133 Returns:
134 SavePlan: Optimized local plan.
135 """
136 return plan
138 def optimize_global_plan(self, plan: SavePlan) -> SavePlan:
139 """
140 Optimize this rank's save plan.
142 Args:
143 plan (SavePlan): This rank's save plan, with storage indices assigned.
145 Returns:
146 SavePlan: Optimized save plan.
147 """
148 return plan
151 def _serialize_bytes_item(self, item: WriteItem, planner: SavePlanner) -> bytes:
152 """Serialize a BYTE_IO item payload while preserving current behavior."""
153 data = planner.get_data(item)
154 if isinstance(data, bytes):
155 return data
156 return pickle.dumps(data)
159 @dcp_timer_decorator
160 def _write_bytes_items(self, plan: SavePlan, planner: SavePlanner) -> list[WriteResult]:
161 """
162 Write all BYTE_IO items into one per-rank bytes file.
164 Args:
165 plan (SavePlan): Save plan containing WriteItems.
166 planner (SavePlanner): Save planner used to resolve runtime data.
168 Returns:
169 list[WriteResult]: Write results for BYTE_IO items.
170 """
171 byte_items = [item for item in plan.items if item.type.value == "byte_io"]
172 if not byte_items:
173 return []
175 file_name = f"_rank{self.rank}_.bytes"
176 file_path = self.checkpoint_dir / file_name
178 results: list[WriteResult] = []
180 with open(file_path, "wb") as f:
181 for item in byte_items:
182 payload = self._serialize_bytes_item(item, planner)
183 offset = f.tell()
184 f.write(payload)
185 length = len(payload)
186 storage_info = StorageInfo(
187 relative_path=file_name,
188 offset=offset,
189 length=length,
190 )
191 results.append(
192 WriteResult(
193 index=item.index,
194 storage_data=storage_info,
195 )
196 )
198 return results
200 def _collect_tensors(
201 self, plan: SavePlan, planner: SavePlanner
202 ) -> tuple[dict[str, Any], dict[MetadataIndex, str]]:
203 """
204 Collect tensor data from planner runtime lookup.
206 Args:
207 plan (SavePlan): Save plan containing WriteItems.
208 planner (SavePlanner): Save planner.
210 Returns:
211 tuple[dict[str, Any], dict[MetadataIndex, str]]: Tensor data keyed by
212 physical safetensors key, plus logical index-to-key mapping.
214 Raises:
215 RuntimeError: If tensor data cannot be resolved for an item.
216 """
217 tensor_items = [
218 item for item in plan.items
219 if item.type.value == "tensor" and item.tensor_data
220 ]
221 fqn_counts = Counter(item.index.fqn for item in tensor_items)
222 reserved_keys = set(fqn_counts)
223 used_keys: set[str] = set()
224 next_chunk_index: dict[str, int] = {}
225 tensor_dict: dict[str, Any] = {}
226 tensor_keys: dict[MetadataIndex, str] = {}
228 for item in tensor_items:
229 tensor = planner.get_data(item)
230 if tensor is None:
231 raise RuntimeError(
232 f"Tensor data could not be resolved for index {item.index}. "
233 f"FQN: {item.index.fqn}"
234 )
236 fqn = item.index.fqn
237 tensor_key = fqn
238 if fqn_counts[fqn] > 1:
239 chunk_index = next_chunk_index.get(fqn, 0)
240 next_chunk_index[fqn] = chunk_index + 1
241 tensor_key = f"{fqn}.__dcp_chunk_{chunk_index}"
242 while tensor_key in reserved_keys or tensor_key in used_keys:
243 tensor_key += "_"
245 used_keys.add(tensor_key)
246 tensor_dict[tensor_key] = tensor
247 tensor_keys[item.index] = tensor_key
248 return tensor_dict, tensor_keys
250 @dcp_timer_decorator
251 def _write_tensors(
252 self,
253 plan: SavePlan,
254 tensor_dict: dict[str, Any],
255 tensor_keys: dict[MetadataIndex, str],
256 ) -> list[WriteResult]:
257 """
258 Write all tensors to safetensors file and create WriteResults.
260 Args:
261 plan (SavePlan): Save plan containing WriteItems.
262 tensor_dict (dict[str, Any]): Dictionary mapping physical keys to tensor data.
263 tensor_keys (dict[MetadataIndex, str]): Logical index-to-key mapping.
265 Returns:
266 list[WriteResult]: List of write results for tensor items.
267 """
268 if not tensor_dict:
269 return []
271 file_name = f"_rank{self.rank}_.safetensors"
272 file_path = self.checkpoint_dir / file_name
273 save_checkpoint_file(tensor_dict, str(file_path))
275 # Record StorageInfo for each tensor
276 # Note: we don't know per-tensor byte offsets, so offset=0, length=-1
277 results: list[WriteResult] = []
278 for item in plan.items:
279 if item.type.value == "tensor" and item.tensor_data:
280 storage_info = StorageInfo(
281 relative_path=file_name,
282 offset=0,
283 length=-1,
284 tensor_key=tensor_keys[item.index],
285 )
286 results.append(
287 WriteResult(
288 index=item.index,
289 storage_data=storage_info,
290 )
291 )
292 return results
294 @dcp_timer_decorator
295 def execute_write(self, plan: SavePlan, planner: SavePlanner) -> list[WriteResult]:
296 """
297 Write data to storage and return per-item storage metadata.
299 Group tensors into safetensors files and bytes into separate files, recording StorageInfo for each item.
301 Args:
302 plan (SavePlan): Save plan containing WriteItems.
303 planner (SavePlanner): Save planner.
305 Returns:
306 list[WriteResult]: List of write results with storage metadata.
307 """
308 results: list[WriteResult] = []
310 # Write all BYTE_IO items into one file per rank
311 results.extend(self._write_bytes_items(plan, planner))
313 # Collect and write tensors
314 tensor_dict, tensor_keys = self._collect_tensors(plan, planner)
315 results.extend(self._write_tensors(plan, tensor_dict, tensor_keys))
317 return results
319 def finalize_checkpoint(self, metadata: Metadata, results: list[list[WriteResult]]) -> None:
320 """
321 Finish writing checkpoint and populate metadata.storage_data.
323 When use_collectives=True: only coordinator saves global metadata to .metadata.
324 When use_collectives=False: each rank saves its own metadata to .rank{rank}_metadata,
325 no cross-rank interaction.
327 Args:
328 metadata (Metadata): Checkpoint metadata to update.
329 results (list[list[WriteResult]]): Write results from all ranks (or single rank when use_collectives=False).
330 """
331 should_save = not self.use_collectives or (self.use_collectives and self.is_coordinator)
332 if not should_save:
333 return
335 # Build storage_data: map MetadataIndex -> StorageInfo
336 storage_md: dict[MetadataIndex, StorageInfo] = {}
337 for wr_list in results:
338 for wr in wr_list:
339 storage_md[wr.index] = wr.storage_data
340 metadata.storage_data = storage_md
342 # Save metadata file
343 if self.use_collectives:
344 metadata_file = self.checkpoint_dir / METADATA_FILE_NAME
345 else:
346 metadata_file = self.checkpoint_dir / f"{self.rank}{METADATA_FILE_NAME}"
347 with open(metadata_file, "wb") as f:
348 pickle.dump(metadata, f)
351def _copy_tensor_to_target(
352 req: ReadItem, tensor: Any, target_tensor: Any, planner: LoadPlanner
353) -> None:
354 """
355 Copy tensor data to target tensor and commit.
357 Args:
358 req (ReadItem): ReadItem request.
359 tensor (Any): Source tensor (tensor-like object).
360 target_tensor (Any): Target tensor (tensor-like object).
361 planner (LoadPlanner): Load planner for committing.
362 """
363 if hasattr(target_tensor, "copy_"):
364 # for torch and ms, call 'copy_' to copy data to target tensor
365 target_tensor.copy_(tensor)
366 else:
367 target_tensor[...] = tensor
368 planner.apply_tensor(req, target_tensor)
371def _fetch_bytes_file(
372 path: str,
373 reqs: list[ReadItem],
374 storage_data: dict[MetadataIndex, StorageInfo],
375) -> list[tuple[ReadItem, Any]]:
376 """
377 Read the payload of each item out of a bytes file, leaving it packed.
379 Args:
380 path (str): Path to the bytes file.
381 reqs (list[ReadItem]): List of ReadItems for this file.
382 storage_data (dict[MetadataIndex, StorageInfo]): Physical storage mapping.
384 Returns:
385 list[tuple[ReadItem, Any]]: Each item with the bytes it asked for. Unpacking them
386 into objects needs the planner, so it is left to :func:`_apply_fetched`.
387 """
388 fetched: list[tuple[ReadItem, Any]] = []
389 with open(path, "rb") as f:
390 for req in reqs:
391 storage_info = _get_storage_info(req, storage_data)
392 f.seek(storage_info.offset)
393 fetched.append((req, f.read(storage_info.length)))
394 return fetched
397def _get_tensor_size(tensor: Any) -> Optional[tuple]:
398 """
399 Get size/shape of a tensor.
401 Args:
402 tensor (Any): Tensor object (tensor-like with shape/size attribute).
404 Returns:
405 Optional[tuple]: Tuple of tensor size or None if not available.
406 """
407 if hasattr(tensor, "size") and callable(tensor.size):
408 return tuple(tensor.size())
409 return getattr(tensor, "shape", None)
412def _get_storage_info(
413 req: ReadItem,
414 storage_data: dict[MetadataIndex, StorageInfo],
415) -> StorageInfo:
416 """Return physical storage metadata for one read request."""
417 storage_info = storage_data.get(req.storage_index)
418 if storage_info is None:
419 raise KeyError(f"StorageInfo not found for index {req.storage_index}")
420 return storage_info
423def _validate_and_copy_tensor(
424 req: ReadItem,
425 tensor: Any,
426 planner: LoadPlanner,
427) -> None:
428 """Validate a loaded tensor slice and copy it to its planner destination."""
429 target_tensor = planner.acquire_tensor(req)
431 target_size = _get_tensor_size(target_tensor)
432 tensor_size = _get_tensor_size(tensor)
433 if target_size is not None and tensor_size is not None and target_size != tensor_size:
434 raise AssertionError(
435 f"req {req.storage_index} mismatch sizes "
436 f"{target_size} vs {tensor_size}"
437 )
438 _copy_tensor_to_target(req, tensor, target_tensor, planner)
441def _fetch_tensor_file(
442 tensor_file: Any,
443 reqs: list[ReadItem],
444 storage_data: dict[MetadataIndex, StorageInfo],
445) -> list[tuple[ReadItem, Any]]:
446 """
447 Slice the region each item asked for off an already open safetensors file.
449 Nothing but the file is touched: no planner, no device, which is what leaves the caller
450 free to decide when what was read is put in place.
452 A name the file does not hold is left to safetensors to report - "File does not contain
453 tensor <name>", which says as much as a check here could. Checking first meant listing
454 the file, and listing brings every name in it across from the reader whether one shard is
455 being read or all of them: 6.5 ms on a file of nine thousand tensors, paid again every
456 time the file is opened, to say ahead of time what the next line says anyway.
458 Args:
459 tensor_file (Any): The open file, as :class:`_OpenFiles` hands it over.
460 reqs (list[ReadItem]): List of ReadItems for this file.
461 storage_data (dict[MetadataIndex, StorageInfo]): Physical storage mapping.
463 Returns:
464 list[tuple[ReadItem, Any]]: Each item with the region it asked for.
465 """
466 fetched: list[tuple[ReadItem, Any]] = []
467 for req in reqs:
468 storage_info = _get_storage_info(req, storage_data)
469 tensor_key = storage_info.tensor_key or req.storage_index.fqn
470 tensor_slices = tuple(
471 slice(int(off), int(off) + int(length))
472 for off, length in zip(req.storage_offsets, req.lengths)
473 )
474 if tensor_slices:
475 tensor = tensor_file.get_slice(tensor_key)[tensor_slices]
476 else:
477 # Scalar entries (rank-0 tensors such as the AdamW ``step``) have no slices to
478 # narrow by, and safetensors before 0.4.3 rejects the empty index with
479 # "too many indices for tensor of dimension 0" - read the whole tensor instead.
480 tensor = tensor_file.get_tensor(tensor_key)
481 fetched.append((req, tensor))
482 return fetched
485def _apply_fetched(
486 fetched: list[tuple[ReadItem, Any]],
487 planner: LoadPlanner,
488) -> None:
489 """
490 Put what was read where the load wants it.
492 The other half of a read, and the half that has to stay with whoever owns the device and
493 the planner: it copies into the state dict and hands each item back to the planner.
495 Args:
496 fetched (list[tuple[ReadItem, Any]]): Items with what was read for them, as
497 :func:`_fetch_tensor_file` and :func:`_fetch_bytes_file` return them.
498 planner (LoadPlanner): Load planner for resolving and committing.
499 """
500 for req, payload in fetched:
501 if req.type is LoadItemType.BYTE_IO:
502 planner.apply_bytes(req, payload)
503 else:
504 _validate_and_copy_tensor(req, payload, planner)
507# How many checkpoint files a read keeps open. Going through the shards in the order every
508# rank agrees on means walking the files over and over, where the old file-at-a-time read
509# went through each one once, and opening one costs far more than the slice read it wraps --
510# 0.70 ms against 0.045 ms, measured on a file holding four hundred tensors. So the files
511# stay open. A safetensors reader is a memory map over the file, so several of them cost
512# next to nothing to hold.
513_FILES_KEPT = 8
516class _OpenFiles:
517 """The checkpoint files a read is holding open, least recently reached for first."""
519 def __init__(self, capacity: int, open_one: Any, close_one: Any) -> None:
520 """
521 Args:
522 capacity (int): How many files to hold before letting the oldest go.
523 open_one (Any): Called with a path, returns a reader for it.
524 close_one (Any): Called with a reader that is being let go.
525 """
526 self._capacity = capacity
527 self._open_one = open_one
528 self._close_one = close_one
529 self._readers: dict[str, Any] = {}
531 def reader(self, path: str) -> Any:
532 """
533 The reader for this file, opening it when it is not already held.
535 Args:
536 path (str): Path of the checkpoint file.
538 Returns:
539 Any: The open reader, now the most recently reached for.
540 """
541 reader = self._readers.pop(path, None)
542 if reader is None:
543 while len(self._readers) >= self._capacity:
544 self._close_one(self._readers.pop(next(iter(self._readers))))
545 reader = self._open_one(path)
546 self._readers[path] = reader
547 return reader
549 def close(self) -> None:
550 """Let go of every file still held."""
551 for reader in self._readers.values():
552 self._close_one(reader)
553 self._readers.clear()
556def _open_checkpoint_files() -> _OpenFiles:
557 """
558 A place to keep the checkpoint files of one read open.
560 Returns:
561 _OpenFiles: Opens through safetensors, holding as many as such a reader is cheap
562 to hold.
563 """
564 return _OpenFiles(
565 _FILES_KEPT,
566 lambda path: safe_open(path, framework="pt", device="cpu"),
567 lambda reader: reader.__exit__(None, None, None),
568 )
571class FileSystemReader(StorageReader):
572 """
573 File system storage reader implementation.
575 Reads checkpoint data from the local file system, loading tensors
576 from safetensors files and bytes from separate files.
577 """
579 def __init__(self, checkpoint_dir: Union[Path, str]) -> None:
580 """
581 Args:
582 checkpoint_dir (Union[Path, str]): Directory the checkpoint is read out of.
583 """
584 self.checkpoint_dir = Path(checkpoint_dir) if isinstance(checkpoint_dir, str) else checkpoint_dir
585 # Cached storage layout: MetadataIndex -> StorageInfo (torch-aligned)
586 self.storage_data: dict[MetadataIndex, StorageInfo] = {}
587 self.rank: int = 0
588 self.is_coordinator: bool = False
590 def initialize_reader(self, checkpoint_id: Optional[Union[Path, str]] = None) -> None:
591 """
592 Initialize storage reader with new checkpoint directory.
594 Args:
595 checkpoint_id (Optional[Union[Path, str]]): New checkpoint directory path. Default None.
596 """
597 if checkpoint_id:
598 self.checkpoint_dir = Path(checkpoint_id) if isinstance(checkpoint_id, str) else checkpoint_id
600 @dcp_timer_decorator
601 def load_metadata(self, **kwargs: Any) -> Metadata:
602 """
603 Load checkpoint metadata from file.
605 When rank is provided in kwargs: load rank-local metadata from .rank{rank}_metadata
606 (for checkpoints saved with use_collectives=False).
607 Otherwise: load global metadata from .metadata.
609 Args:
610 **kwargs: Optional arguments (e.g., rank for rank-local metadata).
612 Returns:
613 Metadata: Metadata object loaded from file.
614 """
615 rank = kwargs.get("rank")
616 if rank is not None:
617 metadata_file = self.checkpoint_dir / f"{rank}{METADATA_FILE_NAME}"
618 else:
619 metadata_file = self.checkpoint_dir / METADATA_FILE_NAME
621 if not metadata_file.exists():
622 raise FileNotFoundError(f"Metadata file not found: {metadata_file}")
623 with open(metadata_file, "rb") as f:
624 metadata = _MetadataUnpickler(f).load()
625 return metadata
627 def configure_reader(self, metadata: Metadata, is_coordinator: bool,
628 **kwargs: Any) -> None:
629 """Configure storage reader."""
630 # Cache storage_data separately for quick lookup in execute_read.
631 # This mirrors torch.filesystem, where reader keeps a storage_data dict.
632 self.storage_data = getattr(metadata, "storage_data", None)
633 self.is_coordinator = is_coordinator
634 # Do not evaluate get_rank() when an offline caller supplies rank. The
635 # default process group is intentionally not initialized by converters.
636 self.rank = kwargs["rank"] if "rank" in kwargs else dist.get_rank()
638 def optimize_local_plan(self, plan: LoadPlan) -> LoadPlan:
639 """
640 Optimize local plan.
642 Args:
643 plan (LoadPlan): Local load plan.
645 Returns:
646 LoadPlan: Optimized local plan.
647 """
648 return plan
650 def optimize_global_plan(self, plan: LoadPlan) -> LoadPlan:
651 """
652 Optimize this rank's load plan.
654 Args:
655 plan (LoadPlan): This rank's load plan.
657 Returns:
658 LoadPlan: Optimized load plan.
659 """
660 return plan
662 def _get_storage_path(self, read_item: ReadItem) -> str:
663 """
664 Get storage file path for a read item.
666 Args:
667 read_item (ReadItem): ReadItem to get path for.
669 Returns:
670 str: Absolute path to the storage file.
671 """
672 if self.storage_data is None:
673 raise KeyError("Checkpoint metadata.storage_data is required for filesystem read")
674 storage_info = self.storage_data.get(read_item.storage_index)
675 if storage_info is None:
676 raise KeyError(f"StorageInfo not found for index {read_item.storage_index}")
677 return str(self.checkpoint_dir / storage_info.relative_path)
679 def _group_items_by_file(self, items: list) -> dict[str, list]:
680 """
681 Group ReadItems by the checkpoint file they read from.
683 Which items a rank goes to storage for at all is settled by :meth:`execute_read`,
684 which only passes on the ones this rank reads itself.
686 Args:
687 items (list[ReadItem]): ReadItems to load from storage.
689 Returns:
690 dict[str, list[ReadItem]]: Dictionary mapping file paths to lists of ReadItems.
691 """
692 per_file: dict[str, list] = {}
693 for read_item in items:
694 path = self._get_storage_path(read_item)
695 per_file.setdefault(path, []).append(read_item)
696 return per_file
698 def _private_batches(self, by_shard: dict) -> dict[str, list]:
699 """
700 Gather the shards no other rank shares into one batch per checkpoint file.
702 Nothing waits on these - they are read after the last broadcast has been handed
703 over - so unlike the shared shards they are under no obligation to be read in the
704 order every rank agreed on, and are free to be read in whatever order costs least.
705 Asking for a file's worth at once is that order: the file is opened, read for every
706 shard that wants it, and left, however far apart in the plan those shards sit.
708 What this is for above all is the pickled entries. Those go through the save side's
709 dedup like everything else, so the one that writes an entry is whichever rank was
710 carrying the least at the time, but the load side never broadcasts them - each rank
711 rebuilds its own object - so every rank reads every entry it wants, out of whatever
712 file the dedup put it in. Nothing holds a bytes file open between reads either, so
713 one batch per shard opened a file for each: eighty entries spread over three files
714 cost eighty opens where three would do.
716 Args:
717 by_shard (dict[MetadataIndex, list[ReadItem]]): The plan's items, gathered by
718 the shard they land in.
720 Returns:
721 dict[str, list[ReadItem]]: The items of every file, keyed by its path, in the
722 order the files first come up in the plan.
723 """
724 private: dict[str, list] = {}
725 for items in by_shard.values():
726 if items[0].source is not None:
727 continue
728 for item in items:
729 private.setdefault(self._get_storage_path(item), []).append(item)
730 return private
732 def _fetch_from_storage(self, items: list, open_files: _OpenFiles) -> list[tuple[ReadItem, Any]]:
733 """
734 Take a set of items off disk, a file at a time, and hand back what was read.
736 Only the checkpoint files are touched: putting any of it in place, which is what
737 needs the planner and the device, is :func:`_apply_fetched`.
739 A file that is not there is left to the open to report. Both of them raise
740 FileNotFoundError naming the path - safetensors as well as the builtin open - so a
741 check ahead of them said no more than they do, and said it with a stat per file, on
742 storage where a stat is a round trip to another machine.
744 Args:
745 items (list[ReadItem]): The items to read, in any order.
746 open_files (_OpenFiles): The files this read is holding open.
748 Returns:
749 list[tuple[ReadItem, Any]]: Each item with what was read for it.
751 Raises:
752 FileNotFoundError: If a file the items name is not there.
753 """
754 fetched: list[tuple[ReadItem, Any]] = []
755 for path, reqs in self._group_items_by_file(items).items():
756 if path.endswith(".bytes"):
757 # BYTE_IO: one bytes file per rank with per-item offsets.
758 fetched.extend(_fetch_bytes_file(path, reqs, self.storage_data))
759 else:
760 # TENSOR: one safetensors file per rank
761 reader = open_files.reader(path)
762 fetched.extend(_fetch_tensor_file(reader, reqs, self.storage_data))
763 return fetched
765 @dcp_timer_decorator
766 def execute_read(
767 self,
768 plan: LoadPlan,
769 planner: LoadPlanner,
770 broadcast_groups: Optional[dict[tuple, Any]] = None,
771 broadcast_batch_bytes: int = 0,
772 ) -> None:
773 """
774 Read data from storage, overlapping the reads with the copies and sends they feed.
776 Shards several ranks hold are dealt with first, in the order the global plan put
777 them in, which every rank arrives at the same way. The rank that was chosen to read
778 one starts its broadcast the moment it has it and goes straight on to the next shard,
779 so a send runs while the one after it is still being read. Only so many sends are
780 left going at once; past that the oldest is waited on to make room.
782 What no other rank shares comes last. Nobody is waiting on it, so reading it earlier
783 would only hold up the broadcasts, and reading it here fills the time the final sends
784 are still in flight. Nobody waiting on it also means nothing holds it to the order the
785 ranks agreed on, so it is read a file at a time rather than a shard at a time: each
786 file is opened once and read for every shard that wants it.
788 Args:
789 plan (LoadPlan): Load plan containing ReadItems.
790 planner (LoadPlanner): Load planner for resolving and committing tensors.
791 broadcast_groups (Optional[dict]): Communication group of every shard the plan
792 marked, keyed by rank tuple, as :func:`broadcast_groups_for_load` yields them.
793 Building one is collective, so it is done before the read rather than here.
794 broadcast_batch_bytes (int): Shards smaller than this are gathered and sent
795 together instead of one at a time. Zero sends every shard on its own.
797 Raises:
798 KeyError: If the plan marked a shard whose group is not among the ones given.
799 """
800 groups = broadcast_groups or {}
802 by_shard: dict[MetadataIndex, list] = {}
803 for item in plan.items:
804 by_shard.setdefault(item.dest_index, []).append(item)
805 shared = sorted(
806 (index for index, items in by_shard.items() if items[0].source is not None),
807 key=lambda index: (index.fqn, index.offset, index.index),
808 )
809 # What this rank goes to storage for, in the order it will ask for it: the shards it
810 # was picked to read for its group, a shard at a time so that a group waiting on this
811 # rank waits for its own shard and no more, then what nobody else shares, a file at a
812 # time.
813 batches = [by_shard[index] for index in shared
814 if by_shard[index][0].source.src_rank == self.rank]
815 batches.extend(self._private_batches(by_shard).values())
817 in_flight: deque = deque()
818 batcher = BroadcastBatcher(broadcast_batch_bytes, groups)
819 open_files = _open_checkpoint_files()
820 try:
821 read = (self._fetch_from_storage(batch, open_files) for batch in batches)
822 for index in shared:
823 items = by_shard[index]
824 if items[0].source.src_rank == self.rank:
825 _apply_fetched(next(read), planner)
826 batcher.add(in_flight, planner.state_dict, items[0])
828 # Whatever is still waiting for company goes now, before the shards nobody
829 # shares: the ranks holding them are waiting on these, and reading what
830 # nobody wants first would keep them waiting through all of it.
831 batcher.flush(in_flight)
832 for batch in read:
833 _apply_fetched(batch, planner)
834 wait_broadcasts(in_flight)
835 finally:
836 open_files.close()
838 if shared:
839 logger.info(
840 "[rank=%d] >>> %d replicated shards sent in %d broadcasts, %d of them gathered",
841 self.rank, len(shared), batcher.sent, batcher.batched,
842 )