Restores state from TensorStoreStateStorage. Args: state: Uninitialized state. Returns: The restored state.
(self, state: Builder.State)
| 419 | return cls.default_config().set(dir=spec) |
| 420 | |
| 421 | def __call__(self, state: Builder.State) -> Builder.State: |
| 422 | """Restores state from TensorStoreStateStorage. |
| 423 | |
| 424 | Args: |
| 425 | state: Uninitialized state. |
| 426 | |
| 427 | Returns: |
| 428 | The restored state. |
| 429 | """ |
| 430 | cfg: TensorStoreStateStorageBuilder.Config = self.config |
| 431 | if cfg.base_dir: |
| 432 | ckpt_dir = build_step_dir(cfg.base_dir, step=cfg.step) |
| 433 | step = cfg.step |
| 434 | else: |
| 435 | ckpt_dir = cfg.dir |
| 436 | step = parse_step_from_dir(cfg.dir) |
| 437 | cfg.storage.max_concurrent_restore_gb = cfg.concurrent_gb |
| 438 | storage = cfg.storage.instantiate() |
| 439 | restored_state = storage.restore_from_dir( |
| 440 | step=step, |
| 441 | state=state.trainer_state, |
| 442 | ckpt_dir=ckpt_dir, |
| 443 | validation=cfg.validation, |
| 444 | ) |
| 445 | built_keys = state.built_keys.union({key for key, _ in flatten_items(restored_state)}) |
| 446 | return Builder.State(step=step, trainer_state=restored_state, built_keys=built_keys) |
| 447 | |
| 448 | |
| 449 | # pylint: disable=abstract-method |
nothing calls this directly
no test coverage detected