(*, step: int, ckpt_dir: str)
| 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. |
nothing calls this directly
no test coverage detected