Forward worker exceptions and exit with it

test/dataloader2/test_dataloader2.py

@@ -62,6 +62,8 @@
 
 mp_ctx_parametrize = parametrize("ctx", mp.get_all_start_methods())
 
+EXCEPTION_ITERATION_NUM = 7
+
 
 class _ReadingServiceWrapper:
     def __init__(self, dp):
@@ -79,6 +81,18 @@ def return_one():
         return 1
 
 
+class MakeMistakeDataPipe(IterDataPipe):
+    def __init__(self, source_datapipe, exc_iteration=EXCEPTION_ITERATION_NUM):
+        self.source_datapipe = source_datapipe
+        self.exc_iteration = exc_iteration
+
+    def __iter__(self):
+        for i, x in enumerate(self.source_datapipe):
+            if i == self.exc_iteration:
+                raise Exception("oops")
+            yield x
+
+
 class TestReadingService(ReadingServiceInterface):
     def initialize(self, dp: DataPipe) -> DataPipe:
         return _ReadingServiceWrapper(dp)  # type: ignore[return-value]
@@ -99,6 +113,19 @@ def test_dataloader2_shutdown(self) -> None:
         data_loader: DataLoader2 = DataLoader2(datapipe=test_data_pipe)
         data_loader.shutdown()
 
+    def test_worker_exception_raised(self):
+        dp = IterableWrapper(range(100)).sharding_filter()
+        dp = MakeMistakeDataPipe(dp)
+        for worker_prefetch_cnt in [0, 5, 10]:
+            for num_workers in [1, 4]:
+                rs = MultiProcessingReadingService(num_workers=num_workers, worker_prefetch_cnt=worker_prefetch_cnt)
+                dl = DataLoader2(dp, reading_service=rs)
+                it = iter(dl)
+                for i in range(EXCEPTION_ITERATION_NUM * num_workers):
+                    next(it)
+                with self.assertRaises(communication.iter.WorkerException):
+                    next(it)
+
     def test_dataloader2_state_dict(self) -> None:
         test_data_pipe = IterableWrapper(range(3))
         data_loader: DataLoader2 = DataLoader2(datapipe=test_data_pipe)
@@ -171,7 +198,6 @@ def test_dataloader2_iterates_correctly(self) -> None:
             self.assertEqual(list(range(10)), actual)
 
     def test_dataloader2_reset(self) -> None:
-
         test_data_pipe = IterableWrapper(range(10))
         reading_services = [None, TestReadingService(), MultiProcessingReadingService(num_workers=1)]
 
@@ -264,7 +290,6 @@ def test_dataloader2_shuffle(self) -> None:
     "fork is not supported. Dying (set die_after_fork=0 to override)",
 )
 class TestDataLoader2EventLoop(TestCase):
-
     # TODO: This needs fixing, see issue 624
     # @skipIfNoDill
     # def test_basic_threading(self):

torchdata/dataloader2/communication/iter.py

@@ -59,6 +59,12 @@ class TerminateRequired(Exception):
     pass
 
 
+class WorkerException(Exception):
+    """
+    Returned by DataPipe when there is a failure/exception from a worker process
+    """
+
+
 class NonBlocking(IterDataPipe):
     not_available_hook = default_not_available_hook
 
@@ -212,6 +218,9 @@ def DataPipeBehindQueues(source_datapipe, protocol, blocking_request_get=False,
                     protocol.response_invalid_state()
                     yield True
                     break
+                except Exception as e:
+                    protocol.response_worker_exception(e)
+                    return
                 protocol.response_next(value)
                 yield True  # Returns control
                 break
@@ -322,6 +331,8 @@ def __iter__(self):
                         raise communication.iter.InvalidStateResetRequired
                     if isinstance(response, communication.messages.TerminateResponse):
                         raise communication.iter.TerminateRequired
+                    if isinstance(response, communication.messages.WorkerExceptionResponse):
+                        raise communication.iter.WorkerException(f"Exception from worker {idx}") from response.exception
                     if len(self.res_buffers[idx]) == 0:  # Only request if buffer is empty
                         self.datapipes[idx].protocol.request_next()
                     yield response.value

torchdata/dataloader2/communication/messages.py

@@ -108,3 +108,8 @@ class InvalidStateResponse(Response):
     """
 
     pass
+
+
+class WorkerExceptionResponse(Response):
+    def __init__(self, exception):
+        self.exception = exception

torchdata/dataloader2/communication/protocol.py

@@ -231,6 +231,12 @@ def response_invalid_state(self):
         self.response_queue.put(communication.messages.InvalidStateResponse())
         self._req_received = None
 
+    def response_worker_exception(self, exception):
+        if not self.have_pending_request():
+            raise Exception("Attempting to reply with pending request")
+        self.response_queue.put(communication.messages.WorkerExceptionResponse(exception))
+        self._req_received = None
+
 
 class IterDataPipeQueueProtocolClient(ProtocolClient):
     def request_reset_iterator(self):

torchdata/datapipes/iter/util/prefetcher.py

@@ -76,6 +76,9 @@ def thread_worker(prefetch_data: _PrefetchData):
                     except communication.iter.TerminateRequired:
                         prefetch_data.run_prefetcher = False
                         prefetch_data.stop_iteration = True
+                    except Exception as e:
+                        prefetch_data.prefetch_buffer.append(e)
+                        break
                 elif prefetch_data.stop_iteration and len(prefetch_data.prefetch_buffer) == 0:
                     prefetch_data.run_prefetcher = False
                 else:  # Buffer is full, waiting for main thread to consume items
@@ -93,7 +96,11 @@ def __iter__(self):
 
             while prefetch_data.run_prefetcher:
                 if len(prefetch_data.prefetch_buffer) > 0:
-                    yield prefetch_data.prefetch_buffer.popleft()
+                    item = prefetch_data.prefetch_buffer.popleft()
+                    if isinstance(item, Exception):
+                        prefetch_data.run_prefetcher = False
+                        raise item
+                    yield item
                 else:
                     # TODO: Calculate sleep interval based on previous availability speed
                     if not prefetch_data.stop_iteration: