-
Notifications
You must be signed in to change notification settings - Fork 3k
fix!: resume a pipeline snapshot taken on a later loop visit #12162
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
b6d51a9
c6a1517
10efc59
7b862e9
a8be660
d19ee3f
eb50372
63bedbd
ca3e3f2
38c8d55
c9a9bb2
11e2ab6
3e5d6fa
e37c7f3
02cb23f
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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`. | ||
|
Comment on lines
+241
to
+243
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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 |
||
|
|
||
| :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, | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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 |
||
|
|
||
| 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: | ||
|
|
||
There was a problem hiding this comment.
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_fallbackdoes not support serializing mixed-type lists.