def __init__(
self,
url: str,
rank: Rank = SINGLE_TRAINER,
num_machines: int = 1,
partition_client: Optional[PartitionClient] = None,
subprocess_name: Optional[str] = None,
subprocess_init: Optional[Callable[[], None]] = None,
) -> None:
self.storage: AbstractCheckpointStorage = CHECKPOINT_STORAGES.make_instance(url)
self.dirty: Set[Tuple[EntityName, Partition]] = set()
self.rank: Rank = rank
self.num_machines: int = num_machines
if self.rank == 0:
self.storage.prepare()
self.checkpoint_version = self.storage.load_version()
self.partition_client = partition_client
self.metadata_providers: List[MetadataProvider] = []
def __init__(
self,
url: str,
rank: Rank = 0,
num_machines: int = 1,
background: bool = False,
partition_client: Optional[PartitionClient] = None,
subprocess_name: Optional[str] = None,
subprocess_init: Optional[Callable[[], None]] = None,
) -> None:
self.storage: AbstractCheckpointStorage = CHECKPOINT_STORAGES.make_instance(
url)
self.dirty: Set[Tuple[EntityName, Partition]] = set()
self.rank: Rank = rank
self.num_machines: int = num_machines
if self.rank == 0:
self.storage.prepare()
self.checkpoint_version = self.storage.load_version()
self.background: bool = background
if self.background:
self.pool: mp.pool.Pool = create_pool(
1,
subprocess_name=subprocess_name,
subprocess_init=subprocess_init,
)
# FIXME In py-3.7 switch to typing.OrderedDict[str, AsyncResult].
self.outstanding: OrderedDict = OrderedDict()
self.prefetched: Dict[str,
Tuple[FloatTensorType,
Optional[OptimizerStateDict]]] = {}
self.partition_client = partition_client
self.metadata_providers: List[MetadataProvider] = []