十亿行挑战赛 (One Billion Row Challenge) 是一个有趣的基准测试,旨在展示基本的数据处理操作。它最初是作为纯 Java 竞赛发起的,现已聚集了其他语言(包括 Python、Rust、Go、Swift 等)的开发者社区。对于许多有兴趣探索文本文件阅读细节、基于哈希的算法和 CPU 优化的软件工程师来说,这项挑战赛非常有用。截至 2024 年年中,One Billion Row Challenge GitHub 存储库已经吸引了超过 1.8K 个分叉,获得了超过 6K 颗星,并启发了数十篇博客文章和视频。
本文将展示如何使用 RAPIDS cuDF pandas 加速器模式完成处理十亿行数据的挑战。具体而言,我们将展示 cuDF pandas 加速器版本 24.08 中的两项新功能 —— 大字符串支持和带预取的托管内存 —— 如何借助 GPU 加速的数据处理工作流程提高大数据的性能。
使用 RAPIDS cuDF pandas 加速器模式进行数据处理
pandas 是一个基于 Python 构建的开源软件库,专门用于数据处理和分析。这是一个灵活的数据处理工具,支持完成一亿行挑战赛所需的操作,包括解析文本文件、按组聚合数值数据以及对表格进行排序。
RAPIDS cuDF 是一个 GPU DataFrame 库,可提供类似于 pandas 的 API,用于加载、过滤和操作数据。RAPIDS cuDF pandas 加速器模式通过统一的 CPU/GPU 用户体验,为 pandas 工作流带来加速计算,且无需更改代码。如需了解更多信息,请参阅 RAPIDS cuDF 将 pandas 的速度提升近 150 倍,无需更改代码。
以下 pandas 脚本足以完成一亿行挑战:
import pandas as pd
df = pd.read_csv(
“measurements.txt”,
sep=';',
header=None,
names=["station", "measure"]
)
df = df.groupby("station").agg(["min", "max", "mean"])
df = df.sort_values("station")
在 cuDF pandas 加速器模式下运行此脚本需要在 Python 或 Jupyter Notebooks 中导入 pandas 之前添加单行命令,该命令为 %load_ext cudf.pandas
。
%load_ext cudf.pandas
import pandas as pd
详细了解 RAPIDS cuDF pandas 加速器模式,包括在 Python 中使用此模式的不同方式,例如使用 Python 模块标志,或通过导入显式启用。
RAPIDS cuDF pandas 加速器模式下的新大型数据处理功能 24.08
24.08 版本的 RAPIDS cuDF pandas 加速器模式包括两个关键功能,可实现更高效的数据处理:大字符串支持和带预取的托管内存池。这些功能共同支持大型 DataFrame 处理 – 多达 2.1亿行的数据,即使在总 GPU 内存为 2-3 倍的情况下也能保持出色的性能。请注意,适用于 Linux 的 Windows 子系统 (WSL2) 对 GPU 超额认购的支持有限,本文中介绍的结果收集于 Ubuntu 22.04。
大字符串支持
大字符串支持使 RAPIDS cuDF 能够在 32 位和 64 位索引之间动态切换。与 PyArrow 中存在字符串和大字符串类型不同,cuDF 中的字符串只有当列数据超过 2.1亿个字符时才会切换到 64 位索引。这使得 cuDF 能够保持较低的内存占用,并提高对小于 2.1亿个字符的列的处理速度,同时仍然支持对大字符串列进行高效处理。
以前,在版本 24.06 中,当 GPU 上的字符串列超过 2.1亿个字符时,会出现字符串溢出,而由此产生的溢出错误会导致数据复制回主机并回退至 pandas 处理。现在,在版本 24.08 中,DataFrames 可能包含大大小小的字符串列,并且每列都会正确处理 cuDF 中的字符串列类型。
带预提取的托管内存池
具有预取功能的受管理内存池使 cuDF 能够同时使用 GPU 和主机内存来存储数据并避免内存不足错误。受管理内存(也称为 CUDA 统一虚拟内存)可维护由 GPU 和主机内存支持的单个地址空间。当 GPU 内核启动时,任何数据如果不可被 GPU 访问,将从主机内存分页(迁移)到 GPU 内存。使用具有受管理内存的内存池可减少每次分配的开销,并减少总体执行时间。预取对于使用受管理内存观察良好性能也很重要,因为它有助于确保数据可用于 GPU 内核,而无需在计算时对数据进行分页,这可能“为时已晚”。
以前,在版本 24.06 中,较大的数据集更有可能耗尽总 GPU 内存,而由此产生的内存不足错误还将导致数据复制回主机并返回 cuDF pandas 处理。现在,在版本 24.08 中,cuDF pandas 加速器模式使用启用预取的托管内存池。请注意,大数据的最佳性能可能取决于数据和工作流程,我们欢迎您的反馈。
使用 NVIDIA GPU 运行一亿行挑战赛
您可以在高显存和低显存的 GPU 上运行挑战,这显示了 RAPDIS cuDF 24.08 中大数据功能对性能的影响。在行数为 10 亿的情况下,挑战始于一个 13.1 GB 的文本文件。在 cuDF 中,这将转化为 16 GB 的字符串数据和 8 GB 的 float64 数据。read_csv
操作的峰值显存占用率约为 76 GB,groupby
聚合操作的峰值显存占用率约为 56 GB。
请注意,输入文件是根据 One Billion Row Challenge GitHub 存储库生成的,存储在本地 NVMe SSD 上,并在基准测试期间由操作系统缓存。Wall 时间是在 Python 中执行完成 subprocess.run
命令的时间,包括所有初始化和加载步骤。
NVIDIA A100 Tensore Core GPU
在使用 RAPIDS cuDF pandas 加速器模式和具有足够内存的 GPU 运行挑战时,大字符串支持对于保持出色性能至关重要。第一个硬件集使用 NVIDIA A100 Tensor Core 80 GB PCIe GPU、 Arm Neoverse-N1 CPU(RAM 为 500 GiB)和 Samsung MZ1L23T8HBLA SSD。对于 200 million 行的行数,cuDF 显示的运行时间约为 6 秒,而 pandas 仅 CPU 的运行时间约为 50 秒。但是,对于 300 million 行,cuDF 24.06 中的字符串溢出会导致 pandas 回退,并将运行时间增加到约 240 秒。
借助 cuDF 24.08 中的大字符串支持,我们观察到运行时间为 17 秒,达到十亿行,远快于 pandas 运行时间(260 秒)和 cuDF 24.06 运行时间(800 秒)。如图 1 中的绿色所示。请注意,由于 GPU 显存不足,正常回退到 CPU 会导致 cuDF 24.06 上的运行时间更长。
NVIDIA Tesla T4 GPU
我们还可以评估 2018 年在 Colab 和 Kaggle 等笔记本电脑平台上广泛使用的旧代 GPU 的性能。在这种情况下,带预取的托管内存池对于良好的性能至关重要。第二个硬件集使用 NVIDIA Tesla T4 14 GB PCIe GPU、具有 376 GiB RAM 的 Intel Xeon Gold 6130 CPU 和 Dell Express Flash PM1725a SSD。对于 200 million 行的行数,RAPIDS cuDF pandas 加速器模式显示的运行时间约为 10 秒,而 pandas 运行时间约为 130 秒。当我们扩展到 1 billion 行时,T4 GPU 的超额认购率约为 5 倍,运行时间仍为 200 秒,而 pandas 的运行时间为 660 秒(图 2)。
总体而言,大字符串支持和带管理内存池与 RAPIDS cuDF pandas 加速器模式 24.08 中的预取相结合,消除了在版本 24.06 中影响性能的数据大小限制。在 24.08 中,NVIDIA A100 GPU 的内存容量更大,导致运行速度更快,主机到 GPU 的数据移动也更少。您可以根据成本效率和性能来决定适合您工作流程的 GPU。现在,使用 cuDF 扩展数据大小的障碍更少,运行时间更可预测。
优化 libcudf 中的挑战
如果您要构建 GPU 加速的数据处理应用程序,并且需要更低的开销和更快的运行时,我们建议使用 RAPIDS libcudf(cuDF 的 CUDA C++ 计算核心)。libcudf 可加速数据库和 DataFrame 操作,从数据摄取和解析到连接、聚合等。
我们发布了一个名为 billion_rows
的新 C++ 示例模块。这些示例展示了使用 libcudf 进行单线程数据分块和多线程数据流水线处理。brc
示例展示了对 13 GB “One Billion Row Challenge” 输入文件的简单单批处理。brc_chunks
示例展示了一种分块模式,该模式读取来自输入文件的字节范围,计算部分聚合,然后合并最终结果。brc_pipeline
示例展示了一种流水线模式,该模式使用多个主机线程和设备 CUDA 流完成分块工作,同时使复制带宽和计算能力饱和。
通过比较 NVIDIA A100 GPU 上的这些方法,我们发现 brc_pipeline
实现了最快的运行时间,使用 256 个块和四个线程,运行时间约为 5.2 秒(图 3)。在 A100 GPU 上的 80 GB GPU 内存中,这三种方法都可以以一亿行的速度完成挑战。分块和流水线可提供更快的运行时间,并大幅降低峰值内存占用。对于一亿行,brc
使用 55 GB 的峰值内存,而 brc_chunks
和 brc_pipeline
使用不到 1 GB 的峰值内存。libcudf billion_rows
示例模块中的技术展示了如何在高效饱和 GPU 资源和 PCIe 带宽的同时完成大型工作流程。
转移到 NVIDIA T4 GPU 后,我们发现 brc_pipeline
方法也实现了最快的运行时间,使用 256 个块和 4 个线程,运行时间约为 5.7 秒(图 4)。对于优化的 brc_pipeline
案例,由于通过 PCIe 从主机到 GPU 的数据传输的限制,T4 和 A100 GPU 的结果看起来都很相似。T4 具有 16 GB 的内存容量,因此 brc
示例在执行 2000 万行后显存不足,无法完成挑战。即使在 T4 等内存容量较低的 GPU 中,也可以有效地完成挑战。
开始使用
在 RAPIDS cuDF pandas 加速器模式下,处理大型数据集比以往更轻松。借助大字符串支持,您的字符串处理工作流可以扩展到超过先前的 2.1亿个字符限制。借助新的托管内存池,您的数据处理内存占用可以超出 GPU 内存限制。要开始使用 cuDF pandas 加速器模式,请查看 RAPIDS cuDF 在 Google Colab 上即时将 pandas 加速 50 倍。
RAPIDS cuDF pandas 加速器模式使用 RAPIDS libcudf (用于 GPU 数据处理的 CUDA C++ 库) 构建。要开始使用 RAPIDS libcudf,请尝试构建和运行一些 C++ 示例。RAPIDS Docker 容器也可用于发行版和夜间构建,以便更轻松地进行测试和部署。