(
self,
step: int,
state: Union[NestedTensor, NestedTensorSpec],
*,
ckpt_dir: str,
validation: CheckpointValidationType = CheckpointValidationType.EXACT,
)
| 611 | self._manager.wait_until_finished() |
| 612 | |
| 613 | def restore_from_dir( |
| 614 | self, |
| 615 | step: int, |
| 616 | state: Union[NestedTensor, NestedTensorSpec], |
| 617 | *, |
| 618 | ckpt_dir: str, |
| 619 | validation: CheckpointValidationType = CheckpointValidationType.EXACT, |
| 620 | ) -> NestedTensor: |
| 621 | spec = self._get_spec(step, state, ckpt_dir) |
| 622 | logging.info("Restoring checkpoint from directory %s", ckpt_dir) |
| 623 | check_state_structure( |
| 624 | read_index_file(ckpt_dir), target_structure=spec.index, validation=validation |
| 625 | ) |
| 626 | restore_tf_savables(spec.tf_ckpt_map, dir=ckpt_dir) |
| 627 | |
| 628 | # TODO(jtian22): Below code needs to change when we work on supporting |
| 629 | # Elastic training in grain input pipeline. |
| 630 | maybe_restore_python_savables( |
| 631 | spec.python_ckpt_map, dir=os.path.join(ckpt_dir, f"python_{jax.process_index()}") |
| 632 | ) |
| 633 | return self._restore_tensorstore_state(state, ckpt_dir=ckpt_dir, spec=spec) |
| 634 | |
| 635 | def _restore_tensorstore_state( |
| 636 | self, state, *, ckpt_dir: str, spec: CheckpointSpec, sync: bool = True |
nothing calls this directly
no test coverage detected