cuDF 基准测试#

此存储库中基准测试的目标是衡量各种 cuDF API 的性能。cuDF 中的基准测试是使用 pytest-benchmark 插件编写的,该插件基于 pytest Python 测试框架。使用 pytest-benchmark 为熟悉 pytest 的开发者提供了无缝体验。我们包含了公共 API 和内部函数的基准测试。前者提供了我们性能的宏观视图,尤其是在与 pandas 相比时。后者帮助我们量化并最大程度地减少 Python 绑定的开销。

注意

我们当前的基准测试完全侧重于测量运行时间。然而,在某些情况下,最小化内存占用同样重要。未来,我们可能会更新基准测试,以包含内存使用量测量。

基准测试组织结构#

在顶层,基准测试被分为 internalAPI 目录。API 基准测试针对我们期望用户使用的公共特性。内部基准测试捕获 cuDF 内部实现的性能,这些内部实现不提供稳定性保证。

在每个目录中,基准测试按函数类型组织。cuDF 中的函数通常分为两类

  1. 诸如 DataFrameSeries 之类的类的方法。

  2. 对上述类进行操作的自由函数,例如 cudf.merge

前者应组织到名为 bench_class.py 的文件中。例如,DataFrame.eval 的基准测试应放在 API/bench_dataframe.py 中。基准测试应在类层次结构的最高通用级别编写。例如,所有类都支持 take 方法,因此这些基准测试应放在 API/bench_frame_or_index.py 中。如果某个方法在不同类中具有略微不同的 API,除非开发者期望某些参数会触发性能特征差异很大的代码路径,否则基准测试应使用最小的通用 API。例如,DataFrame.where 支持其他类不支持的广泛输入(如其他 DataFrame)。因此,除了针对所有 FrameIndex 类的通用基准测试外,我们还有针对 DataFrame 的单独基准测试。

注意

pytest 不支持存在两个同名基准测试文件,即使它们位于不同的目录中。因此,公共类的内部方法的基准测试位于名称以 _internal 为后缀的文件中。例如,DataFrame._apply_boolean_mask 的基准测试属于 internal/bench_dataframe_internal.py

自由函数具有更大的灵活性。广义地说,它们应该分组到包含相似功能的基准测试文件中。例如,I/O 基准测试都可以放在 bench_io.py 中。目前,这些分组由开发者自行决定。

运行基准测试#

默认情况下,pytest 会发现以 test_ 为前缀的测试文件和函数。对于基准测试,我们将 pytest 配置为改用 bench_ 前缀进行搜索。安装 pytest-benchmark 后,运行基准测试就像运行 pytest 一样简单。

运行基准测试时,默认行为是将结果以表格形式输出到终端。一个常见需求是比较更改前后基准测试的性能。我们可以使用 pytest 的 --benchmark-autosave 选项保存输出来生成这些比较。使用此选项时,基准测试运行后,输出将包含一行

Saved benchmark data in: /path/to/XXXX_*.json

XXXX 是一个四位数,用于标识基准测试。如果需要,用户还可以使用 --benchmark-save=NAME 选项,这可以更好地控制生成的文件名。给定两次基准测试运行 XXXXYYYY,可以使用以下命令比较基准测试

pytest-benchmark compare XXXX YYYY

请注意,比较使用的是 pytest-benchmark 命令,而不是 pytest 命令。pytest-benchmark 有许多额外的选项可用于自定义输出。下一行包含一个有用的示例,但开发者应该进行尝试以找到有用的输出

pytest-benchmark compare XXXX YYYY --sort="name" --columns=Mean --name=short --group-by=param

更多详细信息,请参阅 pytest-benchmark 文档

基准测试内容#

基准测试配置#

基准测试必须支持与 pandas 比较作为测试运行。为了满足这些要求,编写基准测试时必须遵循以下规则

  1. 从 config 模块导入 cudfcupy

        from ..common.config import cudf, cupy # Do this
        import cudf, cupy # Not this
    

    这允许分别替换为 pandasnumpy

  2. 避免硬编码基准测试数据集大小,而应使用 config.py 中声明的大小。这使得能够在小型数据集上以“测试”模式运行基准测试,这将快得多。

编写基准测试#

就像基准测试应以层次结构中的最高级别类来编写一样,它们也应尽可能少地假设数据的性质。例如,除非存在有意义的功能差异,否则基准测试不应关心数据的数据类型或可空性。在这些方面不同的对象对于大多数基准测试来说应该是可互换的。以这种方式编写基准测试的目标是然后自动对具有不同属性的对象进行基准测试。我们使用 benchmark_with_object 装饰器支持这种用例。

使用此装饰器的最佳演示是示例

@benchmark_with_object(cls="dataframe", dtype="int", cols=6)
def bench_foo(benchmark, dataframe):
    benchmark(dataframe.foo)

在上面的示例中,bench_foo 将针对包含六列整数数据的 DataFrames 运行。该装饰器允许自动参数化以下对象属性

  • cls: 特定类的对象,例如 DataFrame

  • dtype: 特定数据类型的对象。

  • nulls: 包含和不包含空值的对象。

  • cols: 具有一定列数的对象。

  • rows: 具有一定行数的对象。

在示例中,由于我们没有指定行数或可空性,它将针对每个有效的行数以及可空和不可空数据各运行一次。所有参数的有效集合(例如行数)存储在 common/config.py 文件中。此装饰器允许开发者编写一个适用于多种类型对象的通用基准测试,然后让该基准测试自动运行所有感兴趣的对象。

