Repository: facebookresearch/spdl

Issue:
Add mode to propagate the error from Pipeline to the foreground thread
By default the errors happening in Pipeline is not propagated to the main thread. This is by design, but it is annoying when building a pipeline. (It is not straightforward to catch issue in tests unless you are watching the log.)

An idea to solve this is to add a sort of "strict" mode people use while development. In this mode, the Pipeline reacts to a task error and send some signal to the foreground thread.

Discussion:

Solution patch:
diff --git a/src/spdl/dataloader/_cache_dataloader.py b/src/spdl/dataloader/_cache_dataloader.py
index a2cadec..bd183dc 100644
--- a/src/spdl/dataloader/_cache_dataloader.py
+++ b/src/spdl/dataloader/_cache_dataloader.py
@@ -26,12 +26,7 @@ class CacheDataLoader(Generic[T]):
 
     Args:
         dl: Source iterator. Expected to be a data loader object.
-        num_caches: The number of items (batches) to cache.
-        return_caches_after: The number of iterations to use the original
-            iterator. By default, it uses the same value as ``num_caches``.
-        delete_src: When this iterator starts returning the cached value,
-            call ``del`` on the original data loader so that resources are
-            released.
+        num_caches,return_caches_after,stop_after: See :py:func:`spdl.pipeline.cache_iterator`.
 
     Returns:
         The new iterator.
@@ -43,7 +38,11 @@ class CacheDataLoader(Generic[T]):
     """
 
     def __init__(
-        self, dl: Iterable[T], num_caches: int, return_caches_after: int
+        self,
+        dl: Iterable[T],
+        num_caches: int,
+        return_caches_after: int,
+        stop_after: int | None = None,
     ) -> None:
         log_api_usage_once("spdl.dataloader.CacheDataLoader")
 
@@ -51,19 +50,15 @@ class CacheDataLoader(Generic[T]):
 
         self.num_caches = num_caches
         self.return_caches_after = return_caches_after
+        self.stop_after = stop_after
 
     def __iter__(self) -> Iterator[T]:
         """See :py:func:`spdl.pipeline.cache_iterator` for the detail."""
-        try:
-            stop_after = len(self)
-        except Exception:
-            stop_after = None
-
         return cache_iterator(
             self.dl,
             num_caches=self.num_caches,
             return_caches_after=self.return_caches_after,
-            stop_after=stop_after,
+            stop_after=self.stop_after,
         )
 
     def __len__(self) -> int:
diff --git a/src/spdl/pipeline/_pipeline.py b/src/spdl/pipeline/_pipeline.py
index 8af2786..a478bca 100644
--- a/src/spdl/pipeline/_pipeline.py
+++ b/src/spdl/pipeline/_pipeline.py
@@ -52,6 +52,7 @@ class _EventLoop:
 
         self._task_started = SyncEvent()
         self._task_completed = SyncEvent()
+        self._task_exception: BaseException | None = None
         self._stop_requested = SyncEvent()
 
         self._thread: Thread | None = None
@@ -100,6 +101,11 @@ class _EventLoop:
         _LG.debug("The background task is completed.")
         _LG.debug("The event loop is now shutdown.")
 
+        try:
+            self._task_exception = task.exception()
+        except asyncio.CancelledError:
+            pass
+
     def start(self, *, timeout: float | None = None, daemon: bool = False) -> None:
         """Start the thread and block until the loop is initialized."""
         if self._thread is not None:
@@ -329,6 +335,9 @@ class Pipeline(Generic[T]):
             self._event_loop.join(timeout=timeout)
             self._event_loop_state = _EventLoopState.STOPPED
 
+        if self._event_loop._task_exception is not None:
+            raise self._event_loop._task_exception
+
     @contextmanager
     def auto_stop(self, *, timeout: float | None = None) -> Iterator[None]:
         """Context manager to start/stop the background thread automatically.
