Lightning-AI / Lightning-AI/pytorch-lightning

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

Open
#17,768 2 comments 2 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug help wanted logger: wandb ver: 2.1.x
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.