Dask Working Notes - Posts by Matthew Rocklin (Coiled)

Faster Scheduling

2020-07-21T00:00:00+00:00

This post discusses Dask overhead costs for task scheduling, and then lays out a rough plan for acceleration.

This post is written for other maintainers, and often refers to internal details. It is not intended for broad readability.

How does this problem present?

When we submit large graphs there is a bit of a delay between us calling .compute() and work actually starting on the workers. In some cases, that delay can affect usability and performance.

Additionally, in far fewer cases, the gaps in between tasks can be an issue, especially if those tasks are very short and for some reason can not be made longer.

Who cares?

First, this is a problem that affects about 1-5% of Dask users. These are people who want to process millions of tasks relatively quickly. Let’s list a few use cases:

Xarray/Pangeo workloads at the 10-100TB scale
NVIDIA RAPIDS workloads on large tabular data (GPUs make computing fast, so other costs become relatively larger)
Some mystery use cases inside of some hedge funds

It does not affect the everyday user, who processes 100GB to a few TB of data, and doesn’t mind waiting 10s for things to start running.

Coarse breakdown of costs

When you call x.sum().compute() a few things happen:

Graph generation: Some Python code in a Dask collection library, like dask array, calls the sum function, which generates a task graph on the client side.
Graph Optimization: We then optimize that graph, also on the client side, in order to remove unnecessary work, fuse tasks, apply important high level optimizations, and more.
Graph Serializtion: We now pack up that graph in a form that can be sent over to the scheduler.
Graph Communication: We fire those bytes across a wire over to the scheduler
Scheduler.update_graph: The scheduler receives these bytes, unpacks them, and then updates its own internal data structures
Scheduling: The scheduler then assigns ready tasks to workers
Communicate to workers: The scheduler sends out lots of smaller messages to each of the workers with the tasks that they can perform
Workers work: The workers then perform this work, and start communicating back and forth with the scheduler to receive new instructions

Generally most people today are concerned with steps 1-6. Once things get out to the workers and progress bars start moving people tend to care a bit less (but not zero).

What do other people do?

Let’s look at a few things that other projects do, and see if there are things that we can learn. These are commonly suggested, but there are challenges with most of them.

Rewrite the scheduler it in C++/Rust/C/Cython

Proposal: Python is slow. Want to make it faster? Don’t use Python. See academic projects.

Challenge: This makes sense for some parts of the pipeline above, but not for others. It also makes it harder to attract maintainers.

What we should consider: Some parts of the scheduler and optimization algorithms could be written in a lower level language, maybe Cython. We’ll need to be careful about maintainability.
Distributed scheduling

Proposal: The scheduler is slow, maybe have many schedulers? See Ray.

Challenge: It’s actually really hard to make the kinds of decisions that Dask has to make if scheduling state is spread on many computers. Distributed scheduling works better when the workload is very either uniform or highly decoupled. Distributed scheduling is really attractive to people who like solving interesting/hard problems.

What we should consider: We can move some simple logic down to the workers. We’ve already done this with the easy stuff though. It’s not clear how much additional benefit there is here.
Build specialty scheduling around collections

Proposal: If Dask were to become just a dataframe library or just an array computing library then it could special-case things more effectively. See Spark, Mars, and others.

Challenge: Yes, but Dask is not a dataframe library or an array library. The three use cases we mention above are all very different.

What we should consider: modules like dask array and dask dataframe should develop high level query blocks, and we should endeavor to communicate these subgraphs over the wire directly so that they are more compact.

What should we actually do?

Because our pipeline has many stages, each of which can be slow for different reasons, we’ll have to do many things. Additionally, this is a hard problem because changing one piece of the project at this level has repurcussions for many other pieces. The rest of this post tries to lay out a consistent set of changes. Let’s start with a summary:

