Build a Task that will be run once after the job's `init_group` is run. This task will determine which blobs need to be checkpointed. If retrieve_from_epoch is not None, then the checkpoint metadata is retrieved from a previously saved checkpoint.
(
self,
nodes=None,
retrieve_from_epoch=None,
path_prefix=None,
path_type=None
)
| 190 | path_type: Indicate the type of path where checkpoint files are stored. |
| 191 | """ |
| 192 | def init( |
| 193 | self, |
| 194 | nodes=None, |
| 195 | retrieve_from_epoch=None, |
| 196 | path_prefix=None, |
| 197 | path_type=None |
| 198 | ): |
| 199 | """ |
| 200 | Build a Task that will be run once after the job's `init_group` is run. |
| 201 | This task will determine which blobs need to be checkpointed. |
| 202 | If retrieve_from_epoch is not None, then the checkpoint metadata is |
| 203 | retrieved from a previously saved checkpoint. |
| 204 | """ |
| 205 | assert nodes is None or len(nodes) == 1, ( |
| 206 | 'CheckpointManager only supports single node.') |
| 207 | |
| 208 | with Task(outputs=[self._blob_names]) as task: |
| 209 | if retrieve_from_epoch is None: |
| 210 | ops.GetAllBlobNames( |
| 211 | [], |
| 212 | self._blob_names, |
| 213 | include_shared=False) |
| 214 | else: |
| 215 | full_db_name = db_name(retrieve_from_epoch, |
| 216 | self._node_name, self._db_prefix, path_prefix) |
| 217 | db_type = path_type or self._db_type |
| 218 | logger.info("Initializing checkpoints from = %s" |
| 219 | % full_db_name) |
| 220 | ops.Load( |
| 221 | [], self._blob_names, |
| 222 | db=full_db_name, |
| 223 | db_type=db_type, |
| 224 | absolute_path=True, |
| 225 | keep_device=True, |
| 226 | ) |
| 227 | self._names_output = task.outputs()[0] |
| 228 | return task |
| 229 | |
| 230 | def blob_list(self): |
| 231 | assert self._names_output |