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

[ashleve]

🚀 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

[tchaton]

Hey there,

It is a good idea. Would you mind making a PR for it ?

Best regards,
T.C

[ashleve]

Hey @tchaton.
Sure :)
I can try making a PR, but I haven't made any contributions here before so let me know if I think correctly:

• I don't see any hooks in LightningLoggerBase, that would allow for state saving and restoring (like "on_checkpoint_save" in base Callback class). I assume we should first somehow implement those hooks?
• it seems that CheckpointConnector currently implements state restoriation of callbacks, schedulers, optimizer, etc. but not loggers. I assume we can just extend this functionality to loggers also?

[awaelchli]

@hobogalaxy

Are you still interested? I can help too.

I don't see any hooks in LightningLoggerBase, that would allow for state saving and restoring (like "on_checkpoint_save" in base Callback class). I assume we should first somehow implement those hooks?

because loggers are not callbacks.
but of course we could think of a way to restore loggers. This will be very different between the loggers.

So I suggest the following:

1. One PR that implements dumping a state_dict of some sort for loggers that goes into the checkpoint dict. Implement that first only for the TensorBoardLogger (no restore).
2. A PR that adds the functionality to restore the TensorBoardLogger from the state dict out of a checkpoint
3. Then follow up with other PRs to add this also to the other loggers.

[ashleve]

@awaelchli Thanks for the suggestions. I will try to make a PR this week.

[Queuecumber]

Pretty sure this is a (sort of) dupe of #6205 and #7104
Those dealt with the HPC checkpoints but it should really apply to both HPC checkpoints and regular checkpoints.

[yuvalkirstain]

Hey, any news about this? that's a very useful feature.

[awaelchli]

Hello, the status is that we are currently investigating adding a consistent interface for stateful components, #11429. And loggers would also fall into this category imo and could inherit from this too. I believe this issue can move forward once we have that in place, which is likely for 1.7.

[borisdayma]

Just for info, you can easily know which run created or use a specific artifact with artifact.used_by() and artifact.logged_by()