Memory leak when submitting futures

[sofsse]

I'm new to dask and trying to parallel some calculations on HPC cluster using SLURMCluster. Apparently, I faced a memory leak similar to that discussed in #5960. There is a reproducer of code I run (the only difference is that in real workflow I submit scipy.optimize.minimize function with some args, then get result and save it to file).

from memory_profiler import profile

from dask_jobqueue import SLURMCluster
from dask.distributed import Client, as_completed

@profile
def big_arr():
    return [5]*(10**5)

@profile
def go():

    extra_directives = [f"-J check_mem", f"-o daskout", f"-e daskerr"]
    cluster = SLURMCluster(
        queue='mpi', # your queue
        account="tguskov", # your account
        cores=1,
        processes=1,
        memory="1GB",
        job_name='check_mem',
        walltime='06:00:00',
        job_extra_directives=extra_directives, #_directives
    )
    client = Client(cluster)

    client.cluster.scale(5)
    futures = []
    for i in range(10):
        futures.append(client.submit(big_arr))

    for future in as_completed(futures):
        sol = future.result()
        del future, sol

    client.cluster.scale(0)

    del futures, client, cluster

for i in range(4):
    go()

There is also a graph of mprof memory usage by the script above:

It's clear for me from that graph that smth's wrong with submitting a new future (10 peaks corresponding to 10 submitted futures on each iteration). CreatingSLURMCluster outside of the function go() doesn't make any difference. The lines with aggressive deletion are also doesn't really affect the memory usage.

Here I used dask 2022.1.1 and distributed 2022.1.1 according to temporary solution from #5960, but that didn't work out for me either. The same situation with the latest(2024.8.0) version of dask and distributed.

The memory leak is around a few GB per iteration in my original case where I try to submit around ~500 futures on 100 workers.

I'm also providing my bash script just in case:

#!/bin/sh
#SBATCH -p mpi
#SBATCH -N 1
#SBATCH -n 1
#SBATCH --mem=1GB 
#SBATCH -o output
#SBATCH -e errors

ls -l|grep mprofile|xargs rm -f --

conda activate dask
mprof run --multiprocess --python python check.py # --include-children

Env

dask==2022.01.1/2024.8.0
distributed==2022.01.1/2024.8.0
dask_jobqueue==0.8.5

[jacobtomlinson]

I'm curious whether this is specific to SLURM? Are you able to repdocue this memory leak with LocalCluster? If not then we may want to transfer this issue to dask-jobqueue.

[sofsse]

Sorry for late response, just replaced SlurmCluster with LocalCluster and observed the same behaviour

[guillaumeeb]

Hi @sofsse, could you precise what do you call a memory leak here? Are these the pikes at each submission, or the slow growth of overall memory usage?

[sofsse]

Hi @guillaumeeb, the overall memory usage

[guillaumeeb]

Thanks! Sorry for coming back late.

Also, you are checking the memory consumption on the main script/process right? The one creating the Cluster and Client? If so this is a memory growth from Client or Scheduler side. And on each iteration you are creating a new Scheduler.

Did you try to do some GC cleaning? This example is a bit small on modern system, and doesn't show for me enough proof of a real leak, it could just come from Python GC or the OS.

Playing a bit with your code also rings some bells:

• Submitting the same pure function in a loop is just a one op, it will execute only once.
• What takes time and probably memory is the gathering back of the data, that involves each time some serialization of your array back and forth.
• Since your measuring Scheduler memory growth, there might be some leakage here, or just not freed memory by Python.