title: Parallel processing using the Dask packge in Python

date: 2025-04-22

format:

html:

theme: cosmo

css: assets/styles.css

toc: true

code-copy: true

code-block-bg: true

code-block-border-left: "#31BAE9"

engine: knitr

ipynb-shell-interactivity: all

code-overflow: wrap

<script

src="https://cdn.mathjax.org/mathjax/latest/MathJax.js?config=TeX-AMS-MML_HTMLorMML"

type="text/javascript">

</script>

The Dask package provides a variety of tools for managing parallel computations.

In particular, some of the key ideas/features of Dask are:

- Separate what to parallelize from how and where the parallelization is actually carried out.

- Different users can run the same code on different computational resources (without touching the actual code that does the computation).

- Dask provides distributed data structures that can be treated as a single data structures when runnig operations on them (like Spark and pbdR).

The idea of a 'future' or 'delayed' operation is to tag operations such that they run lazily. Multiple operations can then be pipelined together and Dask can figure out how best to compute them in parallel on the computational resources available to a given user (which may be different than the resources available to a different user).

Let's import dask to get started.

One specifies a "scheduler" to control how parallelization is done, including what machine(s) to use and how many cores on each machine to use.

For example,

This table gives an overview of the different scheduler.

|Type|Description|Multi-node|Copies of objects made?|

|----|-----------|----------|-----------------------|

|synchronous|not in parallel|no|no|

|threaded|threads within current Python session|no|no|

|processes|background Python sessions|no|yes|

|distributed|Python sessions across multiple nodes|yes or no|yes|

Note that because of Python's Global Interpreter Lock (GIL), many computations done in pure Python code won't be parallelized using the 'threaded' scheduler; however computations on numeric data in numpy arrays, Pandas dataframes and other C/C++/Cython-based code will parallelize.

For the next section (Section 3), we'll just assume use of the 'processes' schduler and will provide more details on the other schedulers in the following section (Section 4).

Dask has a large variety of patterns for how you might parallelize a computation.

We'll simply parallelize computation of the mean of a large number of random numbers across multiple replicates as can be seen in `calc_mean.py`.

(Note the code in calc_mean.py is not safe in terms of parallel random number generation - see Section 8 later in this document.)

The basic pattern for setting up a parallel loop is:

You could set the scheduler in the `compute` call:

However, it is best practice to separate what is parallelized from where the parallelization is done, specifying the scheduler at the start of your code.

We can do parallel map operations (i.e., a *map* in the map-reduce or functional programming sense, akin to `lapply` in R).

For this we need to use the *distributed* scheduler, which we'll discuss more later.

Note that the distributed scheduler can work on one node or on multiple nodes.

The map operation appears to cache results. If you rerun the above with the same inputs, you get the same result back essentially instantaneously (even if one removes the setting of the seed from `calc_mean_vargs`). HOWEVER, that means that if there is randomness in the results of your function for a given input, Dask will just continue to return the original output.

You can use `delayed` in more complicated situations than the simple iterations shown above.

`visualize()` uses the `graphviz` package to illustrate the task graph (similar to a directed acyclic graph in a statistical model and to how Tensorflow organizes its computations).

One can also tell Dask to always delay evaluation of a given function:

