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