diff --git a/src/spdl/pipeline/_utils.py b/src/spdl/pipeline/_utils.py
index 710aaca..87d62cb 100644
--- a/src/spdl/pipeline/_utils.py
+++ b/src/spdl/pipeline/_utils.py
@@ -285,7 +285,7 @@ def cache_iterator(
     # Note - Design choice
     # When these optional values are provided, we could choose to not validate.
     # But the purpose of this function is to make sure you are using cache,
-    # so we raise an error if these parameters do not make ogical sense.
+    # so we raise an error if these parameters do not make logical sense.
     if return_caches_after is not None:
         if return_caches_after < num_caches:
             raise ValueError(


Test patch:
diff --git a/tests/spdl_unittest/dataloader/cache_dataloader_test.py b/tests/spdl_unittest/dataloader/cache_dataloader_test.py
index f6a51a6..531f832 100644
--- a/tests/spdl_unittest/dataloader/cache_dataloader_test.py
+++ b/tests/spdl_unittest/dataloader/cache_dataloader_test.py
@@ -85,7 +85,7 @@ def test_CacheDataLoader():
             return self.n
 
     N = 8
-    dl = CacheDataLoader(DL(N), num_caches=2, return_caches_after=3)
+    dl = CacheDataLoader(DL(N), num_caches=2, return_caches_after=3, stop_after=N)
 
     assert dl.n == len(dl) == N
 
diff --git a/tests/spdl_unittest/dataloader/pipeline_test.py b/tests/spdl_unittest/dataloader/pipeline_test.py
index 4ab7fb9..d24c56a 100644
--- a/tests/spdl_unittest/dataloader/pipeline_test.py
+++ b/tests/spdl_unittest/dataloader/pipeline_test.py
@@ -22,6 +22,7 @@ from typing import TypeVar
 import pytest
 from spdl.pipeline import (
     PipelineBuilder,
+    PipelineFailure,
     run_pipeline_in_subprocess,
     TaskHook,
     TaskStatsHook,
@@ -568,12 +569,15 @@ def test_pipeline_hook_failure_enter_stage():
         .build(num_threads=1, task_hook_factory=lambda _: [_enter_stage_fail()])
     )
 
-    with pipeline.auto_stop():
-        assert [] == list(pipeline.get_iterator(timeout=3))
+    with pytest.raises(PipelineFailure):
+        with pipeline.auto_stop():
+            vals = list(pipeline.get_iterator(timeout=3))
+
+    assert vals == []
 
 
 def test_pipeline_hook_failure_exit_stage():
-    """If exit_stage fails, the pipeline does not fail."""
+    """If exit_stage fails, the error is propagated to the front end."""
 
     class _exit_stage_fail(TaskHook):
         @asynccontextmanager
@@ -592,8 +596,10 @@ def test_pipeline_hook_failure_exit_stage():
         .add_sink(1000)
         .build(num_threads=1, task_hook_factory=lambda _: [_exit_stage_fail()])
     )
-    with pipeline.auto_stop():
-        assert list(range(10)) == list(pipeline.get_iterator(timeout=3))
+    with pytest.raises(PipelineFailure):
+        with pipeline.auto_stop():
+            vals = list(pipeline.get_iterator(timeout=3))
+    assert vals == list(range(10))
 
 
 def test_pipeline_hook_failure_enter_task():
@@ -1082,7 +1088,11 @@ def test_pipeline_disaggregate():
 
 
 def test_pipeline_source_failure():
-    """AsyncPipeline continues when source fails."""
+    """AsyncPipeline continues when source fails.
+
+    note: the front end will propagate the error at the end of the `stop`.
+    before that, the pipeline should continue functioning.
+    """
 
     def failing_range(i):
         yield from range(i)
@@ -1097,9 +1107,11 @@ def test_pipeline_source_failure():
         .build(num_threads=1)
     )
 
