Repository: PrefectHQ/prefect
Pull request: Fix error when returning from task and flows with `cache_result_in_memory=False`
<!-- 
Thanks for opening a pull request to Prefect! We've got a few requests to help us review contributions:

- Make sure that your title neatly summarizes the proposed changes.
- Provide a short overview of the change and the value it adds.
- Share an example to help us understand the change in user experience.
- Confirm that you've done common tasks so we can give a timely review.
- Review our contribution guidelines: https://docs.prefect.io/latest/contributing/overview/

Happy engineering!
-->

<!-- Include an overview here -->
We were relying to heavily on results to get the final return value for a task or flow. If the return value wasn't cached in memory on the created result, then we can't return anything. This PR updates the task and flow engines to hold their own reference to the user code return value, so that in most cases we can can return a value without relying on the result holding the return value in memory.

Closes https://github.com/PrefectHQ/prefect/issues/14332

### Example
<!-- 
Share an example of the change in action.

A code blurb is best. Changes to features should include an example that is executable by a new user.
If changing documentation, a link to a preview of the page is great.
 -->

### Checklist
<!-- These boxes may be checked after opening the pull request. -->

- [x] This pull request includes a label categorizing the change e.g. `maintenance`, `fix`, `feature`, `enhancement`, `docs`.
- [x] This pull request references any related issue by including "closes `<link to issue>`"
  <!-- If you do not have permission to add a label, a maintainer will add one for you -->
  - If no issue exists and your change is not a small fix, please [create an issue](https://github.com/PrefectHQ/prefect/issues/new/choose) first.
- [ ] If this pull request adds new functionality, it includes unit tests that cover the changes
- [ ] If this pull request removes docs files, it includes redirect settings in `mint.json`.
- [ ] If this pull request adds functions or classes, it includes helpful docstrings.


Problem statement:
Turning off result caching in memory crashes tasks
MRE:
```python
from prefect import flow, task


@task(cache_result_in_memory=False)
def my_task():
    return 42

@flow
def my_flow():
    my_task()


if __name__ == "__main__":
    my_flow()
```

Output:
```
13:49:25.313 | INFO    | Task run 'my_task' - Created task run 'my_task' for task 'my_task'
13:49:25.375 | INFO    | Task run 'my_task' - Finished in state Completed()
Traceback (most recent call last):
  File "/Users/alexander/dev/PrefectHQ/prefect/my_files/flow2.py", line 10, in <module>
    my_task()
  File "/Users/alexander/dev/PrefectHQ/prefect/src/prefect/tasks.py", line 826, in __call__
    return run_task(
  File "/Users/alexander/dev/PrefectHQ/prefect/src/prefect/task_engine.py", line 862, in run_task
    return run_task_sync(**kwargs)
  File "/Users/alexander/dev/PrefectHQ/prefect/src/prefect/task_engine.py", line 675, in run_task_sync
    return engine.state if return_type == "state" else engine.result()
  File "/Users/alexander/dev/PrefectHQ/prefect/src/prefect/task_engine.py", line 307, in result
    _result = self.state.result(raise_on_failure=raise_on_failure, fetch=True)
  File "/Users/alexander/dev/PrefectHQ/prefect/src/prefect/client/schemas/objects.py", line 262, in result
    return get_state_result(
  File "/Users/alexander/dev/PrefectHQ/prefect/src/prefect/states.py", line 68, in get_state_result
    return _get_state_result(
  File "/Users/alexander/dev/PrefectHQ/prefect/src/prefect/utilities/asyncutils.py", line 389, in coroutine_wrapper
    return run_coro_as_sync(ctx_call())
  File "/Users/alexander/dev/PrefectHQ/prefect/src/prefect/utilities/asyncutils.py", line 248, in run_coro_as_sync
    return call.result()
  File "/Users/alexander/dev/PrefectHQ/prefect/src/prefect/_internal/concurrency/calls.py", line 312, in result
    return self.future.result(timeout=timeout)
  File "/Users/alexander/dev/PrefectHQ/prefect/src/prefect/_internal/concurrency/calls.py", line 182, in result
    return self.__get_result()
  File "/Users/alexander/.pyenv/versions/3.9.18/lib/python3.9/concurrent/futures/_base.py", line 391, in __get_result
    raise self._exception
  File "/Users/alexander/dev/PrefectHQ/prefect/src/prefect/_internal/concurrency/calls.py", line 383, in _run_async
    result = await coro
  File "/Users/alexander/dev/PrefectHQ/prefect/src/prefect/utilities/asyncutils.py", line 231, in coroutine_wrapper
    return await task
  File "/Users/alexander/dev/PrefectHQ/prefect/src/prefect/utilities/asyncutils.py", line 379, in ctx_call
    result = await async_fn(*args, **kwargs)
  File "/Users/alexander/dev/PrefectHQ/prefect/src/prefect/states.py", line 130, in _get_state_result
    result = await _get_state_result_data_with_retries(
  File "/Users/alexander/dev/PrefectHQ/prefect/src/prefect/states.py", line 94, in _get_state_result_data_with_retries
    return await state.data.get()
  File "/Users/alexander/dev/PrefectHQ/prefect/src/prefect/utilities/asyncutils.py", line 379, in ctx_call
    result = await async_fn(*args, **kwargs)
  File "/Users/alexander/dev/PrefectHQ/prefect/src/prefect/results.py", line 498, in get
    raise MissingResult("The result was not persisted and is no longer available.")
prefect.exceptions.MissingResult: The result was not persisted and is no longer available.
```

Hints from issue discussion:
Resolution of this issue likely involves not relying on the result from the task run state when returning a value to the user from the task engine.

Commit message: Add missing change

Patch to review:
diff --git a/src/prefect/flow_engine.py b/src/prefect/flow_engine.py
index a1e5378909f6..a240a0696d5f 100644
--- a/src/prefect/flow_engine.py
+++ b/src/prefect/flow_engine.py
@@ -16,6 +16,7 @@
     Literal,
     Optional,
     Tuple,
+    Type,
     TypeVar,
     Union,
     cast,
@@ -50,6 +51,7 @@
     exception_to_failed_state,
     return_value_to_state,
 )
