# Efficient MPI-based Communication for GPU-Accelerated Dask Applications
††thanks: This research is supported in part by NSF grants #1818253, #1854828,
#1931537, #2007991, #2018627, and XRAC grant #NCR-130002.

Aamir Shafi, Jahanzeb Maqbool Hashmi, Hari Subramoni and Dhabaleswar K. (DK)
Panda
The Ohio State University
{shafi.16, hashmi.29, subramoni.1<EMAIL_ADDRESS>
###### Abstract

Dask is a popular parallel and distributed computing framework, which rivals
Apache Spark to enable task-based scalable processing of big data. The Dask
Distributed library forms the basis of this computing engine and provides
support for adding new communication devices. It currently has two
communication devices: one for TCP and the other for high-speed networks using
UCX-Py—a Cython wrapper to UCX. This paper presents the design and
implementation of a new communication backend for Dask—called MPI4Dask—that is
targeted for modern HPC clusters built with GPUs. MPI4Dask exploits mpi4py
over MVAPICH2-GDR, which is a GPU-aware implementation of the Message Passing
Interface (MPI) standard. MPI4Dask provides point-to-point asynchronous I/O
communication coroutines, which are non-blocking concurrent operations defined
using the async/await keywords from the Python’s asyncio framework. Our
latency and throughput comparisons suggest that MPI4Dask outperforms UCX by
$6\times$ for 1 Byte message and $4\times$ for large messages (2 MBytes and
beyond) respectively. We also conduct comparative performance evaluation of
MPI4Dask with UCX using two benchmark applications: 1) sum of cuPy array with
its transpose, and 2) cuDF merge. MPI4Dask speeds up the overall execution
time of the two applications by an average of $3.47\times$ and $3.11\times$
respectively on an in-house cluster built with NVIDIA Tesla V100 GPUs for
$1-6$ Dask workers. We also perform scalability analysis of MPI4Dask against
UCX for these applications on TACC’s Frontera (GPU) system with upto $32$ Dask
workers on $32$ NVIDIA Quadro RTX 5000 GPUs and $256$ CPU cores. MPI4Dask
speeds up the execution time for cuPy and cuDF applications by an average of
$1.71\times$ and $2.91\times$ respectively for $1-32$ Dask workers on the
Frontera (GPU) system.

###### Index Terms:

Python, Dask, MPI, MVAPICH2-GDR, Coroutines

##  I Introduction

With the end of Moore’s Law [1, 2] in sight, the performance advances in the
computer industry are likely to be driven by the “Top”—1) hardware
architecture, 2) software, and 3) algorithms—as noted by Leiserson recently
[3]. This is in stark comparison to the vision “There’s Plenty of Room at the
Bottom” laid out by Feynman [4] in 1959 referring to semiconductor physics and
fabrication technology. Leiserson [3] argues that the post-Moore generation of
software will focus on reducing software engineering bloat as well as
exploiting parallelism, locality, and heterogeneous hardware. A representative
of the “Top” is the Python programming language, which is a clear winner on
the landscape of data science. Python is a classic example of a language
benefiting from Moore’s law by traditionally focusing on programmer
productivity and reduced development time. In the context of data science
applications, Python’s popularity is due to rich set of free and open-source
libraries that enable the end-to-end data processing pipeline. These
libraries/packages include: core (SciPy), data preparation (NumPy, Pandas),
data visualization (Matlibplot), machine learning (Scikit-learn), and deep
learning (PyTorch, TensorFlow). However in the post-Moore era, it is vital
that Python is able to support parallel and distributed computing as well as
exploit emerging architectures especially accelerators.

Two popular big data computing frameworks that enable high-performance data
science in Python include Dask [5] and Apache Spark [6]. This paper focuses on
Dask, which provides support for natively extending popular data processing
libraries like numPy and Pandas. This allows incremental development of user
applications and lesser modifications to legacy codes when executed with Dask.
Traditionally Dask has mainly supported execution on hosts (CPUs) only, which
means that it was not able to exploit massive parallelism offered by Graphical
Processing Units (GPUs). However this has recently changed as part of the
development of the NVIDIA RAPIDS library, which aims to enable parallel and
distributed computation on clusters of GPUs. RAPIDS adopted a similar
approach—as taken by Dask—of extending already existing Python data processing
libraries for the GPU ecosystem. For instance, RAPIDS support processing of
distributed data stored in cuPy (numPy-like) and cuDF (Pandas-like) formats.
An overarching goal for RAPIDS project is to hide the low-level programming
complexities of the CUDA compute environment from Python developers and make
it easy to deploy and execute full end-to-end processing pipeline on GPUs. An
example of a machine learning library provided by RAPIDS is cuML [7], which is
the GPU-counterpart for Scikit-learn.

In order to support execution of Dask programs on cluster of GPUs, an
efficient communication layer is required. This is extended by the Dask
Distributed library that provides essentials for distributed execution of Dask
programs on parallel hardware. It is an asynchronous I/O application, which
means that it supports non-blocking and concurrent execution of its
routines/functions including communication primitives. This mandates the
restriction that any communication backend—aiming to be part of Dask
Distributed library—must implement coroutines that are non-blocking methods
defined and awaited by using the async and await keywords respectively.
Coroutines support concurrent mode of execution and hence cannot be invoked
like regular Python functions. The main reason behind this non-blocking
asynchronous style of programming is to avoid blocking, or delaying,
networking applications for completing I/O operations.

The Dask Distributed library currently provides two communication devices: one
for TCP and the other for UCX-Py that is a Cython-based wrapper library to UCX
[8]. The UCX-Py communication device is capable of communicating data to/from
GPU device memory directly. However, it fails to deliver high-performance—as
revealed by our performance evaluation detailed in Section V—since it delays
progressing the communication engine by assigning it to a separate coroutine
that executes periodically. We address this performance overhead by advocating
an efficient design where communication coroutines also progress the
communication engine.

