dask clusters with context manager

Question

Consider a simple workflow like this:

from dask.distributed import Client
import time

with Client() as client:
    futs = client.map(time.sleep, list(range(10)))

The above code will submit and almost immediately cancel the futures since the context manager will close. It's possible to keep the context manager open until tasks are completed with client.gather, however that will block further execution in the current process.

I am interested in submitting tasks to multiple clusters (e.g. local and distributed) within the same process, ideally without blocking the current process. It's straightforward to do with explicit definition of different clients and clusters, but is it also possible with context managers (one for each unique client/cluster)?

It might sound like a bit of an anti-pattern, but maybe there is way to close the cluster only after computations all futures run. I tried fire_and_forget and also tried passing shutdown_on_close=False, but that doesn't seem to be implemented.

I think we can agree it's an antipattern - why do you want specifically to use with contexts? — mdurant
The issue is that I'm operating on a SLURM cluster, where if I don't close the cluster properly, e.g. because my main process is killed due to out of memory error, I have jobs (to launch dask workers) left hanging in the queue which need to cleared manually... it's a minor nuisance, so I thought maybe it can be fixed with context manager, but I understand that it's an antipattern now. — SultanOrazbayev
Perhaps use weakref.finalize or atexit? Also, I think you can attach a callback to your future, to have something happen when its done. — mdurant

elukem elukem · Accepted Answer · 2021-05-08T08:38:03

For some Dask cluster/scheduler types, such as the dask-cloudprovider ECSCluster, the approach described above using the with block and shutdown_on_close=False would work fine.

Both ECSCluster and SLURMCluster are derived from SpecCluster. However, ECSCluster passes its **kwargs (including shutdown_on_close) down to the SpecCluster constructor via this call:

super().__init__(**kwargs)

(see the ECSCluster code here)

SLURMCluster does not: it calls the JobQueueCluster constructor which in turn instantiates SpecCluster with only a subset of its parameters:

super().__init__(
    scheduler=scheduler,
    worker=worker,
    loop=loop,
    security=security,
    silence_logs=silence_logs,
    asynchronous=asynchronous,
    name=name,
)

See the JobQueueCluster code here

Therefore SLURMCluster/JobQueueCluster is ignoring shutdown_on_close (and other optional parameters). Looks like an update to JobQueueCluster would be required for your use case.

dask clusters with context manager

1 Answers