dask: Regression: poor memory management in threads and sync schedulers

Also cc @mrocklin @jrbourbeau . IMHO this should be treated as very high priority as it is something that is very likely to hit in the face a newbie user when he plays around with dask.array or dask.dataframe for the very first time.

Got it! See my edits above for the path to debugging, here’s the current diff:

diff --git a/dask/local.py b/dask/local.py
index 9f035ef2..913a4706 100644
--- a/dask/local.py
+++ b/dask/local.py
@@ -463,11 +463,14 @@ def get_async(
                     ntasks = nready
                     chunksize = -(ntasks // -num_workers)
                 else:
-                    ntasks = min(nready, chunksize * num_workers)
+                    used_workers = -(len(state["running"]) // -chunksize)
+                    avail_workers = max(num_workers - used_workers, 0)
+                    ntasks = min(nready, chunksize * avail_workers)
+                    print(nready, ntasks, num_workers, len(state["running"]), avail_workers, chunksize)

                 # Prep all ready tasks for submission
                 args = []
-                for i, key in zip(range(ntasks), state["ready"]):
+                for key in state["ready"][nready-ntasks:nready]:
                     # Notify task is running
                     state["running"].add(key)
                     for f in pretask_cbs:
@@ -487,7 +490,7 @@ def get_async(
                             pack_exception,
                         )
                     )
-                del state["ready"][:ntasks]
+                del state["ready"][nready-ntasks:nready]

                 # Batch submit
                 for i in range(-(len(args) // -chunksize)):
@@ -499,7 +502,11 @@ def get_async(

             # Main loop, wait on tasks to finish, insert new ones
             while state["waiting"] or state["ready"] or state["running"]:
+                print(len(state["waiting"]), len(state["ready"]), len(state["running"]))
                 fire_tasks(chunksize)
+                if not len(state["running"]):
+                    print("Nothing is running")
+                    continue
                 for key, res_info, failed in queue_get(queue).result():
                     if failed:
                         exc, tb = loads(res_info)

I can make a pull request later today if this seems reasonable and we can discuss pros/cons of this approach (and see what tests fail). Here’s the profile as proof:

So in summary:

1. The batch logic was submitting all tasks regardless of how many workers were available.
2. The reverse ordering of the ready tasks is needed otherwise we run into the situation in my last comment where all the data generation tasks are executed before the reduction/dependent tasks which fills up the memory. This is likely similar to the distributed memory back pressure discussions that have been going on elsewhere on this repository.

I wanted to see if @sjperkins was on to something with the reverse changes so I made this script:

$ cat test_dask_memory_regression.py
from dask.diagnostics import ResourceProfiler, Profiler, CacheProfiler, visualize
import dask.array as da
from datetime import datetime

if __name__ == "__main__":
    with Profiler() as prof, ResourceProfiler(dt=0.5) as rprof, CacheProfiler() as cprof:
        d = da.random.random((5000, 250000), chunks=(1, -1)).sum()
        d.compute(optimize_graph=False)
    visualize(
            [rprof, prof, cprof],
            file_path=f'/tmp/profile_{datetime.utcnow():%Y%m%d_%H%M%S}.html',
            show=False
            )

Note I used a smaller number of rows to make the computation go faster for easier testing. Ran it locally and here’s a screenshot of the results (github doesn’t let me attach HTML files, I can put them somewhere if desired).

Version 2021.3.1 - profile_20210813_132924.html

Peak: ~220MB

Version 2021.6.2 - profile_20210813_131334.html

Peak: ~10GB

I then hacked my 2021.6.2 install to undo the reverse changes (without knowing why the changes were made or what purpose they serve) and did NOT see any differences in the profile.

Edit: I know I used 2021.6.2 instead of 2021.4.0. I just tested with 2021.4.0 and same result as 2021.6.2 (same appearance in the profile).

Thanks all for bumping this issue.

@quasiben can you poke @jakirkham about this when he gets back? This seems to be a priority issue that arose from his work on swapping in concurrent.futures.

Also cc @eriknw in case he is interested (although this may not be ordering related)