+from prefect.utilities.annotations import NotSet
 from prefect.utilities.asyncutils import run_coro_as_sync
 from prefect.utilities.callables import (
     call_with_parameters,
@@ -95,6 +97,7 @@ class FlowRunEngine(Generic[P, R]):
     flow_run_id: Optional[UUID] = None
     logger: logging.Logger = field(default_factory=lambda: get_logger("engine"))
     wait_for: Optional[Iterable[PrefectFuture]] = None
+    _result: Union[R, Exception, Type[NotSet]] = NotSet
     _is_started: bool = False
     _client: Optional[SyncPrefectClient] = None
     short_circuit: bool = False
@@ -208,6 +211,11 @@ def set_state(self, state: State, force: bool = False) -> State:
         return state
 
     def result(self, raise_on_failure: bool = True) -> "Union[R, State, None]":
+        if self._result is not NotSet and not isinstance(self._result, State):
+            if isinstance(self._result, Exception) and raise_on_failure:
+                raise self._result
+            else:
+                return self._result
         _result = self.state.result(raise_on_failure=raise_on_failure, fetch=True)  # type: ignore
         # state.result is a `sync_compatible` function that may or may not return an awaitable
         # depending on whether the parent frame is sync or not
@@ -219,13 +227,15 @@ def handle_success(self, result: R) -> R:
         result_factory = getattr(FlowRunContext.get(), "result_factory", None)
         if result_factory is None:
             raise ValueError("Result factory is not set")
+        resolved_result = resolve_futures_to_states(result)
         terminal_state = run_coro_as_sync(
             return_value_to_state(
-                resolve_futures_to_states(result),
+                resolved_result,
                 result_factory=result_factory,
             )
         )
         self.set_state(terminal_state)
+        self._result = resolved_result
         return result
 
     def handle_exception(
@@ -252,6 +262,7 @@ def handle_exception(
                 ),
             )
             state = self.set_state(Running())
+        self._result = exc
         return state
 
     def handle_timeout(self, exc: TimeoutError) -> None:
@@ -572,7 +583,7 @@ def run_context(self):
             except TimeoutError as exc:
                 self.handle_timeout(exc)
             except Exception as exc:
-                self.logger.exception(f"Encountered exception during execution: {exc}")
+                self.logger.exception("Encountered exception during execution: %r", exc)
                 self.handle_exception(exc)
 
     def call_flow_fn(self) -> Union[R, Coroutine[Any, Any, R]]:
diff --git a/src/prefect/task_engine.py b/src/prefect/task_engine.py
index e8c1d17b2dc6..b85cbc863d60 100644
--- a/src/prefect/task_engine.py
+++ b/src/prefect/task_engine.py
@@ -17,6 +17,7 @@
     Optional,
     Sequence,
     Set,
+    Type,
     TypeVar,
     Union,
 )