The Message Passing Interface (MPI) standard [9] is considered as the defacto
programming model for writing parallel applications on modern GPU clusters.
The MVAPICH2-GDR  [10] library provides high-performance support for
communicating data to/from GPU devices via optimized point-to-point and
collective routines. As part of this paper, we design, implement, and evaluate
a new communication backend, called MPI4Dask, based on the MVAPICH2-GDR
library for Dask. MPI4Dask implements point-to-point communication coroutines
using mpi4py  [11] over MVAPICH2-GDR. To the best of our knowledge, this is
the first attempt to use an GPU-enabled MPI library to handle communication
requirements of Dask. This is challenging especially because the Dask
Distributed library is an asynchronous I/O application. Hence any point-to-
point communication operations—implemented via MPI—must be integrated as
communication coroutines that typically have performance penalties over
regular Python functions. An additional goal here is to avoid making
unnecessary modifications to the Dask ecosystem.

Another challenge that MPI4Dask addresses is to provide communication
isolation and support for dynamic connectivity between Dask components
including scheduler, workers, and client. Communication between Dask
entities—as provided by TCP and UCX backends—is based on the abstraction of
endpoints that are established when processes connect with one another. An
endpoint essentially represent a direct point-to-point connection. In MPI4Dask
, we devise a strategy, detailed in Section IV, that relies on sub-
communicator mechanism provided by MPI—coupled with communicator
duplication—to handle this. This allows us to build sub-communicators between
processes to mimic a direct endpoint-based connection.

In order to motivate the need for MPI4Dask, we present latency comparison
between Python communication coroutines implemented using MVAPICH2-GDR and
mpi4py with UCX-Py polling and event-based modes in Figure 1. The latency
graphs—Figure 1(a) and Figure 1(b) depict that the performance of Python
communication coroutines using MVAPICH2-GDR (with mpi4py) are roughly
$5\times$ better for small messages ($1$ Byte to $16$ Bytes) and $2-3\times$
better for small-medium messages ($32$ Bytes to $4$ KBytes). For large
messages ($2$ MBytes to $128$ MBytes) as shown in Figure 1(c), MVAPICH2-GDR
(with mpi4py)-based coroutines are better than UCX-Py point-to-point
coroutines by a factor of $3-4\times$. These performance benefits form the
basis for our motivation to design and implement MPI-based communication
backend for the Dask framework.

(a) Latency (Small)

(b) Latency (Medium)

(c) Throughput (Large)

Figure 1: Latency and Bandwidth comparison of MVAPICH2-GDR (using mpi4py) with
UCX using Ping-pong Benchmark—based on Python Coroutines—on the RI2 Cluster
(V100 GPUs). This demonstrates the performance benefits of of using
MVAPICH2-GDR as compared to UCX at the Python layer and forms the major
motivation of this paper.

We present a detailed performance evaluation of MPI4Dask against TCP and UCX
communication backends using a number of micro-benchmarks and application
benchmarks—these are presented in Section V. For the micro-benchmark
evaluation, MPI4Dask outperforms the UCX communication device in latency and
throughput comparison by ping-pong benchmarks by $5\times$ for small messages,
$2-3\times$ better for small-medium messages, and $3-4\times$ for large
messages. We used these two application benchmarks: 1) sum of cuPy array and
its transpose, and 2) cuDF merge on an in-house cluster and the TACC’s
Frontera (GPU) cluster. On the in-house cluster, we are witnessing an average
speedup of $3.47\times$ for $1-6$ Dask workers for the cuPy application.
Communication time for this application has been reduced by $6.92\times$
compared to UCX communication device. For the cuDF application on the in-house
cluster, there is an average speedup of $3.11\times$ for $2-6$ Dask workers
and $3.22\times$ reduction in communication time. On the Frontera (GPU)
cluster, MPI4Dask outperforms UCX by an average factor of $1.71\times$ for the
cuPy application with $1-32$ Dask workers. For the cuDF application, MPI4Dask
reduces the total execution time by an average factor of $2.91\times$ when
compared to UCX for $1-32$ Dask workers. Reasons for better overall
performance of MPI4Dask against its counterparts include better point-to-point
performance of MVAPICH2-GDR and efficient coroutine implementation. Unlike
UCX-Py that implements a separate coroutine to make progress for UCX worker,
MPI4Dask ensures cooperative progression where every communication coroutine
triggers the communication progression engine.

###  I-A Contributions

Main contributions of this paper are summarized below:

  1. 1.

Design and implementation of MPI4Dask that provides high-performance point-to-
point communication coroutines for Python-based HPC applications. To the best
of our knowledge, MPI4Dask is the first library that implements MPI-based
Python common coroutines that work with the asyncio framework on cluster of
GPUs.

  2. 2.

Integration of MPI4Dask with asynchronous Dask Distributed library. This is a
pioneering effort that enables MPI-based communication for the Dask ecosystem.

  3. 3.

Demonstrate the performance benefits of using MPI4Dask compared to TCP and UCX
devices using basic micro-benchmarks and two application benchmarks (based on
cuPy and cuDF) using an in-house cluster comprising of V100/K80 GPUs.

  4. 4.

Perform scalability evaluation of MPI4Dask against TCP and UCX communication
devices using two application benchmarks (based on cuPy and cuDF) on TACC’s
Frontera (GPU) system with upto $32$ Dask workers on $32$ NVIDIA Quadro RTX
5000 GPUs and $256$ CPU cores.

Rest of the paper is organized as follows. Section II presented relevant
background followed by the design approach of the MPI4Dask library in Section
III. Section IV presents implementation details. Later, we evaluate MPI4Dask
using a communication micro-benchmark and two application-level benchmarks.
Section VII concludes the paper.

