Lightning-AI/pytorch-lightning

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

Open

#17.768 aberto em 7 de jun. de 2023

Ver no GitHub
 (2 comments) (2 reactions) (0 assignees)Python (3.233 forks)batch import
bughelp wantedlogger: wandbver: 2.1.x

Métricas do repositório

Stars
 (26.687 stars)
Métricas de merge de PR
 (Mesclagem média 9d 15h) (3 fundiu PRs em 30d)

Description

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

Guia do colaborador