Repository: mosaicml/streaming

Issue:
deadlock when process exits before end of epoch
## Environment
<!-- Provide the below as necessary -->

- OS: AlmaLinux 8.8
- Hardware (GPU, or instance type): no gpu


## To reproduce
1. take any StreamingDataset
2. create an iterator
3. do one iteration of the iterator, but do NOT finish an epoch.
4. try to exit the program.

The following script will deadlock on exit:
```python
from streaming import MDSWriter
from streaming import StreamingDataset

### generate a sample dataset...
# directory for a sample dataset
out_root = 'deadlock_data'

# A dictionary of input fields to an Encoder/Decoder type
columns = {
    'anumber': 'int'
}

# some sample data
samples = [
    {
        'anumber': 123123
    }
    for _ in range(1000)
]

# write out the sample data
print("writing data...")
with MDSWriter(out=out_root, columns=columns) as out:
    for sample in samples:
        out.write(sample)
### finished generating the sample dataset


print("creating dataset object...")
remote_dir = 'deadlock_data'
local_dir = 'local_cache'
dataset = StreamingDataset(remote=remote_dir, local_dir=local_dir, split=None)
i = iter(dataset)
print("accessing the iterator...")

print(next(i))

print("exiting...")

# process will now hang on exit...

```

## Expected behavior
The program should exit.

## Additional context
I believe this is caused by the threads spawned by the iterator (`_prepare_thread`  and `_ready_thread`) never exiting. When the interpreter exits, it waits for these threads and so will hang. One hack to fix it is to add the following:
```python
import threading
from streaming.base.dataset import _IterState
def kill_iterators(dataset):
    if hasattr(dataset, "_iterator"):
        with dataset._iterator._lock:
            if dataset._iterator._state == _IterState.ITERATING:
                dataset._iterator._state = _IterState.EXITING
threading._register_atexit(kill_iterators, dataset)
```
Note that we need to use `threading._register_atexit` rather than the `atexit` module because the code that waits on the threads happens before `atexit`. Also, it does not work to call `dataset._iterator.exit()` inside `kill_iterators` because this will wait for two threads to exit, which causes its own deadlock if the process has been forked because one fork will not have any threads to exit.





Discussion:
Hi @acutkosky , thanks for finding out the bug. Not sure how much bandwidth do you have, but you can open a PR with the change, would love to review it. Thanks!
Thanks a lot @acutkosky. I was about to open a similar issue. Any ideas on the timeline for a bugfix? Seems pretty important to me!

Resolving patch:
diff --git a/streaming/base/dataset.py b/streaming/base/dataset.py
--- a/streaming/base/dataset.py
+++ b/streaming/base/dataset.py
@@ -6,6 +6,7 @@
 import json
 import logging
 import os
+import sys
 import warnings
 from concurrent.futures import ThreadPoolExecutor, wait
 from concurrent.futures._base import Future
@@ -103,6 +104,29 @@ def __init__(self, sample_ids: NDArray[np.int64]) -> None:
         self._state = 0
         self._num_exited = 0
 
+        # python will attempt to join all threads on shutdown.
+        # Here, we register a call to self.non_blocking_exit to run
+        # at shutdown to prevent a deadlock.
+        # In python version >=3.9 this can be accomplished via
+        # threading._register_atexit but not with the atexit module.
+        # In older python versions, the atexit module can be used, and
+        # threading._register_atexit does not exist.
+        if sys.version_info[1] <= 8:  # check if python version <=3.8
+            import atexit
+            atexit.register(self.non_blocking_exit)
+        else:
+            from threading import _register_atexit  # pyright: ignore
+            _register_atexit(self.non_blocking_exit)
+
+    def non_blocking_exit(self) -> None:
+        """Signal threads to exit without blocking.
+
+        This will be called at process exit.
+        """
+        with self._lock:
+            if self._state == _IterState.ITERATING:
+                self._state = _IterState.EXITING
+
     def exit(self) -> None:
         """Signal threads to exit, wait until they have all exited, then return.
 


Test patch:
diff --git a/tests/test_streaming.py b/tests/test_streaming.py
--- a/tests/test_streaming.py
+++ b/tests/test_streaming.py
@@ -4,6 +4,7 @@
 import math
 import os
 import shutil
+from multiprocessing import Process
 from typing import Any, Tuple
 
 import pytest
@@ -483,6 +484,32 @@ def test_dataloader_fixed_balanced_sampling(local_remote_dir: Any, batch_size: i
             assert samples_seen != first_samples_seen
 
 
+@pytest.mark.parametrize('num_samples', [9867])
+@pytest.mark.parametrize('seed', [1234])
+@pytest.mark.usefixtures('local_remote_dir')
+def test_dataloader_mid_epoch_exit(local_remote_dir: Tuple[str, str], num_samples: int, seed: int):
+    local, remote = local_remote_dir
+    convert_to_mds(out_root=remote, dataset_name='sequencedataset', num_samples=num_samples)
+
+    def run_one_iter(local: str, remote: str, seed: int) -> None:
+        # Build a StreamingDataset
+        dataset = StreamingDataset(local=local, remote=remote, shuffle_seed=seed)
+
+        # Do one iteration
+        it = iter(dataset)
+        next(it)
+        # Test if we can exit...
+
+    p = Process(target=run_one_iter, args=(local, remote, seed))
+    p.start()
+    p.join(5)
+    p.terminate()
+
+    result = p.exitcode
+
+    assert result == 0
+
+
 @pytest.mark.parametrize('batch_size', [128])
 @pytest.mark.parametrize('drop_last', [False, True])
 @pytest.mark.parametrize('shuffle', [False, True])