##  II Background

This section covers the background for this paper, which includes the
fundamentals of the MPI standard, the Dask framework, and the asyncio package.
This section also presents related work.

###  II-A Message Passing Interface (MPI)

The Message Passing Interface (MPI) API is considered the defacto standard for
writing parallel applications. The MPI API defines a set of point-to-point and
collective communication routines that are provided as convenience functions
to application developers. In the context of the Dask Distributed library, the
most relevant set of functions are the non-blocking point-to-point
communication functions like MPI_Isend() and MPI_Irecv(). Both of these
functions return an MPI_Request object that can be used to invoke MPI_Test()
function to check if the non-blocking communication operation has completed or
not. In this paper, we make use of a GPU-aware MPI library called MVAPICH2-GDR
[10] that provides optimized point-to-point and collective communication
support for GPU devices. GPU-awareness here means that the MPI library is
capable of directly communication data efficiently to/from GPU memory instead
of staging it through the host system. The Dask ecosystem is implemented in
the Python programming language. This implies that the MPI communication
backend must also be implemented in Python. For this reason, we use the GPU-
aware mpi4py library that provides Cython [12] wrappers to native MPI
library—MVAPICH2-GDR in this case. For performance reasons, it is important
that the Python wrapper library is capable of communicating data to/from GPU
memory directly without incurring the overhead of serialization/de-
serialization. mpi4py supports efficient exchange of GPU data stored in cuPy
and cuDF format.

###  II-B Dask

Dask is a popular data science framework for Python programmers. It converts
user application into a task-graph, which is later executed lazily on
distributed hardware. This execution requires data exchange supported through
implicit communication by the Dask runtime. The Dask ecosystem is a suite of
Python packages. One such package is Dask Distributed that provides various
components like scheduler, workers, and client. This library also supports
point-to-point functionality between these Dask components. Figure 2 depicts
the distributed execution model of a Dask program. There are three types of
connections between Dask entities: 1) data connections (solid lines) for
exchanging application data, 2) control connections (dotted lines) for
exchanging heart-beat messages to maintain status of workers and detect any
failures, and 3) dynamic connections (dashed lines) between workers to resolve
dependencies, work-stealing, or achieving higher throughput. Currently Dask
has two communication devices: 1) the TCP device that exploits the
asynchronous Tornado library, 2) the UCX backend that uses UCX-Py [13]—a
Cython wrapper library—on top of the native UCX [8] communication library.

Figure 2: Dask Execution Model. The scheduler and workers form a Dask cluster.
The client executes user program by connecting with the Dask cluster.

###  II-C The asyncio Package

Since version $3.5$, Python has introduced a new package called asyncio that
allows writing concurrent non-blocking I/O applications. The main idea behind
this package is to allow networking (or other I/O) applications to efficiently
utilize CPU without unnecessarily getting blocked for long-running I/O
operations. This is possible because as the Python program executes, it keeps
defining tasks that are stored in a task queue and execute concurrently as
soon as it is possible to execute them. These tasks are defined through
coroutines, which are functions defined using the async keyword and invoked
later using the await keyword.

##  III Design Overview of MPI4Dask

This section presents the design of the MPI4Dask library. We first discuss
communication requirements of the Dask framework. This is followed by coverage
of layered architecture of Dask with focus on communication devices.

We begin with communication requirements of the Dask framework.

  1. 1.

Scalability: Provide scalability by exploiting low-latency and high-throughput
for cluster of GPUs typically deployed in modern data centers.

  2. 2.

Coroutines: The communication backend for Dask Distributed library is
asynchronous and executes within an event loop, as part of an asyncio
application. Hence, the communication backend needs to support non-blocking
point-to-point send and receive operations through asyncio coroutines defined
using the async/await syntax.

  3. 3.

Elasticity: The Dask cluster (e.g., scheduler and a set of workers) is an
elastic entity meaning that the workers dynamically join or leave the cluster.
In this context, the communication backend needs to support such dynamic
connectivity for clients and workers. Dask, however, also executes with static
number of processes using the Dask-MPI library. In this paper we adopt this
approach.

  4. 4.

Serialization/De-serialization: Dask applications support distributed
processing of GPU-based buffers including cuPy and cuDF DataFrames. The
communication engine should be able to support communication to/from GPU-based
Python objects/data-structures supported by Dask.

Figure 3 presents a layered architecture for Dask. The top layer represents
the Dask framework, which includes support for various data structures and
storage formats including Dask Bags, Arrays, and DataFrames. Also there is
support for asynchronous execution through Delayed and Future objects. The
user application is internally converted to a task graph by Dask. The next
layer shows various packages that are part of the Dask ecosystem. These
include Dask-MPI, Dask-CUDA, and Dask-Jobqueue. The next layer represents the
Dask Distributed library. This package supports distributed computation
through scheduler, worker, and client. This library also contains the
distributed.comm module, shown here as “Comm Layer”. This layer provides the
API that all Dask communication backends must implement. A subset of this API
can be seen in Listing 1. As shown, the Dask Distributed library provides TCP
and UCX-Py backends represented in this layer by tcp.py and ucx.py. MPI4Dask
is implemented as part of the Dask Distributed “Comm Layer” as a multi-layered
software sitting over mpi4py, which in turn exploits MVAPICH2-GDR for GPU-
aware MPI communication. The bottom layer in Figure 3 represents cluster of
GPUs as parallel hardware. Dask is also capable of running on shared memory
system like a laptop or a desktop.

Figure 3: Dask Layered Architecture with Communication Backends. Yellow boxes
are designed and evaluated as a part of this paper.

⬇

1class Comm(ABC):

2 @abstractmethod

3 def read(self, deserializers=None):

4 @abstractmethod

5 def write(self, msg, serializers=None,

6 on_error=None):

