core.dist_checkpointing.strategies.cached_metadata_filesystem_reader#

FS Reader with metadata cached support.

Module Contents#

Classes#

CachedMetadataFileSystemReader

Extends FileSystemReader to cache metadata for improved performance.

API#

class core.dist_checkpointing.strategies.cached_metadata_filesystem_reader.CachedMetadataFileSystemReader(
path: Union[str, os.PathLike],
cache_metadata: bool = True,
)#

Bases: torch.distributed.checkpoint.FileSystemReader

Extends FileSystemReader to cache metadata for improved performance.

Metadata is shared across all reader instances that use the same checkpoint directory (same path), since the loaded metadata is identical.

.. attribute:: _metadata_cache

Class-level cache keyed by checkpoint path.

Type:

Dict[str, Metadata]

Initialization

Initialize with file system path.

Parameters:

path (Union[str, os.PathLike]) – Path to the checkpoint directory or file.

_metadata_cache: Dict[str, torch.distributed.checkpoint.Metadata]#

None

read_metadata() torch.distributed.checkpoint.Metadata#

Read metadata from file system, caching for subsequent calls. Shared across instances when the checkpoint directory is the same.

Returns:

Checkpoint metadata.

Return type:

Metadata

classmethod clear_metadata_cache()#

Clear the metadata cache.