Don't delete checkpoint when other steps fail

Author: zprobst

nodestream/pipeline/extractors/extractor.py

@@ -24,8 +24,11 @@ async def start(self, context: StepContext):
             await self.resume_from_checkpoint(checkpoint)
 
     async def finish(self, context: StepContext):
-        context.debug("Clearing checkpoint for extractor since extractor is finished.")
-        context.object_store.delete(CHECKPOINT_OBJECT_KEY)
+        if not context.pipeline_encountered_fatal_error:
+            context.debug(
+                "Clearing checkpoint for extractor since extractor is finished."
+            )
+            context.object_store.delete(CHECKPOINT_OBJECT_KEY)
 
     async def make_checkpoint(self) -> T:
         return None

nodestream/pipeline/object_storage.py

@@ -309,7 +309,7 @@ def delete(self, key: str):
         self.store.delete(key)
 
 
-class S3ObjectStore(ObjectStore):
+class S3ObjectStore(ObjectStore, alias="s3"):
     __slots__ = ("client", "bucket_name")
 
     def __init__(self, bucket_name: str, **client_factory_args):

nodestream/pipeline/progress_reporter.py

@@ -37,6 +37,11 @@ class PipelineProgressReporter:
     on_start_callback: Callable[[], None] = field(default=no_op)
     on_finish_callback: Callable[[Metrics], None] = field(default=no_op)
     on_fatal_error_callback: Callable[[Exception], None] = field(default=no_op)
+    encountered_fatal_error: bool = field(default=False)
+
+    def on_fatal_error(self, exception: Exception):
+        self.encountered_fatal_error = True
+        self.on_fatal_error_callback(exception)
 
     @classmethod
     def for_testing(cls, results_list: list) -> "PipelineProgressReporter":

nodestream/pipeline/step.py

@@ -26,6 +26,11 @@ def __init__(
         self.index = index
         self.object_store = object_store
 
+    @property
+    def pipeline_encountered_fatal_error(self) -> bool:
+        """Whether the pipeline has encountered a fatal error."""
+        return self.reporter.encountered_fatal_error
+
     def report_error(
         self,
         message: str,
@@ -51,7 +56,7 @@ def report_error(
             stack_info=True,
         )
         if fatal:
-            self.reporter.on_fatal_error_callback(exception)
+            self.reporter.on_fatal_error(exception)
         else:
             Metrics.get().increment(Metric.NON_FATAL_ERRORS)
 

tests/data/test_snapshot_handling_during_errors.py

@@ -0,0 +1,38 @@
+import pytest
+
+from nodestream.pipeline import Pipeline
+from nodestream.pipeline.extractors import IterableExtractor
+from nodestream.pipeline.object_storage import ObjectStore
+from nodestream.pipeline.progress_reporter import PipelineProgressReporter
+from nodestream.pipeline.step import Step
+
+
+class ExplodesAfter(Step):
+    def __init__(self, explode_after: int):
+        self.count = 0
+        self.explode_after = explode_after
+
+    async def process_record(self, record, context):
+        self.count += 1
+        if self.count > self.explode_after:
+            raise Exception("Exploded!")
+        yield record
+
+
+@pytest.mark.asyncio
+async def test_snapshot_handling_during_errors(mocker):
+    extractor = IterableExtractor.range(0, 100000000)
+    kaboom = ExplodesAfter(5000)
+    storage = mocker.Mock(spec=ObjectStore)
+    storage.namespaced.return_value = storage
+    pipeline = Pipeline((extractor, kaboom), 1000, storage)
+
+    def raise_fatal_error(e):
+        # raise e
+        pass
+
+    await pipeline.run(
+        PipelineProgressReporter(on_fatal_error_callback=raise_fatal_error)
+    )
+    storage.delete.assert_not_called()
+    assert storage.put_picklable.call_count == 6