Efficient scalable shuffle - P2P shuffle extension

[fjetter]

Motivation

Shuffles are an integral part of many distributed data manipulation algorithms. Common DataFrame operations relying on shuffling include sort, merge, set_index, or various groupby operations (e.g. groupby().apply(), groupby(split_out>1)) whereas the most stereotypical array workload is the rechunk. There are many other applications for an efficient shuffle implementation which justifies taking a dedicated approach to solve this issue.

Shuffling is a poor fit for centralized graph-based scheduling, since the graph is all-to-all (naive O(N²), dask O(N logN); in size where N is the number of partitions), yet the core logic of a shuffle is so simple, it benefits little from centralized coordination, while suffering significant overhead from it. With task-based shuffles, the amount of data we can shuffle effectively (before workers run out of memory, or the scheduler crashes or bottlenecks) is severely limited. Allowing workers to autonomously exchange data with their peers and manage disk and memory usage in a more granular way allows us to push that limit significantly higher.

See https://coiled.io/blog/better-shuffling-in-dask-a-proof-of-concept/ for more background.

This issue tracks the current implementation progress and highlights various milestones. We intend to update the top-level description of this issue continuously such that this issue can serve as an always up-to-date overview of the current efforts.

Goals

• Can reliably shuffle orders-of-magnitude larger datasets (in total size and number of partitions) than the current task-based shuffle
• Can shuffle larger-than-memory datasets by spilling to disk
• Constant, predictable memory footprint per worker, which scales linearly with partition size, not total number of partitions
• Just works, without users needing to tune parameters (buffer sizes, etc.)
• Graceful restarting when possible, and quick failure when not
• All state is cleaned up on success, failure, or cancellation
• Shuffle performance is IO-bound (network, disk)
• Resilience to worker failures via restart of computation

Roadmap

1 - Foundations and dask.DataFrame ✅

The implementation effort so far focused on creating a stable foundation for the things to come and is deriving from the early prototype. This stage mostly focused on a consistent concurrency model that supports off-band, direct peer to peer communication between workers and integrates well with the existing task based scheduling logic.

This was developed at the example of a DataFrame based shuffle and we consider this now ready to use!

For detailed instructions, known issues and feedback, please see #7509. We encourage all users of dask.DataFrame to try this and report with feedback.

2 - dask.Array rechunking

The new shuffle extension is currently build to handle pandas DataFrames and is using pyarrow behind the scenes. It's architecture is built with generic types in mind and will be suited just as well for array workloads. One of the most popular many-to-many problems is the array rechunking which we will implement next using this extension.

Basic functionality is being set up in #7534

This approach already provides constant time array rechunking but sometimes falls short in terms of walltime performance compared to old style task based shuffling.

3 - Misc

This next stage is not as refined as the intial ones. There are many smaller to medium sized issues that will either expand adoption of the P2P algorithm or make it run faster and smoother. This section will become more refined over time.

• Enable P2P shuffling by default dask#9991
• Optional Disk for P2P shuffling #7572
• P2P shuffling and queuing combined may cause high memory usage with dask.dataframe.merge #7496
• string[pyarrow] dtype does not roundtrip in P2P shuffling #7420
• Asynchronous disk IO for shuffle
• Hook up instrumentation to primary dashboard (e.g. task stream)
• Extend usage of new shuffle algorithm to other APIs like map_overlap
• Add support for Bags
• Performance tuning
• ...

[dcherian]

Hello, this is very cool to see, and I'm excited to see it being applied to arrays!

In my Xarray-centric experience the shuffle is usually expressed as a out-of-order indexing of a dask.Array (for e.g. index out all elements belonging to a group in a groupby operation; and present that to the user as a new array). See this comment and below for some discussion. Currently this just results in an explosion of tasks, and is only workable for a limited set of problems.

For built-in reductions, we can use a dask.array.reduction style approach (see flox)but we still need this kind of shuffling for GroupBy.map much like a dataframe.

Is it feasible for the indexing/slicing code to use the shuffle for such cases.

[fjetter]

Is it feasible for the indexing/slicing code to use the shuffle for such cases.

So far, we implemented this only for the API call arr.rechunk and haven't considered more generic slicing. However, the biggest lift was not the rechunk API but rather figuring out how to chop up the array and piece it back together properly but that's basically slicing :)

We are currently generating a plan (this all is subject to change) that calculates the "target chunk" for every slice of an input chunk

distributed/distributed/shuffle/_rechunk.py

