[WandbLogger] Call `wandb.finish()` to ensure all artifacts are uploaded before training ends.
#17.768 aberto em 7 de jun. de 2023
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.
@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