From d15575fe36553f7a3f58901356e43939ca1a0ebb Mon Sep 17 00:00:00 2001 From: Eilen6316 Date: Tue, 18 Aug 2026 11:41:05 +0800 Subject: [PATCH] Make streaming failures fail closed by default --- burr/core/__init__.py | 2 + burr/core/application.py | 54 ++++++++-- docs/concepts/streaming-actions.rst | 25 ++++- tests/core/test_application.py | 159 ++++++++++++++++++++++++++-- 4 files changed, 219 insertions(+), 21 deletions(-) diff --git a/burr/core/__init__.py b/burr/core/__init__.py index ddc1f53e6..fee031aad 100644 --- a/burr/core/__init__.py +++ b/burr/core/__init__.py @@ -31,6 +31,7 @@ ApplicationBuilder, ApplicationContext, ApplicationGraph, + StreamingFailurePolicy, ) from burr.core.graph import Graph, GraphBuilder from burr.core.state import State @@ -42,6 +43,7 @@ "ApplicationBuilder", "ApplicationGraph", "ApplicationContext", + "StreamingFailurePolicy", "Condition", "default", "expr", diff --git a/burr/core/application.py b/burr/core/application.py index 25bce4a10..d192b86f4 100644 --- a/burr/core/application.py +++ b/burr/core/application.py @@ -19,6 +19,7 @@ import contextvars import dataclasses +import enum import functools import inspect import logging @@ -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 @@ -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.""" @@ -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. " @@ -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) @@ -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. " @@ -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 @@ -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. " @@ -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) @@ -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. " @@ -1396,6 +1408,7 @@ 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. @@ -1403,6 +1416,8 @@ def stream_result( :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. @@ -1424,8 +1439,9 @@ def stream_result( halt_after condition. The :py:class:`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: @@ -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 @@ -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 @@ -1647,6 +1665,7 @@ 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. @@ -1654,6 +1673,8 @@ async def astream_result( :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 `, 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. @@ -1675,8 +1696,9 @@ async def astream_result( halt_after condition. The :py:class:`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: @@ -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 @@ -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 @@ -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, @@ -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_ """ @@ -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 @@ -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, @@ -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) """ @@ -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 diff --git a/docs/concepts/streaming-actions.rst b/docs/concepts/streaming-actions.rst index 95e3a06e0..14ef8e13e 100644 --- a/docs/concepts/streaming-actions.rst +++ b/docs/concepts/streaming-actions.rst @@ -221,6 +221,24 @@ 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... @@ -228,9 +246,10 @@ 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 diff --git a/tests/core/test_application.py b/tests/core/test_application.py index 9313cefc9..d2069d5be 100644 --- a/tests/core/test_application.py +++ b/tests/core/test_application.py @@ -46,6 +46,7 @@ Application, ApplicationBuilder, ApplicationContext, + StreamingFailurePolicy, _adjust_single_step_output, _arun_function, _arun_multi_step_streaming_action, @@ -62,10 +63,14 @@ ) from burr.core.graph import Graph, GraphBuilder, Transition from burr.core.persistence import ( + AsyncInMemoryPersister, AsyncDevNullPersister, BaseStateLoader, BaseStatePersister, DevNullPersister, + InMemoryPersister, + PersisterHook, + PersisterHookAsync, PersistedStateData, SQLLitePersister, ) @@ -1620,13 +1625,29 @@ def update(self, result: dict, state: State) -> State: return state.update(**result).append(tracker=result["count"]) -def test__run_single_step_streaming_action_graceful_exception(): - """When the generator raises but yields a final state in finally, stream completes gracefully.""" +def test__run_single_step_streaming_action_fails_closed_by_default(): action = SingleStepStreamingCounterWithException().with_name("counter") state = State({"count": 0, "tracker": []}) generator = _run_single_step_streaming_action( action, state, inputs={}, sequence_id=0, partition_key="pk", app_id="app" ) + + with pytest.raises(RuntimeError, match="simulated failure"): + list(generator) + + +def test__run_single_step_streaming_action_commit_partial(): + action = SingleStepStreamingCounterWithException().with_name("counter") + state = State({"count": 0, "tracker": []}) + generator = _run_single_step_streaming_action( + action, + state, + inputs={}, + sequence_id=0, + partition_key="pk", + app_id="app", + failure_policy=StreamingFailurePolicy.COMMIT_PARTIAL, + ) results = list(generator) intermediate = [(r, s) for r, s in results if s is None] final = [(r, s) for r, s in results if s is not None] @@ -1650,8 +1671,7 @@ def test__run_single_step_streaming_action_exception_propagates(): list(generator) -async def test__run_single_step_streaming_action_graceful_exception_async(): - """Async: when the generator raises but yields a final state in finally, stream completes.""" +async def test__run_single_step_streaming_action_fails_closed_by_default_async(): action = SingleStepStreamingCounterWithExceptionAsync().with_name("counter") state = State({"count": 0, "tracker": []}) generator = _arun_single_step_streaming_action( @@ -1663,6 +1683,25 @@ async def test__run_single_step_streaming_action_graceful_exception_async(): partition_key="pk", lifecycle_adapters=LifecycleAdapterSet(), ) + + with pytest.raises(RuntimeError, match="simulated failure"): + async for _ in generator: + pass + + +async def test__run_single_step_streaming_action_commit_partial_async(): + action = SingleStepStreamingCounterWithExceptionAsync().with_name("counter") + state = State({"count": 0, "tracker": []}) + generator = _arun_single_step_streaming_action( + action=action, + state=state, + inputs={}, + sequence_id=0, + app_id="app", + partition_key="pk", + lifecycle_adapters=LifecycleAdapterSet(), + failure_policy=StreamingFailurePolicy.COMMIT_PARTIAL, + ) results = [] async for item in generator: results.append(item) @@ -1695,13 +1734,29 @@ async def test__run_single_step_streaming_action_exception_propagates_async(): pass -def test__run_multi_step_streaming_action_graceful_exception(): - """When the generator raises but yields a final result in finally, stream completes.""" +def test__run_multi_step_streaming_action_fails_closed_by_default(): action = MultiStepStreamingCounterWithException().with_name("counter") state = State({"count": 0, "tracker": []}) generator = _run_multi_step_streaming_action( action, state, inputs={}, sequence_id=0, partition_key="pk", app_id="app" ) + + with pytest.raises(RuntimeError, match="simulated failure"): + list(generator) + + +def test__run_multi_step_streaming_action_commit_partial(): + action = MultiStepStreamingCounterWithException().with_name("counter") + state = State({"count": 0, "tracker": []}) + generator = _run_multi_step_streaming_action( + action, + state, + inputs={}, + sequence_id=0, + partition_key="pk", + app_id="app", + failure_policy=StreamingFailurePolicy.COMMIT_PARTIAL, + ) results = list(generator) intermediate = [(r, s) for r, s in results if s is None] final = [(r, s) for r, s in results if s is not None] @@ -1725,8 +1780,25 @@ def test__run_multi_step_streaming_action_exception_propagates(): list(generator) -async def test__run_multi_step_streaming_action_graceful_exception_async(): - """Async: when the generator raises but yields a final result in finally, stream completes.""" +async def test__run_multi_step_streaming_action_fails_closed_by_default_async(): + action = MultiStepStreamingCounterWithExceptionAsync().with_name("counter") + state = State({"count": 0, "tracker": []}) + generator = _arun_multi_step_streaming_action( + action=action, + state=state, + inputs={}, + sequence_id=0, + app_id="app", + partition_key="pk", + lifecycle_adapters=LifecycleAdapterSet(), + ) + + with pytest.raises(RuntimeError, match="simulated failure"): + async for _ in generator: + pass + + +async def test__run_multi_step_streaming_action_commit_partial_async(): action = MultiStepStreamingCounterWithExceptionAsync().with_name("counter") state = State({"count": 0, "tracker": []}) generator = _arun_multi_step_streaming_action( @@ -1737,6 +1809,7 @@ async def test__run_multi_step_streaming_action_graceful_exception_async(): app_id="app", partition_key="pk", lifecycle_adapters=LifecycleAdapterSet(), + failure_policy=StreamingFailurePolicy.COMMIT_PARTIAL, ) results = [] async for item in generator: @@ -1770,6 +1843,76 @@ async def test__run_multi_step_streaming_action_exception_propagates_async(): pass +def test_stream_result_failure_does_not_commit_state(): + persister = InMemoryPersister() + action = MultiStepStreamingCounterWithException().with_name("counter") + app = ( + ApplicationBuilder() + .with_actions(action) + .with_transitions() + .with_entrypoint("counter") + .with_state(count=0, tracker=[]) + .with_identifiers(app_id="app", partition_key="pk") + .with_hooks(PersisterHook(persister)) + .build() + ) + + _, streaming_result = app.stream_result(halt_after=["counter"]) + with pytest.raises(RuntimeError, match="simulated failure"): + streaming_result.get() + + persisted_state = persister.load("pk", "app") + assert app.state["count"] == 0 + assert persisted_state["state"]["count"] == 0 + assert persisted_state["status"] == "failed" + + +def test_stream_result_commit_partial_requires_explicit_policy(): + action = MultiStepStreamingCounterWithException().with_name("counter") + app = ( + ApplicationBuilder() + .with_actions(action) + .with_transitions() + .with_entrypoint("counter") + .with_state(count=0, tracker=[]) + .build() + ) + + _, streaming_result = app.stream_result( + halt_after=["counter"], + failure_policy=StreamingFailurePolicy.COMMIT_PARTIAL, + ) + result, state = streaming_result.get() + + assert result == {"count": 1} + assert state["count"] == 1 + assert app.state["count"] == 1 + + +async def test_astream_result_failure_does_not_commit_state(): + persister = AsyncInMemoryPersister() + action = MultiStepStreamingCounterWithExceptionAsync().with_name("counter") + app = await ( + ApplicationBuilder() + .with_actions(action) + .with_transitions() + .with_entrypoint("counter") + .with_state(count=0, tracker=[]) + .with_identifiers(app_id="app", partition_key="pk") + .with_hooks(PersisterHookAsync(persister)) + .abuild() + ) + + _, streaming_result = await app.astream_result(halt_after=["counter"]) + with pytest.raises(RuntimeError, match="simulated failure"): + await streaming_result.get() + + persisted_state = await persister.load("pk", "app") + assert app.state["count"] == 0 + assert persisted_state["state"]["count"] == 0 + assert persisted_state["status"] == "failed" + + class SingleStepActionWithDeletionAsync(SingleStepActionWithDeletion): async def run_and_update(self, state: State, **run_kwargs) -> Tuple[dict, State]: return {}, state.wipe(delete=["to_delete"])