Lightning-AI/pytorch-lightning

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

Open

#17,768 创建于 2023年6月7日

在 GitHub 查看
 (2 评论) (2 反应) (0 负责人)Python (3,233 fork)batch import
bughelp wantedlogger: wandbver: 2.1.x

仓库指标

Star
 (26,687 star)
PR 合并指标
 (平均合并 9天 15小时) (30 天内合并 3 个 PR)

描述

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

贡献者指南