Lightning-AI / Lightning-AI/pytorch-lightning

Trainer does not wait for neptune logger completion and logger connection stays open unless explicitly closed

Open
#19,750 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug logger: neptune
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

I'm performing a naive hyperparameter sweep using the PL Trainer and NeptuneLogger. After the successful completion of `Trainer.fit()` I see that the neptune run is still not complete on Neptune App until the Jupyter Kernel is killed. I also see odd behavior where the next run will start and the training will be terminated almost immediately (what I suspect to be the NeptuneLogger instance synchronizing with the server and then stopping training?).

A snippet of the code is below:

```python
def train(hparams):
model = ImageClassifier(hparams["model_name"], num_classes=hparams["num_classes"], lr=hparams["lr"])

neptune_logger = NeptuneLogger(
project="project_name",
api_token=neptune_token
)

neptune_logger.log_hyperparams(params=hparams)

trainer = Trainer(
callbacks=[checkpoint_callback, early_stopping_callback],
max_epochs=hparams["max_epochs"],
accelerator=hparams["training_device"],
logger=neptune_logger,
)

trainer.fit(model, train_dataloader, val_dataloader)
neptune_logger.log_model_summary(model=model, max_depth=-1)

model = ImageClassifier.load_from_checkpoint("checkpoints/best-checkpoint.ckpt", model_name=hparams["model_name"], num_classes=num_classes, lr=hparams["lr"])
script = model.to_torchscript()
torch.jit.save(script, "traced_model.pt")
neptune_logger.run["model"].track_files("traced_model.pt")

max_epochs = [30, 40, 50]
lrs = [5e-2, 1e-3, 5e-3]
batch_sizes = [(32, 32, 32), (64, 64, 64), (128, 128, 128)]
val_split_sizes = [0.2, 0.3, 0.4]

combinations = list(itertools.product(max_epochs, lrs, batch_sizes, val_split_sizes))

for hp in combinations:
hparams = {"num_classes": num_classes, "max_epochs": hp[0], "lr": hp[1], "batch_sizes": hp[2], "val_split_size": hp[3]}
print("=============Training============")
print(f"Parameters: {run_params}")
train(run_params)
print("=============Complete============")
```

### What version are you seeing the problem on?

v2.2

### How to reproduce the bug

_No response_

### Error messages and logs

```
# Error messages and logs here please
```

### Environment

Current environment

