Repository: PrefectHQ/prefect
Issue category: Performance Issue

Task input persisted leading to memory not being released (same for output).
### First check

- [X] I added a descriptive title to this issue.
- [X] I used the GitHub search to find a similar request and didn't find it.
- [X] I searched the Prefect documentation for this feature.

### Prefect Version

2.x

### Describe the current behavior

Hi,

I really like the design of Prefect. To me Prefect is close to perfect. There is just the issue that **passing large sets of data to and from tasks quickly eats up all memory**. This eaten up memory cannot be released, or at least, I could not find any option to do so. Actually, being able to pass data into tasks is one of the advertised advantages of Prefect over Airflow:

```Python
from prefect import flow, task
import os
import psutil
import sys
import gc


@task(persist_result=False, cache_result_in_memory=False)  # <----- Remove this line, and the memory is released -----
def my_task(df):
    pass


@flow
def my_sub_flow_1():

    print(f"Memory before: {psutil.Process(os.getpid()).memory_info()[0] / float(1024*1024)}MiB")
    df = bytearray(1024*1024*1024)  # 1024MiB of memory

    my_task(df)

    print(f'{sys.getrefcount(df)} references to df')
    del df  # release memory
    gc.collect()  # garbage collection not needed, just be certain
    print(f"Memory after: {psutil.Process(os.getpid()).memory_info()[0] / float(1024*1024)}MiB")


my_sub_flow_1()
```

output is:

```
Memory before: 163MiB
4 references to df
Memory after: 1187MiB
```

then removing the `@task` decorator, all memory gets released by the `del` and everything works fine. Of course, this is a minimal example. In real live, data engineers like me want to run flows with lots of sub_flows passing large pandas data frames around without running out of memory.


### Describe the proposed behavior

It would be a really great enhancement if we could run sub_flows like the one below with large dataframes without running out of memory quickly once having more subflows or tasks.

```Python
@flow(empty_cache_after_flowrun=True):  # <----- something like this here
def my_sub_flow_1():
    df=task_load_data1()
    df=task_modify_data(df)
    task_upload_data(df, release_)  
	
    prefect.release_task_cache  # <----- or something like this here
   
    df=task_load_data2()
    df=task_modify_data(df)
    task_upload_data(df)   	
```

### Example Use

see above. Currently, all big objects passed from one task to another need to go via disk or any other store to be able to release memory. This somehow defeats the great idea of Prefect.

### Additional context

Debugging showed that the class `class UnpersistedResult(BaseResult)` persists the _outputs_ of task '_forever_' by design in `self._cache`. Surprisingly _inputs_ are also persisted somewhere. I spend quite some time debugging all the options, reading documentation and googling. In fact, the documentation is to me misleading as it suggests that things can be non-persistent. persist_result=False, cache_result_in_memory=False are not of help either, especially not to the problem of input persistence.

I also tried to modify the source to release the `self._cache` after it has been read to contribute, but sadly, I did not find a non destructive way to solve the issue to help also others that might be affected. Actually, I am wondering that I am the only one having this issue.

Ps.: Other options I found are not using tasks at all so that no data is passed around. I also experimented with using a wrapper classes to encapsulate each object passed around and then deleting the large object from the class (leaving the class itself in memory), but that somehow defeats the simplicity of prefect.




Functions to edit:
src/prefect/_internal/concurrency/calls.py:get_current_call
src/prefect/_internal/concurrency/calls.py:set_current_call
src/prefect/_internal/concurrency/cancellation.py:AsyncCancelScope.__exit__
Functions added: src/prefect/_internal/concurrency/calls.py:Future._invoke_callbacks

Patch:
diff --git a/src/prefect/_internal/concurrency/calls.py b/src/prefect/_internal/concurrency/calls.py
index 13e9e9bbfc61..6ca4dc45567e 100644
--- a/src/prefect/_internal/concurrency/calls.py
+++ b/src/prefect/_internal/concurrency/calls.py
@@ -10,6 +10,7 @@
 import dataclasses
 import inspect
 import threading
+import weakref
 from concurrent.futures._base import (
     CANCELLED,
     CANCELLED_AND_NOTIFIED,
@@ -32,21 +33,30 @@
 T = TypeVar("T")
 P = ParamSpec("P")
 
-
-# Tracks the current call being executed
-current_call: contextvars.ContextVar["Call"] = contextvars.ContextVar("current_call")
+# Tracks the current call being executed. Note that storing the `Call`
+# object for an async call directly in the contextvar appears to create a
+# memory leak, despite the fact that we `reset` when leaving the context
+# that sets this contextvar. A weakref avoids the leak and works because a)
+# we already have strong references to the `Call` objects in other places
+# and b) this is used for performance optimizations where we have fallback
+# behavior if this weakref is garbage collected. A fix for issue #10952.
+current_call: contextvars.ContextVar["weakref.ref[Call]"] = contextvars.ContextVar(
+    "current_call"
+)
 
 # Create a strong reference to tasks to prevent destruction during execution errors
 _ASYNC_TASK_REFS = set()
 
 
 def get_current_call() -> Optional["Call"]:
-    return current_call.get(None)
+    call_ref = current_call.get(None)
+    if call_ref:
+        return call_ref()
 
 
 @contextlib.contextmanager
 def set_current_call(call: "Call"):
-    token = current_call.set(call)
+    token = current_call.set(weakref.ref(call))
     try:
         yield
     finally:
@@ -181,6 +191,29 @@ def result(self, timeout=None):
             # Break a reference cycle with the exception in self._exception
             self = None
 
+    def _invoke_callbacks(self):
+        """
+        Invoke our done callbacks and clean up cancel scopes and cancel
+        callbacks. Fixes a memory leak that hung on to Call objects,
+        preventing garbage collection of Futures.
+
+        A fix for #10952.
+        """
+        if self._done_callbacks:
+            done_callbacks = self._done_callbacks[:]
+            self._done_callbacks[:] = []
+
+            for callback in done_callbacks:
+                try:
+                    callback(self)
+                except Exception:
+                    logger.exception("exception calling callback for %r", self)
+
+        self._cancel_callbacks = []
+        if self._cancel_scope:
+            self._cancel_scope._callbacks = []
+            self._cancel_scope = None
+
 
 @dataclasses.dataclass
 class Call(Generic[T]):
diff --git a/src/prefect/_internal/concurrency/cancellation.py b/src/prefect/_internal/concurrency/cancellation.py
index 8dcb4e6519de..25c2c2b5ad9b 100644
--- a/src/prefect/_internal/concurrency/cancellation.py
+++ b/src/prefect/_internal/concurrency/cancellation.py
@@ -270,6 +270,11 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             # Mark as cancelled
             self.cancel(throw=False)
 
+        # TODO: Can we also delete the scope?
+        # We have to exit this scope to prevent leaking memory. A fix for
+        # issue #10952.
+        self._anyio_scope.__exit__(exc_type, exc_val, exc_tb)
+
         super().__exit__(exc_type, exc_val, exc_tb)
 
         if self.cancelled() and exc_type is not CancelledError: