我们使用 cuDF 和 Dask 实验了单节点多 GPU 连接。我们发现 GPU 内的计算比通信更快。我们还介绍了近期工作的背景和计划，包括使用 UCX 改进 Dask 中的高性能通信。
在最近的一篇文章中，我们展示了 Dask + cuDF 如何利用多个 GPU 并行加速读取 CSV 文件。在添加了几个 GPU 后，该操作很快受限于磁盘速度。现在我们尝试一种非常不同的操作：多 GPU 连接。
这种工作负载可能通信量很大，尤其是当连接列未良好排序时，因此提供了与解析 CSV 相反的另一个极端情况下的良好示例。
在这里，我们使用 Dask 数组和 Dask 数据框构建了两个具有共享 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)
我们有两个 Dask 数据框，由我们随机数据的许多 Pandas 数据框组成。现在我们在这上面映射 cudf.from_pandas 函数，以创建一个由 cuDF 数据框组成的 Dask 数据框。
import dask
import cudf
gleft = left.map_partitions(cudf.from_pandas)
gright = right.map_partitions(cudf.from_pandas)
gleft, gright = dask.persist(gleft, gright) # 将数据持久化到设备内存中
这里的妙处在于没有任何特殊的 dask_pandas_dataframe_to_dask_cudf_dataframe 函数。Dask 与 cuDF 很好地组合在一起。我们无需做任何特殊的事情来支持它。
我们还将数据持久化到了设备内存中。
在此之后，简单的操作变得轻松快速，并利用了我们的八个 GPU。
>>> gleft.x.sum().compute() # 这需要 250ms
500004719.254711
我们将使用标准的 Pandas 语法合并数据集，将结果持久化到 RAM 中，然后等待
out = gleft.merge(gright, on=['id']) # 这是惰性的
现在我们查看此计算的 Dask 诊断图。
当我们查看 Dask 的任务流图时，我们看到我们的八个线程（每个线程管理一个 GPU）中的每一个都将大部分时间花在通信上（红色表示通信时间）。实际的合并和连接任务相对于数据传输时间来说相当快。
这并不太令人惊讶。对于此计算，我禁用了设备之间的任何通信尝试（更多内容见下文），因此数据正在从 GPU 移动到 CPU 内存，然后进行序列化并放到 TCP 套接字上。我们在一台机器上移动了数十 GB 的数据，因此系统总吞吐量约为 1GB/s，这对于 Python 中的 TCP-on-localhost 来说是典型的。
我们还可以更深入地查看 Dask 火焰图风格的计算成本图。这显示了我们函数中的哪些行占用了大部分时间（至少到 Python 级别）。
这个 火焰图 显示了我们在计算时（不包括上面提到的主要通信成本）花费时间的 cudf 代码行。对于那些试图进一步优化性能的人来说，这可能会很有趣。它显示我们的大部分开销都在内存分配上。与通信一样，这个问题实际上也已在 RAPIDS 的可选内存管理池中得到修复，只是它还不是默认设置，所以我在这里没有使用它。
cuDF 库实际上有一个不错的单节点多 GPU 通信方法，我在本次实验中特意关闭了它。该方法巧妙地利用 Dask 使用 Dask 的常规通道（这很小且快速）来传递设备指针信息，然后使用该信息启动侧通道通信来传输大部分数据。这种方法有效，但有些脆弱。我倾向于放弃它，转而选择……
UCX。 UCX 项目提供了一个单一 API，它封装了多种传输方式，如 TCP、Infiniband、共享内存以及 GPU 特定的传输方式。UCX 声称可以根据可用的硬件找到在两个点之间传输数据的最佳方式。如果 Dask 能够使用它进行通信，那么它将提供单机上高效的 GPU 到 GPU 通信，以及在存在 Infiniband 等高效网络硬件时（甚至在 GPU 上下文之外）高效的机器间通信。
在这方面我们还有一些工作要做
这项工作正在由研究 UCX 的 Akshay Vekatesh 和 Dask/Pandas 核心开发者 Tom Augspurger 进行中。我猜他们很快就会写关于它的文章。我期待着看到这项工作的成果，无论对于 Dask 还是对于普遍的高性能 Python 来说。
值得指出的是，这项工作不仅会帮助 GPU 用户。它应该也会帮助使用高级网络硬件的任何人，包括主流的科学 HPC 社区。
单节点多 GPU 连接潜力巨大。事实上，早期的 RAPIDS 开发者通过我简要提及的巧妙通信技巧，使其运行速度比我上面实现的要快得多。本文的主要目的是提供一个可用于未来连接的基准测试，并强调在并行计算中何时通信至关重要。
现在 GPU 加速了我们每个工作块的计算时间，我们越来越发现其他系统成为了瓶颈。以前我们不太关心通信，因为计算成本是可比的。现在计算成本降低了一个数量级，我们技术栈的其他方面变得更加重要。
我期待着看到这项工作的发展。
如果上面的工作听起来让你感兴趣，那就来帮忙吧！有很多容易上手且影响巨大的工作可以做。
如果您有兴趣通过专注于这些主题获得报酬，那么请考虑申请一份工作。NVIDIA 的 RAPIDS 团队正在招聘使用 GPU 进行 Dask 开发以及其他数据分析库开发项目的工程师。
版权所有 © 2022 Dask 核心开发者。采用 New-BSD 许可证。
