Repository: PrefectHQ/prefect
Pull request: Fix three memory leaks
<!-- 
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!
-->

Fix ~two~ three memory leaks. Fixes #10952.

<!-- Include an overview here -->

### 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. -->

- [ ] This pull request references any related issue by including "closes `<link to issue>`"
	- 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.
- [ ] This pull request includes tests or only affects documentation.
- [ ] This pull request includes a label categorizing the change e.g. `maintenance`, `fix`, `feature`, `enhancement`, `docs`.
  <!-- If you do not have permission to add a label, a maintainer will add one for you -->

For documentation changes:

- [ ] This pull request includes redirect settings in `netlify.toml` for files that are removed or renamed.

For new functions or classes in the Python SDK:

- [ ] This pull request includes helpful docstrings.
- [ ] If a new Python file was added, this pull request contains a stub page in the Python SDK docs and an entry in `mkdocs.yml` navigation.

Problem statement:
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.

Hints from issue 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_
Hi @sibbiii, I just merged a fix that will be in the next release, but if you could give it a try on `prefect@main`, that would be amazing!  I was able to address the _parameter_ issue as you described in the original writeup, but the _result_ issue is a bit more invasive of an issue.  I'll bring this back to the team to talk about how we might improve the performance/memory use around result caching between flow runs.
> Hi @sibbiii, I just merged a fix that will be in the next release, but if you could give it a try on prefect@main, that would be amazing!


Hi @chrisguidry,

I imported main head (2.14.4+12.g6d7acc78b9) and reran the example at the beginning of this issue. The output is still the same, that is, memory is _not_ released if the `@task` decorator is in place and is released if the decorator is removed.

Sebastian
Ah yes, @sibbiii, that is likely to still be the case during the flow run.  My fix corrected an issue where memory wasn't freed _between_ flow runs, but it wouldn't have addressed freeing memory _during_ a flow run.  Check out this example for what my change fixed:

```python
import gc
import os
import sys

import psutil
from prefect import flow, task


@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 task: {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 task: {psutil.Process(os.getpid()).memory_info()[0] / float(1024*1024)}MiB"
    )


if __name__ == "__main__":
    print(
        f"Memory before flow: {psutil.Process(os.getpid()).memory_info()[0] / float(1024*1024)}MiB"
    )

    my_sub_flow_1()

    gc.collect()  # garbage collection not needed, just be certain
    print(
        f"Memory after flow: {psutil.Process(os.getpid()).memory_info()[0] / float(1024*1024)}MiB"
    )
```

Before the fix:

```
Memory before flow: 80.4453125MiB
...
Memory before task: 136.1328125MiB
09:26:15.670 | INFO    | Flow run 'spiked-ostrich' - Created task run 'my_task-0' for task 'my_task'
09:26:15.671 | INFO    | Flow run 'spiked-ostrich' - Executing 'my_task-0' immediately...
09:26:16.353 | INFO    | Task run 'my_task-0' - Finished in state Completed()
4 references to df
Memory after task: 1162.3828125MiB
09:26:16.566 | INFO    | Flow run 'spiked-ostrich' - Finished in state Completed('All states completed.')
Memory after flow: 1163.3203125MiB

```

After the fix:

```
Memory before flow: 84.30078125MiB
...
Memory before task: 99.55078125MiB
09:21:21.617 | INFO    | Flow run 'uncovered-trogon' - Created task run 'my_task-0' for task 'my_task'
09:21:21.618 | INFO    | Flow run 'uncovered-trogon' - Executing 'my_task-0' immediately...
09:21:22.285 | INFO    | Task run 'my_task-0' - Finished in state Completed()
3 references to df
Memory after task: 1165.80078125MiB <---- what you're observing
09:21:22.531 | INFO    | Flow run 'uncovered-trogon' - Finished in state Completed('All states completed.')
Memory after flow: 141.91796875MiB <---- what I was able to fix in the first pass
```

I'm going to re-open this for further investigation.
Hi @chrisguidry,

I checked: you are right. It works when the flow is finished. 

By the way, now that the input issue is fixed, 
there is a _hack_ to release memory using something like this:

```
class SelfDestruct:

    class AlreadyDestructedException(Exception):
        pass

    class NoObjectStored:
        pass  # Singleton for defining that there is no object in the stored

    def __init__(self, obj):
        self._obj_store = obj

    def get_destruct(self, destruct=True):
        if self._obj_store is self.NoObjectStored:
            raise self.AlreadyDestructedException

        result = self._obj_store
        if destruct:
            self._obj_store = self.NoObjectStored
        return result
```

and then

```
@task()  
def my_task(df):
    return myHelper.SelfDestruct(df)

@flow
def my_sub_flow():
    df = ...
    result = my_task(df).get_destruct()
```

which actually works fine, that is, the memory is released. It just does not look like "clean code". Still I could actually add this logic to the `@task` and `@flow` decorators ...    

Other than this, I had no smart Idea so far how to fix the issue.

Commit message: Fix two memory leaks

Patch to review:
diff --git a/src/prefect/_internal/concurrency/calls.py b/src/prefect/_internal/concurrency/calls.py
index fbd74ba211e9..9d7e5b4c0f99 100644
--- a/src/prefect/_internal/concurrency/calls.py
+++ b/src/prefect/_internal/concurrency/calls.py
@@ -182,6 +182,22 @@ def result(self, timeout=None):
             # Break a reference cycle with the exception in self._exception
             self = None
 
+    def _invoke_callbacks(self):
+        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..e8c5392ac87e 100644
--- a/src/prefect/_internal/concurrency/cancellation.py
+++ b/src/prefect/_internal/concurrency/cancellation.py
@@ -270,6 +270,8 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             # Mark as cancelled
             self.cancel(throw=False)
 
+        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:


Code review comments:
diff_hunk: @@ -270,6 +270,8 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             # Mark as cancelled
             self.cancel(throw=False)
 
+        self._anyio_scope.__exit__(exc_type, exc_val, exc_tb)
path: src/prefect/_internal/concurrency/cancellation.py
text: @user1:
OHHHHH

@user1:
Can probably also `del self._anyio_scope` afterward too

@author:
I didn't have time to test that out but added a TODO and will get back to it.