单节点多 GPU Dataframe 连接

摘要

我们使用 cuDF 和 Dask 对单节点多 GPU 连接进行了实验。我们发现 GPU 内计算比通信更快。我们还介绍了近期工作的背景和计划，包括使用 UCX 改进 Dask 中的高性能通信。

这是本文实验的 Notebook

引言

在最近的一篇文章中，我们展示了 Dask + cuDF 如何利用多个 GPU 并行加速 CSV 文件读取。在我们添加了一些 GPU 后，该操作很快受限于磁盘速度。现在我们尝试一种非常不同的操作：多 GPU 连接。

这个工作负载可能会涉及大量通信，特别是如果用于连接的列没有很好的排序。因此，它提供了一个很好的例子，与解析 CSV 形成另一个极端。

基准测试

使用 CPU 构建随机数据

这里我们使用 Dask array 和 Dask dataframe 来构建两个共享 id 列的随机表。我们可以调整每个表的行数和键的数量，以各种方式使连接具有挑战性。

import dask.array as da
import dask.dataframe as dd

n_rows = 1000000000
n_keys = 5000000

left = dd.concat([
    da.random.random(n_rows).to_dask_dataframe(columns='x'),
    da.random.randint(0, n_keys, size=n_rows).to_dask_dataframe(columns='id'),
], axis=1)

n_rows = 10000000

right = dd.concat([
    da.random.random(n_rows).to_dask_dataframe(columns='y'),
    da.random.randint(0, n_keys, size=n_rows).to_dask_dataframe(columns='id'),
], axis=1)

发送到 GPU

我们有两个 Dask dataframe，由许多包含我们随机数据的 Pandas dataframe 组成。现在，我们对这些数据应用 cudf.from_pandas 函数，生成一个由 cuDF dataframe 组成的 Dask dataframe。

import dask
import cudf

gleft = left.map_partitions(cudf.from_pandas)
gright = right.map_partitions(cudf.from_pandas)

gleft, gright = dask.persist(gleft, gright)  # persist data in device memory

这里的优点在于，没有特殊的 dask_pandas_dataframe_to_dask_cudf_dataframe 函数。Dask 与 cuDF 很好地结合了。我们无需做任何特别的事情来支持它。

我们还将数据持久化在设备内存中。

之后，简单的操作变得容易且快速，并使用了我们的八个 GPU。

>>> gleft.x.sum().compute()  # this takes 250ms
500004719.254711

连接

我们将使用标准的 Pandas 语法合并数据集，将结果持久化在 RAM 中，然后等待。

out = gleft.merge(gright, on=['id'])  # this is lazy

分析和结果

现在我们查看 Dask 为此计算生成的诊断图。

任务流和通信

当我们查看 Dask 的任务流图时，我们看到我们的八个线程（每个线程管理一个 GPU）将大部分时间花在了通信上（红色表示通信时间）。实际的合并（merge）和连接（concat）任务相对于数据传输时间来说非常快。

这并不太令人惊讶。对于这次计算，我关闭了设备之间的任何通信尝试（下面会详细介绍），因此数据正在从 GPU 移动到 CPU 内存，然后序列化并放入 TCP 套接字。我们在单机上移动了几十 GB 数据，系统的总吞吐量约为 1GB/s，这对于 Python 中的 TCP-on-localhost 来说是典型的。

计算火焰图

我们还可以更深入地查看 Dask 火焰图样式的计算成本图。这显示了我们函数中哪行代码花费的时间最多（至少在 Python 级别）。

这个火焰图显示了我们在计算时（不包括上述主要通信成本）在哪些 cuDF 代码行上花费了时间。对于那些试图进一步优化性能的人来说，这可能很有趣。它显示我们大部分成本都在内存分配上。像通信一样，这实际上也已经在 RAPIDS 的可选内存管理池中得到修复，但它目前还不是默认设置，所以我这里没有使用它。

高效通信计划

cuDF 库实际上有一种不错的单节点多 GPU 通信方法，我在本次实验中有意关闭了它。该方法巧妙地利用 Dask 的常规通道（这部分很小且快速）来传递设备指针信息，然后利用这些信息启动侧通道通信进行大部分数据传输。这种方法很有效，但有些脆弱。我倾向于放弃它，转而采用……

UCX。UCX 项目提供了一个单一 API，它封装了多种传输方式，如 TCP、Infiniband、共享内存以及 GPU 特定的传输。UCX 声称能够根据可用的硬件找到在两点之间传输数据的最佳方式。如果 Dask 能够使用它进行通信，那么它将不仅在单机上提供高效的 GPU 到 GPU 通信，而且在存在 Infiniband 等高效网络硬件时，即使在 GPU 上下文之外，也能提供高效的跨机器通信。

我们在这里需要做一些工作：

1. 我们需要构建一个 UCX 的 Python 封装库
2. 我们需要围绕这个 ucx-py 库构建一个可选的 Dask 通信模块 (Comm)，允许用户指定端点，例如 ucx://path-to-scheduler
3. 我们需要构建引用设备内存的 Python 内存视图（memoryview）类似对象
4. ...

这项工作已经在进行中，由 UCX 开发者 Akshay Vekatesh 和核心 Dask/Pandas 开发者 Tom Augspurger 负责。我猜他们很快就会写文章介绍它。我很期待看到它的成果，无论是对 Dask 还是对高性能 Python 总体而言。

值得指出的是，这项努力不仅仅会帮助 GPU 用户。它应该能帮助任何使用高级网络硬件的用户，包括主流的科学 HPC 社区。

摘要

单节点多 GPU 连接前景光明。事实上，早期的 RAPIDS 开发者通过我简短提及的巧妙通信技巧，使得运行速度比我上面实现的快得多。本文的主要目的是为连接提供一个我们将来可以使用的基准，并强调在并行计算中通信何时至关重要。

既然 GPU 已经加速了我们每个工作块的计算时间，我们发现其他系统越来越成为瓶颈。以前我们不太关心通信，因为计算成本与之相当。现在计算成本降低了一个数量级，我们技术栈的其他方面变得更加重要。

我很期待看到事情的发展。

来帮忙！

如果上面的工作让你感兴趣，那就来帮忙吧！有很多容易实现且影响深远的工作要做。

如果您有兴趣受薪专注于这些课题，请考虑申请工作。NVIDIA 的 RAPIDS 团队正在招聘工程师，负责 Dask 与 GPU 的开发以及其他数据分析库开发项目。

• 高级库软件工程师 - RAPIDS