You can also control evaluation of tasks using the [Futures interface for managing tasks](https://docs.dask.org/en/latest/futures.html). Unlike use of `delayed`, the evaluation occurs immediately instead of via lazy evaluation.

Dask provides the ability to work on data structures that are split (sharded/chunked) across workers. There are two big advantages of this:

- You can do calculations in parallel because each worker will work on a piece of the data.

- When the data is split across machines, you can use the memory of multiple machines to handle much larger datasets than would be possible in memory on one machine. That said, Dask processes the data in chunks, so one often doesn't need a lot of memory, even just on one machine.

Because computations are done in external compiled code (e.g., via numpy) it's effective to use the threaded scheduler when operating on one node to avoid having to copy and move the data.

Dask dataframes are Pandas-like dataframes where each dataframe is split into groups of rows, stored as smaller Pandas dataframes.

One can do a lot of the kinds of computations that you would do on a Pandas dataframe on a Dask dataframe, but many operations are not possible, as discussed in the [Dask dataframe API](http://docs.dask.org/en/latest/dataframe-api.html).

By default dataframes are handled by the threads scheduler.

Here's an example of reading from a dataset of flight delays (about 11 GB data).

You can get [the flight delay data](https://www.stat.berkeley.edu/share/paciorek/1987-2008.csvs.tgz).

You should see this:

Bags are like lists but there is no particular ordering, so it doesn't make sense to ask for the i'th element.

You can think of operations on Dask bags as being like parallel map operations on lists in Python or R.

By default bags are handled via the multiprocessing scheduler.

Let's see some basic operations on a large dataset of Wikipedia log files.

You can get a [subset of the Wikipedia data](https://www.stat.berkeley.edu/share/paciorek/wikistats_example.tar.gz).

That should look like this:

Note that it is quite inefficient to do the `find()` (and therefore necessarily reading

the data in) and then compute on top of that intermediate result in

two separate calls to `compute()`. More in Section 6.

::: {.callout-important title="Memory use danger"}

Be careful about running `compute` such that it produces a large object in the main Python session.

E.g., above you would not want to do `data = wiki.compute()` as that would pull the entire dataset into your main Python session as a single (very large) list.

:::

Dask arrays are numpy-like arrays where each array is split up by both rows and columns into smaller numpy arrays.

One can do a lot of the kinds of computations that you would do on a numpy array on a Dask array, but many operations are not possible, as discussed in the [Dask array API](http://docs.dask.org/en/latest/array-api.html).

By default arrays are handled via the threads scheduler.

Let's first see operations on a single node, using a single 13 GB 2-d array. Note that Dask uses lazy evaluation, so creation of the array doesn't happen until an operation requiring output is done.

Here we specify that the chunks (the sub-arrays) are 10000 by 10000.

For a row-based operation, we would presumably only want to chunk things up by row,

but this doesn't seem to actually make a difference, presumably because the

mean calculation can be done in pieces and only a small number of summary statistics

moved between workers.

Of course, given the lazy evaluation, this timing comparison is not just timing the

actual row mean calculations.

But this doesn't really clarify the story...

Dask will avoid storing all the chunks in memory. (It appears to just generate them on the fly.)

Here we have an 80 GB array but we never use more than a few GB of memory (based on `top` or `free -h`).

This should be straightforward based on using Dask distributed. However, one would want to

be careful about creating arrays by distributing the data from a single Python process

as that would involve copying between machines.

The computation here effectively parallelizes. However, if instead of using the `default_rng` Generator constructor, one uses the old numpy syntax of `np.random.normal(size = n)`, one would see it takes four times as long, so is not parallelizing.

The problem is presumably occurring because of Python's Global Interpreter Lock (GIL): any computations done in pure Python code can not be parallelized using the 'threaded' scheduler. However, computations on numeric data in numpy arrays, pandas dataframes and other C/C++/Cython-based code would parallelize.

Exactly why one form of numpy code encounters the GIL and the other is not clear to me.

We can effectively parallelize regardless of the GIL by using multiple Python processes.

::: {.callout-important title="Interactive vs. background use"}

The above code will work when run in an interactive Python

session. However, if you want to run it within a Python script (i.e., in a background/batch job),

you'll need to configure the scheduler and run the code within an `if __name__ == '__main__'` block:

:::

According to the Dask documentation, using `Distributed` on a local machine

has advantages over multiprocessing, including the diagnostic dashboard

(see Section 7) and better handling of when copies need to be made. As we saw previously,

using Distributed also allows us to use the handy `map()` operation.

::: {.callout-important title="Interactive vs. background use"}

The above code will work when run in an interactive Python

session. However, if you want to run it within a Python script (i.e., in a background/batch job),

you'll need to configure the scheduler and run the code within an `if __name__ == '__main__'` block:

:::

We need to set up a scheduler on one machine (possibly the machine we are on)

and workers on whatever machines we want to do the computation on.

One option is to use the `dask-ssh` command to start up the scheduler and workers.

(Note that for this to work we need to have password-less SSH working to connect

to the various machines.)

export SCHED=$(hostname)

dask-ssh --scheduler ${SCHED} radagast.berkeley.edu radagast.berkeley.edu arwen.berkeley.edu arwen.berkeley.edu &

Then in Python, connect to the cluster via the scheduler.

Alternatively, you can start the workers from within Python:

To run within a Slurm job we can use `dask-ssh` or a combination of `dask-scheduler`

and `dask-worker`..

Provided that we have used --ntasks or --nodes and --ntasks-per-node to set the

number of CPUs desired (and not --cpus-per-task), we can use `srun` to

enumerate where the workers should run.

First we'll start the scheduler and the workers.

export SCHED=$(hostname):8786

dask-scheduler&

sleep 50

srun dask-worker tcp://${SCHED} --local-directory /tmp & # might need machinename.berkeley.edu:8786

sleep 100 # might need even more time to make sure workers start up fully

Then in Python, connect to the cluster via the scheduler.

Let's process the 500 GB of Wikipedia log data. (Note that

when I tried this on our local department cluster, I had some errors that might be related

to the SCF not being set up for fast parallel I/O.) You'll need to change the path to

the data if using this code yourself.

Alternatively, we can use dask-ssh, but I've had problems sometimes with

using SSH to connect between nodes of a Slurm job, so the approach above

is likely to be more robust as it relies on Slurm itself to connect

between nodes.

export SCHED=$(hostname)

srun hostname > .hosts

dask-ssh --scheduler ${SCHED} --hostfile .hosts

We can set up nested parallelization (or an arbitrary set of computations) and just have Dask's delayed functionality figure out how to do the parallelization, provided there is a single call to the compute() method.

When using `delayed`, Dask starts up each delayed evaluation separately (i.e., dynamic allocation).

This is good for load-balancing, but each task induces

some overhead (a few hundred microseconds).

Even with a distributed `map()` it doesn't appear possible to ask that the tasks

be broken up into batches.

So if you have many quick tasks, you probably

want to break them up into batches manually, to reduce the impact of the overhead.

As far as I can tell, Dask avoids keeping all the pieces of a distributed object or

computation in memory. However, in many cases this can mean repeating

computations or re-reading data if you need to do multiple operations on a dataset.

For example, if you are create a Dask distributed dataset from data on disk, I think this means that

every distinct set of computations (each computational graph) will involve reading

the data from disk again.

One implication is that if you can include all computations

on a large dataset within a single computational graph (i.e., a call to `compute`)

that may be much more efficient than making separate calls.

Here's an example with Dask dataframe on the airline delay data, where we make sure to

do all our computations as part of one graph:

However, I also tried the above where I added `air.count()` to the `dask.compute` call and something went wrong - the computation time increased a lot and there was a lot of memory use. I'm not sure what is going on.

Note that when reading from disk, disk caching by the operating system (saving files that are used repeatedly in memory) can also greatly speed up I/O. (Note this can very easily confuse you in terms of timing your code..., e.g., simply copying the data to your machine can put them in the cache, so subsequent reading into Python can take advantage of that.)

Except for the 'threads' scheduler, copies will be made of all objects passed to the workers.

In general you want to delay the input objects. There are a couple reasons why:

- Dask hashes the object to create a name, and if you pass the same object as an argument multiple times, it will repeat that hashing.

- When using the distributed scheduler *only*, delaying the inputs will prevent sending the data separately for every task (rather it should send the data separately for each worker).

In this example, most of the "computational" time is actually spent transferring the data

rather than computing the mean. Whether there is a copy per task or a copy per process seems to depend on exactly what the parallelized code is doing (perhaps based on whether there is random number generation in the code).

Here's an example of using the Distributed (local) scheduler.

Copies have to be made, but if we delay the data object, there is only one copy per worker.

Note that with the Dask distributed scheduler, it is complicated to assess what is going on, because the scheduler seems to try to optimize assignment of tasks to workers in a way that may cause an imbalance in the number of tasks assigned to each worker. In this example, all the tasks are assigned to a single worker.

That took a few seconds if we delay the data but takes ~20 seconds if we don't.

Also, Dask gives a warning about sending the data to the workers in advance. I’m not sure of the distinction between what it is recommending and use of `dask.delayed(x)`. When I tried to use `scatter()` in various ways, I wasn't able to silence the warning.

Note that in either case, we incur the memory usage of the original 'x' plus copies of 'x' on the workers.

For this to make the most sense we want to be on a system where we can read multiple files

without having the bottleneck of accessing a single spinning hard disk. For example the

UC Berkeley Savio cluster filesystem is set up for fast parallel I/O.

On systems with a single spinning hard disk or a single SSD, you might experiment

to see how effectively things scale as you read (or write) multiple files in parallel.

Here we'll demo this. Ideally you would do this on a system with fast, parallel I/O.

I'm not sure why that didn't scale perfectly (i.e., that 21 files on 21 or more workers would take only 28 seconds),

but we do see that it was quite a bit faster than sequentially reading the data would be.

With a resource manager like Kubernetes, Dask can scale the number of workers up and down to adapt to the computational needs of a workflow. Similarly, if submitting jobs to Slurm via Dask, it will scale up and down automatically - see Section 9.

Dask distributed provides a web interface showing the status of your work.

(Unfortunately I've been having some problems viewing the interface on SCF machines, but hopefully this will work for you.)

By default Dask uses port 8787 for the web interface.

If your Python session is not running on your local machine, you can set up port forwarding to view the web interface in the browser on your local machine, e.g.,

ssh -L 8787:localhost:8787 name_of_remote_machine

Then go to `localhost:8787` in your local browser.

In the code above, I was cavalier about the seeds for the random number generation in the different parallel computations.

The general problem is that we want the random numbers generated on each worker to not overlap with the random numbers generated on other workers. But random number generation involves numbers from a periodic sequence. Simply setting different seeds on different workers does not guarantee non-overlapping blocks of random numbers (though in most cases they probably would not overlap).

Using the basic numpy RNG, one can simply set different seeds for each task, but as mentioned above that doesn't guarantee non-overlapping random numbers.

We can use functionality with numpy's PCG64 or MT19937 generators to be completely safe in our parallel random number generation. Each provide a `jumped()` function that moves the RNG ahead as if one had generated a very large number of random variables ($2^{128}$ for the Mersenne Twister and nearly that for the PCG64).

Here’s how we can set up the use of the PCG64 generator:

Now let’s see how to jump forward. And then verify that jumping forward two increments is the same as making two separate jumps.

We can also use `jumped()` with the Mersenne Twister.

So the strategy to parallelize across tasks (or potentially workers if random number generation is done sequentially for tasks done by a single worker) is to give each task the same seed and use `jumped(i)` where `i` indexes the tasks (or workers).

One caution is that it appears that the period for PCG64 is $2^{128}$ and that `jumped(1)` jumps forward by nearly that many random numbers. That seems quite strange, and I don’t understand it.

Alternatively as [recommended in the docs](https://numpy.org/doc/stable/reference/random/bit_generators/pcg64.html):

One can submit jobs to a scheduler such as Slurm from within Python. In general I don't recommend

this as it requires you to be running Python within a stand-alone server while the Slurm

job is running, but here is how one can do it.

Note that the SLURMCluster() call defines the parameters of a single Slurm job, but

`scale()` is what starts one or more jobs. If you ask for more workers than there are

processes defined in your job definition, more than one Slurm job will be launched.

Be careful to request as many processes as cores; if you leave out `processes`, it will