For Dask array/dataframe let’s use high level graphs more aggressively so that we can communicate only abstract representations between the client and scheduler.
But this breaks low level graph optimizations, fuse, cull, and slice fusion in particular. We can make these unnecessary with two changes:
- We can make high level graphs considerably smarter to handle cull and slice fusion
- We can move a bit more of the scheduling down to the workers to replicate the advantages of low-level fusion there
Then, once all of the graph manipulation happens on the scheduler, let’s try to accelerate it, hopefully in a language that the current dev community can understand, like Cython
At the same time in parallel, let’s take a look at our network stack

We’ll go into these in more depth below

Graph Generation

High Level Graph History

A year or two ago we moved graph generation costs from user-code-typing time to graph-optimization-time with high level graphs

y = x + 1                 # graph generation used to happen here
(y,) = dask.optimize(y,)  # now it happens here

This really improved usability, and also let us do some high level optimizations which sometimes allowed us to skip some lower-level optimization costs.

Can we push this further?

The first four stages of our pipeline happen on the client:

Graph generation: Some Python code in a Dask collection library, like dask array, calls the sum function, which generates a task graph on the client side.
Graph Optimization: We then optimize that graph, also on the client side, in order to remove unnecessary work, fuse tasks, apply important high level optimizations, and more.
Graph Serializtion: We now pack up that graph in a form that can be sent over to the scheduler.
Graph Communication: We fire those bytes across a wire over to the scheduler

If we’re able to stay with the high level graph representation through these stages all the way until graph communication, then we can communicate a far more compact representation up to the scheduler. We can drop a lot of these costs, at least for the high level collection APIs (delayed and client.submit would still be slow, client.map might be ok though).

This has a couple of other nice benefits:

User’s code won’t block, and we can alert the user immediately that we’re on the job
We’ve centralized costs in just the scheduler, so there is now only one place where we might have to think about low-level code

