1# mypy: allow-untyped-defs 2# Copyright (c) Meta Platforms, Inc. and affiliates 3from typing import Any, Protocol, runtime_checkable 4 5import torch 6 7 8@runtime_checkable 9class _Checkpointable(Protocol): # noqa: PYI046 10 """ 11 Interface for checkpointable objects. 12 Implemented as a protocol, implicit subtyping is supported so subclasses do not need to inherit this explicitly. 13 This is to allow arbitrary objects/tensor subclasses to hook into DCP seamlessly through implementing the interface. 14 """ 15 16 def __create_write_items__(self, fqn: str, object: Any): 17 """ 18 Return a list of WriteItems based on object's contents. 19 """ 20 raise NotImplementedError( 21 "_Checkpointable._create_write_items is not implemented" 22 ) 23 24 def __create_chunk_list__(self): 25 """ 26 Return a list of `ChunkStorageMetadata` based on object's contents. 27 """ 28 raise NotImplementedError( 29 "_Checkpointable._create_chunk_list is not implemented" 30 ) 31 32 def __get_tensor_shard__(self, index) -> torch.Tensor: 33 """ 34 Return a 'torch.Tensor' shard based on 'MetadataIndex'. 35 """ 36 raise NotImplementedError( 37 "_Checkpointable._get_tensor_shard is not implemented" 38 ) 39