Store logger experiment id in checkpoint to enable correct resuming of experiments

🚀 Feature

Enable each logger a possibility to attach custom data to checkpoint, like id of experiment defined by logger.

Motivation

Currently, if you use loggers like Wandb, Comet or Nepture, you have to manually restore correct experiment id when resuming from checkpoint in order to resume logging of previous experiment, instead of creating a new experiment in logger.

It would be better if id of experiment/run could be attached by logger to every checkpoint, and that id could be then automatically passed to logger when resuming from those checkpoints.

This was recently mentioned in #4935 and it seems to me like a very useful feature.

Solution

So my suggestion is the following: Currently existing loggers could accept parameters like

WandbLogger(store_id_in_cktp=True)

which would make them attach experiment id to checkpoint. Or alternatively, all loggers should always attached their id without specifying it.

Then, when we resume experiment from checkpoint, trainer could accept parameter like resume_logger_experiment

logger = WandbLogger()
trainer = Trainer(resume_from_checkpoint="last.ckpt", logger=logger, resume_logger_experiment=True)

which would make trainer modify logger experiment id.

I've so far only used wandb, but I assume other loggers like Neptune or Tensorboard, also have equivalent "id" parameter than enables them to resume experiment? If so this change should affect every one of them.

cc @borda @awaelchli @edward-io @ananthsub @rohitgr7 @kamil-kaczmarek @Raalsky @Blaizzy @ninginthecloud