fix!: resume a pipeline snapshot taken on a later loop visit (#12162)

Co-authored-by: David S. Batista <dsbatista@gmail.com>

Author: sjrl

haystack/core/component/types.py

@@ -73,7 +73,8 @@ def is_variadic(self) -> bool:
     @property
     def is_mandatory(self) -> bool:
         """Check if the input is mandatory."""
-        return self.default_value == _empty
+        # Identity, so a default with a custom `__eq__`, such as a DataFrame, is never compared to the sentinel.
+        return self.default_value is _empty
 
     def __post_init__(self) -> None:
         try:

haystack/core/pipeline/base.py

@@ -27,7 +27,7 @@
     PipelineValidationError,
 )
 from haystack.core.pipeline.component_checks import (
-    _NO_OUTPUT_PRODUCED,
+    _NoOutputProduced,
     all_predecessors_executed,
     are_all_sockets_ready,
     can_component_run,
@@ -1223,7 +1223,9 @@ def _consume_component_inputs(
         greedy_inputs_to_remove = set()
         for socket_name, socket in component["input_sockets"].items():
             socket_inputs = component_inputs.get(socket_name, [])
-            socket_inputs_values = [sock["value"] for sock in socket_inputs if sock["value"] is not _NO_OUTPUT_PRODUCED]
+            socket_inputs_values = [
+                sock["value"] for sock in socket_inputs if not isinstance(sock["value"], _NoOutputProduced)
+            ]
 
             # if we are resuming a component, the inputs are already consumed, so we just return the first input
             if is_resume:
@@ -1546,12 +1548,12 @@ def _write_component_outputs(
         :param include_outputs_from: Set of component names that should always return an output from the pipeline.
         """
         for receiver_name, sender_socket, receiver_socket, conversion_strategy in receivers:
-            # We either get the value that was produced by the actor or we use the _NO_OUTPUT_PRODUCED class to indicate
+            # We either get the value that was produced by the actor or we use a _NoOutputProduced marker to indicate
             # that the sender did not produce an output for this socket.
             # This allows us to track if a predecessor already ran but did not produce an output.
-            value = component_outputs.get(sender_socket.name, _NO_OUTPUT_PRODUCED)
+            value = component_outputs.get(sender_socket.name, _NoOutputProduced())
 
-            if value is not _NO_OUTPUT_PRODUCED and conversion_strategy:
+            if not isinstance(value, _NoOutputProduced) and conversion_strategy:
                 try:
                     value = _convert_value(value=value, conversion_strategy=conversion_strategy)
                 except Exception as e:
@@ -1588,7 +1590,8 @@ def _write_component_outputs(
                 )
             else:
                 # If the receiver socket is not lazy variadic, it is greedy variadic or non-variadic.
-                # We overwrite with the new input if it's not _NO_OUTPUT_PRODUCED or if the current value is None.
+                # We overwrite with the new input if it's not a _NoOutputProduced marker, or if the current value
+                # is None.
                 _write_to_standard_socket(
                     inputs=inputs,
                     receiver_name=receiver_name,
@@ -1861,5 +1864,5 @@ def _write_to_standard_socket(
     current_value = inputs[receiver_name].get(receiver_socket_name)
 
     # Only overwrite if there's no existing value, or we have a new value to provide
-    if current_value is None or value is not _NO_OUTPUT_PRODUCED:
+    if current_value is None or not isinstance(value, _NoOutputProduced):
         inputs[receiver_name][receiver_socket_name] = [{"sender": component_name, "value": value}]

haystack/core/pipeline/breakpoint.py

@@ -13,7 +13,8 @@
 
 from haystack import logging
 from haystack.core.errors import PipelineInvalidPipelineSnapshotError
-from haystack.dataclasses.breakpoints import Breakpoint, PipelineSnapshot, PipelineState
+from haystack.dataclasses.breakpoints import INTERNAL_INPUTS_FORMAT, Breakpoint, PipelineSnapshot, PipelineState
+from haystack.utils import _deserialize_value_with_schema
 from haystack.utils.base_serialization import _serialize_with_field_fallback
 
 logger = logging.getLogger(__name__)
@@ -211,6 +212,48 @@ def _save_pipeline_snapshot(
     return str(full_path)
 
 
+def _serialize_internal_inputs(inputs: dict[str, Any]) -> dict[str, Any]:
+    """
+    Serialize the pipeline's internal inputs state, keeping the `sender` of every input intact.
+
+    The `sender` tells a resumed run which inputs came from a predecessor and which from outside the pipeline, and
+    therefore whether the paused component can be triggered again.
+
+    The inputs a socket received are stored keyed by position rather than as a list, because a list gets a single
+    schema derived from its first item. A socket with several senders can hold values of different types, for
+    example a value from one sender and a `_NoOutputProduced` marker from another, and those need one schema
+    each to survive the round trip.
+
+    :param inputs: The pipeline's internal inputs state.
+    :returns: A dict of the form `{"serialization_schema": ..., "serialized_data": ...}`.
+    """
+    positional_inputs = {
+        component_name: {
+            socket_name: {str(position): entry for position, entry in enumerate(socket_inputs)}
+            for socket_name, socket_inputs in socket_dict.items()
+        }
+        for component_name, socket_dict in inputs.items()
+    }
+
+    return _serialize_with_field_fallback(positional_inputs, description="the inputs of the current pipeline state")
+
+
+def _deserialize_internal_inputs(serialized_inputs: dict[str, Any]) -> dict[str, Any]:
+    """
+    Restore the pipeline's internal inputs state from a snapshot written by `_serialize_internal_inputs`.
+
+    :param serialized_inputs: The `inputs` of a `PipelineState` whose `inputs_format` is `INTERNAL_INPUTS_FORMAT`.
+    :returns: The pipeline's internal inputs state.
+    """
+    return {
+        component_name: {
+            socket_name: [socket_inputs[position] for position in sorted(socket_inputs, key=int)]
+            for socket_name, socket_inputs in socket_dict.items()
+        }
+        for component_name, socket_dict in _deserialize_value_with_schema(serialized_inputs).items()
+    }
+
+
 def _create_pipeline_snapshot(
     *,
     inputs: dict[str, Any],
@@ -226,7 +269,8 @@ def _create_pipeline_snapshot(
     Create a snapshot of the pipeline at the point where the breakpoint was triggered.
 
     :param inputs: The current pipeline snapshot inputs.
-    :param component_inputs: The inputs to the component that triggered the breakpoint.
+    :param component_inputs: The inputs of the component that triggered the breakpoint, as they were before the
+        component consumed them.
     :param break_point: The breakpoint that triggered the snapshot.
     :param component_visits: The visit count of the component that triggered the breakpoint.
     :param original_input_data: The original input data.
@@ -239,11 +283,8 @@ def _create_pipeline_snapshot(
     component_name = break_point.component_name
 
     transformed_original_input_data = _transform_json_structure(original_input_data)
-    transformed_inputs = _transform_json_structure({**inputs, component_name: component_inputs})
 
-    serialized_inputs = _serialize_with_field_fallback(
-        transformed_inputs, description="the inputs of the current pipeline state"
-    )
+    serialized_inputs = _serialize_internal_inputs({**inputs, component_name: component_inputs})
     serialized_original_input_data = _serialize_with_field_fallback(
         transformed_original_input_data, description="original input data for `pipeline.run`"
     )
@@ -253,7 +294,10 @@ def _create_pipeline_snapshot(
 
     return PipelineSnapshot(
         pipeline_state=PipelineState(
-            inputs=serialized_inputs, component_visits=component_visits, pipeline_outputs=serialized_pipeline_outputs
+            inputs=serialized_inputs,
+            component_visits=component_visits,
+            pipeline_outputs=serialized_pipeline_outputs,
+            inputs_format=INTERNAL_INPUTS_FORMAT,
         ),
         timestamp=datetime.now(),
         break_point=break_point,

haystack/core/pipeline/component_checks.py

@@ -4,9 +4,38 @@
 
 from typing import Any
 
-from haystack.core.component.types import InputSocket, _empty
+from haystack.core.component.types import InputSocket
 
-_NO_OUTPUT_PRODUCED = _empty
+
+class _NoOutputProduced:
+    """
+    Marker a socket input holds when its sender ran without producing a value for that socket.
+
+    Test for it with `isinstance(value, _NoOutputProduced)`, never with `==`: comparing runs `__eq__` on the value
+    being checked, which for a value such as a DataFrame does not return a bool. It carries no state, so instances
+    are interchangeable and every one of them means the same thing.
+
+    It is serializable because it reaches a pipeline snapshot as part of the pipeline's inputs.
+    """
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, _NoOutputProduced)
+
+    def __hash__(self) -> int:
+        return hash(_NoOutputProduced)
+
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Serialize the marker. It carries no state, so the payload is empty.
+        """
+        return {}
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "_NoOutputProduced":  # noqa: ARG003
+        """
+        Deserialize the marker.
+        """
+        return cls()
 
 
 def can_component_run(component: dict, inputs: dict) -> bool:
@@ -103,7 +132,7 @@ def any_socket_value_from_predecessor_received(socket_inputs: list[dict[str, Any
     :param socket_inputs: Inputs for the component's socket.
     """
     # When sender is None, the input was provided from outside the pipeline.
-    return any(inp["value"] is not _NO_OUTPUT_PRODUCED and inp["sender"] is not None for inp in socket_inputs)
+    return any(not isinstance(inp["value"], _NoOutputProduced) and inp["sender"] is not None for inp in socket_inputs)
 
 
 def has_user_input(inputs: dict) -> bool:
@@ -142,7 +171,7 @@ def any_socket_input_received(socket_inputs: list[dict]) -> bool:
 
     :param socket_inputs: Inputs for the socket.
     """
-    return any(inp["value"] is not _NO_OUTPUT_PRODUCED for inp in socket_inputs)
+    return any(not isinstance(inp["value"], _NoOutputProduced) for inp in socket_inputs)
 
 
 def has_lazy_variadic_socket_received_all_inputs(socket: InputSocket, socket_inputs: list[dict]) -> bool:
@@ -156,7 +185,7 @@ def has_lazy_variadic_socket_received_all_inputs(socket: InputSocket, socket_inp
     actual_senders = {
         sock["sender"]
         for sock in socket_inputs
-        if sock["value"] is not _NO_OUTPUT_PRODUCED and sock["sender"] is not None
+        if not isinstance(sock["value"], _NoOutputProduced) and sock["sender"] is not None
     }
     return expected_senders == actual_senders
 
@@ -176,7 +205,7 @@ def has_socket_received_all_inputs(socket: InputSocket, socket_inputs: list[dict
     if (
         socket.is_variadic
         and socket.is_greedy
-        and any(sock["value"] is not _NO_OUTPUT_PRODUCED for sock in socket_inputs)
+        and any(not isinstance(sock["value"], _NoOutputProduced) for sock in socket_inputs)
     ):
         return True
 
@@ -185,7 +214,7 @@ def has_socket_received_all_inputs(socket: InputSocket, socket_inputs: list[dict
         return True
 
     # The socket is not variadic and the only expected input is complete.
-    return not socket.is_variadic and socket_inputs[0]["value"] is not _NO_OUTPUT_PRODUCED
+    return not socket.is_variadic and not isinstance(socket_inputs[0]["value"], _NoOutputProduced)
 
 
 def all_predecessors_executed(component: dict, inputs: dict) -> bool:

haystack/core/pipeline/pipeline.py

@@ -21,13 +21,14 @@
 from haystack.core.pipeline.breakpoint import (
     SnapshotCallback,
     _create_pipeline_snapshot,
+    _deserialize_internal_inputs,
     _save_pipeline_snapshot,
     _validate_break_point_against_pipeline,
     _validate_pipeline_snapshot_against_pipeline,
 )
 from haystack.core.pipeline.utils import _deepcopy_with_exceptions
 from haystack.dataclasses import AsyncStreamingCallbackT, StreamingCallbackT, StreamingChunk, select_streaming_callback
-from haystack.dataclasses.breakpoints import Breakpoint, PipelineSnapshot
+from haystack.dataclasses.breakpoints import INTERNAL_INPUTS_FORMAT, Breakpoint, PipelineSnapshot
 from haystack.dataclasses.streaming_chunk import _invoke_streaming_callback
 from haystack.telemetry import pipeline_running
 from haystack.utils import _deserialize_value_with_schema
@@ -347,6 +348,9 @@ def run(  # noqa: PLR0915, PLR0912, C901
             include_outputs_from = set()
 
         pipeline_outputs: dict[str, Any] = {}
+        # Set when resuming from a snapshot that predates `INTERNAL_INPUTS_FORMAT` and therefore lost the sender of
+        # each input. Cleared as soon as the paused component has run.
+        legacy_resume_component: str | None = None
 
         if not pipeline_snapshot:
             # normalize `data`
@@ -362,14 +366,26 @@ def run(  # noqa: PLR0915, PLR0912, C901
             # We track component visits to decide if a component can run.
             component_visits = dict.fromkeys(ordered_component_names, 0)
 
+            inputs = self._convert_to_internal_format(pipeline_inputs=data)
+
         else:
             # Validate the pipeline snapshot against the current pipeline graph
             _validate_pipeline_snapshot_against_pipeline(pipeline_snapshot, self.graph)
 
             # Handle resuming the pipeline from a snapshot
             component_visits = pipeline_snapshot.pipeline_state.component_visits
             ordered_component_names = pipeline_snapshot.ordered_component_names
-            data = _deserialize_value_with_schema(pipeline_snapshot.pipeline_state.inputs)
+            data = _deserialize_value_with_schema(pipeline_snapshot.original_input_data)
+
+            if pipeline_snapshot.pipeline_state.inputs_format == INTERNAL_INPUTS_FORMAT:
+                inputs = _deserialize_internal_inputs(pipeline_snapshot.pipeline_state.inputs)
+            else:
+                # A legacy snapshot lost the sender of each input, so the only thing we can do is treat them as if
+                # they came from outside the pipeline. The paused component then has to be let through once.
+                inputs = self._convert_to_internal_format(
+                    pipeline_inputs=_deserialize_value_with_schema(pipeline_snapshot.pipeline_state.inputs)
+                )
+                legacy_resume_component = pipeline_snapshot.break_point.component_name
 
             # include_outputs_from from the snapshot when resuming
             include_outputs_from = pipeline_snapshot.include_outputs_from
@@ -391,7 +407,6 @@ def run(  # noqa: PLR0915, PLR0912, C901
                 "haystack.pipeline.execution_mode": "sync",
             },
         ) as span:
-            inputs = self._convert_to_internal_format(pipeline_inputs=data)
             priority_queue = self._fill_queue(ordered_component_names, inputs, component_visits)
 
             # check if pipeline is blocked before execution
@@ -431,10 +446,15 @@ def run(  # noqa: PLR0915, PLR0912, C901
                         component_name, component_visits[component_name]
                     )
 
-                if pipeline_snapshot:
-                    is_resume = pipeline_snapshot.break_point.component_name == component_name
-                else:
-                    is_resume = False
+                is_resume = legacy_resume_component == component_name
+                if is_resume:
+                    # Only the first execution of the paused component needs the legacy handling. Later visits of the
+                    # same component inside a loop receive regular inputs and must consume them normally.
+                    legacy_resume_component = None
+
+                # A snapshot has to store the component's inputs as they were before the component consumed them, so
+                # that resuming re-triggers the component the same way this run did.
+                component_inputs_before_consume = inputs.get(component_name, {})
                 component_inputs = self._consume_component_inputs(
                     component_name=component_name, component=component, inputs=inputs, is_resume=is_resume
                 )
@@ -444,13 +464,6 @@ def run(  # noqa: PLR0915, PLR0912, C901
                 # initialization
                 component_inputs = self._add_missing_input_defaults(component_inputs, component["input_sockets"])
 
-                # Scenario 1: Pipeline snapshot is provided to resume the pipeline at a specific component
-                # Deserialize the component_inputs if they are passed in the pipeline_snapshot.
-                # this check will prevent other component_inputs generated at runtime from being deserialized
-                if pipeline_snapshot and component_name in pipeline_snapshot.pipeline_state.inputs.keys():
-                    for key, value in component_inputs.items():
-                        component_inputs[key] = _deserialize_value_with_schema(value)
-
                 try:
                     component_outputs = self._run_component(
                         component_name=component_name,
@@ -474,7 +487,7 @@ def run(  # noqa: PLR0915, PLR0912, C901
                     # Create a snapshot of the state of the pipeline before the error occurred.
                     pipeline_snapshot = _create_pipeline_snapshot(
                         inputs=_deepcopy_with_exceptions(inputs),
-                        component_inputs=_deepcopy_with_exceptions(component_inputs),
+                        component_inputs=_deepcopy_with_exceptions(component_inputs_before_consume),
                         break_point=saved_break_point,
                         component_visits=component_visits,
                         original_input_data=data,

haystack/dataclasses/breakpoints.py

@@ -8,6 +8,10 @@
 
 from haystack.utils.dataclasses import _warn_on_inplace_mutation
 
+# Value of `PipelineState.inputs_format` used by every snapshot Haystack creates now. See that field for the shape
+# each format implies.
+INTERNAL_INPUTS_FORMAT = "internal"
+
 
 @dataclass(frozen=True)
 class Breakpoint:
@@ -53,11 +57,17 @@ class PipelineState:
     :param component_visits: A dictionary mapping component names to their visit counts.
     :param inputs: The inputs processed by the pipeline at the time of the snapshot.
     :param pipeline_outputs: Dictionary containing the final outputs of the pipeline up to the breakpoint.
+    :param inputs_format: Which format `inputs` are stored in. `"internal"` means each input records the component
+        that sent it, keyed by the position it arrived in, as in
+        `{component: {socket: {"0": {"sender": ..., "value": ...}}}}`. `None` marks snapshots taken before Haystack
+        recorded the sender, which hold one flattened value per socket, `{component: {socket: value}}`, and can only
+        be resumed on a component's first visit.
     """
 
     inputs: dict[str, Any]
     component_visits: dict[str, int]
     pipeline_outputs: dict[str, Any]
+    inputs_format: str | None = None
 
     def to_dict(self) -> dict[str, Any]:
         """