From 4449a89e0585428d2a31c4317a2b01eb1b807123 Mon Sep 17 00:00:00 2001 From: Afloat16 <132809511+Afloat16@users.noreply.github.com> Date: Fri, 2 Oct 2026 03:08:49 -0700 Subject: [PATCH] fix: keep pipeline snapshots reusable after resume --- .../pipelines/pipeline-breakpoints.mdx | 2 + haystack/core/pipeline/pipeline.py | 3 +- ...napshot-run-counters-037e24bfd9a7fb82.yaml | 7 +++ test/core/pipeline/test_breakpoint.py | 43 +++++++++++++++++-- 4 files changed, 51 insertions(+), 4 deletions(-) create mode 100644 releasenotes/notes/preserve-pipeline-snapshot-run-counters-037e24bfd9a7fb82.yaml diff --git a/docs-website/docs/concepts/pipelines/pipeline-breakpoints.mdx b/docs-website/docs/concepts/pipelines/pipeline-breakpoints.mdx index 80fa18a2bcc..f3838e0c5bc 100644 --- a/docs-website/docs/concepts/pipelines/pipeline-breakpoints.mdx +++ b/docs-website/docs/concepts/pipelines/pipeline-breakpoints.mdx @@ -148,6 +148,8 @@ Resume from memory: result = pipeline.run(data={}, pipeline_snapshot=snapshot) ``` +The saved component visit counts stay unchanged after a resumed run, including when a downstream component fails. Resuming the same snapshot again therefore uses the original visit counts for `max_runs_per_component`. + Resume from disk: ```python diff --git a/haystack/core/pipeline/pipeline.py b/haystack/core/pipeline/pipeline.py index c68d7e08d6b..b46633890cd 100644 --- a/haystack/core/pipeline/pipeline.py +++ b/haystack/core/pipeline/pipeline.py @@ -376,7 +376,8 @@ def run( # noqa: PLR0915, PLR0912, C901 _validate_pipeline_snapshot_against_pipeline(pipeline_snapshot, self.graph) # Handle resuming the pipeline from a snapshot - component_visits = pipeline_snapshot.pipeline_state.component_visits + # Advancing or failing a resumed run must not consume the snapshot's saved run budget. + component_visits = pipeline_snapshot.pipeline_state.component_visits.copy() ordered_component_names = pipeline_snapshot.ordered_component_names data = _deserialize_value_with_schema(pipeline_snapshot.original_input_data) diff --git a/releasenotes/notes/preserve-pipeline-snapshot-run-counters-037e24bfd9a7fb82.yaml b/releasenotes/notes/preserve-pipeline-snapshot-run-counters-037e24bfd9a7fb82.yaml new file mode 100644 index 00000000000..a6605634e52 --- /dev/null +++ b/releasenotes/notes/preserve-pipeline-snapshot-run-counters-037e24bfd9a7fb82.yaml @@ -0,0 +1,7 @@ +--- +fixes: + - | + ``Pipeline.run()`` no longer changes the visit counts stored in a + ``PipelineSnapshot`` when resuming it. Replaying the same snapshot or + retrying after a downstream failure now preserves the saved state and + does not incorrectly exhaust ``max_runs_per_component``. diff --git a/test/core/pipeline/test_breakpoint.py b/test/core/pipeline/test_breakpoint.py index e7663cd2c4d..77b7e8db04f 100644 --- a/test/core/pipeline/test_breakpoint.py +++ b/test/core/pipeline/test_breakpoint.py @@ -13,7 +13,7 @@ from haystack.components.joiners import BranchJoiner, ListJoiner from haystack.components.routers import ConditionalRouter from haystack.components.routers.conditional_router import Route -from haystack.core.errors import BreakpointException, PipelineInvalidPipelineSnapshotError +from haystack.core.errors import BreakpointException, PipelineInvalidPipelineSnapshotError, PipelineRuntimeError from haystack.core.pipeline import Pipeline from haystack.core.pipeline.breakpoint import ( HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED, @@ -171,8 +171,8 @@ def run(self, value: int) -> dict[str, Any]: return {"done": f"finished at {value}"} -def _three_component_pipeline() -> Pipeline: - pipeline = Pipeline() +def _three_component_pipeline(*, max_runs_per_component: int = 100) -> Pipeline: + pipeline = Pipeline(max_runs_per_component=max_runs_per_component) pipeline.add_component("comp1", _AppendingComponent()) pipeline.add_component("comp2", _AppendingComponent()) pipeline.add_component("comp3", _AppendingComponent()) @@ -191,6 +191,42 @@ def _looping_pipeline() -> Pipeline: class TestResumeFromPipelineSnapshot: + def test_resuming_snapshot_preserves_saved_state_and_run_budget(self): + pipeline = _three_component_pipeline(max_runs_per_component=1) + + with pytest.raises(BreakpointException) as exc_info: + pipeline.run(data={"comp1": {"input_value": "test"}}, break_point=Breakpoint(component_name="comp2")) + snapshot = exc_info.value.pipeline_snapshot + assert snapshot is not None + saved_state = snapshot.to_dict() + expected = {"comp3": {"result": "test_processed_processed_processed"}} + + assert pipeline.run(data={}, pipeline_snapshot=snapshot) == expected + assert pipeline.run(data={}, pipeline_snapshot=snapshot) == expected + assert snapshot.to_dict() == saved_state + + def test_failed_resume_does_not_modify_snapshot_or_prevent_retry(self, monkeypatch): + pipeline = _three_component_pipeline(max_runs_per_component=1) + + with pytest.raises(BreakpointException) as exc_info: + pipeline.run(data={"comp1": {"input_value": "test"}}, break_point=Breakpoint(component_name="comp2")) + snapshot = exc_info.value.pipeline_snapshot + assert snapshot is not None + saved_state = snapshot.to_dict() + + def fail(input_value: str) -> dict[str, str]: + raise RuntimeError(f"injected downstream failure for {input_value}") + + with monkeypatch.context() as patch: + patch.setattr(pipeline.get_component("comp3"), "run", fail) + with pytest.raises(PipelineRuntimeError, match="injected downstream failure"): + pipeline.run(data={}, pipeline_snapshot=snapshot) + + assert pipeline.run(data={}, pipeline_snapshot=snapshot) == { + "comp3": {"result": "test_processed_processed_processed"} + } + assert snapshot.to_dict() == saved_state + def test_break_point_with_pipeline_snapshot_steps_through_pipeline(self): pipeline = _three_component_pipeline() @@ -207,6 +243,7 @@ def test_break_point_with_pipeline_snapshot_steps_through_pipeline(self): second_snapshot = exc_info.value.pipeline_snapshot assert second_snapshot is not None assert second_snapshot.pipeline_state.component_visits == {"comp1": 1, "comp2": 1, "comp3": 0} + assert first_snapshot.pipeline_state.component_visits == {"comp1": 1, "comp2": 0, "comp3": 0} # resume from the second snapshot and run to completion result = pipeline.run(data={}, pipeline_snapshot=second_snapshot)