def _save_checkpoints()

in src/hyperpod_nemo_adapter/utils/callbacks/checkpoint.py [0:0]

• 12 lines of code
• 5 McCabe index (conditional complexity)

    def _save_checkpoints(self, trainer: "pl.Trainer"):
        """
        At each step, we check if we should save Peft checkpoint.
        """
        checkpoint_io = trainer.strategy.checkpoint_io
        assert isinstance(checkpoint_io, SageMakerCheckpointIO)
        checkpoint_dir = self.checkpoint_dir
        checkpoint_info = []
        monitor_candidates = self._monitor_candidates(trainer)

        if self._should_save_sharded(trainer, monitor_candidates):
            checkpoint_info.append(self._save_peft_sharded(trainer, checkpoint_dir, monitor_candidates))
        # Note that PEFT FULL checkpoints reuse the configs for saving regular FULL checkpoints
        if self._should_save_full(trainer):
            checkpoint_info.append(self._save_peft_full(trainer, checkpoint_dir))

        for checkpoint_type, checkpoint_dir in checkpoint_info:
            self._save(trainer, checkpoint_io, checkpoint_type, checkpoint_dir)