Recommended way to load previous checkpoints from trainer callbacks?

[xsys-technology]

I'd like to reload a previous checkpoint when some condition is met within a callback that triggers when training epoch begins:

class MyCallback(Callback):

        def on_train_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
            # Check some condition 'my_cond'
            #       Reload a previous checkpoint 'some_checkpoint.ckpt' if 'my_cond' == True

What is the recommended way to do this given that the model is sharded across multiple 'gpu's when on_train_epoch_start is called? I have a hunch that TorchCheckpointIO probably can handle this but I'm not sure about getting the data mapping correct.

I'm using DDPPlugin strategy:

trainer = pl.Trainer.from_argparse_args(args,
                                            strategy=DDPPlugin(find_unused_parameters=False),
                                            callbacks=[checkpoint_callback,...])

I believe the trainer's checkpoint connector can use TorchCheckpointIO, via restore(), something like:

class MyCallback(Callback):

        def on_train_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
            if my_cond:
                    trainer.checkpoint_connector.restore('some_checkpoint.ckpt')

Is this safe to do considering the state of the optimizer(s), lr_scheduler(s), etc...?

I imagine the right way to do this looks at least something like this.