Bug: explicit dataloader iterator `_shutdown_workers()` call causes lightning to crash when pytorch ignores the error
bugneeds triagever: 2.6.x
### Bug description
When using `CombinedLoader` with persistent workers and large tensors, Lightning calls `_shutdown_workers()` explicitly, which exposes a PyTorch race condition that PyTorch normally suppresses via `__del__`. This causes jobs to exit with non-zero code.
**File:** [`src/lightning/pytorch/utilities/combined_loader.py:400`](https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/pytorch/utilities/combined_loader.py#L400)
## Expected Behavior
Lightning should match PyTorch's behavior and suppress worker shutdown errors (exit code 0).
## Actual Behavior
Lightning explicitly calls `_shutdown_workers()`, causing the error to propagate and the process to terminate non zero exit code.
## Error Output
```
Predicting ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1/1 0:00:00 • 0:00:00 0.00it/s terminate called without an active exception
Predicting ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1/1 0:00:00 • 0:00:00 0.00it/s
Traceback (most recent call last):
File "/homes/michaelva/personal_repos/pytorch-lightning/src/lightning/pytorch/trainer/call.py", line 49, in _call_and_handle_interrupt
return trainer_fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/homes/michaelva/personal_repos/pytorch-lightning/src/lightning/pytorch/trainer/trainer.py", line 996, in _predict_impl
results = self._run(model, ckpt_path=ckpt_path, weights_only=weights_only)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/homes/michaelva/personal_repos/pytorch-lightning/src/lightning/pytorch/trainer/trainer.py", line 1091, in _run
self._teardown()
File "/homes/michaelva/personal_repos/pytorch-lightning/src/lightning/pytorch/trainer/trainer.py", line 1112, in _teardown
loop.teardown()
File "/homes/michaelva/personal_repos/pytorch-lightning/src/lightning/pytorch/loops/prediction_loop.py", line 210, in teardown
self._data_fetcher.teardown()
File "/homes/michaelva/personal_repos/pytorch-lightning/src/lightning/pytorch/loops/fetchers.py", line 82, in teardown
self._combined_loader.reset()
File "/homes/michaelva/personal_repos/pytorch-lightning/src/lightning/pytorch/utilities/combined_loader.py", line 367, in reset
_shutdown_workers_and_reset_iterator(iterable)
File "/homes/michaelva/personal_repos/pytorch-lightning/src/lightning/pytorch/utilities/combined_loader.py", line 400, in _shutdown_workers_and_reset_iterator
dataloader._iterator._shutdown_workers()
File "/home/michaelva/.conda/envs/lightning-dev/lib/python3.11/site-packages/torch/utils/data/dataloader.py", line 1618, in _shutdown_workers
w.join(timeout=_utils.MP_STATUS_CHECK_INTERVAL)
File "/home/michaelva/.conda/envs/lightning-dev/lib/python3.11/multiprocessing/process.py", line 149, in join
res = self._popen.wait(timeout)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/michaelva/.conda/envs/lightning-dev/lib/python3.11/multiprocessing/popen_fork.py", line 40, in wait
if not wait([self.sentinel], timeout):
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/michaelva/.conda/envs/lightning-dev/lib/python3.11/multiprocessing/connection.py", line 930, in wait
ready = selector.select(timeout)
^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/michaelva/.conda/envs/lightning-dev/lib/python3.11/selectors.py", line 415, in select
fd_event_list = self._selector.poll(timeout)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/michaelva/.conda/envs/lightning-dev/lib/python3.11/site-packages/torch/utils/data/_utils/signal_handling.py", line 73, in handler
_error_if_any_worker_fails()
RuntimeError: DataLoader worker (pid 1931486) is killed by signal: Aborted.
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "/homes/michaelva/personal_repos/pytorch-lightning/pytorch_lightning_reproduce.py", line 53, in <module>
trainer.predict(model=TestModel(), dataloaders=dataloader)
File "/homes/michaelva/personal_repos/pytorch-lightning/src/lightning/pytorch/trainer/trainer.py", line 947, in predict
return call._call_and_handle_interrupt(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/homes/michaelva/personal_repos/pytorch-lightning/src/lightning/pytorch/trainer/call.py", line 70, in _call_and_handle_interrupt
trainer._teardown()
File "/homes/michaelva/personal_repos/pytorch-lightning/src/lightning/pytorch/trainer/trainer.py", line 1112, in _teardown
loop.teardown()
File "/homes/michaelva/personal_repos/pytorch-lightning/src/lightning/pytorch/loops/prediction_loop.py", line 210, in teardown
self._data_fetcher.teardown()
File "/homes/michaelva/personal_repos/pytorch-lightning/src/lightning/pytorch/loops/fetchers.py", line 80, in teardown
self.reset()
File "/homes/michaelva/personal_repos/pytorch-lightning/src/lightning/pytorch/loops/fetchers.py", line 142, in reset
super().reset()
File "/homes/michaelva/personal_repos/pytorch-lightning/src/lightning/pytorch/loops/fetchers.py", line 76, in reset
self.length = sized_len(self.combined_loader)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/homes/michaelva/personal_repos/pytorch-lightning/src/lightning/fabric/utilities/data.py", line 52, in sized_len
length = len(dataloader) # type: ignore [arg-type]
^^^^^^^^^^^^^^^
File "/homes/michaelva/personal_repos/pytorch-lightning/src/lightning/pytorch/utilities/combined_loader.py", line 358, in __len__
raise RuntimeError("Please call `iter(combined_loader)` first.")
RuntimeError: Please call `iter(combined_loader)` first.
```
## Root Cause
This is a timing-dependent race condition in PyTorch when shutting down workers with large tensors in memory. PyTorch suppresses this error in `__del__`, but Lightning's explicit `_shutdown_workers()` call exposes it.
The issue occurs when:
- Using `num_workers > 1`
- Using `multiprocessing_context="spawn"`
- Using `persistent_workers=True`
- Processing large tensors (e.g., multi camera high-resolution images)
## Reproduction
```python
import os
import torch
from torch.utils.data import Dataset, DataLoader
from lightning import Trainer
from lightning import LightningModule
class LargeTensorDataset(Dataset):
def __init__(self, size=100):
self.size = size
def __len__(self):
return self.size
def __getitem__(self, idx):
# Create large tensors (98.3 MB per tensor)
images = torch.randn(8, 640, 640, 3, dtype=torch.float64)
return {'images': images, 'index': idx}
def get_dataloader():
dataset = LargeTensorDataset(size=50)
num_workers = os.cpu_count()
print(f"Creating DataLoader with {num_workers} workers...")
dataloader = DataLoader(
dataset,
batch_size=8,
num_workers=num_workers,
multiprocessing_context='spawn',
persistent_workers=True
)
return dataloader
class TestModel(LightningModule):
def predict_step(self, batch, batch_idx, dataloader_idx=0):
return batch
if __name__ == "__main__":
trainer = Trainer(max_epochs=1, limit_predict_batches=1)
dataloader = get_dataloader()
trainer.predict(model=TestModel(), dataloaders=dataloader)
```
**Note:** This is timing-dependent and may not reproduce every time. More likely with higher worker counts and larger tensors.
## Proposed Solution
Change `_shutdown_workers_and_reset_iterator()` in combined_loader.py:
**Current (line 397-402):**
```python
def _shutdown_workers_and_reset_iterator(dataloader: object) -> None:
if hasattr(dataloader, "_iterator"):
if isinstance(dataloader._iterator, _MultiProcessingDataLoaderIter):
dataloader._iterator._shutdown_workers()
dataloader._iterator = None
```
**Proposed:**
```python
def _shutdown_workers_and_reset_iterator(dataloader: object) -> None:
if hasattr(dataloader, "_iterator"):
if isinstance(dataloader._iterator, _MultiProcessingDataLoaderIter):
del dataloader._iterator
dataloader._iterator = None
```
## Environment
- **Lightning:** 2.6.1
- **PyTorch:** 2.9.1
- **OS:** Ubuntu 22.04.5 LTS
- **CPU:** Intel(R) Xeon(R) Gold 6354 CPU @ 3.00GHz (4 cores)
### What version are you seeing the problem on?
v2.6
### Reproduced in studio
_No response_
### How to reproduce the bug
```python
```
### Error messages and logs
### Environment
_No response_
### More info
_No response_
cc @ethanwharris
0 条评论