Fix task linking

[desertaxle]

Closes #15391

In #14651 we started tracking task run states with a WeakValueDictionary to avoid memory bloat in flow with many tasks. Using a WeakValueDictionary can lead task run states to be garbage collected too soon which remove the information necessary to consistently show data dependencies between tasks.

This PR goes back to using a dict to track task run states, but drops the return value from the user's code to avoid memory bloat.

Using the following script, we can see that memory usage stays roughly consistent (or even improves a bit) before and after this change:

from memory_profiler import profile

from prefect import flow, task


@task
def with_task():
    return [{"abc": "123"} for _ in range(10_000)]

@flow
@profile
def some_flow(n: int = 100):
    for _ in range(n):
        with_task()
    print("DONE")


@profile
def main():
    some_flow(n=100)


if __name__ == "__main__":
    main()

Here are the results:

Before
n=100 w/ WeakValueDictionary

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
    10    133.5 MiB    133.5 MiB           1   @flow
    11                                         @profile
    12                                         def some_flow(n: int = 100):
    13    138.6 MiB  -1302.8 MiB         101       for _ in range(n):
    14    138.6 MiB  -1297.7 MiB         100           with_task()
    15    125.7 MiB    -12.9 MiB           1       print("DONE")

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
    18    129.2 MiB    129.2 MiB           1   @profile
    19                                         def main():
    20    125.8 MiB     -3.3 MiB           1       some_flow(n=100)

n=1000 w/ WeakValueDictionary

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
    10    133.0 MiB    133.0 MiB           1   @flow
    11                                         @profile
    12                                         def some_flow(n: int = 100):
    13    141.2 MiB -19463.3 MiB        1001       for _ in range(n):
    14    141.2 MiB -19455.5 MiB        1000           with_task()
    15    120.8 MiB    -20.5 MiB           1       print("DONE")

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
    18    128.5 MiB    128.5 MiB           1   @profile
    19                                         def main():
    20    121.1 MiB     -7.5 MiB           1       some_flow(n=1000)

After
n=100 w/ dict

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
    10    133.5 MiB    133.5 MiB           1   @flow
    11                                         @profile
    12                                         def some_flow(n: int = 100):
    13    140.5 MiB  -5924.6 MiB         101       for _ in range(n):
    14    140.5 MiB  -5917.7 MiB         100           with_task()
    15     67.2 MiB    -73.3 MiB           1       print("DONE")

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
    18    129.5 MiB    129.5 MiB           1   @profile
    19                                         def main():
    20     69.4 MiB    -60.1 MiB           1       some_flow(n=100)

n=1000 w/ dict

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
    10    133.9 MiB    133.9 MiB           1   @flow
    11                                         @profile
    12                                         def some_flow(n: int = 100):
    13    133.9 MiB -42578.9 MiB        1001       for _ in range(n):
    14    130.2 MiB -48355.7 MiB        1000           with_task()
    15     64.5 MiB    -69.4 MiB           1       print("DONE")

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
    18    131.6 MiB    131.6 MiB           1   @profile
    19                                         def main():
    20     66.9 MiB    -64.6 MiB           1       some_flow(n=1000)

Alternatives considered

I explored using weakref.finalizer to try to intelligently remove tracked task run states when the associate return value is garbage collection, but weak references cannot be created for all objects (like int or str) so unfortunately, that approach didn't work.

[codspeed-hq]

CodSpeed Performance Report

Merging #15418 will not alter performance

_{Comparing fix-task-linkingwq (c8192e0) with main (9d78f77)}

Summary

✅ 3 untouched benchmarks