7

8class Listener(ABC):

9 @abstractmethod

10 async def start(self):

11 @abstractmethod

12 def stop(self):

13

14class Connector(ABC):

15 @abstractmethod

16 def connect(self, address, deserialize=True)

Listing 1: Subset of the Communication Backend API as Mandated by the Dask
Distributed Library

##  IV Implementation Details of the MPI4Dask Library

This section outlines the implementation details of the MPI4Dask communication
library.

###  IV-A Bootstrapping the Dask ecosystem and initializing MPI

There are many ways to start execution of Dask programs. The manual way is to
start the scheduler followed by multiple workers on the command line. After
this, the client is ready to connect with the Dask cluster and execute the
user program. In order to automate this process, Dask provides a number of
utility “Cluster” classes. An example of this is the LocalCUDACluster from the
Dask CUDA package. Once an instance of LocalCUDACluster object has been
initiated, a client can connect to execute a user program. Other example of
utility cluster classes include SLURMCluster, SGECluster, PBSCluster, and
others.

MPI4Dask currently requires that user initiates the execution of Dask program
using the Dask-MPI package. Dask-MPI uses the bootstrapping mechanism provided
by MPI libraries to start Dask scheduler, client, and one or more workers. We
use the mpirun_rsh utility provided by MVAPICH2-GDR . Figure 2 illustrates
this where the mpirun_rsh utility was used to start $4$ MPI processes. Here
processes $0$ and $1$ assume the role of scheduler and client, while all the
other processes—in this case $2-3$ become worker processes. This bootstrapping
is done by the Dask-MPI package and this is the point where mpi4py and
MVAPICH2-GDR are also initialized. The CUDA_VISIBLE_DEVICES environment
variable is used to map Dask worker processes on a particular GPU in a node.
This is particularly important when a node has multiple GPUs since we would
want multiple Dask workers to utilize distinct GPUs on the system. Note that
Dask-MPI does not provide any communication between Dask components and this
is what we address in this paper.

###  IV-B Point-to-point Communication Coroutines

A challenge that we tackle as part of this work is to implement asynchronous
communication coroutines using mpi4py over MVAPICH2-GDR that can be
incorporated inside the Dask Distributed layer. To the best of our knowledge,
this is first such effort to exploit MPI-based communication inside an asyncio
application in Python.

mpi4py provides two variants of point-to-point functions. The first option is
to use the lowercase methods like Comm.send()/Comm.isend() to communicate
data/to from Python objects. This involves picking/unpickling
(serialization/de-serialization) of Python objects. The second option is to
use methods like Comm.Send()/Comm.Isend() that start with uppercase letter and
communicate data to/from directly from user-specified buffer. The second
option is efficient and preferred for achieving high-performance in the Dask
Distributed library. Since we are implementing MPI4Dask for GPU buffers, the
specified buffer to mpi4py communication routines must support the
__cuda_array_interface__.

As previously mentioned, the Dask Distributed library is an asynchronous
application due to which it executes coroutines in a non-blocking and
concurrent manner. Also for this reason, using blocking MPI point-to-point
methods will result in deadlock for the application. In order to tackle this,
our implementation of MPI4Dask makes use of Comm.Isend()/Comm.Irecv() methods
that return Request objects. Later, MPI4Dask checks for completion of pending
communication using the Request.Test() method. Instead of checking the
completion status in a busy-wait loop—that will result in a deadlock
situation—MPI4Dask calls the asyncio.sleep() method that allows other
coroutines to make progress while waiting for communication to complete.
Listings 2 and 3 provide outlines of send and receive communication coroutines
in MPI4Dask.

⬇

1request = comm.Isend([buf, size], dest, tag)

2status = request.Test()

3

4while status is False:

5 await asyncio.sleep(0)

6 status = request.Test()

Listing 2: The Comm.Isend()-based Send Communication Coroutine Implemented by
MPI4Dask.

⬇

1request = comm.Irecv([buf, size], src, tag)

2status = request.Test()

3

4while status is False:

5 await asyncio.sleep(0)

6 status = request.Test()

Listing 3: The Comm.Irecv()-based Receive Communication Coroutine Implemented
by MPI4Dask.

###  IV-C Handling Large Messages

Dask is a programming framework for writing data science applications and it
is common for such libraries to handle and exchange large amounts of data. As
a consequence it is possible that the higher layers of the Dask ecosystem
attempt to communicate large messages using the underlying communication
infrastructure. We have experienced this with MPI4Dask that relies on the
point-to-point non-blocking primitives provided by MPI—Comm.Isend() and
Comm.Irecv() methods. These functions accept an argument int count that is
used to specify the size of the message being sent or received. The maximum
value that this parameter can hold is $2^{31}-1$ bytes, which corresponds a
message size of $2$ GB$-1$. On the other hand, the Dask Distributed library
attempts to communicate messages larger than this value including upto $64$
GB. We have catered for this requirement by dividing the large message into
several chunks of $1$ GB for the actual communication using the Comm.Isend()
and Comm.Irecv() methods. Typically the buffers specified to these
communication functions are Python objects and hence subscriptable using the
slice notation array[start:end]. This approach works for numPy and cuPy arrays
and hence the buffer argument can be subscripted in a loop to implement
chunking. But this strategy does not work for cuDF and RAPIDS Memory Manager
(RMM) buffers rmm.DeviceBuffer. For these, we implement chunking by
incrementing the offset argument in the communication loop.

###  IV-D Providing Communication Isolation

