Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion haystack/core/component/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
17 changes: 10 additions & 7 deletions haystack/core/pipeline/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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}]
58 changes: 51 additions & 7 deletions haystack/core/pipeline/breakpoint.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)
Expand Down Expand Up @@ -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.
Comment on lines +222 to +225

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This explains why this function is needed as a workaround since _serialize_with_field_fallback does not support serializing mixed-type lists.


: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`.
Comment on lines +241 to +243

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Related to comment https://github.com/deepset-ai/haystack/pull/12162/changes#r3656582668 that this function is only needed to work around _deserialize_value_with_schema not supporting deserializing mixed type lists.


: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],
Expand All @@ -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.
Expand All @@ -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`"
)
Expand All @@ -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,
Expand Down
43 changes: 36 additions & 7 deletions haystack/core/pipeline/component_checks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Comment on lines +10 to +12

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was upgraded to a proper standalone class b/c it must support serde since it can be included in the serialized inputs in PipelineState.inputs


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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand All @@ -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

Expand All @@ -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

Expand All @@ -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:
Expand Down
43 changes: 28 additions & 15 deletions haystack/core/pipeline/pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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`
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
)
Expand All @@ -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,
Expand All @@ -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,
Expand Down
10 changes: 10 additions & 0 deletions haystack/dataclasses/breakpoints.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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]:
"""
Expand Down
Loading
Loading