feat: pipeline checkpoints crash allow to resume from snapshot

haystack/core/pipeline/breakpoint.py

@@ -140,7 +140,11 @@ def load_pipeline_snapshot(file_path: Union[str, Path]) -> PipelineSnapshot:
 
 
 def _save_pipeline_snapshot_to_file(
-    *, pipeline_snapshot: PipelineSnapshot, snapshot_file_path: Union[str, Path], dt: datetime
+    *,
+    pipeline_snapshot: PipelineSnapshot,
+    snapshot_file_path: Union[str, Path],
+    dt: datetime,
+    file_name: Optional[str] = None,
 ) -> None:
     """
     Save the pipeline snapshot dictionary to a JSON file.
@@ -158,17 +162,18 @@ def _save_pipeline_snapshot_to_file(
 
     snapshot_file_path.mkdir(exist_ok=True)
 
-    # Generate filename
-    # We check if the agent_name is provided to differentiate between agent and non-agent breakpoints
-    if isinstance(pipeline_snapshot.break_point, AgentBreakpoint):
-        agent_name = pipeline_snapshot.break_point.agent_name
-        component_name = pipeline_snapshot.break_point.break_point.component_name
-        visit_nr = pipeline_snapshot.pipeline_state.component_visits.get(component_name, 0)
-        file_name = f"{agent_name}_{component_name}_{visit_nr}_{dt.strftime('%Y_%m_%d_%H_%M_%S')}.json"
-    else:
-        component_name = pipeline_snapshot.break_point.component_name
-        visit_nr = pipeline_snapshot.pipeline_state.component_visits.get(component_name, 0)
-        file_name = f"{component_name}_{visit_nr}_{dt.strftime('%Y_%m_%d_%H_%M_%S')}.json"
+    if not file_name:
+        # Generate filename
+        # We check if the agent_name is provided to differentiate between agent and non-agent breakpoints
+        if isinstance(pipeline_snapshot.break_point, AgentBreakpoint):
+            agent_name = pipeline_snapshot.break_point.agent_name
+            component_name = pipeline_snapshot.break_point.break_point.component_name
+            visit_nr = pipeline_snapshot.pipeline_state.component_visits.get(component_name, 0)
+            file_name = f"{agent_name}_{component_name}_{visit_nr}_{dt.strftime('%Y_%m_%d_%H_%M_%S')}.json"
+        else:
+            component_name = pipeline_snapshot.break_point.component_name
+            visit_nr = pipeline_snapshot.pipeline_state.component_visits.get(component_name, 0)
+            file_name = f"{component_name}_{visit_nr}_{dt.strftime('%Y_%m_%d_%H_%M_%S')}.json"
 
     try:
         with open(snapshot_file_path / file_name, "w") as f_out:
@@ -219,7 +224,7 @@ def _create_pipeline_snapshot(
     return pipeline_snapshot
 
 
-def _save_pipeline_snapshot(pipeline_snapshot: PipelineSnapshot) -> PipelineSnapshot:
+def _save_pipeline_snapshot(pipeline_snapshot: PipelineSnapshot, f_name: Optional[str] = None) -> PipelineSnapshot:
     """
     Save the pipeline snapshot to a file.
 
@@ -244,7 +249,7 @@ def _save_pipeline_snapshot(pipeline_snapshot: PipelineSnapshot) -> PipelineSnap
     if snapshot_file_path is not None:
         dt = pipeline_snapshot.timestamp or datetime.now()
         _save_pipeline_snapshot_to_file(
-            pipeline_snapshot=pipeline_snapshot, snapshot_file_path=snapshot_file_path, dt=dt
+            pipeline_snapshot=pipeline_snapshot, snapshot_file_path=snapshot_file_path, dt=dt, file_name=f_name
         )
 
     return pipeline_snapshot

haystack/core/pipeline/pipeline.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from copy import deepcopy
+from datetime import datetime
 from typing import Any, Mapping, Optional, Union
 
 from haystack import logging, tracing
