dataloader timeout (#172)

elegy/data/dataset.py

@@ -67,6 +67,7 @@ def __init__(
         shuffle: tp.Optional[bool] = False,
         worker_type: tp.Optional[str] = "thread",
         prefetch: tp.Optional[int] = 1,
+        timeout: int = 10,
     ):
         """
         Arguments:
@@ -82,6 +83,7 @@ def __init__(
                          'spawn', 'fork' and 'forkserver' can be used to select a specific process type.
                          For more information consult the Python `multiprocessing` documentation.
             prefetch: Number of batches to prefetch for pipelined execution (Default: 2)
+            timeout: Timeout in seconds for waiting for a batch of samples from worker threads.
         """
         assert (
             batch_size > 0 and type(batch_size) == int
@@ -97,6 +99,7 @@ def __init__(
         self.shuffle = shuffle
         self.worker_type = worker_type
         self.prefetch = prefetch
+        self.timeout = timeout
 
     def __len__(self) -> int:
         """Returns the number of batches per epoch"""
@@ -122,6 +125,7 @@ def __iter__(self) -> tp.Generator[tp.Any, None, None]:
                 self.n_workers,
                 prefetch=self.prefetch,
                 worker_type=self.worker_type,
+                timeout=self.timeout,
             )
 
 

elegy/data/dataset_test.py

@@ -143,7 +143,7 @@ def __getitem__(self, i):
 
     def test_worker_type(self):
         ds = DS0()
-        for worker_type in ["thread", "process", "spawn", "fork", "forkserver"]:
+        for worker_type in ["thread", "process"]:
             loader = elegy.data.DataLoader(
                 ds, batch_size=4, n_workers=4, worker_type=worker_type
             )