Skip to content

Make streaming failures fail closed by default #889

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.

Sign up for GitHub

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

Open
Eilen6316 wants to merge 1 commit into
base: main
Choose a base branch
from
+219 −21
Open

Make streaming failures fail closed by default #889

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
2 changes: 2 additions & 0 deletions burr/core/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
ApplicationBuilder,
ApplicationContext,
ApplicationGraph,
StreamingFailurePolicy,
)
from burr.core.graph import Graph, GraphBuilder
from burr.core.state import State
Expand All @@ -42,6 +43,7 @@
"ApplicationBuilder",
"ApplicationGraph",
"ApplicationContext",
"StreamingFailurePolicy",
"Condition",
"default",
"expr",
Expand Down
54 changes: 44 additions & 10 deletions burr/core/application.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

import contextvars
import dataclasses
import enum
import functools
import inspect
import logging
Expand Down Expand Up @@ -89,6 +90,13 @@
StateTypeToSet = TypeVar("StateTypeToSet")


class StreamingFailurePolicy(str, enum.Enum):
"""Controls how streaming actions handle exceptions after producing output."""

FAIL_CLOSED = "fail_closed"
COMMIT_PARTIAL = "commit_partial"


def _validate_result(result: Any, name: str, schema: ActionSchema = DEFAULT_SCHEMA) -> None:
# TODO -- split out the action schema into input/output schema types
# Currently they're tied together, but this doesn't make as much sense for single-step actions
Expand Down Expand Up @@ -328,6 +336,7 @@ def _run_single_step_streaming_action(
app_id: str,
partition_key: Optional[str],
lifecycle_adapters: LifecycleAdapterSet = LifecycleAdapterSet(),
failure_policy: StreamingFailurePolicy = StreamingFailurePolicy.FAIL_CLOSED,
) -> Generator[Tuple[dict, Optional[State]], None, None]:
"""Runs a single step streaming action. This API is internal-facing.
This normalizes + validates the output."""
Expand Down Expand Up @@ -367,7 +376,7 @@ def _run_single_step_streaming_action(
)
yield result, None
except Exception as e:
if state_update is None:
if failure_policy != StreamingFailurePolicy.COMMIT_PARTIAL or state_update is None:
raise
logger.warning(
"Streaming action '%s' raised %s after yielding %d items. "
Expand Down Expand Up @@ -397,6 +406,7 @@ async def _arun_single_step_streaming_action(
app_id: str,
partition_key: Optional[str],
lifecycle_adapters: LifecycleAdapterSet = LifecycleAdapterSet(),
failure_policy: StreamingFailurePolicy = StreamingFailurePolicy.FAIL_CLOSED,
) -> AsyncGenerator[Tuple[dict, Optional[State]], None]:
"""Runs a single step streaming action in async. See the synchronous version for more details."""
action.validate_inputs(inputs)
Expand Down Expand Up @@ -435,7 +445,7 @@ async def _arun_single_step_streaming_action(
)
yield result, None
except Exception as e:
if state_update is None:
if failure_policy != StreamingFailurePolicy.COMMIT_PARTIAL or state_update is None:
raise
logger.warning(
"Streaming action '%s' raised %s after yielding %d items. "
Expand Down Expand Up @@ -467,6 +477,7 @@ def _run_multi_step_streaming_action(
app_id: str,
partition_key: Optional[str],
lifecycle_adapters: LifecycleAdapterSet = LifecycleAdapterSet(),
failure_policy: StreamingFailurePolicy = StreamingFailurePolicy.FAIL_CLOSED,
) -> Generator[Tuple[dict, Optional[State[ApplicationStateType]]], None, None]:
"""Runs a multi-step streaming action. E.G. one with a run/reduce step.
This API is internal-facing. Note that this converts the shape of a
Expand Down Expand Up @@ -505,7 +516,7 @@ def _run_multi_step_streaming_action(
count += 1
yield next_result, None
except Exception as e:
if result is None:
if failure_policy != StreamingFailurePolicy.COMMIT_PARTIAL or result is None:
raise
logger.warning(
"Streaming action '%s' raised %s after yielding %d items. "
Expand All @@ -531,6 +542,7 @@ async def _arun_multi_step_streaming_action(
app_id: str,
partition_key: Optional[str],
lifecycle_adapters: LifecycleAdapterSet = LifecycleAdapterSet(),
failure_policy: StreamingFailurePolicy = StreamingFailurePolicy.FAIL_CLOSED,
) -> AsyncGenerator[Tuple[dict, Optional[State[ApplicationStateType]]], None]:
"""Runs a multi-step streaming action in async. See the synchronous version for more details."""
action.validate_inputs(inputs)
Expand Down Expand Up @@ -563,7 +575,7 @@ async def _arun_multi_step_streaming_action(
count += 1
yield next_result, None
except Exception as e:
if result is None:
if failure_policy != StreamingFailurePolicy.COMMIT_PARTIAL or result is None:
raise
logger.warning(
"Streaming action '%s' raised %s after yielding %d items. "
Expand Down Expand Up @@ -1396,13 +1408,16 @@ def stream_result(
halt_after: list[str],
halt_before: Optional[list[str]] = None,
inputs: Optional[Dict[str, Any]] = None,
failure_policy: StreamingFailurePolicy = StreamingFailurePolicy.FAIL_CLOSED,
) -> Tuple[Action, StreamingResultContainer[ApplicationStateType, Union[dict, Any]]]:
"""Streams a result out.

:param halt_after: The list of actions/tags to halt after execution of. It will halt on the first one.
:param halt_before: The list of actions/tags to halt before execution of. It will halt on the first one. Note that
if this is met, the streaming result container will be empty (and return None) for the result, having an empty generator.
:param inputs: Inputs to the action -- this is if this action requires an input that is passed in from the outside world
:param failure_policy: Controls whether an exception after yielding output fails the action or
explicitly commits the last partial result. Defaults to fail closed.
:return: A streaming result container, which is a generator that will yield results as they come in, as well as cache/give you the final result,
and update state accordingly.

Expand All @@ -1424,8 +1439,9 @@ def stream_result(
halt_after condition.

The :py:class:`StreamingResultContainer <burr.core.action.StreamingResultContainer>` is meant as a convenience -- specifically this allows for
hooks, callbacks, etc... so you can take the control flow and still have state updated afterwards. Hooks/state update will be called after an exception
is thrown during streaming, or the stream is completed. Note that it is undefined behavior to attempt to execute another action while a stream is in progress.
hooks, callbacks, etc... so you can take control flow while retaining the final state update. Hooks are called after an exception is thrown
during streaming, or the stream is completed; state is updated only after successful completion or when ``COMMIT_PARTIAL`` is selected.
Note that it is undefined behavior to attempt to execute another action while a stream is in progress.


To see how this works, let's take the following action (simplified as a single-node workflow) as an example:
Expand Down Expand Up @@ -1607,6 +1623,7 @@ def callback(
app_id=self._uid,
partition_key=self._partition_key,
lifecycle_adapters=self._adapter_set,
failure_policy=failure_policy,
)
return next_action, StreamingResultContainer(
generator, self._state, process_result, callback
Expand All @@ -1621,6 +1638,7 @@ def callback(
app_id=self._uid,
partition_key=self._partition_key,
lifecycle_adapters=self._adapter_set,
failure_policy=failure_policy,
)
except Exception as e:
# We only want to raise this in the case of an exception
Expand All @@ -1647,13 +1665,16 @@ async def astream_result(
halt_after: list[str],
halt_before: Optional[list[str]] = None,
inputs: Optional[Dict[str, Any]] = None,
failure_policy: StreamingFailurePolicy = StreamingFailurePolicy.FAIL_CLOSED,
) -> Tuple[Action, AsyncStreamingResultContainer[ApplicationStateType, Union[dict, Any]]]:
"""Streams a result out in an asynchronous manner.

:param halt_after: The list of actions/tags to halt after execution of. It will halt on the first one.
:param halt_before: The list of actions/tags to halt before execution of. It will halt on the first one. Note that
if this is met, the streaming result container will be empty (and return None) for the result, having an empty generator.
:param inputs: Inputs to the action -- this is if this action requires an input that is passed in from the outside world
:param failure_policy: Controls whether an exception after yielding output fails the action or
explicitly commits the last partial result. Defaults to fail closed.
:return: An asynchronous :py:class:`AsyncStreamingResultContainer <burr.core.action.AsyncStreamingResultContainer>`, which is a generator that will yield results as they come in, as well as cache/give you the final result,
and update state accordingly.

Expand All @@ -1675,8 +1696,9 @@ async def astream_result(
halt_after condition.

The :py:class:`AsyncStreamingResultContainer <burr.core.action.AsyncStreamingResultContainer>` is meant as a convenience -- specifically this allows for
hooks, callbacks, etc... so you can take the control flow and still have state updated afterwards. Hooks/state update will be called after an exception
is thrown during streaming, or the stream is completed. Note that it is undefined behavior to attempt to execute another action while a stream is in progress.
hooks, callbacks, etc... so you can take control flow while retaining the final state update. Hooks are called after an exception is thrown
during streaming, or the stream is completed; state is updated only after successful completion or when ``COMMIT_PARTIAL`` is selected.
Note that it is undefined behavior to attempt to execute another action while a stream is in progress.


To see how this works, let's take the following action (simplified as a single-node workflow) as an example:
Expand Down Expand Up @@ -1864,6 +1886,7 @@ async def callback(
app_id=self._uid,
partition_key=self._partition_key,
lifecycle_adapters=self._adapter_set,
failure_policy=failure_policy,
)
return next_action, AsyncStreamingResultContainer(
generator, self._state, process_result, callback
Expand All @@ -1886,6 +1909,7 @@ async def callback(
app_id=self._uid,
partition_key=self._partition_key,
lifecycle_adapters=self._adapter_set,
failure_policy=failure_policy,
)
except Exception as e:
# We only want to raise this in the case of an exception
Expand Down Expand Up @@ -1920,6 +1944,7 @@ def stream_iterate(
halt_after: Optional[Union[str, List[str]]] = None,
halt_before: Optional[Union[str, List[str]]] = None,
inputs: Optional[Dict[str, Any]] = None,
failure_policy: StreamingFailurePolicy = StreamingFailurePolicy.FAIL_CLOSED,
) -> Generator[
Tuple[Action, StreamingResultContainer[ApplicationStateType, Union[dict, Any]]],
None,
Expand All @@ -1940,6 +1965,7 @@ def stream_iterate(
:param halt_after: Action names/tags to halt before
:param halt_before: _description_, defaults to None
:param inputs: _description_, defaults to None
:param failure_policy: Streaming exception policy passed to each action. Defaults to fail closed.
:return: _description_
:yield: _description_
"""
Expand All @@ -1951,7 +1977,10 @@ def stream_iterate(
while self.has_next_action():
next_action = self.get_next_action()
_, streaming_result = self.stream_result(
halt_after=[next_action.name], halt_before=None, inputs=inputs
halt_after=[next_action.name],
halt_before=None,
inputs=inputs,
failure_policy=failure_policy,
)
yield next_action, streaming_result
# We need to ensure it's fully exhausted before going to the next action
Expand All @@ -1965,6 +1994,7 @@ async def astream_iterate(
halt_after: Optional[Union[str, List[str]]] = None,
halt_before: Optional[Union[str, List[str]]] = None,
inputs: Optional[Dict[str, Any]] = None,
failure_policy: StreamingFailurePolicy = StreamingFailurePolicy.FAIL_CLOSED,
) -> AsyncGenerator[
Tuple[
Action,
Expand All @@ -1978,6 +2008,7 @@ async def astream_iterate(
:param halt_after: Action names/tags to halt after the action completes
:param halt_before: Action names/tags to halt after
:param inputs: Inputs to the first action run
:param failure_policy: Streaming exception policy passed to each action. Defaults to fail closed.
:return: Async generator yielding tuples of (action, streaming_result_container)
:yield: Tuples of (action, streaming_result_container)
"""
Expand All @@ -1988,7 +2019,10 @@ async def astream_iterate(
while self.has_next_action():
next_action = self.get_next_action()
_, streaming_result = await self.astream_result( # Use astream_result
halt_after=[next_action.name], halt_before=None, inputs=inputs
halt_after=[next_action.name],
halt_before=None,
inputs=inputs,
failure_policy=failure_policy,
)
yield next_action, streaming_result
# We need to ensure it's fully exhausted before going to the next action
Expand Down
25 changes: 22 additions & 3 deletions docs/concepts/streaming-actions.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -221,16 +221,35 @@ This object is effectively a cached iterator. You can use it as follows:
result, state = await async_streaming_result.get()
print(result) # all at once

Streaming failures fail closed by default. If a streaming action raises, Burr propagates the original
exception and does not apply the action's state update, even if the action already produced partial
results. Applications that intentionally accept partial results must opt in for that execution:

.. code-block:: python

from burr.core import StreamingFailurePolicy

action, streaming_result = application.stream_result(
halt_after="streaming_response",
inputs={"prompt": prompt},
failure_policy=StreamingFailurePolicy.COMMIT_PARTIAL,
)

The same policy is available on ``astream_result()``, ``stream_iterate()``, and
``astream_iterate()``. ``COMMIT_PARTIAL`` runs the reducer with the last available result, so use it
only when partial output is valid application state.


Thus you can run this in a web-service, a streamlit app, etc...

--------------
Considerations
--------------

All hooks/state update will be called once the iterator completes, or an exception interrupts the iterator
and it has to be cleaned up. You can call ``.stream_result()`` or ``.astream_result()`` on non-streaming
results, and it will return a ``StreamingResultContainer`` with an empty iterator that returns the result.
Hooks are called once the iterator completes, or an exception interrupts the iterator and it has to be
cleaned up. State is updated only after successful completion, unless ``COMMIT_PARTIAL`` was explicitly
selected. You can call ``.stream_result()`` or ``.astream_result()`` on non-streaming results, and it
will return a ``StreamingResultContainer`` with an empty iterator that returns the result.
If streaming items are run as intermediate nodes in the graph, they will be run as normal actions
(effectively fully exhausted), and the result will be returned as a single item. Currently
you cannot use synchronous streaming actions as asynchronous streaming actions, but we will likely be
Expand Down
Loading
Loading