* CUDA:
- GPU:
- NVIDIA GeForce RTX 2060 SUPER
- available: True
- version: 12.1
* Lightning:
- lightning-utilities: 0.11.2
- pytorch-lightning: 2.2.1
- torch: 2.1.0
- torchaudio: 2.1.0
- torchmetrics: 1.3.2
- torchvision: 0.16.0
* Packages:
- aiohttp: 3.9.3
- aiosignal: 1.3.1
- anyio: 4.0.0
- argon2-cffi: 23.1.0
- argon2-cffi-bindings: 21.2.0
- arrow: 1.3.0
- asttokens: 2.4.1
- async-lru: 2.0.4
- attrs: 23.1.0
- babel: 2.13.1
- backports.functools-lru-cache: 1.6.5
- beautifulsoup4: 4.12.2
- bleach: 6.1.0
- boto3: 1.34.81
- botocore: 1.34.81
- bottleneck: 1.3.5
- bravado: 11.0.3
- bravado-core: 6.1.1
- brotli: 1.1.0
- cached-property: 1.5.2
- certifi: 2023.7.22
- cffi: 1.16.0
- charset-normalizer: 3.3.2
- click: 8.1.7
- comm: 0.1.4
- contourpy: 1.1.1
- cycler: 0.12.1
- datasets: 2.18.0
- debugpy: 1.8.0
- decorator: 5.1.1
- defusedxml: 0.7.1
- dill: 0.3.8
- entrypoints: 0.4
- exceptiongroup: 1.1.3
- executing: 2.0.1
- fastjsonschema: 2.18.1
- filelock: 3.13.1
- fonttools: 4.43.1
- fqdn: 1.5.1
- frozenlist: 1.4.1
- fsspec: 2024.2.0
- future: 1.0.0
- gitdb: 4.0.11
- gitpython: 3.1.43
- gmpy2: 2.1.2
- huggingface-hub: 0.22.2
- idna: 3.4
- importlib-metadata: 6.8.0
- importlib-resources: 6.1.0
- ipykernel: 6.26.0
- ipython: 8.17.2
- ipython-genutils: 0.2.0
- ipywidgets: 8.1.1
- isoduration: 20.11.0
- jedi: 0.19.1
- jinja2: 3.1.2
- jmespath: 1.0.1
- json5: 0.9.14
- jsonpointer: 2.4
- jsonref: 1.1.0
- jsonschema: 4.19.2
- jsonschema-specifications: 2023.7.1
- jupyter: 1.0.0
- jupyter-client: 7.4.9
- jupyter-console: 6.6.3
- jupyter-contrib-core: 0.4.0
- jupyter-contrib-nbextensions: 0.7.0
- jupyter-core: 5.5.0
- jupyter-events: 0.8.0
- jupyter-highlight-selected-word: 0.2.0
- jupyter-latex-envs: 1.4.6
- jupyter-lsp: 2.2.0
- jupyter-nbextensions-configurator: 0.6.1
- jupyter-server: 2.9.1
- jupyter-server-terminals: 0.4.4
- jupyterlab: 4.0.7
- jupyterlab-pygments: 0.2.2
- jupyterlab-server: 2.25.0
- jupyterlab-widgets: 3.0.9
- kiwisolver: 1.4.5
- lightning-utilities: 0.11.2
- lxml: 4.9.2
- markupsafe: 2.1.3
- matplotlib: 3.8.1
- matplotlib-inline: 0.1.6
- mistune: 3.0.1
- monotonic: 1.6
- mpmath: 1.3.0
- msgpack: 1.0.8
- multidict: 6.0.5
- multiprocess: 0.70.16
- munkres: 1.1.4
- nbclassic: 1.0.0
- nbclient: 0.8.0
- nbconvert: 7.10.0
- nbformat: 5.9.2
- neptune: 1.10.2
- nest-asyncio: 1.5.8
- networkx: 3.2.1
- notebook: 6.5.6
- notebook-shim: 0.2.3
- numexpr: 2.8.7
- numpy: 1.26.0
- oauthlib: 3.2.2
- overrides: 7.4.0
- packaging: 23.2
- pandas: 2.1.1
- pandocfilters: 1.5.0
- parso: 0.8.3
- pexpect: 4.8.0
- pickleshare: 0.7.5
- pillow: 9.4.0
- pip: 23.3.1
- pkgutil-resolve-name: 1.3.10
- platformdirs: 3.11.0
- ply: 3.11
- prometheus-client: 0.18.0
- prompt-toolkit: 3.0.39
- psutil: 5.9.5
- ptyprocess: 0.7.0
- pure-eval: 0.2.2
- pyarrow: 15.0.2
- pyarrow-hotfix: 0.6
- pycparser: 2.21
- pygments: 2.16.1
- pyjwt: 2.8.0
- pyparsing: 3.1.1
- pyqt5: 5.15.9
- pyqt5-sip: 12.12.2
- pysocks: 1.7.1
- python-dateutil: 2.8.2
- python-json-logger: 2.0.7
- pytorch-lightning: 2.2.1
- pytz: 2023.3.post1
- pyyaml: 6.0.1
- pyzmq: 24.0.1
- qtconsole: 5.4.4
- qtpy: 2.4.1
- referencing: 0.30.2
- requests: 2.31.0
- requests-oauthlib: 2.0.0
- rfc3339-validator: 0.1.4
- rfc3986-validator: 0.1.1
- rpds-py: 0.10.6
- s3transfer: 0.10.1
- safetensors: 0.4.2
- send2trash: 1.8.2
- setuptools: 68.2.2
- simplejson: 3.19.2
- sip: 6.7.12
- six: 1.16.0
- smmap: 5.0.1
- sniffio: 1.3.0
- soupsieve: 2.5
- stack-data: 0.6.2
- swagger-spec-validator: 3.0.3
- sympy: 1.12
- terminado: 0.17.1
- timm: 0.9.16
- tinycss2: 1.2.1
- toml: 0.10.2
- tomli: 2.0.1
- torch: 2.1.0
- torchaudio: 2.1.0
- torchmetrics: 1.3.2
- torchvision: 0.16.0
- tornado: 6.3.3
- tqdm: 4.66.2
- traitlets: 5.13.0
- triton: 2.1.0
- types-python-dateutil: 2.8.19.14
- typing-extensions: 4.8.0
- typing-utils: 0.1.0
- tzdata: 2023.3
- uri-template: 1.3.0
- urllib3: 2.0.7
- wcwidth: 0.2.9
- webcolors: 1.13
- webencodings: 0.5.1
- websocket-client: 1.6.4
- wheel: 0.41.3
- widgetsnbextension: 4.0.9
- xxhash: 3.4.1
- yarl: 1.9.4
- zipp: 3.17.0
* System:
- OS: Linux
- architecture:
- 64bit
- ELF
- processor: x86_64
- python: 3.11.6
- release: 5.4.0-166-generic
- version: #183-Ubuntu SMP Mon Oct 2 11:28:33 UTC 2023

### More info

_No response_

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

No repository files, tests, or concrete reproduction steps are identified in the report. Start by reproducing the Jupyter sweep with Trainer.fit() and NeptuneLogger, then trace the logger lifecycle around training completion; done means the run finishes and the connection closes before the next sweep iteration without killing the kernel.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning, observability
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.