Resuming WandB runs in PL

[rahulvigneswaran]

When I start a run, I always generate a wandb id using wandb.util.generate_id() and save it alongside the ckpt.

When a run crashes, I try to resume the trainer by providing the appropriate ckpt_path in trainer.fit and try to resume the wandb logger by doing the following (cfg.LOGGER.WANDB_ID is the same wandb id that i saved in the earlier step),

wandb_logger = WandbLogger(project=cfg.LOGGER.PROJECT_NAME, 
                        name=LOGGER_NAME,
                        log_model=cfg.LOGGER.LOG_MODEL,
                        id=cfg.LOGGER.WANDB_ID,
                        resume="allow",
                        reinit=True)

While the trainer resumes without any hiccup, the wandb logger, instead of resuming the previous run, creates a new run and logs from the resumed epoch.

Where am I going wrong and how to fix it?

[rahulvigneswaran]

Hey @williamFalcon , Any idea about this?

[daniellepintz]

@rahulvigneswaran could this be because you set reinit=True?

[rahulvigneswaran]

@daniellepintz The servers that use to run my experiments are down. Give me a couple of days, I will get back.

[rahulvigneswaran]

@daniellepintz I gave it a try but unfortunately removing it/setting it to False didn't help.

[konstantinjdobler]

@rahulvigneswaran I believe the W&B docs state that you should set resume="must" instead of "allow".

[talhaanwarch]

We can resume it this way

from pytorch_lightning.loggers import WandbLogger
wandb_logger = WandbLogger(project="vitals-{}".format(vital),id='w5pelznp',resume='must')

i got the id from folder name, which look like run-20220812_220817-w5pelznp

[icekang]

This one works for me, thanks!

[samils7]

Any other solutions?

[kabouzeid]

You can also set WANDB_RUN_ID=foo env variable.

[ivegner]

I've solved this in LightningCLI the following way:

# main.py
class MyLightningCLI(LightningCLI):
    def add_arguments_to_parser(self, parser):
        parser.add_argument(
            "--resume_run_id", default="", type=str, help="W&B run ID to resume from"
        )
        parser.link_arguments("resume_run_id", "trainer.logger.init_args.resume", compute_fn=lambda x: "must" if x else "never")

    def before_instantiate_classes(self):
        subcommand = self.config.subcommand
        c = self.config[subcommand]
        run_id = None

        if c.resume_run_id:
            run_id = c.resume_run_id
            api = wandb.Api()
            artifact = api.artifact(f'{c.trainer.logger.init_args.project}/model-{run_id}:latest', type="model")
            artifact_dir = artifact.download()
            c.ckpt_path=str(Path(artifact_dir) / "model.ckpt")
        else:
            run_id = wandb.util.generate_id()

        # also make sure that ModelCheckpoints go to the right place
        c.trainer.logger.init_args.id = run_id
        for callback in c.trainer.callbacks:
            if callback.class_path == "pytorch_lightning.callbacks.ModelCheckpoint":
                callback.init_args.dirpath = f"checkpoints/{run_id}"

if __name__ == "__main__":
    _cli = MyLightningCLI(
        ...
    )

and the relevant part of my config.yaml is

trainer:
  logger:
    class_path: pytorch_lightning.loggers.WandbLogger
    init_args:
      project: MyProject
      log_model: all
  callbacks:
  - class_path: pytorch_lightning.callbacks.ModelCheckpoint
    init_args:
      save_top_k: 1
      save_last: True
      monitor: val/total_loss
      filename: '{epoch}-{step}-{validation_loss:.3f}'

Now, I can just do

python main.py fit --config config.yaml --resume_run_id <some_run_id>

and it will download the latest checkpoint and restart appropriately using ckpt_path, and will start a new run if --resume_run_id is not provided.