After the initial bootstrapping, MPI4Dask has full connectivity between all
processes. This is provided by the default communicator MPI_COMM_WORLD
initialized by the MPI library. However using this communicator naively might
lead to message interference. This is possible when a particular worker might
want to use same tag for both data and control message exchanges. There are
multiple ways to tackle this. The approach that we choose in this paper is to
rely on MPI sub-communicators. Figure 4 outlines this approach. The $4$ MPI
processes here engage with one another to form five sub-communicators
represented by colored bi-directional edges between a pair of processes. The
pseudocode for this is shown in Listing 4. Each process stores a handle to
these sub-communicators in a table called comm_table. Creating a new MPI sub-
communicator is a costly operation and for this reason all of this is done at
the startup. The MPI.Group.Incl() and MPI.COMM_WORLD.Create() functions (from
mpi4py) are used for creating sub-communicators that only contain two
processes. Complexity of making new sub-communicators is $O(\frac{n^{2}}{2})$
where $n$ is the total number of MPI processes. Our approach has negligible
overhead because these sub-communicators are initialized at the startup and
re-used later during the program execution.

Figure 4: Building new Sub-Communicators from MPI.COMM_WORLD at Startup. The
bi-directional edges between processes represent newly built sub-
communicators.

⬇

1for i in range(size):

2 for j in range(i+1, size):

3 incls = [i, j]

4 new_group = MPI.Group.Incl(group, incls)

5 new_comm = MPI.COMM_WORLD.Create(new_group)

6 if rank == i:

7 comm_table.update({j : new_comm})

8 else if rank == j:

9 comm_table.update({i : new_comm})

Listing 4: Nested Initialization Loop Executed at Each MPI Process during
Startup to Build new Sub-communicators from MPI.COMM_WORLD. The newly built
sub-communicators encapsulate only two processes and are stored in the
comm_table.

###  IV-E Handling Dynamic Connectivity

The Dask Distributed library support connectivity between Dask components
including scheduler, workers, and client. The UCX and TCP communication
devices maintain information for remote endpoints since this information is
required for the actual communication. In MPI4Dask, we have replaced the
abstraction of endpoint with sub-communicator. This enables us to utilize the
existing communication infrastructure of the Dask ecosystem. Apart from
communications being setup at the startup, Dask allows dynamic connections
between workers as well. In this context, new communication channels are only
established in the Dask Distributed library through traditional client/server
semantics of connect/accept. In MPI4Dask we handle this requirement by
starting a server that listens for incoming connections—and invokes a
connection handler callback function—using the asyncio.start_server() method.
Dask scheduler and worker processes use this function since they act as
listeners and other processes are allowed to connect with them. Later any
process (scheduler, client, or worker) that attempts connecting to a listener
process do so by using the asyncio.open_connection() function. When two
processes connect with one another—and this is something that happens
frequently in Dask at startup—each process gets the relevant sub-communicator
from the comm_table by doing a lookup based on the destination process rank in
the MPI.COMM_WORLD communicator. The corresponding sub-communicator from the
comm_table is duplicated by using the Comm.Dup() function. This ensures that a
new sub-communicator is built for every new dynamic connection. For
performance reasons, we are maintaining a configurable cache of these sub-
communicators.

##  V Performance Evaluation

This section presents performance evaluation of MPI4Dask against UCX and TCP
(using IPoIB) communication devices using 1) Ping Pong micro-benchmark
(Section V-A), 2) Two application benchmarks (cuPy and cuDF) on an in-house
cluster (RI2) with two types of GPU nodes with NVIDIA Tesla V100s and NVIDIA
Tesla K80s (Section V-B), and 3) Scalability results for the same two
application benchmarks on TACC’s Frontera (GPU) cluster with NVIDIA Quadro RTX
5000 GPUs (Section V-C). The hardware specifications for the RI2 cluster and
Frontera (GPU) subsystem are shown in Table I. The following versions of the
software were used: UCX v1.8.0, UCX-Py v0.17.0, Dask Distributed v2.30,
MVAPICH2-GDR v2.3.4, and mpi4py v3.0.3.

TABLE I: Hardware specification of the in-house RI2 and TACC’s Frontera (GPU) clusters. Columns 1 and 2 titled RI2-V100 and RI2-K80 provide details for two types of nodes on the RI2 cluster. Column 3 titled RI2-V100 provide details for the nodes used on the Frontera (GPU) cluster. Specification | RI2-V100 | RI2-K80 | Frontra (GPU)
---|---|---|---
Number of Nodes | 16 | 16 | 90
Processor Family | Xeon Broadwell | Xeon Broadwell | Xeon Broadwell
Processor Model | E5-2680 v4 | E5-2680 v4 | E5-2620 v4
Clock Speed | 2.4 GHz | 2.4 GHz | 2.1 GHz
Sockets | 2 | 2 | 2
Cores Per socket | 14 | 14 | 16
RAM (DDR4) | 128 GB | 128 GB | 128
GPU Family | Tesla V100 | Tesla K80 | Quadro RTX 5000
GPUs | 1 | 2 | 4
GPU Memory | 32 GB | 12 GB | 16 GB
Interconnect | IB-EDR (100G) | IB-EDR (100G) | IB-FDR (56G)

###  V-A Latency and Throughput Comparison

We first perform latency and throughput comparisons between MPI4Dask and other
communication backends—in particular UCX-Py—using a Ping Pong benchmark. Also
we add UCX (Tag API) and MVAPICH2-GDR to our comparisons for baseline
performance. UCX v1.8.0 was used alongwith ucx_perftest [14] to evaluate
latency and throughput for UCX Tag API. For MVAPICH2-GDR, we used the
osu_latency test that is part of the OSU Micro-Benchmark Suite (OMB) [15].
UCX-Py was also evaluated using its own Ping Pong benchmark [16]. UCX-Py has
two modes of execution: 1) polling-based, and 2) event-based. The polling-
based mode is latency-bound and hence more efficient than the event-based
mode. Comparisons also include mpi4py running over MVAPICH2-GDR—this is
labeled as MV2-GDR (mpi4py). The benchmark used for mpi4py is an in-house
Python version of OMB. Lastly we plot MPI4Dask that essentially represents the
MPI4Dask communication device integrated into the Dask Distributed library. We
always execute the TCP device using IP over InfiniBand (IPoIB) protocol that
provides best performance for this backend. For this reason, we use TCP and
IPoIB interchangeably in this section.

