Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/trainer/connectors/checkpoint_connector.py (Python)

Code before:
        checkpoint_path = self.resume_checkpoint_path
        if not checkpoint_path:
            return

        # clear cache before restore
        if self.trainer._device_type == DeviceType.GPU:
            torch.cuda.empty_cache()

        # Try to read the checkpoint file at `checkpoint_path`. If not exist, do not restore checkpoint.
        fs = get_filesystem(checkpoint_path)
        if not fs.exists(checkpoint_path):
            raise FileNotFoundError(f"Checkpoint at {checkpoint_path} not found. Aborting training.")

        rank_zero_info(f"Restoring states from the checkpoint file at {checkpoint_path}")
        checkpoint, load_optimizer_states = self.trainer.training_type_plugin.restore_model_state_from_ckpt_path(
            checkpoint_path, map_location=lambda storage, loc: storage
        )
        self.loaded_checkpoint = checkpoint
        self._load_optimizer_states = load_optimizer_states

    def resume_end(self) -> None:
        """ Signal the connector that all states have resumed and memory for the checkpoint object can be released. """
        rank_zero_info(f"Restored all states from the checkpoint file at {self.resume_checkpoint_path}")
        self.resume_checkpoint_path = None
        self.loaded_checkpoint = dict()

        # clear cache after restore
        if self.trainer._device_type == DeviceType.GPU:
            torch.cuda.empty_cache()

        # wait for all to catch up
        self.trainer.training_type_plugin.barrier("CheckpointConnector.resume_end")

    def restore_weights(self) -> None:
        """
        Attempt to restore a checkpoint (e.g. weights) in this priority:
        1. from HPC weights
        2. from `resume_from_checkpoint` file
        3. don't restore
        """
        self.restore(self.resume_checkpoint_path)

    def restore(self, checkpoint_path: str) -> bool:
        """
        Load model/training states from a 'PyTorch-Lightning checkpoint' file through file-read and state-restore.
        All restored states are listed in return value description of `dump_checkpoint`.
        """
        self.resume_checkpoint_path = checkpoint_path
        self.resume_start()
        model = self.trainer.lightning_module


Review comment (suggestion):
reviewers: this part has moved to `resume_start()`

Code after:

        Raises:
            FileNotFoundError: If the path to the checkpoint file is provided but the file does not exist.
        """
        self.resume_checkpoint_path = self.hpc_resume_path or self.resume_checkpoint_path
        checkpoint_path = self.resume_checkpoint_path
        if not checkpoint_path:
            return

        # clear cache before restore
        if self.trainer._device_type == DeviceType.GPU:
            torch.cuda.empty_cache()

        # Try to read the checkpoint file at `checkpoint_path`. If not exist, do not restore checkpoint.
        fs = get_filesystem(checkpoint_path)
        if not fs.exists(checkpoint_path):
            raise FileNotFoundError(f"Checkpoint at {checkpoint_path} not found. Aborting training.")

        rank_zero_info(f"Restoring states from the checkpoint file at {checkpoint_path}")
        checkpoint, load_optimizer_states = self.trainer.training_type_plugin.restore_model_state_from_ckpt_path(
            checkpoint_path, map_location=lambda storage, loc: storage
        )
        self._loaded_checkpoint = checkpoint
        self._load_optimizer_states = load_optimizer_states

    def resume_end(self) -> None:
        """ Signal the connector that all states have resumed and memory for the checkpoint object can be released. """
        rank_zero_info(f"Restored all states from the checkpoint file at {self.resume_checkpoint_path}")
        self.resume_checkpoint_path = None
        self._loaded_checkpoint = dict()

        # clear cache after restore
        if self.trainer._device_type == DeviceType.GPU:
            torch.cuda.empty_cache()

        # wait for all to catch up
        self.trainer.training_type_plugin.barrier("CheckpointConnector.resume_end")

    def restore(self, checkpoint_path: Optional[Union[Path, str]] = None) -> bool:
        """
        Attempt to restore model/training states from a 'PyTorch-Lightning checkpoint' file
        through file-read and state-restore, in this priority:

        1. from HPC weights if found
        2. from `resume_from_checkpoint` file if provided
        3. don't restore

        All restored states are listed in return value description of `dump_checkpoint`.
        """
        self.resume_checkpoint_path = checkpoint_path or self.resume_checkpoint_path
        self.resume_start()