`ThreadPrefetchIterDataset` memory leaks / zombie threads that hang exiting?
- Dominant language
- Python
- Stars
- 779
- Forks
- 86
- Avg merge
- 2d 6h
- Merged PRs (30d)
- 6
Description
(This seems related to #1021 , but I don't _think_ it's the same issue since `grain_nightly` presumably has the fix to #1021 from #1051.)
This minimal reproducer:
```python
import grain
import jax
import numpy as np
from grain.experimental import device_put as grain_device_put
from grain.sources import RandomAccessDataSource
class FakeDataSource(RandomAccessDataSource):
def __len__(self):
return 1000
def __getitem__(self, idx):
return {"x": np.ones((32,), dtype=np.float32) * idx}
# has to be in main() to hang
def main():
source = FakeDataSource()
dataset = grain.MapDataset.source(source)
# has to have this to hang---something to do with nested threaded pre-fetching?
dataset = grain_device_put(dataset, device=jax.devices()[0])
# if this is for _ in iter(dataset) it doesn't hang
data_iter = iter(dataset)
for _ in data_iter:
raise RuntimeError
if __name__ == "__main__":
main()
```
will raise the `RuntimeError`, print the traceback, and then instead of exiting it will just hang indefinitely. I am running it with `uv run --python 3.13 --with grain --with jax repro.py`.
If everything is moved from `main()` to the top level of the script, then it fails with an additional exception and the process exits fully and returns to the shell:
```
Exception ignored in:
Traceback (most recent call last):
File ".../.venv/lib/python3.13/site-packages/grain/_src/python/dataset/dataset.py", line 1602, in __del__
File ".../.venv/lib/python3.13/site-packages/grain/_src/python/dataset/transformations/prefetch.py", line 989, in close
File ".../.venv/lib/python3.13/site-packages/grain/_src/python/dataset/transformations/prefetch.py", line 1006, in _stop_prefetch
File ".../.venv/lib/python3.13/site-packages/grain/_src/python/dataset/transformations/prefetch.py", line 995, in _clear_buffer
AttributeError: 'NoneType' object has no attribute 'Empty'
```
Running with Python 3.14 instead (`uv run --python 3.14 --with grain --with jax repro.py`) prints a different extra traceback and exits back to the shell:
```
Exception ignored while calling deallocator :
Traceback (most recent call last):
File ".../.cache/uv/archive-v0/86c0NR2RDIQ-nixYPZcBP/lib/python3.14/site-packages/grain/_src/python/dataset/dataset.py", line 1551, in __del__
self.close()
File ".../.cache/uv/archive-v0/86c0NR2RDIQ-nixYPZcBP/lib/python3.14/site-packages/grain/_src/python/dataset/transformations/prefetch.py", line 959, in close
self._stop_prefetch()
File ".../.cache/uv/archive-v0/86c0NR2RDIQ-nixYPZcBP/lib/python3.14/site-packages/grain/_src/python/dataset/transformations/prefetch.py", line 977, in _stop_prefetch
self._prefetch_thread.join()
File ".../.local/share/uv/python/cpython-3.14.2-linux-x86_64-gnu/lib/python3.14/threading.py", line 1133, in join
self._os_thread_handle.join(timeout)
PythonFinalizationError: cannot join thread at interpreter shutdown
```
Running the `uv` commands with `--with grain_nightly` instead gives same results as far as I can tell. I ran with `grain_nightly==0.2.16.dev20260112` and current `grain==0.2.15`. CUDA-enabled JAX or CPU only also don't seem to matter, although I ran both tests on a Linux machine with a GPU.
While the example may seem contrived, I distilled it from a real-world failure in my training code where issues in my main training loop were causing jobs to hang indefinitely instead of exiting. I can't use `for _ in iter(dataset)` as a workaround because I need a reference to the `data_iter` object for checkpointing.
Contributor guide
Assessment
This issue has not been assessed yet.