Coverage for haystack/core/pipeline/breakpoint.py: 90%
110 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-12 09:42 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import json
6import os
7from collections.abc import Callable
8from datetime import datetime
9from pathlib import Path
10from typing import Any
12from networkx import MultiDiGraph
14from haystack import logging
15from haystack.core.errors import PipelineInvalidPipelineSnapshotError
16from haystack.dataclasses.breakpoints import INTERNAL_INPUTS_FORMAT, Breakpoint, PipelineSnapshot, PipelineState
17from haystack.utils.base_serialization import _serialize_with_field_fallback
19logger = logging.getLogger(__name__)
21# Environment variable to control pipeline snapshot file saving (disabled by default)
22HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED = "HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED"
24# Type alias for snapshot callback function
25# The callback receives a PipelineSnapshot and optionally returns a file path string
26SnapshotCallback = Callable[[PipelineSnapshot], str | None]
29def _is_snapshot_save_enabled() -> bool:
30 """
31 Check if pipeline snapshot file saving is enabled via environment variable.
33 The environment variable HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED controls whether
34 pipeline snapshots are saved to files. By default (when the variable is not set),
35 saving is disabled. Only "true" and "1" (case-insensitive) enable saving; any other value disables it.
37 :returns: True if snapshot saving is enabled, False otherwise.
38 """
39 value = os.environ.get(HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED, "false").lower()
40 return value in ("true", "1")
43def _validate_break_point_against_pipeline(break_point: Breakpoint, graph: MultiDiGraph) -> None:
44 """
45 Validates the breakpoints passed to the pipeline.
47 Makes sure the breakpoint contains a valid components registered in the pipeline.
49 :param break_point: a breakpoint to validate
50 """
51 if break_point.component_name not in graph.nodes:
52 raise ValueError(f"break_point {break_point} is not a registered component in the pipeline")
55def _validate_pipeline_snapshot_against_pipeline(pipeline_snapshot: PipelineSnapshot, graph: MultiDiGraph) -> None:
56 """
57 Validates that the pipeline_snapshot contains valid configuration for the current pipeline.
59 Raises a PipelineInvalidPipelineSnapshotError if any component in pipeline_snapshot is not part of the
60 target pipeline.
62 :param pipeline_snapshot: The saved state to validate.
63 """
65 pipeline_state = pipeline_snapshot.pipeline_state
66 valid_components = set(graph.nodes.keys())
68 # Check if the ordered_component_names are valid components in the pipeline
69 invalid_ordered_components = set(pipeline_snapshot.ordered_component_names) - valid_components
70 if invalid_ordered_components:
71 raise PipelineInvalidPipelineSnapshotError(
72 f"Invalid pipeline snapshot: components {invalid_ordered_components} in 'ordered_component_names' "
73 f"are not part of the current pipeline."
74 )
76 # A snapshot must describe the complete pipeline it resumes. If the current pipeline contains components that
77 # were not present when the snapshot was created, they would never be added to the resume queue and could be
78 # silently skipped.
79 missing_ordered_components = valid_components - set(pipeline_snapshot.ordered_component_names)
80 if missing_ordered_components:
81 raise PipelineInvalidPipelineSnapshotError(
82 f"Invalid pipeline snapshot: components {missing_ordered_components} in the current pipeline are not "
83 f"present in 'ordered_component_names'."
84 )
86 # Check if the original_input_data is valid components in the pipeline
87 serialized_input_data = pipeline_snapshot.original_input_data["serialized_data"]
88 invalid_input_data = set(serialized_input_data.keys()) - valid_components
89 if invalid_input_data:
90 raise PipelineInvalidPipelineSnapshotError(
91 f"Invalid pipeline snapshot: components {invalid_input_data} in 'input_data' "
92 f"are not part of the current pipeline."
93 )
95 # Validate 'component_visits'
96 invalid_component_visits = set(pipeline_state.component_visits.keys()) - valid_components
97 if invalid_component_visits:
98 raise PipelineInvalidPipelineSnapshotError(
99 f"Invalid pipeline snapshot: components {invalid_component_visits} in 'component_visits' "
100 f"are not part of the current pipeline."
101 )
103 component_name = pipeline_snapshot.break_point.component_name
104 visit_count = pipeline_snapshot.pipeline_state.component_visits[component_name]
106 logger.info(
107 "Resuming pipeline from {component} with visit count {visits}", component=component_name, visits=visit_count
108 )
111def load_pipeline_snapshot(file_path: str | Path) -> PipelineSnapshot:
112 """
113 Load a saved pipeline snapshot.
115 :param file_path: Path to the pipeline_snapshot file.
116 :returns:
117 Dict containing the loaded pipeline_snapshot.
118 """
120 file_path = Path(file_path)
122 try:
123 with open(file_path, encoding="utf-8") as f:
124 pipeline_snapshot_dict = json.load(f)
125 except FileNotFoundError as e:
126 raise FileNotFoundError(f"File not found: {file_path}") from e
127 except json.JSONDecodeError as e:
128 raise json.JSONDecodeError(f"Invalid JSON file {file_path}: {str(e)}", e.doc, e.pos) from e
129 except OSError as e:
130 raise OSError(f"Error reading {file_path}: {str(e)}") from e
132 try:
133 pipeline_snapshot = PipelineSnapshot.from_dict(pipeline_snapshot_dict)
134 except ValueError as e:
135 raise ValueError(f"Invalid pipeline snapshot from {file_path}: {str(e)}") from e
137 logger.info("Successfully loaded the pipeline snapshot from: {file_path}", file_path=file_path)
138 return pipeline_snapshot
141def _save_pipeline_snapshot(
142 pipeline_snapshot: PipelineSnapshot,
143 raise_on_failure: bool = True,
144 snapshot_callback: SnapshotCallback | None = None,
145) -> str | None:
146 """
147 Save the pipeline snapshot dictionary to a JSON file, or invoke a custom callback.
149 If a `snapshot_callback` is provided, it will be called with the pipeline snapshot instead of saving to a file.
150 This allows users to customize how snapshots are handled (e.g., saving to a database, sending to a remote service).
152 When no callback is provided, the default behavior saves to a JSON file:
153 - The filename is generated based on the component name, visit count, and timestamp.
154 - The component name is taken from the break point's `component_name`.
155 - The visit count is taken from the pipeline state's `component_visits` for the component name.
156 - The timestamp is taken from the pipeline snapshot's `timestamp` or the current time if not available.
157 - The file path is taken from the break point's `snapshot_file_path`.
158 - If the `snapshot_file_path` is None, the function will return without saving.
160 The default file saving behavior is disabled. To enable it, set the environment variable
161 `HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED` to "true" or "1". When disabled,
162 the function will return None without saving to a file (custom callbacks are still invoked).
164 :param pipeline_snapshot: The pipeline snapshot to save.
165 :param raise_on_failure: If True, raises an exception if saving fails. If False, logs the error and returns None.
166 :param snapshot_callback: Optional callback function that receives the PipelineSnapshot.
167 If provided, the callback is invoked instead of the default file-saving behavior.
168 The callback should return an optional string (e.g., a file path or identifier) or None.
170 :returns:
171 The full path to the saved JSON file (or the value returned by the callback), or None if
172 `snapshot_file_path` is None, no callback is provided, snapshot saving is disabled, or saving fails and
173 `raise_on_failure` is False.
174 :raises:
175 Exception: If saving the JSON snapshot fails (when raise_on_failure is True).
176 """
177 # If a callback is provided, use it instead of the default file-saving behavior
178 if snapshot_callback is not None:
179 try:
180 result = snapshot_callback(pipeline_snapshot)
181 logger.info("Pipeline snapshot handled by custom callback.")
182 return result
183 except Exception as error:
184 logger.exception("Failed to handle pipeline snapshot with custom callback. Error: {error}", error=error)
185 if raise_on_failure:
186 raise
187 return None
189 # Check if snapshot saving is enabled via environment variable (disabled by default)
190 if not _is_snapshot_save_enabled():
191 logger.debug("Pipeline snapshot file saving is disabled via HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED env var.")
192 return None
194 break_point = pipeline_snapshot.break_point
195 snapshot_file_path = break_point.snapshot_file_path
197 if snapshot_file_path is None:
198 return None
200 dt = pipeline_snapshot.timestamp or datetime.now()
201 snapshot_dir = Path(snapshot_file_path)
203 component_name = break_point.component_name
204 visit_nr = pipeline_snapshot.pipeline_state.component_visits.get(component_name, 0)
205 timestamp = dt.strftime("%Y_%m_%d_%H_%M_%S")
206 file_name = f"{component_name}_{visit_nr}_{timestamp}.json"
207 full_path = snapshot_dir / file_name
209 try:
210 snapshot_dir.mkdir(parents=True, exist_ok=True)
211 with open(full_path, "w") as f_out:
212 json.dump(pipeline_snapshot.to_dict(), f_out, indent=2)
213 logger.info(
214 "Pipeline snapshot saved to '{full_path}'. You can use this file to debug or resume the pipeline.",
215 full_path=full_path,
216 )
217 except Exception as error:
218 logger.exception("Failed to save pipeline snapshot to '{full_path}'. Error: {e}", full_path=full_path, e=error)
219 if raise_on_failure:
220 raise
221 return None
223 return str(full_path)
226def _create_pipeline_snapshot(
227 *,
228 inputs: dict[str, Any],
229 component_inputs: dict[str, Any],
230 break_point: Breakpoint,
231 component_visits: dict[str, int],
232 original_input_data: dict[str, Any],
233 ordered_component_names: list[str],
234 include_outputs_from: set[str],
235 pipeline_outputs: dict[str, Any],
236) -> PipelineSnapshot:
237 """
238 Create a snapshot of the pipeline at the point where the breakpoint was triggered.
240 :param inputs: The current pipeline snapshot inputs.
241 :param component_inputs: The inputs of the component that triggered the breakpoint, as they were before the
242 component consumed them.
243 :param break_point: The breakpoint that triggered the snapshot.
244 :param component_visits: The visit count of the component that triggered the breakpoint.
245 :param original_input_data: The original input data.
246 :param ordered_component_names: The ordered component names.
247 :param include_outputs_from: Set of component names whose outputs should be included in the pipeline results.
248 :param pipeline_outputs: The current outputs of the pipeline.
249 :returns:
250 A PipelineSnapshot containing the state of the pipeline at the point of the breakpoint.
251 """
252 component_name = break_point.component_name
254 transformed_original_input_data = _transform_json_structure(original_input_data)
256 serialized_inputs = _serialize_with_field_fallback(
257 {**inputs, component_name: component_inputs}, description="the inputs of the current pipeline state"
258 )
259 serialized_original_input_data = _serialize_with_field_fallback(
260 transformed_original_input_data, description="original input data for `pipeline.run`"
261 )
262 serialized_pipeline_outputs = _serialize_with_field_fallback(
263 pipeline_outputs, description="outputs of the current pipeline state"
264 )
266 return PipelineSnapshot(
267 pipeline_state=PipelineState(
268 inputs=serialized_inputs,
269 component_visits=component_visits,
270 pipeline_outputs=serialized_pipeline_outputs,
271 inputs_format=INTERNAL_INPUTS_FORMAT,
272 ),
273 timestamp=datetime.now(),
274 break_point=break_point,
275 original_input_data=serialized_original_input_data,
276 ordered_component_names=ordered_component_names,
277 include_outputs_from=include_outputs_from,
278 )
281def _transform_json_structure(data: dict[str, Any] | list[Any] | Any) -> Any:
282 """
283 Transforms a JSON structure by removing the 'sender' key and moving the 'value' to the top level.
285 For example:
286 "key": [{"sender": null, "value": "some value"}] -> "key": "some value"
288 :param data: The JSON structure to transform.
289 :returns: The transformed structure.
290 """
291 if isinstance(data, dict):
292 # If this dict has both 'sender' and 'value', return just the value
293 if "value" in data and "sender" in data:
294 return data["value"]
295 # Otherwise, recursively process each key-value pair
296 return {k: _transform_json_structure(v) for k, v in data.items()}
298 if isinstance(data, list):
299 # First, transform each item in the list.
300 transformed = [_transform_json_structure(item) for item in data]
301 # If the original list has exactly one element and that element was a dict
302 # with 'sender' and 'value', then unwrap the list.
303 if len(data) == 1 and isinstance(data[0], dict) and "value" in data[0] and "sender" in data[0]:
304 return transformed[0]
305 return transformed
307 # For other data types, just return the value as is.
308 return data