Scaling Python: dipping my toes into dask

As folks might know, I’m finishing off my PhD in voice and speech data at ANU’s School of Cybernetics at the moment, and am using a lot of data science tools like jupyter and pandas as part of my methods. I recently came across a task that needed parallel processing capabilities - where code is written to take advantage of multiple CPUs (or GPUs) on a host. To tackle this, I used dask- a library focused on scaling Python libraries like pandas and scikit-learn to different infrastructures - from a laptop to a large cluster. Although Python has built-in multiprocessing libraries, I chose to focus on learning about dask because of its close association with pandas, which I’ve been using extensively in my PhD data work.

However, dask required a steep learning curve - it’s been decades since I covered parallel or concurrent programming in my undergrad - and I’m noting my learnings here for anyone else following a similar learning path.

The problem

The specific problem I was trying to solve with dask was to append a value to a column in a pandas dataframe, where another column value matched a value from a tsv file. This sounds really straightforward doesn’t it?

for rows in dataframe:
    for rows in tsv file:
        if match:
            append value to column
            break # we don’t need to keep searching if matched

And it would be.

Except for the problem of scale.

In my case, the pandas dataframe was over 1 million rows long, and the tsv file was over 2 million rows long. This meant that for every row in the dataframe, Python had to search up to 2 million rows in the tsv - meaning an upper bound of 1 million x 2 million = 2 trillion (2 x 10¹²) operations to match all the rows in the dataframe (and a lower bound, assuming that each row in the dataframe matched the first row in the tsv file, which is, admittedly, unrealistic, of 1 million or 1 x 10⁶ operations). In Big O notation, used to show the computational complexity of algorithms, this approach would have linear complexity - $O(n)$ - that is, the computation scales linearly with an increase in either the dataframe or tsv file size (I think - algorithms are not my forte. I’m sure if I’m wrong I will be corrected 😉 ).

I initially set this to run on jupyter, thinking, naïvely, that it would automagically find a way to parallelise this operation over my laptop’s 16 cores. I was wrong, and jupyter, using the tqdmlibrary, estimated that this operation would take 21 days to complete. I confirmed this by running htop, which showed that only one of my CPU cores was heavily loaded.

Selecting dask

I first considered some optimisations to reduce the computation required.

Could I remove the value in the tsv file once a match was found, thus reducing the search space of subsequent operations? This would have been possible if I was certain that there was one and only match for each row in the dataframe with a row in the tsv file, but that wasn’t the case here.

Could I use Python’s inbuilt multiprocessing library to help?

I tried this, and my jupyter notebook ran for several hours , but I had done something wrong with the start() and join() calls, and manual verification of the matching failed.

So, I turned to dask, primarily because it appeared to be specifically addressing the problem of scaling data science-based workloads in Python, and I had also seen Genevieve Buckley’s talk using dask for image processing at PyConline AU 2020, and thought it might be the right tool for the job.

First steps with dask

But I hadn’t done concurrency or parallelisation work since undergrad, and needed to start from the basics. I tried the dask tutorial, which I had to run in binder because I couldn’t find the right set of dependencies to run it under pip on my local machine (no, I’m not running conda and pip together, sorry).

binder, to be blunt, sucked. It refreshed my notebook every waves hands 15 minutes, and I lost the work that I had done. So, a frustrating start. But I was grateful that there was at least a tutorial available.

This gave me at least some basic concepts of dask:

The Client - which communicates with the Scheduler, whose job it is to schedule tasks across the workers, processes, threads and processors.
The concept of collections - data structures such as Arrays, Dataframes and Bags, which mimic data structures familiar to those who use pandas or numpy.
Dask dataframes - which parallelise pandas dataframes using partitioning
The concept of lazy evaluation, in which computation is delayed until needed, which provides efficiencies in task sharing. The meta keyword, when passed to a function that undertakes some sort of computation, tells daskwhat the resulting output data structure should be. daskhas some handy built-ins for generating values for meta. From what I could understand, metaassistsdask in optimisation, because the resulting data structure is known a priori.
The concept of intermediate results, for example entire dataframes that are calculated before being reduced in some way, and how dask handles them.
Using the persist() function to store commonly accessed data in memory, as long as that object was small enough to fit into available memory.

That was enough to get started on the problem.

Setting up my environment to run dask