Lines 134 to 162 in f4328bb

	def rechunk_slicing(
	old: ChunkedAxes, new: ChunkedAxes
	) -> dict[NIndex, list[tuple[ShardID, NSlice]]]:
	"""Calculate how to slice the old chunks to create the new chunks
	Returns
	-------
	Mapping of each old chunk to a list of tuples defining where each slice
	of the old chunk belongs. Each tuple consists of the index
	of the new chunk, the index of the slice within the composition of slices
	creating the new chunk, and the slice to be applied to the old chunk.
	"""
	from dask.array.rechunk import intersect_chunks
	ndim = len(old)
	intersections = intersect_chunks(old, new)
	new_indices = product(*(range(len(c)) for c in new))
	slicing = defaultdict(list)
	for new_index, new_chunk in zip(new_indices, intersections):
	sub_shape = [len({slice[dim][0] for slice in new_chunk}) for dim in range(ndim)]
	shard_indices = product(*(range(dim) for dim in sub_shape))
	for shard_index, sliced_chunk in zip(shard_indices, new_chunk):
	old_index, nslice = zip(*sliced_chunk)
	slicing[old_index].append((ShardID(new_index, shard_index), nslice))
	return slicing

effectively this is just a mapping like old_chunk_ix -> list of (new_chunk_ix, slice), i.e. if you can express an operating like this we can use this P2P backend.

The only difference between rechunk and generic slicing (IIUC) is that rechunk typically slices contiguously. The extension currently doesn't care (or optimize for) this so I am pretty confident we can express generic indexing as well using this. There are very likely cases where the old style indexing is still better so a challenge will be figuring out which algorithm to choose.

@dcherian is arr.rechunk also useful for you or will you only be able to try this once we implement indexing/slicing?

[dcherian]

how to chop up the array and piece it back together properly but that's basically slicing :)

That's what I was thinking ;)

@dcherian is arr.rechunk also useful for you or will you only be able to try this once we implement indexing/slicing?

Yes it is useful but I'll have to make up a use-case. It has been such a problem for so long, we try to avoid it as much as possible, or use rechunker to create a second copy of the dataset.

The other common place this bites us is arr.reshape (or stack with xarray) Example: dask/dask#5544 (EDIT: there are so many issues pinging this one!)

cc @TomNicholas @jbusecke @rabernat

[mrocklin]

I think that the right approach here is to start to generalize out the shuffle system to the point where a sharp and adventurous developer, like @dcherian , could try implementing his own operation. This will also likely come up with map_overlap with the imaging people.

My sense is that there is likely some abstraction that could be designed similar to Aggregation but more complex that could enable folks outside of @hendrikmakait and @fjetter to engage effectively here.

[TomNicholas]

This looks very exciting! I have a test problem I would like to try it out on. What's the easiest way for me to do that? Use Coiled with distributed >2023.3.0?

[hendrikmakait]

@TomNicholas: Yes, using distributed>=2023.3.1 with a Coiled cluster should be the easiest way for you to go about that, though I suggest using distributed>=2023.5.0. See the instructions on the Pangeo Discourse (https://discourse.pangeo.io/t/rechunking-large-data-at-constant-memory-in-dask-experimental/3266/2) on how to make sure you're using p2p rechunking. Just note that disabling fusion is not necessary anymore if you're using distributed>=2023.5.0.

On a related note, I'm currently improving the P2P rechunking implementation (distributed#7897). I still have to create a range of benchmarks to test these changes and a decision heuristic between tasks and p2p rechunking (dask/dask#10226), so I might add your example to the list.

Please let me know if you have more common rechunking workloads you care about. Creating an issue on https://github.com/coiled/benchmarks/issues might be the best way to do that.

[dcherian]

I'm not sure why P2P rechunking helps with Tom's test problem. Can someone explain that?

ms = MemorySampler()

ds = xr.Dataset(
    dict(
        anom_u=(["time", "face", "j", "i"], da.random.random((5000, 1, 987, 1920), chunks=(10, 1, -1, -1))),
        anom_v=(["time", "face", "j", "i"], da.random.random((5000, 1, 987, 1920), chunks=(10, 1, -1, -1))),
    )
)

quad = ds**2
quad["uv"] = ds.anom_u * ds.anom_v
mean = quad.mean("time")

with ms.sample():
    mean.compute()

[hendrikmakait]

@dcherian: I don't think it does. I've tested the problem recently, and it appears to lack rechunking of any sort. Hence P2P rechunking won't help. @TomNicholas: Do you have more information on this? Where would you have expected it to help?

[TomNicholas]

Okay hmmm - in that case there must be some other change that's caused me to see a totally different performance between running an old version of distributed locally and the new version with P2P on Coiled. I'll try both on Coiled instead now to eliminate possible causes.

[TomNicholas]

Okay I take it back, 'array.rechunk.method': 'tasks' and 'array.rechunk.method': 'p2p' behaved indistinguishably on this problem when I tried both with distributed main via Coiled.

[fjetter]

Okay I take it back, 'array.rechunk.method': 'tasks' and 'array.rechunk.method': 'p2p' behaved indistinguishably on this problem when I tried both with distributed main via Coiled.

If you ran this on Coiled, can you provide us with the cluster ID?