Fix shared memory queue not shutdown issue and cause load_with_workers hangs forever when exits issue

Author: fangpenlin

pyproject.toml

@@ -11,6 +11,7 @@ dependencies = [
 
 [dependency-groups]
 dev = [
+    "pytest-timeout>=2.4.0",
     "pytest>=8.4.0",
     "tqdm>=4.67.1",
 ]

tests/test_loader.py

@@ -2,6 +2,7 @@
 from multiprocessing.managers import SharedMemoryManager
 
 import numpy as np
+import pytest
 import tinygrad
 import tqdm
 
@@ -96,3 +97,22 @@ def test_share_memory_enabled():
             assert y.numpy().shape == label_size
             count += 1
     assert count == n
+
+
+@pytest.mark.timeout(10)
+def test_generator_early_stops_queue_not_shutdown():
+    data_size = (5,)
+    label_size = (4,)
+    num_worker = 4
+
+    def forever_gen():
+        while True:
+            yield 1
+
+    loader = RandomLoader(data_size=data_size, label_size=label_size)
+    with load_with_workers(
+        loader, forever_gen(), num_worker, shared_memory_enabled=True
+    ) as generator:
+        for i, _ in enumerate(tqdm.tqdm(generator)):
+            if i > 10:
+                break

tinyloader/loader.py

@@ -17,18 +17,41 @@
 class Loader(abc.ABC):
     @abc.abstractmethod
     def make_request(self, item: typing.Any) -> typing.Any:
+        """Called to make data loading request to potentially passing to the workers. Ideally the return value should
+        be easily pickable otherwise it might be very slow.
+
+        :param item: The item to generate the loading request for
+        :return: a pickable value for the worker process or the current process to load
+        """
         raise NotImplementedError
 
     @abc.abstractmethod
     def load(self, request: typing.Any) -> tuple[np.typing.NDArray, ...]:
+        """Called to load data for the given item. Potentially called from a worker process.
+
+        :param request: Request for loading the data
+        :return: The loaded data, should be a tuple of numpy's ndarray
+        """
         raise NotImplementedError
 
     @abc.abstractmethod
     def post_process(
         self, response: tuple[np.typing.NDArray, ...]
     ) -> tuple[tinygrad.Tensor, ...]:
+        """Called to convert numpy's ndarray returned from the `load` method into tinygrad's Tensor for training or
+        testing purpose. This method will be called from the process which invokes the loading generator.
+
+        :param response: Response ndarray values returned by the `load` method
+        :return: A tuple of tinygrad Tensor for training / testing or other purpose
+        """
         raise NotImplementedError
 
+    def shutdown(self):
+        """Called to shutdown resources associated with the loader. Like, abort async operations or release files and
+        etc.
+
+        """
+
 
 @dataclasses.dataclass(frozen=True)
 class SharedBuffer:
@@ -112,6 +135,9 @@ def push(self, shared_buffer: SharedBuffer):
         self._queue.put(shared_buffer.index)
         logger.debug("Pushed shared buffer %s", shared_buffer)
 
+    def shutdown(self):
+        self._queue.shutdown(immediate=True)
+
 
 class SharedMemoryShim(Loader):
     def __init__(
@@ -181,6 +207,10 @@ def __reduce__(self):
         # avoid pickling SharedMemoryManager, only care about the underlying loader in `load` method anyway
         return self.__class__, (self.loader, None, 0), None
 
+    def shutdown(self):
+        for mem_pool in self._memory_pools.values():
+            mem_pool.shutdown()
+
 
 def load(
     loader: Loader, items: typing.Sequence[typing.Any]
@@ -244,4 +274,7 @@ def generate() -> typing.Generator[tuple[tinygrad.Tensor, ...], None, None]:
                 ),
             )
 
-        yield generate()
+        try:
+            yield generate()
+        finally:
+            actual_loader.shutdown()