Repository: PrefectHQ/prefect

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.




Discussion:
Hi @sibbiii, thanks for the well written issue and investigation. I can reproduce the behavior with your MRE. 

I think it is unlikely that the references are maintained as part of results as you're not returning anything from your function.

I think this likely has to do with the dependency tracking or as your said _input persistence_. `prefect` keeps tracks of inputs as a way to track/resolve dependencies between tasks. We already provide some mechanisms like `quote()` to opt out of dependency tracking. This is possibly a bug where a reference is maintained when it shouldn't be OR a an enhancement. Either way it seems like it makes sense to at least have an opt out.
Hi, 

> I think it is unlikely that the references are maintained as part of results as you're not returning anything from your function.

You are right, I checked, they are not maintained in UnpersistedResult.

Thereby invested some time to create another minimal example how to fill up memory by calling a task with some input data. 
When executing the code below, memory usage (e.g. open the task manager) increases with time.

<img src="https://github.com/PrefectHQ/prefect/assets/115080722/89e99b95-a7c2-4bf0-bbc4-604335c2aa5d" width="400">

```Python
from prefect import flow, task
import gc


@task  # <----- Remove this line, and the memory is released -----
def my_task(some_large_object):
    pass


@flow
def my_sub_flow():
    my_task(bytearray(1024 * 1024 * 300))  # 300MiB


for i in range(100):
    my_sub_flow()
    gc.collect()
```
I did not find out yet why the memory cannot be released if the task decorator is added.
Interestingly, calling a task in a loop (similar to the example below) does not fill up the memory.

For the output of a task, filling up the memory is easy to reproduce:
```Python
from prefect import flow, task


@task  # <----- Remove this line, and the memory is released -----
def my_task():
    return bytearray(1024*1024*300)  # 300MiB of memory


@flow
def my_sub_flow():
    for i in range(100):
        my_task()


my_sub_flow()
```
Here a reference to the data is kept in self._cache of UnpersistedResult.

Unfortunately, not being able to release the memory for inputs and outputs does result in memory issues when running flows with many tasks or lots of data being transferred to and from tasks. I hope there are other ways to keeping track of dependencies between tasks rather than "storing" all input data by keeping the reference, same for the outputs.
I think these issues highly impact long running flows/tasks. Where memory is slowly crippling in and got OutOfMemory.

My use case is http calls to a website with hundreds of thousands of sitemaps.


I've been investigating this memory leak and one issue I've found with the parameter case is that after garbage collection, the chain of referrers looks like:
* the task parameters
  * the `kwargs` dict on a `Call` object to `get_task_call_return_value`
    * an intermediate `hamt` object (a `frozenset` somewhere?)
      * a `contextvars.Context` object
        * the `kwargs` dict on another another `Call` for `QueueService.drain` of the `APILogWorker` service

I think what's happening here is that even when the flow and tasks have both finished running, there are lingering references to them in the service that streams logs back to the Prefect API, so I'm going to try to figure out how we can break those references.

One possible solution is that the `Call` object can probably drop its `args`/`kwargs` immediately after it has finished running, as they shouldn't be used after that.  Need to experiment some more....
For a task that returns a large value, it's also the `Call`, but this time it's the `Call.future` that's holding a long-term reference to the result of the call.  It's a bit trickier to know when we're done with the `future` of a `Call` since a caller may theoretically call for the `result` multiple times.  Looking to see if there's a way we can just drop the `Call` object itself when we're done with it...
In the case of task return values, at the end of a flow run there are three references:

1: a `Call` object's `.future` <-- this is something we may be able to address
2: an `UnpersistedResult` in the `SequentialTaskRunner._results` cache of `State` results
3: another copy of an `UnpersistedResult` referenced by the list of `State` results in the  `Future._result` of the future for `create_then_begin_flow_run`

That first one seems like something we may be able to address, but the second/third start to interfere with how our result processing works, which might need to be opened as a separate issue.  I do believe I can address the _parameter_ issue by breaking the `Call`'s reference to the task's `args`/`kwargs`, which should help with this issue about task _inputs_

Resolving patch:
diff --git a/src/prefect/_internal/concurrency/calls.py b/src/prefect/_internal/concurrency/calls.py
--- a/src/prefect/_internal/concurrency/calls.py
+++ b/src/prefect/_internal/concurrency/calls.py
@@ -253,7 +253,9 @@ def run(self) -> Optional[Awaitable[T]]:
                 # If an event loop is available, return a task to be awaited
                 # Note we must create a task for context variables to propagate
                 logger.debug(
-                    "Scheduling coroutine for call %r in running loop %r", self, loop
+                    "Scheduling coroutine for call %r in running loop %r",
+                    self,
+                    loop,
                 )
                 task = self.context.run(loop.create_task, self._run_async(coro))
 
@@ -312,7 +314,14 @@ def _run_sync(self):
         try:
             with set_current_call(self):
                 with self.future.enforce_sync_deadline() as cancel_scope:
-                    result = self.fn(*self.args, **self.kwargs)
+                    try:
+                        result = self.fn(*self.args, **self.kwargs)
+                    finally:
+                        # Forget this call's arguments in order to free up any memory
+                        # that may be referenced by them; after a call has happened,
+                        # there's no need to keep a reference to them
+                        self.args = None
+                        self.kwargs = None
 
             # Return the coroutine for async execution
             if inspect.isawaitable(result):
@@ -342,7 +351,14 @@ async def _run_async(self, coro):
         try:
             with set_current_call(self):
                 with self.future.enforce_async_deadline() as cancel_scope:
-                    result = await coro
+                    try:
+                        result = await coro
+                    finally:
+                        # Forget this call's arguments in order to free up any memory
+                        # that may be referenced by them; after a call has happened,
+                        # there's no need to keep a reference to them
+                        self.args = None
+                        self.kwargs = None
         except CancelledError:
             # Report cancellation
             if cancel_scope.timedout():
@@ -383,10 +399,15 @@ async def run_and_return_result():
 
     def __repr__(self) -> str:
         name = getattr(self.fn, "__name__", str(self.fn))
-        call_args = ", ".join(
-            [repr(arg) for arg in self.args]
-            + [f"{key}={repr(val)}" for key, val in self.kwargs.items()]
-        )
+
+        args, kwargs = self.args, self.kwargs
+        if args is None or kwargs is None:
+            call_args = "<dropped>"
+        else:
+            call_args = ", ".join(
+                [repr(arg) for arg in args]
+                + [f"{key}={repr(val)}" for key, val in kwargs.items()]
+            )
 
         # Enforce a maximum length
         if len(call_args) > 100:


Test patch: