See `BaseCheckpointer.restore` docstring for details. A complete checkpoint is one with an "index" file, which is only written after the entire checkpoint has been written.
(
self,
*,
step: Optional[int] = None,
state: Union[NestedTensor, NestedTensorSpec],
)
| 1216 | return fs.exists(ckpt_index) |
| 1217 | |
| 1218 | def restore( |
| 1219 | self, |
| 1220 | *, |
| 1221 | step: Optional[int] = None, |
| 1222 | state: Union[NestedTensor, NestedTensorSpec], |
| 1223 | ) -> tuple[Optional[int], NestedTensor]: |
| 1224 | """See `BaseCheckpointer.restore` docstring for details. |
| 1225 | |
| 1226 | A complete checkpoint is one with an "index" file, which is only written after the entire |
| 1227 | checkpoint has been written. |
| 1228 | """ |
| 1229 | cfg: Checkpointer.Config = self.config |
| 1230 | validation = cfg.validation or CheckpointValidationType.EXACT |
| 1231 | |
| 1232 | # Apply restore_state_filter if configured. |
| 1233 | if cfg.restore_state_filter is not None: |
| 1234 | state = cfg.restore_state_filter.instantiate()(state) |
| 1235 | |
| 1236 | def validate_and_restore(*, step: int, ckpt_dir: str): |
| 1237 | ckpt_index = os.path.join(ckpt_dir, "index") |
| 1238 | if not self._index_exists(ckpt_dir): |
| 1239 | raise ValueError( |
| 1240 | f"Checkpoint {ckpt_dir} is incomplete -- expected {ckpt_index} to be present." |
| 1241 | ) |
| 1242 | restored_state = self._storage.restore_from_dir( |
| 1243 | step=step, |
| 1244 | state=state, |
| 1245 | ckpt_dir=ckpt_dir, |
| 1246 | validation=validation, |
| 1247 | ) |
| 1248 | logging.info("Restored state from ckpt at step %s", step) |
| 1249 | if "summary_writer" in self.children: |
| 1250 | self.summary_writer.log_checkpoint( |
| 1251 | step=step, |
| 1252 | state=state, |
| 1253 | ckpt_dir=ckpt_dir, |
| 1254 | action=CheckpointerAction.RESTORE, |
| 1255 | ) |
| 1256 | return restored_state |
| 1257 | |
| 1258 | if step is not None: |
| 1259 | # For a specified step, we try to load it. |
| 1260 | ckpt_dir = self.ckpt_dir(step) |
| 1261 | return step, validate_and_restore(step=step, ckpt_dir=ckpt_dir) |
| 1262 | |
| 1263 | try: |
| 1264 | # Latest checkpoint path, if it exists, is guaranteed to be complete. |
| 1265 | ckpt_dir = self.latest_checkpoint_path(cfg.dir) |
| 1266 | step = parse_step_from_dir(ckpt_dir) |
| 1267 | restored_state = validate_and_restore(step=step, ckpt_dir=ckpt_dir) |
| 1268 | except IndexError: |
| 1269 | # No checkpoint path exists. Return with input state. |
| 1270 | logging.info("Could not find any completed checkpoints under %s", cfg.dir) |
| 1271 | restored_state = state |
| 1272 | |
| 1273 | return step, restored_state |