@@ -17,6 +18,7 @@
 )
 from haystack.core.pipeline.breakpoint import (
     _create_pipeline_snapshot,
+    _save_pipeline_snapshot,
     _trigger_break_point,
     _validate_break_point_against_pipeline,
     _validate_pipeline_snapshot_against_pipeline,
@@ -175,6 +177,9 @@ def run(  # noqa: PLR0915, PLR0912, C901, pylint: disable=too-many-branches
         :param pipeline_snapshot:
             A dictionary containing a snapshot of a previously saved pipeline execution.
 
+        :param state_persistence_path:
+            The path where the pipeline state should be saved if `state_persistence` is `True`.
+
         :returns:
             A dictionary where each entry corresponds to a component name
             and its output. If `include_outputs_from` is `None`, this dictionary
@@ -390,6 +395,24 @@ def run(  # noqa: PLR0915, PLR0912, C901, pylint: disable=too-many-branches
                 except PipelineRuntimeError as error:
                     # Attach partial pipeline outputs to the error before re-raising
                     error.pipeline_outputs = pipeline_outputs
+                    # Create a snapshot of the last good state of the pipeline before the error occurred.
+                    last_good_state_snapshot = self.create_last_good_state_snapshot(
+                        component_inputs,
+                        component_name,
+                        component_visits,
+                        data,
+                        include_outputs_from,
+                        inputs,
+                        ordered_component_names,
+                        pipeline_outputs,
+                    )
+                    f_name = f"last_good_state_{datetime.now().strftime('%Y_%m_%d_%H_%M_%S')}.json"
+                    _save_pipeline_snapshot(pipeline_snapshot=last_good_state_snapshot, f_name=f_name)
+                    logger.info(
+                        "A snapshot of the last good state of the pipeline has been saved to '{f_name}'. "
+                        "You can inspect this snapshot to fix the error and resume the pipeline from this point.",
+                        f_name=f_name,
+                    )
                     raise error
 
                 # Updates global input state with component outputs and returns outputs that should go to
@@ -416,3 +439,47 @@ def run(  # noqa: PLR0915, PLR0912, C901, pylint: disable=too-many-branches
                 )
 
             return pipeline_outputs
+
+    @staticmethod
+    def create_last_good_state_snapshot(  # pylint: disable=too-many-positional-arguments
+        component_inputs: dict[str, Any],
+        component_name: str,
+        component_visits: dict[str, int],
+        data: dict[str, dict[str, Any]],
+        include_outputs_from: set[str],
+        inputs: dict[str, Any],
+        ordered_component_names: list[str],
+        pipeline_outputs: dict[str, Any],
+    ) -> PipelineSnapshot:
+        """
+        Creates a snapshot of the last good state of the pipeline before an error occurred.
+
+        This snapshot can be used for debugging purposes to understand the state of the pipeline, edited and passed
+        to the `pipeline_snapshot` parameter of the `Pipeline.run()` method to resume execution from this point.
+
+        :param component_inputs: Inputs to the component that was being executed when the error occurred.
+        :param component_name: Name of the component that was being executed when the error occurred.
+        :param component_visits: Current state of component visits.
+        :param data: Original input data provided to the pipeline.
+        :param include_outputs_from: Set of component names whose outputs are included in the pipeline output.
+        :param inputs: Current state of all inputs in the pipeline.
+        :param ordered_component_names: List of component names in the order they were added to the pipeline.
+        :param pipeline_outputs: Current state of the pipeline outputs.
+
+        :return: A PipelineSnapshot representing the last good state of the pipeline.
+        """
+
+        pipeline_snapshot_inputs_serialised = deepcopy(inputs)
+        pipeline_snapshot_inputs_serialised[component_name] = deepcopy(component_inputs)
+        pipeline_snapshot = _create_pipeline_snapshot(
+            inputs=pipeline_snapshot_inputs_serialised,
+            # Dummy breakpoint to pass the component_name and state_persistence_path to the _save_pipeline_snapshot
+            break_point=Breakpoint(component_name=component_name, visit_count=0, snapshot_file_path="debug"),
+            component_visits=component_visits,
+            original_input_data=data,
+            ordered_component_names=ordered_component_names,
+            include_outputs_from=include_outputs_from,
+            pipeline_outputs=pipeline_outputs,
+        )
+
+        return pipeline_snapshot

releasenotes/notes/pipeline-resume-state-13ca79f53277e858.yaml

@@ -0,0 +1,6 @@
+---
+features:
+  - |
+    A snapshot of the last successful step is now saved when an error occurs during a Pipeline run. This allows you
+    to inspect the snapshot, potentially identify and fix the error, and later resume the pipeline from that point
+    onwards. Avoiding to re-run the entire pipeline from the start.

test/core/pipeline/test_pipeline_crash_regular_pipeline_outputs_raised.py

@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+import os
 from unittest.mock import MagicMock, patch
 
 import numpy as np
@@ -232,3 +233,7 @@ async def test_async_hybrid_rag_pipeline_crash_on_embedding_retriever(
         assert "prompt_builder" not in pipeline_outputs, "Prompt builder should not have run due to crash"
         assert "llm" not in pipeline_outputs, "LLM should not have run due to crash"
         assert "answer_builder" not in pipeline_outputs, "Answer builder should not have run due to crash"
+
+        # check that a pipeline snapshot file was created in the "debug" directory
+        snapshot_files = os.listdir("debug")
+        assert any(f.endswith(".json") for f in snapshot_files), "No pipeline snapshot file found in debug directory"