[WandbLogger] Call `wandb.finish()` to ensure all artifacts are uploaded before training ends.
#17,768 建立於 2023年6月7日
倉庫指標
- 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.
@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