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

Code before:
        # 2. Attempt to restore states from `resume_from_checkpoint` file
        elif self.resume_checkpoint_path is not None:
            self.restore(self.resume_checkpoint_path)

        # wait for all to catch up
        self.trainer.training_type_plugin.barrier('TrainerIOMixin.restore_weights')

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

    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`.
        """
        # 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.")

        checkpoint, load_optimizer_states = self.trainer.training_type_plugin.restore_model_state_from_ckpt_path(
            checkpoint_path, map_location=lambda storage, loc: storage
        )

        model = self.trainer.lightning_module

        if self.trainer._device_type == DeviceType.GPU:
            model.cuda(self.trainer.root_gpu)

        # restore training state
        self.restore_training_state(checkpoint, load_optimizer_states)

        rank_zero_info(f"Restored states from the checkpoint file at {checkpoint_path}")
        return True

    def restore_model_state(self, model: LightningModule, checkpoint) -> None:
        """
        Restore model states from a 'PyTorch-Lightning checkpoint' dictionary object
        """

        # restore datamodule states
        if self.trainer.datamodule is not None:
            self.trainer.datamodule.on_load_checkpoint(checkpoint)

        # hook: give user access to checkpoint if needed.
        model.on_load_checkpoint(checkpoint)

        # restore model state_dict
        model.load_state_dict(checkpoint['state_dict'])


Review comment (suggestion):
@awaelchli we shouldn't print this if `not self.resume_checkpoint_path`

Code after:
            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()
        model = self.trainer.lightning_module