(a) Latency (Small)

(b) Latency (Medium)

(c) Throughput (Large)

Figure 5: Latency/Bandwidth comparison of MPI4Dask with UCX-Py (Polling and
Event Modes) using Ping Pong Benchmark on RI2 Cluster with V100 GPUs. UCX (Tag
API) and MVAPICH2-GDR numbers are also presented for baseline performance.

(a) Latency (Small)

(b) Latency (Medium)

(c) Throughput (Large)

Figure 6: Latency/Throughput comparison of MPI4Dask with UCX-Py (Polling and
Event Modes) using Ping Pong Benchmark on RI2 Cluster with K80 GPUs. UCX (Tag
API) and MVAPICH2-GDR numbers are also presented for baseline performance.

Figure 5(a) shows latency for message sizes $1$ Byte to $4$ KByte—here
MPI4Dask outperforms UCX-Py (Polling) by $5\times$. Similarly for medium-sized
messages presented in Figure 5(b)—$8$ KByte to $512$ KByte—MPI4Dask is better
than UCX-Py (Polling) by $2-3\times$. In throughput comparisons for large
messages, $1$ MByte and beyond, MPI4Dask outperforms UCX-Py (Polling) by
$3-4\times$ in throughput as shown in Figure 5(c). It is important to note
that the difference in performance between MV2-GDR (mpi4py) and MV2-GDR is
minimal. This means that regular Python functions over native code do not
introduce much overhead. However when executing communication coroutines—by
employing MPI4Dask or UCX-Py—there is additional overhead due to the asyncio
framework. We observe similar performance patterns for Tesla K80 GPUs in
latency and throughput comparisons between MPI4Dask and UCX-Py (Polling) as
depicted in Figure 6. As noted earlier, we are focusing on performance
comparison of MPI4Dask with UCX-Py (Polling) since it is the more efficient
mode of the UCX-Py library.

###  V-B Application Benchmarks

We also evaluated MPI4Dask against UCX-Py for two application benchmarks: 1)
sum of cuPy array and its transpose, and 2) cuDF merge. The cuPy benchmark
presents strong scaling results as the problem size remains the same as more
Dask workers are used in the computation. On the other hand, the cuDF
benchmark presents weak scaling results as the problem size increases with
increase in the number of Dask workers. This evaluation was done on the in-
house RI2 cluster.

####  V-B1 Sum of cuPy Array and its Transpose

This benchmark [17] creates a cuPy array and distributes its chunks across
Dask workers. The benchmark adds these distributed chunks to their transpose,
forcing the GPU data to move around over the network. The following operations
are performed:
y = x + x.T
y = y.persist()
wait(y)

Performance comparison graphs for the first application benchmark—sum of cuPy
array with its transpose—are shown in Figure 7. Dask follows the One Process
per GPU (OPG) model of execution, which means that a single worker process
with multiple threads is initiated for each GPU. On the RI2 cluster, each node
has a single GPU. For this reason, we instantiate a single worker process with
$28$ worker threads to fully exploit the available CPU cores. Figure 7(a)
shows execution time where we are witnessing an average speedup of
$3.47\times$ for $2-6$ Dask workers. Figure 7(b) shows the communication time
for the benchmark run. This shows that MPI4Dask is better than UCX by
$6.92\times$ on average for $2-6$ workers. Average throughput comparison for
Dask workers is shown in Figure 7(c), which depicts that MPI4Dask outperform
UCX by $5.17\times$ on average for $2-6$ workers. This benchmark application
is presenting strong scaling results. As can be seen in Figure 7(a), the cuPy
application does not exhibit impressive speedups as the number of workers are
increased. This is due to the nature of this benchmark that is designed to
stress and evaluate communication performance. This is suitable for this paper
because we are primarily interested in comparative performance of
communication devices and not demonstrating application-level performance of
the Dask ecosystem.

####  V-B2 cuDF Merge

cuDF DataFrames are table-like data-structure that are stored in the GPU
memory. As part of this application [18], a merge operation is carried out for
multiple cuDF data frames. Performance comparison graphs for the second
application benchmark—cuDF merge operation—are shown in Figure 8. Figure 8(a)
shows execution time where we are witnessing an average speedup of
$3.11\times$ for $2-6$ Dask workers. Figure 8(b) shows time for communication
that is taking place. This shows that MPI4Dask is better than UCX by
$3.22\times$ on average for $2-6$ workers. Average throughput comparison for
Dask workers is presented in Figure 8(c), which depicts that MPI4Dask
outperforms UCX by $3.82\times$ on average for $2-6$ workers. The cuDF merge
benchmark is presenting weak scaling results. As can be observed in Figure 8,
the execution time with MPI4Dask is increasing slightly with an increase in
the number of workers. This is due to efficient communication performance
provided by MPI4Dask to the Dask execution. Note that $3.2$, $6.4$, $9.6$,
$12.8$, $16$, and $19.2$ GB of data is processed for $1$, $2$, $3$, $4$, $5$,
and $6$ Dask workers respectively.

Note that UCX-Py is executed in polling-mode that is the efficient mode.
Reasons for better overall performance of MPI4Dask against other counterparts
include better point-to-point performance of MVAPICH2-GDR , chunking
implemented for messages greater than $1$ GB, and efficient coroutine
implementation for MPI4Dask as compared to UCX-Py. Unlike UCX-Py that
implements a separate coroutine to make communication progress for UCX worker,
MPI4Dask ensures cooperative progression where every communication coroutine
triggers the communication progression engine.