-    with pipeline.auto_stop():
-        results = list(pipeline.get_iterator(timeout=3))
-        assert results == [1 + 2 * i for i in range(10)]
+    with pytest.raises(PipelineFailure):
+        with pipeline.auto_stop():
+            results = list(pipeline.get_iterator(timeout=3))
+
+    assert results == [1 + 2 * i for i in range(10)]
 
 
 def test_pipeline_type_error():
@@ -1116,8 +1128,11 @@ def test_pipeline_type_error():
         .build(num_threads=1)
     )
 
-    with pipeline.auto_stop():
-        assert [] == list(pipeline.get_iterator(timeout=3))
+    with pytest.raises(PipelineFailure):
+        with pipeline.auto_stop():
+            vals = list(pipeline.get_iterator(timeout=3))
+
+    assert vals == []
 
 
 def test_pipeline_task_failure():
@@ -1206,9 +1221,10 @@ def test_pipeline_fail_middle():
     with pytest.raises(RuntimeError):
         apl.get_item(timeout=1)
 
-    with apl.auto_stop():
-        with pytest.raises(EOFError):
-            apl.get_item(timeout=1)
+    with pytest.raises(PipelineFailure):
+        with apl.auto_stop():
+            with pytest.raises(EOFError):
+                apl.get_item(timeout=1)
 
     assert pwc.cache == []
 
@@ -1663,24 +1679,25 @@ def test_pipeline_failures(output_order: str):
             raise ValueError(f"Only evan numbers are allowed. {x}")
         return x
 
-    src = range(10)
-
     builder = (
         PipelineBuilder()
-        .add_source(src)
+        .add_source(range(10))
         .pipe(fail_odd, output_order=output_order, concurrency=10)
         .add_sink(1)
     )
 
     pipeline = builder.build(num_threads=1)
     with pipeline.auto_stop():
-        ite = pipeline.get_iterator(timeout=3)
-        assert list(ite) == [i for i in src if not (i % 2)]
+        vals = list(pipeline.get_iterator(timeout=3))
+
+    assert vals == [0, 2, 4, 6, 8]
 
     pipeline = builder.build(num_threads=1, max_failures=2)
-    with pipeline.auto_stop():
-        ite = pipeline.get_iterator(timeout=3)
-        assert list(ite) == [i for i in range(6) if not (i % 2)]
+    with pytest.raises(PipelineFailure):
+        with pipeline.auto_stop():
+            vals = list(pipeline.get_iterator(timeout=3))
+
+    assert vals == [0, 2, 4]
 
 
 @pytest.mark.parametrize("output_order", ["completion", "input"])
@@ -1710,10 +1727,29 @@ def test_pipeline_failures_multiple_pipeline(output_order: str):
     pipeline1 = builder.build(num_threads=1, max_failures=1)
     pipeline2 = builder.build(num_threads=1, max_failures=2)
 
-    with pipeline2.auto_stop():
-        ite = pipeline2.get_iterator(timeout=3)
-        assert list(ite) == [i for i in range(6) if not (i % 2)]
+    with pytest.raises(PipelineFailure):
+        with pipeline2.auto_stop():
+            vals = list(pipeline2.get_iterator(timeout=3))
+
+    assert vals == [0, 2, 4]
 
-    with pipeline1.auto_stop():
-        ite = pipeline1.get_iterator(timeout=3)
-        assert list(ite) == [i for i in range(4) if not (i % 2)]
+    with pytest.raises(PipelineFailure):
+        with pipeline1.auto_stop():
+            vals = list(pipeline1.get_iterator(timeout=3))
+
+    assert vals == [0, 2]
+
+
+def test_pipeline_propagate_source_failure():
+    """When source itrator fails, the exception is propagated to the front end"""
+
+    def failure_source():
+        raise RuntimeError("Foo")
+        yield None
+
+    pipeline = (
+        PipelineBuilder().add_source(failure_source()).add_sink().build(num_threads=1)
+    )
+    with pytest.raises(RuntimeError):
+        with pipeline.auto_stop():
+            pass