@@ -63,6 +64,7 @@
     return_value_to_state,
 )
 from prefect.transactions import Transaction, transaction
+from prefect.utilities.annotations import NotSet
 from prefect.utilities.asyncutils import run_coro_as_sync
 from prefect.utilities.callables import call_with_parameters, parameters_to_args_kwargs
 from prefect.utilities.collections import visit_collection
@@ -89,6 +91,7 @@ class TaskRunEngine(Generic[P, R]):
     retries: int = 0
     wait_for: Optional[Iterable[PrefectFuture]] = None
     context: Optional[Dict[str, Any]] = None
+    _result: Union[R, Exception, Type[NotSet]] = NotSet
     _initial_run_context: Optional[TaskRunContext] = None
     _is_started: bool = False
     _client: Optional[SyncPrefectClient] = None
@@ -304,6 +307,11 @@ def set_state(self, state: State, force: bool = False) -> State:
         return new_state
 
     def result(self, raise_on_failure: bool = True) -> "Union[R, State, None]":
+        if self._result is not NotSet:
+            if isinstance(self._result, Exception) and raise_on_failure:
+                raise self._result
+            else:
+                return self._result
         _result = self.state.result(raise_on_failure=raise_on_failure, fetch=True)
         # state.result is a `sync_compatible` function that may or may not return an awaitable
         # depending on whether the parent frame is sync or not
@@ -339,6 +347,7 @@ def handle_success(self, result: R, transaction: Transaction) -> R:
         if transaction.is_committed():
             terminal_state.name = "Cached"
         self.set_state(terminal_state)
+        self._result = result
         return result
 
     def handle_retry(self, exc: Exception) -> bool:
@@ -365,9 +374,11 @@ def handle_retry(self, exc: Exception) -> bool:
                 new_state = Retrying()
 
             self.logger.info(
-                f"Task run failed with exception {exc!r} - "
-                f"Retry {self.retries + 1}/{self.task.retries} will start "
-                f"{str(delay) + ' second(s) from now' if delay else 'immediately'}"
+                "Task run failed with exception: %r - " "Retry %s/%s will start %s",
+                exc,
+                self.retries + 1,
+                self.task.retries,
+                str(delay) + " second(s) from now" if delay else "immediately",
             )
 
             self.set_state(new_state, force=True)
@@ -375,7 +386,9 @@ def handle_retry(self, exc: Exception) -> bool:
             return True
         elif self.retries >= self.task.retries:
             self.logger.error(
-                f"Task run failed with exception {exc!r} - Retries are exhausted"
+                "Task run failed with exception: %r - Retries are exhausted",
+                exc,
+                exc_info=True,
             )
             return False
 
@@ -394,6 +407,7 @@ def handle_exception(self, exc: Exception) -> None:
                 )
             )
             self.set_state(state)
