Lightning-AI / Lightning-AI/pytorch-lightning
Trainer does not wait for neptune logger completion and logger connection stays open unless explicitly closed
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
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
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
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