Lightning-AI/pytorch-lightning

[WandbLogger] Call `wandb.finish()` to ensure all artifacts are uploaded before training ends.

Open

#17.768 aperta il 7 giu 2023

Vedi su GitHub
 (2 commenti) (2 reazioni) (0 assegnatari)Python (3233 fork)batch import
bughelp wantedlogger: wandbver: 2.1.x

Metriche repository

Star
 (26.687 star)
Metriche merge PR
 (Merge medio 9g 15h) (3 PR mergiate in 30 g)

Descrizione

Bug description

Currently WandbLogger doesn't explicitly call wandb.finish(), so it's possible that not all artifacts are synced to the server at the time train.fit() returns, especially when you have a large amount of artifacts to upload.

We know that wandb.init() spawns a new background process to log data to a run, and it asynchronously upload the logs and artifacts to the wandb server. The uploading will finish after 1-2 mins or more depending on your artifact size. It's necessary to explicitly call this barrier method to ensure all artifacts are uploaded before trainer.fit() returns.

https://github.com/Lightning-AI/lightning/blob/420eb6f248560c9c721601df791bf8b2f2214fe6/src/lightning/pytorch/loggers/wandb.py#LL566C1-L573C1

@rank_zero_only
def finalize(self, status: str) -> None:
    if status != "success":
        # Currently, checkpoints only get logged on success
        return
    # log checkpoints as artifacts
    if self._checkpoint_callback and self._experiment is not None:
        self._scan_and_log_checkpoints(self._checkpoint_callback)
    
    # Ensure that all artifacts get uploaded before trainer.fit() returns
    wandb.finish()

What version are you seeing the problem on?

master

How to reproduce the bug

Increase save_top_k, train a large model, and set WandbLogger(log_model=True).

Error messages and logs

# Error messages and logs here please

Environment

#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):

More info

No response

cc @awaelchli @morganmcg1 @borisdayma @scottire @parambharat

Guida contributor