Lightning-AI / Lightning-AI/pytorch-lightning
[WandbLogger] Call `wandb.finish()` to ensure all artifacts are uploaded before training ends.
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
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
```python
@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
Current 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
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in src/lightning/pytorch/loggers/wandb.py at the finalize method linked in the issue. Reproduce with a large model, increased save_top_k, and WandbLogger(log_model=True), then verify that training does not return until artifact uploads are complete. Done means the logger explicitly finishes the W&B run after checkpoint artifacts are handled.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 52/100