feat: pipeline checkpoints crash allow to resume from snapshot (#9743)

* initial PoC idea running

* removing test code

* cleaning up

* wip

* cleaning up demos

* adding more pipelines to test persistence saving

* wip

* wip

* working example for logging components inputs in run time

* reverting to a simpler solution for intermediate results

* cleaning up

* testing that in a crash components outputs/inputs up to the crash point are returned

* adding tests for state persistance in a RAG pipeline

* updataing tests for state persistance in a RAG pipeline

* removing use cases of agent tests

* adding LICENSE header

* adding LICENSE header

* adding release notes

* updating tests for mocked components only

* updating release notes

* adapting PipelineRuntimeError

* cleaning up tests

* fixing test pipeline crash components inputs/outputs are saved

* fixing tests for state persistance

* removing code

* removing code

* removing code

* updating release notes

* validating parameters

* cleaning

* wip: debugging

* removing persistance tests

* formatting

* formatting

* cleaning up code

* updating release notes

* adding missing docstrings

* typo in release notes

* Update haystack/core/pipeline/pipeline.py

Co-authored-by: Sebastian Husch Lee <[email protected]>

* PR comments

* handling potential issues with saving the snapshot file

* updating tests

* Update haystack/core/pipeline/pipeline.py

Co-authored-by: Sebastian Husch Lee <[email protected]>

* Update haystack/core/pipeline/pipeline.py

Co-authored-by: Sebastian Husch Lee <[email protected]>

* some more improvements

* fixing exxception

* fixing exception error name conflict

---------

Co-authored-by: Sebastian Husch Lee <[email protected]>

Author: davidsbatista

haystack/core/pipeline/pipeline.py

@@ -17,6 +17,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,
@@ -25,6 +26,7 @@
 from haystack.dataclasses.breakpoints import AgentBreakpoint, Breakpoint, PipelineSnapshot
 from haystack.telemetry import pipeline_running
 from haystack.utils import _deserialize_value_with_schema
+from haystack.utils.misc import _get_output_dir
 
 logger = logging.getLogger(__name__)
 
@@ -390,6 +392,36 @@ 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.
+                    pipeline_snapshot_inputs_serialised = deepcopy(inputs)
+                    pipeline_snapshot_inputs_serialised[component_name] = deepcopy(component_inputs)
+                    out_dir = _get_output_dir("pipeline_snapshot")
+                    break_point = Breakpoint(
+                        component_name=component_name,
+                        visit_count=component_visits[component_name],
+                        snapshot_file_path=out_dir,
+                    )
+                    last_good_state_snapshot = _create_pipeline_snapshot(
+                        inputs=pipeline_snapshot_inputs_serialised,
+                        break_point=break_point,
+                        component_visits=component_visits,
+                        original_input_data=data,
+                        ordered_component_names=ordered_component_names,
+                        include_outputs_from=include_outputs_from,
+                        pipeline_outputs=pipeline_outputs,
+                    )
+                    try:
+                        _save_pipeline_snapshot(pipeline_snapshot=last_good_state_snapshot)
+                        logger.info(
+                            "Saved a snapshot of the pipeline's last valid state to '{out_path}'. "
+                            "Review this snapshot to debug the error and resume the pipeline from here.",
+                            out_path=out_dir,
+                        )
+                    except Exception as save_error:
+                        logger.error(
+                            "Failed to save a snapshot of the pipeline's last valid state with error: {e}", e=save_error
+                        )
                     raise error
 
                 # Updates global input state with component outputs and returns outputs that should go to

haystack/utils/misc.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import mimetypes
+import tempfile
 from pathlib import Path
 from typing import Any, Optional, Union, overload
 
@@ -81,3 +82,36 @@ def _guess_mime_type(path: Path) -> Optional[str]:
     mime_type = mimetypes.guess_type(path.as_posix())[0]
     # lookup custom mappings if the mime type is not found
     return CUSTOM_MIMETYPES.get(extension, mime_type)
+
+
+def _get_output_dir(out_dir: str) -> str:
+    """
+    Find or create a writable directory for saving status files.
+
+    Tries in the following order:
+
+        1. ~/.haystack/{out_dir}
+        2. {tempdir}/haystack/{out_dir}
+        3. ./.haystack/{out_dir}
+
+    :raises RuntimeError: If no directory could be created.
+    :returns:
+        The path to the created directory.
+    """
+
+    candidates = [
+        Path.home() / ".haystack" / out_dir,
+        Path(tempfile.gettempdir()) / "haystack" / out_dir,
+        Path.cwd() / ".haystack" / out_dir,
+    ]
+
+    for candidate in candidates:
+        try:
+            candidate.mkdir(parents=True, exist_ok=True)
+            return str(candidate)
+        except Exception:
+            continue
+
+    raise RuntimeError(
+        f"Could not create a writable directory for output files in any of the following locations: {candidates}"
+    )

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

@@ -0,0 +1,7 @@
+---
+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. Currently, only available in `Pipeline` and not yet
+    in `AsyncPipeline`.

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
@@ -21,6 +22,7 @@
 from haystack.document_stores.in_memory import InMemoryDocumentStore
 from haystack.document_stores.types import DuplicatePolicy
 from haystack.utils.auth import Secret
+from haystack.utils.misc import _get_output_dir
 
 
 def setup_document_store():
@@ -232,3 +234,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 "pipeline_snapshot" directory
+        snapshot_files = os.listdir(_get_output_dir("pipeline_snapshot"))
+        assert any(f.endswith(".json") for f in snapshot_files), "No pipeline snapshot file found in debug directory"