+            self._result = exc
 
     def handle_timeout(self, exc: TimeoutError) -> None:
         if not self.handle_retry(exc):
@@ -539,8 +553,8 @@ def say_hello(name):
 
                             @flow
                             def example_flow():
-                                future = say_hello.submit(name="Marvin)
-                                future.wait()
+                                say_hello.submit(name="Marvin)
+                                say_hello.wait()
 
                             example_flow()
                                       """
diff --git a/tests/test_flows.py b/tests/test_flows.py
index 9808c0bda1c3..11dd56bbe09d 100644
--- a/tests/test_flows.py
+++ b/tests/test_flows.py
@@ -998,6 +998,38 @@ class Foo:
                 def bar():
                     pass
 
+    def test_returns_when_cache_result_in_memory_is_false_sync_flow(self):
+        @flow(cache_result_in_memory=False)
+        def my_flow():
+            return 42
+
+        assert my_flow() == 42
+
+    async def test_returns_when_cache_result_in_memory_is_false_async_flow(self):
+        @flow(cache_result_in_memory=False)
+        async def my_flow():
+            return 42
+
+        assert await my_flow() == 42
+
+    def test_raises_correct_error_when_cache_result_in_memory_is_false_sync_flow(self):
+        @flow(cache_result_in_memory=False)
+        def my_flow():
+            raise ValueError("Test")
+
+        with pytest.raises(ValueError, match="Test"):
+            my_flow()
+
+    async def test_raises_correct_error_when_cache_result_in_memory_is_false_async_flow(
+        self,
+    ):
+        @flow(cache_result_in_memory=False)
+        async def my_flow():
+            raise ValueError("Test")
+
+        with pytest.raises(ValueError, match="Test"):
+            await my_flow()
+
 
 class TestSubflowCalls:
     async def test_subflow_call_with_no_tasks(self):
diff --git a/tests/test_tasks.py b/tests/test_tasks.py
index 52376437ea2c..a6eca829e104 100644
--- a/tests/test_tasks.py
+++ b/tests/test_tasks.py
@@ -487,6 +487,38 @@ class Foo:
                 def bar():
                     pass
 
+    def test_returns_when_cache_result_in_memory_is_false_sync_task(self):
+        @task(cache_result_in_memory=False)
+        def my_task():
+            return 42
+
+        assert my_task() == 42
+
+    async def test_returns_when_cache_result_in_memory_is_false_async_task(self):
+        @task(cache_result_in_memory=False)
+        async def my_task():
+            return 42
+
+        assert await my_task() == 42
+
+    def test_raises_correct_error_when_cache_result_in_memory_is_false_sync_task(self):
+        @task(cache_result_in_memory=False)
+        def my_task():
+            raise ValueError("Test")
+
+        with pytest.raises(ValueError, match="Test"):
+            my_task()
+
+    async def test_raises_correct_error_when_cache_result_in_memory_is_false_async_task(
+        self,
+    ):
+        @task(cache_result_in_memory=False)
+        async def my_task():
+            raise ValueError("Test")
+
+        with pytest.raises(ValueError, match="Test"):
+            await my_task()
+
 
 class TestTaskRun:
     async def test_sync_task_run_inside_sync_flow(self):


Code review comments:
diff_hunk: @@ -95,6 +97,7 @@ class FlowRunEngine(Generic[P, R]):
     flow_run_id: Optional[UUID] = None
     logger: logging.Logger = field(default_factory=lambda: get_logger("engine"))
     wait_for: Optional[Iterable[PrefectFuture]] = None
+    _result: Union[R, Exception, Type[NotSet]] = NotSet
path: src/prefect/flow_engine.py
text: @user1:
Everything in the PR looks awesome except I really want this to not have the word "result" in it haha.  `_flow_return_value`, I'm not sure....

@author:
That word is a tad overloaded. I'll update to `_return_value` and add a comment for clarity.

@user1:
Woot!