这项工作得到 Anaconda Inc 以及 Moore Foundation 的数据驱动发现计划支持

上周在 PyCon DE 会议上，我很幸运地遇到了 Cython 的核心开发者之一 Stefan Behnel。我们一起努力优化了一个小型基准测试，该测试代表了 Dask 的核心任务调度器，这是一个主要受限于数据结构的纯 Python 应用程序。

我们的基准测试是一个玩具问题，它创建了三个数据结构，这些结构使用字典、列表和集合相互索引，然后进行一些简单的算术运算。（您无需深入理解此基准测试即可阅读本文。）

import random
import time

nA = 100
nB = 100
nC = 100

A = {'A-%d' % i: ['B-%d' % random.randint(0, nB - 1)
for i in range(random.randint(0, 5))]
for i in range(nA)}

B = {'B-%d' % i: {'C-%d' % random.randint(0, nC - 1)
for i in range(random.randint(1, 3))}
for i in range(nB)}

C = {'C-%d' % i: i for i in range(nC)}

data = ['A-%d' % i for i in range(nA)]


def f(A, B, C, data)
for a_key in data
b_keys = A[a_key]
for b_key in b_keys
for c_key in B[b_key]
C[c_key] += 1


start = time.time()

for i in range(10000)
f(A, B, C, data)

end = time.time()

print("Duration: %0.3f seconds" % (end - start))

$ python benchmark.py
Duration: 1.12 seconds

这是一个非典型的 Python 优化问题，因为它主要受限于数据结构访问（字典、列表、集合），而不是 Cython 通常优化的数值运算（浮点算术的嵌套 for 循环）。Python 在这方面已经相当快了，通常比 Java 或 C++ 等编译型语言慢 2-5 倍，但我们仍然希望在可能的情况下进行改进。

在本文中，我们结合了两种不同的方法来优化受限于数据结构的工作负载：

1. 使用 Cython 编译 Python 代码，无需其他注解
2. 对字符串进行驻留（interning）以实现更快的字典查找

最后，在文章末尾，我们还在 PyPy 下运行了基准测试以比较性能。

Cython

首先，我们使用 Cython 编译我们的 Python 代码。通常使用 Cython 时，我们会用类型注解变量，为编译器提供足够的信息以完全避免使用 Python。然而，在我们的例子中，没有太多数值运算，而且无论如何我们都会使用 Python 数据结构，所以这帮助不大。我们未经修改地编译了原始 Python 代码。

cythonize -i benchmark.py

然后运行

$ python -c "import benchmark"
Duration: 0.73 seconds

这使得速度从 1.1 秒提高到 0.73 秒，提升显著。相较于典型的 Cython 加速（通常是 10-100 倍），这次提升并不巨大，但这对于我们的调度器来说将是一个*非常受欢迎*的改变，我们已经为 5% 的优化努力了一段时间了。

字符串驻留

我们的第二个技巧是对字符串进行驻留。这意味着我们尝试始终只保留每个字符串的一个副本。这在进行字典查找时可以提高性能，原因如下：

1. Python 只计算一次字符串的哈希值（字符串计算后会缓存其哈希值）
2. Python 在进行值相等性检查（慢）之前，会先检查对象身份（快）

或者，根据经验来说，在 Python 中 `text is text` 比 `text == text` 要快。如果您确保每个字符串只有一个副本，那么只需要进行像 `text is text` 这样的身份比较。

因此，如果在程序中任何时候看到字符串 "abc" 都与所有其他 "abc" 字符串是完全相同的对象，那么字符串字典查找将只需要进行指针/整数相等性检查，而无需进行完整的字符串比较。

在我们的基准测试中添加字符串驻留，代码如下：

inter = {}

def intern(x)
try
return inter[x]
except KeyError
inter[x] = x
return x

A = {intern('A-%d' % i): [intern('B-%d' % random.randint(0, nB - 1))
for i in range(random.randint(0, 5))]
for i in range(nA)}

B = {intern('B-%d' % i): {intern('C-%d' % random.randint(0, nC - 1))
for i in range(random.randint(1, 3))}
for i in range(nB)}

C = {intern('C-%d' % i): i for i in range(nC)}

data = [intern('A-%d' % i) for i in range(nA)]

# 基准测试的其余部分与之前相同

这使得运行时间从 1.1 秒下降到 0.75 秒。请注意，这还不包括上面描述的单独的 Cython 改进。

Cython + 字符串驻留

我们可以结合这两种优化。这将运行时间缩短到约 0.45 秒，比原始时间提高了 2-3 倍。

cythonize -i benchmark2.py

$ python -c "import benchmark2"
Duration: 0.46 seconds

PyPy

或者，我们可以直接在 PyPy 中运行所有代码。

$ pypy3 benchmark1.py # 原始版本
Duration: 0.25 seconds

$ pypy3 benchmark2.py # 包含字符串驻留
Duraiton: 0.20 seconds

因此，PyPy 在这类代码上可能比 Cython 快得多（这并不令人意外）。字符串驻留有一点帮助，但效果没那么大。

这相当令人鼓舞。Dask 调度器可以在 PyPy 下运行，即使 Dask 客户端和工作节点在正常的 CPython 下运行（用于与完整的 PyData 生态系统配合使用）。

Dask 基准测试的初步结果

我们开始这个实验时，假设我们的玩具基准测试在性能特性方面某种程度上代表了 Dask 的调度器。当然，这个假设是错误的。Dask 调度器要复杂得多，很难建立一个简单的玩具示例来代表其性能。

当我们在一个实际使用 Dask 调度器的稍微复杂一些的基准测试上尝试这些技巧时，我们发现以下结果：

• Cython: 几乎没有效果
• 字符串驻留: 几乎没有效果
• PyPy: 几乎没有效果

然而，我只在这上面花费了很短的时间（二十分钟？），所以我希望这里没有性能提升是因为投入的努力不足。

如果有人对此感兴趣，我希望这篇博文包含足够的信息，可以帮助他们进一步研究。