(a) Execution Time Comparison

(b) Communication Time Comparison

(c) Aggregate Throughput Between Workers

Figure 7: Sum of cuPy Array and its Transpose (cuPy Dims: 16K$\times$16K,
Chunk size: 4K, Partitions: 16): Performance Comparison between IPoIB, UCX,
and MPI4Dask on the RI2 Cluster. This benchmark presents strong scaling
results. $28$ threads are started in a single Dask worker.

(a) Execution Time Comparison

(b) Communication Time Comparison

(c) Aggregate Throughput Between Workers

Figure 8: cuDF Merge Operation Benchmark (Chunk Size: 1E8, Shuffle: True,
Fraction Match: 0.3): Performance Comparison between IPoIB, UCX, and MPI4Dask
on the RI2 Cluster. This benchmark presents weak scaling results. $28$ threads
are started in a single Dask worker. $3.2$ GB of data is processed per Dask
worker (or GPU)

###  V-C Scalability Results on the TACC Frontera (GPU) Cluster

This section presents the scalability results for the two application
benchmarks introduced earlier in Section V-B. This evaluation was done on the
Frontera (GPU) system that is equipped with $360$ NVIDIA Quadro RTX 5000 GPUs
in $90$ nodes. Following the OPG model of execution, we started $4$ processes
on a single node—one process per GPU—with $8$ worker threads in each process.
This evaluation is presented in Figure 9. Figure 9(a) shows the execution time
comparison between TCP (using IPoIB), UCX, and MPI4Dask devices for the sum of
cuPy with its transpose application with $1-32$ Dask workers—here MPI4Dask
outperforms UCX by an average factor of $1.71\times$. Figures 9(b) and 9(c)
shows the execution time and throughput comparison between all communication
devices for the cuDF merge operation. Again, MPI4Dask outperforms UCX by an
average factor of $2.91\times$ for overall execution time for $1-32$ Dask
workers. The merge operation throughout depicts the rate at which workers
merge the cuDF input data and it also follows a similar pattern. The
performance results shown in Figure 9(a) present strong scaling results for
the cuPy application, while Figures 9(b) and 9(c) present weak scaling results
for the cuDF merge operation. The performance shown by $1$ Dask worker in
Figures 9(b) and 9(c) is efficient due to small problem size and no
communication overhead. However, the problem size becomes larger and more
realistic with additional Dask workers. Note that $1.6$ GB of data is
processed per GPU for the cuDF application—this means that $51.2$ GB of data
is processed on $32$ GPUs.

(a) Execution Time Comparison for Sum of cuPy Array and its Transpose
Benchmark. (cuPy Dims: 20E3$\times$20E3, Chunk size: 5E2, Partitions: 16E2)

(b) Execution Time Comparison for cuDF Merge Benchmark. (Chunk Size: 5E7,
Shuffle: True, Fraction Match: 0.3)

(c) Average Throughput For Workers for the cuDF Merge Benchmark

Figure 9: Execution Time Comparison for Sum of cuPy Array with its Transpose
Benchmark, cuDF Merge Benchmark, and Average Throughput for Dask Workers for
the cuDF Merge Benchmark. This evaluation was done on the Frontera (GPU)
system at TACC. This comparison is between MPI4Dask , UCX, and TCP (using
IPoIB) communication devices. There are $4$ GPUs in a single node. Each Dask
worker has $8$ threads.

##  VI Related Work

A framework similar to Dask is the Apache Spark software [6]. It supports
programming data science applications using Scala, Java, Python and R.
Traditionally Spark has also only supported execution on hosts (CPUs). This
meant that it was not able to exploit massive parallelism offered by GPUs.
However GPU support has been added by the RAPIDS project recently to the
Apache Spark 3.0 through a Spark-RAPIDS plugin [19]. Vanilla versions of
Apache Spark only support data communication over Ethernet or other high-speed
networks using the TCP backend. There have been efforts [20, 21] to address
this shortcoming for Apache Spark.

The Dask Distributed library currently support two communication backends. The
first device—called TCP—makes use of Python’s Tornado framework [22]. The
second device—called UCX—uses a Cython wrapper called UCX-Py [13] over the UCX
[8] communication library. In an earlier effort [23], we have developed a new
communication device called Blink for Dask. However Blink is limited to CPU-
only execution of Dask programs and hence is not relevant here. Also, recently
there has been another effort [24] to develop an MPI communication layer on
the Dask Distributed github repository. However, this is currently a work in
progress and is not addressing GPU-based data science applications, which is
the focus of this paper. At this time, it is only possible to use TCP and UCX
based devices for Dask on cluster of GPUs. This paper demonstrates that
MPI4Dask —developed as part of this work—delivers better performance than the
current state of the art communication options for Dask on cluster of GPUs.

##  VII Conclusions and Future Work

Python is an emerging language on the landscape of scientific computing with
support for distributed computing engines like Dask. Rivaling the Apache Spark
ecosystem, Dask allows incremental parallelization of Big Data applications.
Recently the Dask software has been extended, as part of the NVIDIA RAPIDS
framework, to support distributed computation on cluster of GPUs. Support for
GPU-based data storage and processing APIs like cuPy and cuDF—GPU-counterparts
of numPy and Pandas respectively—has also been added. As part of this work, we
have extended the Dask Distributed library with a new communication device
called MPI4Dask based on the popular MPI standard. MPI4Dask exploits the
mpi4py wrapper software on top of the MVAPICH2-GDR library to offer an
efficient alternative to existing backends—TCP and UCX—from within a non-
blocking asynchronous I/O framework. This is done by implementing high-
performance communication coroutines using MPI that support the async/await
syntax. The performance evaluation done, on an in-house cluster built with
Tesla V100/K80 GPUs and the Frontera (GPU) cluster equipped with Quadro RTX
5000 GPUs—using the Ping Pong micro-benchmark and two other application
benchmarks—suggest that MPI4Dask clearly outperforms other communication
devices. This is due to efficient implementation of communication coroutines
in MPI4Dask that use the cooperative progress approach with non-blocking
point-to-point MPI calls. In the future we plan to extend MPI4Dask with
support for dynamic process management, which will allow entities like workers
and clients to dynamically join and leave the Dask Cluster. Also we aim to
extend MPI4Dask for host-based (CPU only) distributed computation with Dask.