(some conversation here: https://github.com/dask/distributed/issues/3872)

However, low-level graph optimizations are going to be a problem

In principle changing the distributed scheduler to accept a variety of graph layer types is a tedious but straightforward problem. I’m not concerned.

The bigger concern is what to do with low-level graph optimizations. Today we have three of these that really matter:

Task fusion: this is what keeps your read_parquet task merged with your subsequent blockwise tasks
Culling: this is what makes df.head() or x[0] fast
Slice fusion: this is why x[:100][5] works well

In order for us to transmit abstract graph layers up to the scheduler, we need to remove the need for these low level graph optimizations. I think that we can do this with a combination of two approaches:

More clever high level graph manipulation

We already do this a bit with blockwise, which has its own fusion, and which removes much of the need for fusion generally. But other blockwise-like operations, like read_* will probably have to join the Blockwise family.

Getting culling to work properly may require us to teach each of the individual graph layers how to track dependencies in each layer type and cull themselves. This may get tricky.

Slicing is doable, we just need someone to go in, grok all of the current slicing optimizations, and make high level graph layers for these computations. This would be a great project for a sharp masters student

Send speculative tasks to the workers

High level Blockwise fusion handles many of the use cases for low-level fusion, but not all. For example I/O layers like dd.read_parquet or da.from_zarr aren’t fused at a high level.

We can resolve this either by making them blockwise layers (this requires expanding the blockwise abstraction, which may be hard) or alternatively we can start sending not-yet-ready tasks to workers before all of their dependencies are finished if we’re highly confident that we know where they’re going to go. This would give us some of the same results of fusion, but would keep all of the task types separate (which would be nice for diagnostics) and might still give us some of the same performance benefits that we get from fusion.

Unpack abstract graph layers on the scheduler

So after we’ve removed the need for low level optimizations, and we just send the abstract graph layers up to the scheduler directly, we’ll need to teach the scheduler how to unpack those graph layers.

This is a little tricky because the Scheduler can’t run user Python code (for security reasons). We’ll have to register layer types (like blockwise, rechunk, dataframe shuffle) that the scheduler knows about and trusts ahead of time. We’ll still always support custom layers, and these will be at the same speed that they’ve always been, but hopefully there will be far less need for these if we go all-in on high level layers.

Rewrite scheduler in low-level language

Once most of the finicky bits are moved to the scheduler, we’ll have one place where we can focus on low level graph state manipulation.

Dask’s distributed scheduler is two things:

A Tornado TCP application that receives signals from clients and workers and send signals out to clients and workers

This is async-heavy networking code
A complex state machine internally that responds to those state changes

This is a complex data structure heavy Python code

Networking

Jim has an interesting project here that shows promise: https://github.com/jcrist/ery Reducing latency between workers and the scheduler would be good, and would help to accelerate stages 7-8 in the pipeline listed at the top of this post.

State machine

Rewriting the state machine in some lower level language would be fine. Ideally this would be in a language that was easy for the current maintainer community to maintain, (Cython?) but we may also consider making a more firm interface here that would allow other groups to experiment safely.

There are some advantages to this (more experimentation by different groups) but also some costs (splitting of core efforts and mismatches for users). Also, I suspect that splitting out also probably means that we’ll probably lose the dashboard, unless those other groups are very careful to expose the same state to Bokeh.

There is more exploration to do here. Regardless I think that it probably makes sense to try to isolate the state machine from the networking system. Maybe this also makes it easier for people to profile in isolation.

In speaking with a few different groups most people have expressed reservation about having multiple different state machine codes. This was done in MapReduce and Spark and resulted in difficult to maintain community dynamics.

High Level Graph Optimizations

Once we have everything in smarter high level graph layers, we will also be more ripe for optimization.

We’ll need a better way to write down these optimizations with a separated traversal system and a set of rules. A few of us have written these things before, maybe it’s time we revisit them.

What we need

This would require some effort, but I think that it would hit several high profile problems at once. There are a few tricky things to get right:

A framework for high level graph layers
An optimization system for high level graph layers
Separation of the scheduler into two parts

For this I think that we’ll need people who are fairly familiar with Dask to do this right.

And there there is a fair amount of follow-on work

Build a hierarchy of layers for dask dataframe
Build a hierarchy of layers for dask array
Build optimizations for those to remove the need for low level graph optimizations
Rewrite core parts of the scheduler in Cython
Experiment with the networking layer, maybe with a new Comm

I’ve been thinking about the right way to enact this change. Historically most Dask changes over the past few years have been incremental or peripheral, due to how burdened the maintainers are. There might be enough pressure on this problem though that we can get some dedicated engineering effort from a few organizations though, which might change how possible this is. We’ve gotten 25% time from a few groups. I’m curious if we can gate 100% time for some people for a few months.

Large SVDs

2020-05-13T00:00:00+00:00

We perform Singular Value Decomposition (SVD) calculations on large datasets.

We modify the computation both by using fully precise and approximate methods, and by using both CPUs and GPUs.

In the end we compute an approximate SVD of 200GB of simulated data and using a mutli-GPU machine in 15-20 seconds. Then we run this from a dataset stored in the cloud where we find that I/O is, predictably, a major bottleneck.

SVD - The simple case

Dask arrays contain a relatively sophisticated SVD algorithm that works in the tall-and-skinny or short-and-fat cases, but not so well in the roughly-square case. It works by taking QR decompositions of each block of the array, combining the R matrices, doing another smaller SVD on those, and then performing some matrix multiplication to get back to the full result. It’s numerically stable and decently fast, assuming that the intermediate R matrices of the QR decompositions mostly fit in memory.

The memory constraints here are that if you have an n by m tall and skinny array (n >> m) cut into k blocks then you need to have about m**2 * k space. This is true in many cases, including typical PCA machine learning workloads, where you have tabular data with a few columns (hundreds at most) and many rows.

It’s easy to use and quite robust.

import dask.array as da

x = da.random.random((10000000, 20))
x

	Array	Chunk
Bytes	1.60 GB	100.00 MB
Shape	(10000000, 20)	(625000, 20)
Count	16 Tasks	16 Chunks
Type	float64	numpy.ndarray

20 10000000

u, s, v = da.linalg.svd(x)

This works fine in the short and fat case too (when you have far more columns than rows) but we’re always going to assume that one of your dimensions is unchunked, and that the other dimension has chunks that are quite a bit longer, otherwise, things might not fit into memory.

Approximate SVD

If your dataset is large in both dimensions then the algorithm above won’t work as is. However, if you don’t need exact results, or if you only need a few of the components, then there are a number of excellent approximation algorithms.

Dask array has one of these approximation algorithms implemented in the da.linalg.svd_compressed function. And with it we can compute the approximate SVD of very large matrices.

We were recently working on a problem (explained below) and found that we were still running out of memory when dealing with this algorithm. There were two challenges that we ran into:

The algorithm requires multiple passes over the data, but the Dask task scheduler was keeping the input matrix in memory after it had been loaded once in order to avoid recomputation. Things still worked, but Dask had to move the data to disk and back repeatedly, which reduced performance significantly.

We resolved this by including explicit recomputation steps in the algorithm.
Related chunks of data would be loaded at different times, and so would need to stick around longer than necessary to wait for their associated chunks.

We resolved this by engaging task fusion as an optimization pass.

Before diving further into the technical solution we quickly provide the use case that was motivating this work.

Application - Genomics

Many studies are using genome sequencing to study genetic variation between different individuals within a species. These includes studies of human populations, but also other species such as mice, mosquitoes or disease-causing parasites. These studies will, in general, find a large number of sites in the genome sequence where individuals differ from each other. For example, humans have more than 100 million variable sites in the genome, and modern studies like the UK BioBank are working towards sequencing the genomes of 1 million individuals or more.

In diploid species like humans, mice or mosquitoes, each individual carries two genome sequences, one inherited from each parent. At each of the 100 million variable genome sites there will be two or more “alleles” that a single genome might carry. One way to think about this is via the Punnett square, which represents the different possible genotypes that one individual might carry at one of these variable sites:

In the above there are three possible genotypes: AA, Aa, and aa. For computational genomics, these genotypes can be encoded as 0, 1, or 2. In a study of a species with M genetic variants assayed in N individual samples, we can represent these genotypes as an (M x N) array of integers. For a modern human genetics study, the scale of this array might approach (100 million x 1 million). (Although in practice, the size of the first dimension (number of variants) can be reduced somewhat, by at least an order of magnitude, because many genetic variants will carry little information and/or be correlated with each other.)

These genetic differences are not random, but carry information about patterns of genetic similarity and shared ancestry between individuals, because of the way they have been inherited through many generations. A common task is to perform a dimensionality reduction analysis on these data, such as a principal components analysis (SVD), to identify genetic structure reflecting these differencies in degree of shared ancestry. This is an essential part of discovering genetic variants associated with different diseases, and for learning more about the genetic history of populations and species.

Reducing the time taken to compute an analysis such as SVD, like all science, allows for exploring larger datasets and testing more hypotheses in less time. Practically, this means not simply a fast SVD but an accelerated pipeline end-to-end, from data loading to analysis, to understanding.

We want to run an experiment in less time than it takes to make a cup of tea

Performant SVDs w/ Dask

Now that we have that scientific background, let’s transition back to talking about computation.

To stop Dask from holding onto the data we intentionally trigger computation as we build up the graph. This is a bit atypical in Dask calculations (we prefer to have as much of the computation at once before computing) but useful given the multiple-pass nature of this problem. This was a fairly easy change, and is available in dask/dask #5041.

Additionally, we found that it was helpful to turn on moderately wide task fusion.

import dask
dask.config.set({"optimization.fuse.ave-width": 5})

Then things work fine

We’re going to try this SVD on a few different choices of hardware including:

A MacBook Pro
A DGX-2, an NVIDIA worksation with 16 high-end GPUs and fast interconnect
A twenty-node cluster on AWS

Macbook Pro

We can happily perform an SVD on a 20GB array on a Macbook Pro

import dask.array as da

x = da.random.random(size=(1_000_000, 20_000), chunks=(20_000, 5_000))

u, s, v = da.linalg.svd_compressed(x, k=10, compute=True)
v.compute()

This call is no longer entirely lazy, and it recomputes x a couple times, but it works, and it works using only a few GB of RAM on a consumer laptop.

It takes around 2min 30s time to compute that on a laptop. That’s great! It was super easy to try out, didn’t require any special hardware or setup, and in many cases is fast enough. By working locally we can iterate quickly.

Now that things work, we can experiment with different hardware.

Adding GPUs (a 15 second SVD)

Disclaimer: one of the authors (Ben Zaitlen) works for NVIDIA

We can dramatically increase performance by using a multi-GPU machine. NVIDIA and other manufacturers now make machines with multiple GPUs co-located in the same physical box. In the following section, we will run the calculations on a DGX2, a machine with 16 GPUs and fast interconnect between the GPUs.

Below is almost the same code, running in significantly less same time but we make the following changes:

We increase the size of the array by a factor of 10x
We switch out NumPy for CuPy, a GPU NumPy implementation
We use a sixteen-GPU DGX-2 machine with NVLink interconnects between GPUs (NVLink will dramatically decrease transfer time between workers)

On A DGX2 we can calculate an SVD on a 200GB Dask array between 10 to 15 seconds.

The full notebook is here, but the relevant code snippets are below:

# Some GPU specific setup
from dask_cuda import LocalCluster

cluster = LocalCluster(...)
client = Client(cluster)

import cupy
import dask.array as da
rs = da.random.RandomState(RandomState=cupy.random.RandomState)

# Create the data and run the SVD as normal
x = rs.randint(0, 3, size=(10_000_000, 20_000),
               chunks=(20_000, 5_000), dtype="uint8")
x = x.persist()

u, s, v = da.linalg.svd_compressed(x, k=10, seed=rs)
v.compute()

To see this run, we recommend viewing this screencast:

Read dataset from Disk

While impressive, the computation above is mostly bound by generating random data and then performing matrix calculations. GPUs are good at both of these things.

In practice though, our input array won’t be randomly generated, it will be coming from some dataset stored on disk or increasingly more common, stored in the cloud. To make things more realistic we perform a similar calculation with data stored in a Zarr format in GCS

In this Zarr SVD example, we load a 25GB GCS backed data set onto a DGX2, run a few processing steps, then perform an SVD. The combination of preprocessing and SVD calculations ran in 18.7 sec and the data loading took 14.3 seconds.

Again, on a DGX2, from data loading to SVD we are running in time less than it would take to make a cup of tea. However, the data loading can be accelerated. From GCS we are reading into data into the main memory of the machine (host memory), uncompressing the zarr bits, then moving the data from host memory to the GPU (device memory). Passing data back and forth between host and device memory can significantly decrease performance. Reading directly into the GPU, bypassing host memory, would improve the overall pipeline.

And so we come back to a common lesson of high performance computing:

High performance computing isn’t about doing one thing exceedingly well, it’s about doing nothing poorly.

In this case GPUs made our computation fast enough that we now need to focus on other components of our system, notably disk bandwidth, and a direct reader for Zarr data to GPU memory.

Cloud

Diclaimer: one of the authors (Matthew Rocklin) works for Coiled Computing

We can also run this on the cloud with any number of frameworks. In this case we used the Coiled Cloud service to deploy on AWS

from coiled_cloud import Cloud, Cluster
cloud = Cloud()

cloud.create_cluster_type(
    organization="friends",
    name="genomics",
    worker_cpu=4,
    worker_memory="16 GiB",
    worker_environment={
        "OMP_NUM_THREADS": 1,
        "OPENBLAS_NUM_THREADS": 1,
        # "EXTRA_PIP_PACKAGES": "zarr"
    },
)

cluster = Cluster(
    organization="friends",
    typename="genomics",
    n_workers=20,
)

from dask.distributed import Client
client = Client(cluster)

# then proceed as normal

Using 20 machines with a total of 80 CPU cores on a dataset that was 10x larger than the MacBook pro example we were able to run in about the same amount of time. This shows near optimal scaling for this problem, which is nice to see given how complex the SVD calculation is.

A screencast of this problem is viewable here

Compression

One of the easiest ways for us to improve performance is to reduce the size of this data through compression. This data is highly compressible for two reasons:

The real-world data itself has structure and repetition (although the random play data does not)
We’re storing entries that take on only four values. We’re using eight-bit integers when we only needed two-bit integers

Let’s solve the second problem first.

Compression with bit twiddling

Ideally Numpy would have a two-bit integer datatype. Unfortunately it doesn’t, and this is hard because memory in computers is generally thought of in full bytes.

To work around this we can use bit arithmetic to shove four values into a single value Here are functions that do that, assuming that our array is square, and the last dimension is divisible by four.

import numpy as np

def compress(x: np.ndarray) -> np.ndarray:
    out = np.zeros_like(x, shape=(x.shape[0], x.shape[1] // 4))
    out += x[:, 0::4]
    out += x[:, 1::4] << 2
    out += x[:, 2::4] << 4
    out += x[:, 3::4] << 6
    return out


def decompress(out: np.ndarray) -> np.ndarray:
    back = np.zeros_like(out, shape=(out.shape[0], out.shape[1] * 4))
    back[:, 0::4] = out & 0b00000011
    back[:, 1::4] = (out & 0b00001100) >> 2
    back[:, 2::4] = (out & 0b00110000) >> 4
    back[:, 3::4] = (out & 0b11000000) >> 6
    return back

Then, we can use these functions along with Dask to store our data in a compressed state, but lazily decompress on-demand.

x = x.map_blocks(compress).persist().map_blocks(decompress)

That’s it. We compress each block our data and store that in memory. However the output variable that we have, x will decompress each chunk before we operate on it, so we don’t need to worry about handling compressed blocks.

Compression with Zarr

A slightly more general but probably less efficient route would be to compress our arrays with a proper compression library like Zarr.

The example below shows how we do this in practice.

import zarr
import numpy as np
from numcodecs import Blosc
compressor = Blosc(cname='lz4', clevel=3, shuffle=Blosc.BITSHUFFLE)


x = x.map_blocks(zarr.array, compressor=compressor).persist().map_blocks(np.array)

Additionally, if we’re using the dask-distributed scheduler then we want to make sure that the Blosc compression library doesn’t use additional threads. That way we don’t have parallel calls of a parallel library, which can cause some contention

def set_no_threads_blosc():
    """ Stop blosc from using multiple threads """
    import numcodecs
    numcodecs.blosc.use_threads = False

# Run on all workers
client.register_worker_plugin(set_no_threads_blosc)

This approach is more general, and probably a good trick to have up ones’ sleeve, but it also doesn’t work on GPUs, which in the end is why we ended up going with the bit-twiddling approach one section above, which uses API that was uniformly accessible within the Numpy and CuPy libraries.

Final Thoughts

In this post we did a few things, all around a single important problems in genomics.

We learned a bit of science
We translated a science problem into a computational problem, and in particular into a request to perform large singular value decompositions
We used a canned algorithm in dask.array that performed pretty well, assuming that we’re comfortable going over the array in a few passes
We then tried that algorithm on three architectures quickly
1. A Macbook Pro
2. A multi-GPU machine
3. A cluster in the cloud
Finally we talked about some tricks to pack more data into the same memory with compression

This problem was nice in that we got to dive deep into a technical science problem, and yet also try a bunch of architecture quickly to investigate hardware choices that we might make in the future.

We used several technologies here today, made by several different communities and companies. It was great to see how they all worked together seamlessly to provide a flexible-yet-consistent experience.