Distributed memory management plugin

[mrocklin]

It would be useful to have a plugin that tracked and managed distributed memory.

Currently results are copied around to various workers as necessary for computation. Our current policy is to avoid movement when possible and keep duplicated data around until it is no longer necessary with the thought that data that ends up being duplicated organically is likely in high demand. This can currently be manually controlled by two synchronous operations:

1. rebalance(futures): takes a set of futures and tries to balance them around a set of workers so that all workers store roughly the same number of bytes
2. replicate(futures): Change the number of replicas for results. This is commonly used to spread around commonly needed intermediate results efficiently (this performs a tree-scatter) or to clean up replicated data (n=1) and increase free space

Currently people do this by hand. This is error prone for a number of reasons and generally rare among non-expert users. It would be interesting to consider a scheduler plugin that watched current memory usage and responded dynamically by moving data around in the background. Operations like rebalance, and replicate might then set objectives for this dynamic system to accomplish rather than perform the operation explicitly.

There are a few objectives when moving memory around:

1. Respect replication constraints set by users (came up in Improve recovery time in worker failure scenarios #3184)
2. Redistribute data evenly across workers
3. Free up memory, particularly when we have highly replicated data
4. But still respect replication, particularly when it is cheap or frequently used

I think that this would improve memory usage when we're constrained. I think that it shouldn't impact task scheduling much if we act in a safe manner like verifying that copied data has been moved to one location before deleting it from its origin, etc.. It may be challenging to do this efficiently in linear time, but we could also punt if we notice that the scheduler is otherwise busy.

[sjperkins]

I have a scenario where the align functionality would be useful. Briefly, I compute a function taking many, large, chunked inputs and producing a single Chi-squared value output. Chi-squared's of each individual chunk are summed together to form the final Chi-squared value. All my inputs are stored as dask arrays in an xarray Dataset.

I think that I'm observing a lot of data transfer of the large input chunks (and spilling to disk when workers start running out of memory), mostly because the inputs to each individual function are not automatically aligned on each worker.

As mentioned above, I'm considering persisting each input chunk by hand to avoid data movement and spillage.

I had a couple of thoughts on this:

1. xarray's labelled dimensions provide some information on how to align chunks. e.g. if two arrays both share a dimension dim0, then the dimension provides a trivial alignment for vectorised operations on chunks (but not for transpose-like operations).

2. Persisting/pinning operations to workers is a graph partitioning mechanism. I wonder if it this could be abstracted further by assigning tasks to virtual partitions. Then, tasks within a partition would be constrained to execute on a particular thread/process (in the standalone dask case) or a node (in the distributed dask case).

I'm not sure whether this functionality should end up in dask, dask distributed or xarray. IIRC xarray tries to offload as much decision-making to dask as possible so I'm not sure how appropriate something like a xarray.Dataset.persist(align='vector') would be. I can see graph partitioning being useful to avoid data movement in the dask multiprocessing and distributed cases where you want to avoid transferring large amounts of data by IPC/network.

[mrocklin]

This has come up a bit recently. I've updated the original post (removing align).

cc @fjetter @jacobtomlinson @quasiben @cjnolet

[mrocklin]

Closing in favor of #4982