## References

  * [1] G. E. Moore, “Cramming more components onto integrated circuits,” _Electronics_ , vol. 38, no. 8, April 1965.
  * [2] G. E. Moore, “Progress in digital integrated electronics [technical literaiture, copyright 1975 ieee. reprinted, with permission. technical digest. international electron devices meeting, ieee, 1975, pp. 11-13.],” _IEEE Solid-State Circuits Society Newsletter_ , vol. 11, no. 3, pp. 36–37, 2006.
  * [3] C. E. Leiserson, N. C. Thompson, J. S. Emer, B. C. Kuszmaul, B. W. Lampson, D. Sanchez, and T. B. Schardl, “There’s plenty of room at the top: What will drive computer performance after moore’s law?” _Science_ , vol. 368, no. 6495, 2020. [Online]. Available: https://science.sciencemag.org/content/368/6495/eaam9744
  * [4] R. P. Feynman, “There’s Plenty of Room at the Bottom,” 1959.
  * [5] M. Rocklin, “Dask: Parallel computation with blocked algorithms and task scheduling,” in _Proceedings of the 14th Python in Science Conference_ , K. Huff and J. Bergstra, Eds., 2015, pp. 130 – 136.
  * [6] M. Zaharia, R. S. Xin, P. Wendell, T. Das, M. Armbrust, A. Dave, X. Meng, J. Rosen, S. Venkataraman, M. J. Franklin _et al._ , “Apache Spark: A Unified Engine for Big Data Processing,” _Communications of the ACM_ , vol. 59, no. 11, pp. 56–65, 2016.
  * [7] S. Raschka, J. T. Patterson, and C. Nolet, “Machine learning in python: Main developments and technology trends in data science, machine learning, and artificial intelligence,” _ArXiv_ , vol. abs/2002.04803, 2020.
  * [8] P. Shamis, M. G. Venkata, M. G. Lopez, M. B. Baker, O. Hernandez, Y. Itigin, M. Dubman, G. Shainer, R. L. Graham, L. Liss, Y. Shahar, S. Potluri, D. Rossetti, D. Becker, D. Poole, C. Lamb, S. Kumar, C. Stunkel, G. Bosilca, and A. Bouteiller, “UCX: An Open Source Framework for HPC Network APIs and Beyond,” in _2015 IEEE 23rd Annual Symposium on High-Performance Interconnects_ , 2015, pp. 40–43.
  * [9] “Message Passing Interface (MPI),” http://www.mpi-forum.org. Accessed: .
  * [10] MVAPICH2-GDR Development Team, “MVAPICH2-GDR User Guide,” http://mvapich.cse.ohio-state.edu/userguide/gdr.
  * [11] L. D. Dalcin, R. R. Paz, P. A. Kler, and A. Cosimo, “Parallel distributed computing using python,” _Advances in Water Resources_ , vol. 34, no. 9, pp. 1124 – 1139, 2011, new Computational Methods and Software Tools. [Online]. Available: http://www.sciencedirect.com/science/article/pii/S0309170811000777
  * [12] S. Behnel, R. Bradshaw, C. Citro, L. Dalcin, D. S. Seljebotn, and K. Smith, “Cython: The best of both worlds,” _Computing in Science & Engineering_, vol. 13, no. 2, pp. 31–39, 2011.
  * [13] UCX, “UCX-Py Github,” https://github.com/rapidsai/ucx-py.
  * [14] ——, “Performance Measurement,” https://github.com/openucx/ucx/wiki/Performance-measurement.
  * [15] Network-Based Computing Laboratory, “OSU Micro-Benchmarks 5.6.3,” http://mvapich.cse.ohio-state.edu/benchmarks.
  * [16] UCX-Py, “Benchmarks,” https://github.com/rapidsai/ucx-py/blob/branch-0.15/benchmarks/local-send-recv.py.
  * [17] NVIDIA RAPIDS Team, “Sum of cuPy Array and its Transpose,” https://github.com/rapidsai/dask-cuda/blob/branch-0.16/dask_cuda/benchmarks/local_cupy_transpose_sum.py.
  * [18] ——, “cuDF Merge,” https://github.com/rapidsai/dask-cuda/blob/branch-0.16/dask_cuda/benchmarks/local_cudf_merge.py.
  * [19] NVIDIA, “RAPIDS Accelerator For Apache Spark,” https://github.com/NVIDIA/spark-rapids.
  * [20] X. Lu, D. Shankar, S. Gugnani, and D. K. Panda, “High-performance design of apache spark with RDMA and its benefits on various workloads,” in _2016 IEEE International Conference on Big Data (Big Data)_ , 2016, pp. 253–262.
  * [21] M. Inc., _SparkRDMA_ , 2020 (accessed July 15, 2020). [Online]. Available: https://github.com/Mellanox/SparkRDMA
  * [22] Tornado Web Server, “Dask TCP Backend,” https://www.tornadoweb.org.
  * [23] A. Shafi, J. Hashmi, H. Subramoni, and D. Panda, “ Blink: Towards Efficient RDMA-based Communication Coroutines for Parallel Python Applications ,” December 2020.
  * [24] Ian Thomas, “[WIP]: MPI communication layer,” https://github.com/dask/distributed/pull/4142.