Installing daskwithin Python was reasonably straightforward. It essentially requires installing a Python package, and a range of dependencies. The only one that had to be installed as a binary was graphviz. Because I was using dask with jupyter, I installed the dask-labextension package into my jupyter virtual environment, which provided observability into dask’s operations, such as CPU and memory utilisation, and task execution.

Writing my dataframe iteration as a function

One of design features I really appreciated while learningdaskwas that much of the functionality of pandas dataframes was replicated: that is, I could create a dask dataframe, and for the most part, it would have similar behaviour to that of a pandasdataframe. daskabstracted away the complexity underneath.

However, in order to take full advantage of dask’s capabilities, I needed to write code in a particular way. Iterating over a dataframe in pandas, using something like iterrows(), is an inherently expensive operation, and the same is true in dask.

However, dask provides a way to map a function across partitions of a dataframe, so that the iterations can be processed in parallel.

out_ddf = ddf.map_partitions(myFunction,
myParameter,
meta=columns_meta)

This seemed straightforward at first, but the devil was in the detail.

Finding optimal settings is the hardest part of `dask`

Partitioning the dask dataframe

The key thing I struggled with was understanding how to optimise settings like the number of partitions of the dask dataframe, the number of workers or threads the Client should have and so on. There are several tradeoffs in these decisions: the more partitions your dask dataframe has, the more overhead there is in joining the results of each of the partitions (which are computed in parallel), but the more it takes advantages of multiple cores, if available.

The daskdocumentation has some information on best practices for partitioning dask dataframes, and there was also a useful article on Stack Overflow on the topic by Matthew Rocklin, a core developer of dask, and who answers a lot of the Stack Overflow question on dask (huge thank you, Matthew!).

In essence, daskdefaults to partitioning dataframes so that each partition is around 100Mb in size. What I found in my experiments however, was that each partition engaged a single CPU core. That is, if the dataframe was partitioned into eight partitions by default, then only eight cores were engaged. That meant that I couldn’t have a higher number of partitions than my number of cores. I eventually settled on using 100 partitions, where I had 512 cores (256 cores dual-threaded) available.

Optimal settings for the `Client`

I had similar struggles in finding the optimal settings for Client, too. In dask, the Client is the main interface with the Scheduler. The Scheduler’s job is to optimise execution of the workload, such as by determining the dependency graph, in which the dependency of computation operations is calculated, decomposing operations into smaller operations, and co-ordinating workers that are performing computation, and reporting results back to the Client.

At this point, I decided to read through Holden Karau and Mika Kimmins’ book Scaling Python with Dask, to glean some more pointers (no pun intended).

This explained in more detail the relationship between the Client and Scheduler, but I still remain confused about different parts of the dask ecosystem, such as distributed - the library that allows for lazy evaluation.

In particular, this book provided more information on the optimal choice of Client settings, including balancing the number of workers and threads depending on the type of compute workload you were processing. For my workload - iterating over a large dataframe - I found that a small number of workers (24) each having a large number of threads (64) worked well, as long as I assigned sufficient RAM to each worker (64GB).

In previous iterations, I had tried increasing the number of workers well over 100, but rapidly ran into the situation of workers dying unexpectedly. When digging into this, I found that it was always the “communications” worker that was dying, and I suspect what was happening was that there were too many workers for the communications worker to keep track of inter-worker communications, causing it to die. Again, there was a trade-off here: more workers meant more parallelisation of available compute infrastructure, but it also increased the number of worker inter-communication processes.

Nested delayed objects

Another “gotcha” for me was nested delayed objects.

Above, I had rewritten my function using map_partitions():

out_ddf = ddf.map_partitions(myFunction, 
                             myParameter,  
                             meta=columns_meta)

I then needed to perform lazy evaluation on this, using;

result_df = out_ddf.compute()

However, when I inspected the value of result_df, it was a dataframe, but it contained 100 delayed objects, rather than the processed dataframes I was hoping for. This is because the mapping function had nested delayed objects - so I had to run another level of lazy evaluation:

completed_df = dask.compute(*result_df)

This did result in the evaluated dataframes I needed - and then I simply combined the 100 dataframes together with a merge operation.

merged_df = pd.concat(completed_df, ignore_index=True)

Next steps with `dask`

There’s a lot that I still need to learn with dask - such as different data structures like Arrays and Bags, which I didn’t use in this task. But I’m still pretty pleased at being able to figure out enough dask to compute something that was going to take 21 days on a single CPU in around 12 hours on a CPU cluster!