参数化测试#

benchmark_with_object 装饰器涵盖了大多数用例,并自动保证了基准测试的基本覆盖范围。然而,许多基准测试需要更定制化的对象。在某些情况下,这些对象将是调用其方法的主要目标。例如,基准测试可能需要一个具有特定数据分布的 Series。在其他情况下,这些对象将作为参数传递给其他函数。一个例子是 DataFrame.where,它接受多种类型的对象进行过滤。

在第一种情况下,夹具(fixtures)应遵循 certain rules。编写夹具时,开发者应使数据大小取决于基准测试配置。benchmarks/common/config.py 文件定义了在基准测试中使用的标准数据大小。可以为调试目的调整这些数据大小(参见下面的测试基准测试)。夹具大小应相对于 config 模块中定义的 NUM_ROWS 和/或 NUM_COLS 变量。这些规则确保了这些夹具与 benchmark_with_object 提供的夹具之间的一致性。

与 pandas 比较#

对 cuDF 进行基准测试的一个重要方面是将其与 pandas 进行比较。我们经常希望生成定量比较,因此需要使其尽可能容易。我们的基准测试通过设置环境变量 CUDF_BENCHMARKS_USE_PANDAS 来支持这一点。当检测到此变量时,所有基准测试将自动使用 pandas 运行而不是 cuDF。因此,只需运行两次基准测试(一次设置此变量,一次不设置)即可轻松生成比较。请注意,此变量仅影响 API 基准测试,不影响内部基准测试,因为后者甚至不能保证是有效的 pandas 代码。

注意

CUDF_BENCHMARKS_USE_PANDAS 有效地将 cudf 重新映射到 pandas,并将 cupy 重新映射到 numpy。它通过在 common.config.py 中对这些模块进行别名来实现。这就是为什么开发者从 config.py 导入这些包至关重要的原因。

测试基准测试#

基准测试需要与 cuDF 中的 API 更改保持同步更新。然而,我们不能简单地在 CI 中运行基准测试。这样做会消耗太多资源,并且会显著减慢开发周期

为了平衡这些问题,我们的基准测试也支持以“测试”模式运行。为此,开发者可以设置环境变量 CUDF_BENCHMARKS_DEBUG_ONLY。当在此变量下运行基准测试时,所有数据大小都会被设置为最小值,并且大小的数量也会减少。我们的 CI 测试利用这一点来确保基准测试保持有效的代码。

注意

benchmark_with_object 提供的对象遵循在 common/config.py 中定义的 NUM_ROWSNUM_COLSCUDF_BENCHMARKS_DEBUG_ONLY 通过有条件地重新定义这些值来工作。这就是为什么开发者在定义自定义夹具或用例时使用这些变量至关重要的原因。

性能分析#

虽然不是我们基准测试套件的严格组成部分,但性能分析是一个普遍需求,因此我们在此提供一些指南。以下是分析基准测试的两种简单方法(可能还有其他方法)

  1. pytest-profiling 插件。

  2. py-spy 包。

使用前者就像向 pytest 调用添加 --profile(或 --profile-svg)参数一样简单。后者则需要从 py-spy 调用 pytest,如下所示

py-spy record -- pytest bench_foo.py

每种工具都有不同的优势,并提供略有不同的信息。开发者应该尝试两者,看看哪种适合特定的工作流程。也鼓励开发者分享他们发现的有用的替代方案。

高级主题#

本节讨论 cuDF 基准测试工作原理的一些底层细节。对于普通开发者或基准测试编写者来说,这些细节通常不是必需的。此信息主要针对希望扩展可以轻松进行基准测试的对象类型的开发者。

理解 benchmark_with_object#

在底层,benchmark_with_object 由两个关键部分组成:夹具联合(fixture unions)和一些装饰器魔法。

夹具联合#

夹具联合是 pytest_cases 的一个特性。夹具联合是一个夹具,当用作测试函数参数时,它将触发测试为联合中包含的每个夹具运行一次。由于大多数 cuDF 基准测试可以使用相同的相对较小的对象集运行,因此我们的基准测试生成了可能的夹具的笛卡尔积,然后创建了所有可能的联合。

此特性对于我们的基准测试设计至关重要。对于每个相关的参数组合(大小、可空性等),我们都会以编程方式生成一个新的夹具。生成的夹具根据以下方案明确命名:{classname}_dtype_{dtype}[_nulls_{true|false}][[_cols_{num_cols}]_rows_{num_rows}]。如果夹具名称不包含特定组件,则表示该组件所有值的联合。例如,考虑夹具 dataframe_dtype_int_rows_100。此夹具是具有不同列数的可空和不可空 DataFrame 的联合。

benchmark_with_object 装饰器#

上述联合体的长名称在编写测试时很麻烦。此外,将此信息嵌入名称中意味着,为了更改使用的参数,需要替换整个基准测试的夹具名称。benchmark_with_object 装饰器是解决此问题的方案。当用于测试函数时,它本质上将函数参数名称替换为真实的夹具。在我们上面的原始示例中

@benchmark_with_object(cls="dataframe", dtype="int", cols=6)
def bench_foo(benchmark, dataframe):
    benchmark(dataframe.foo)

在功能上等同于

def bench_foo(benchmark, dataframe_dtype_int_cols_6):
    benchmark(dataframe_dtype_int_cols_6.foo)