Python 并发与 asyncio 7.5 使用线程进行 CPU 密集型工作

7.5 使用线程进行 CPU 密集型工作

全局解释器锁（GIL）是 Python 中一个棘手的话题。经验法则是在阻塞的 I/O 工作中使用多线程，因为 I/O 会释放 GIL。这在大多数情况下是正确的，但并非总是如此。为了正确释放 GIL 并避免任何并发错误，运行的代码需要避免与 Python 对象（如字典、列表、Python 整数等）交互。这可能发生在我们库的大部分工作都在低级 C 代码中完成时。有一些著名的库，如 hashlib 和 NumPy，在纯 C 中执行计算密集型工作并释放 GIL。这使得我们可以使用多线程来提高某些计算密集型工作负载的性能。我们将研究两个实例：使用 hashlib 对敏感文本进行哈希以增强安全性，以及使用 NumPy 解决数据分析问题。

7.5.1 使用 hashlib 进行多线程

在当今的世界里，安全至关重要。确保数据不被黑客读取是避免泄露客户敏感数据（如密码或其他可用于识别或伤害他们的信息）的关键。

哈希算法通过将输入数据转换为一个对人类不可读且不可恢复（如果算法安全）的新数据来解决这个问题。例如，密码 “password” 可能被哈希为一个看起来像 a12bc21df 的字符串。虽然没有人能读取或恢复输入数据，但我们仍然能够检查某段数据是否与哈希匹配。这对于验证用户登录时的密码或检查数据是否被篡改等情况非常有用。

如今有许多不同的哈希算法，如 SHA512、BLAKE2 和 scrypt，尽管对于存储密码而言，SHA 并不是最佳选择，因为它容易受到暴力破解攻击。这些算法中的许多都在 Python 的 hashlib 库中实现。该库中的许多函数在哈希大于 2048 字节的数据时会释放 GIL，因此多线程是提升该库性能的一种选择。此外，scrypt 函数（用于哈希密码）始终会释放 GIL。

让我们引入一个（希望是）假设性的场景，看看在 hashlib 中使用多线程何时可能有用。想象一下，你刚刚开始了一份新工作，担任一家成功企业的首席软件架构师。你的经理分配给你第一个任务，以了解公司的开发流程——一个小的登录系统问题。为了调试这个问题，你开始查看一些数据库表，令你震惊的是，所有客户的密码都以明文形式存储！这意味着如果数据库被攻破，攻击者可以获取所有客户的密码并以他们身份登录，可能暴露诸如保存的信用卡号等敏感数据。你将此问题告知经理，他们要求你尽快找到解决方案。

使用 scrypt 算法哈希明文密码是解决此类问题的好方法。它既安全，原密码也无法恢复，因为它引入了 盐。盐是一个随机数，确保我们为密码获得的哈希是唯一的。为了测试使用 scrypt，我们可以快速编写一个同步脚本来创建随机密码并对其进行哈希，以了解需要多长时间。在这个例子中，我们将测试 10,000 个随机密码。

列表 7.16 使用 scrypt 哈希密码

import hashlibimport osimport stringimport timeimport randomdef random_password(length: int) -> bytes:    ascii_lowercase = string.ascii_lowercase.encode()    return b''.join(bytes(random.choice(ascii_lowercase)) for _ in range(length))passwords = [random_password(10) for _ in range(10000)]def hash(password: bytes) -> str:    salt = os.urandom(16)    return str(hashlib.scrypt(password, salt=salt, n=2048, p=1, r=8))start = time.time()for password in passwords:    hash(password)end = time.time()print(end - start)

我们首先编写一个函数来创建随机的小写字母密码，然后使用该函数创建 10,000 个 10 个字符长的随机密码。然后，我们使用 scrypt 函数哈希每个密码。我们将略过细节（scrypt 函数的 n, p, r 参数），但这些用于调整我们希望哈希的安全程度和内存/处理器使用情况。

在你拥有的服务器上运行这段代码，这些服务器是 2.4 吉赫兹 8 核机器，这段代码在 40 多秒内完成，这不算太差。问题是，你有一个庞大的用户基础，你需要哈希 10 亿个密码。基于这次测试的计算，哈希整个数据库将需要大约 40 多天！我们可以拆分数据集并在多台机器上运行此过程，但鉴于速度如此之慢，我们需要大量的机器。我们能否使用线程来提高速度，从而减少所需的时间和机器数量？让我们运用我们对多线程的了解来尝试一下。我们将创建一个线程池，并在多个线程中哈希密码。

列表 7.17 使用多线程和 asyncio 哈希

import asyncioimport functoolsimport hashlibimport osfrom concurrent.futures.thread import ThreadPoolExecutorimport randomimport stringfrom util import async_timeddef random_password(length: int) -> bytes:    ascii_lowercase = string.ascii_lowercase.encode()    return b''.join(bytes(random.choice(ascii_lowercase)) for _ in range(length))passwords = [random_password(10) for _ in range(10000)]def hash(password: bytes) -> str:    salt = os.urandom(16)    return str(hashlib.scrypt(password, salt=salt, n=2048, p=1, r=8))@async_timed()async def main():    loop = asyncio.get_running_loop()    tasks = []    with ThreadPoolExecutor() as pool:        for password in passwords:            tasks.append(loop.run_in_executor(pool, functools.partial(hash, password)))    await asyncio.gather(*tasks)asyncio.run(main())

这种方法涉及我们创建一个线程池执行器并为每个要哈希的密码创建一个任务。由于 hashlib 释放了 GIL，我们实现了不错的性能提升。这段代码运行时间约为 5 秒，相比之前的 40 秒。我们刚刚将运行时间从 47 天减少到了 5 秒多一点！作为下一步，我们可以将此应用与其他机器并发运行以进一步缩短运行时间，或者获取一台拥有更多 CPU 核心的机器。

7.5.2 使用 NumPy 进行多线程

NumPy 是一个极其流行的 Python 库，广泛应用于数据科学和机器学习项目。它具有大量用于数组和矩阵的数学函数，通常比普通的 Python 数组性能更好。这种性能提升的原因是，该库的大部分底层实现是用低级语言（如 C 和 Fortran）编写的，这些语言通常比 Python 更快。

由于该库的许多操作在 Python 之外的低级代码中进行，这为 NumPy 释放 GIL 并允许我们对部分代码进行多线程处理创造了机会。需要注意的是，此功能并未很好地记录，但通常可以安全地假设矩阵操作可能通过多线程获得性能提升。当然，取决于 numpy 函数的实现方式，收益可能是大是小。如果代码直接调用 C 函数并在那里释放 GIL，可能会有更大的收益；如果在任何低级调用周围有大量的支持性 Python 代码，收益将较小。由于文档记录不足，你可能需要尝试将多线程添加到应用程序中的特定瓶颈处（你可以通过剖析来确定瓶颈），并进行基准测试以了解收益。然后，你需要决定额外的复杂性是否值得任何潜在的收益。

为了在实践中看到这一点，我们将创建一个包含 40 亿个数据点的大型矩阵，共 50 行。我们的任务是为每一行计算均值。NumPy 有一个高效的 mean 函数来计算这个值。该函数有一个 axis 参数，允许我们计算轴上的所有均值而无需编写循环。在我们的例子中，axis=1 将计算每一行的均值。

列表 7.18 使用 NumPy 计算大矩阵的均值

import numpy as npimport timedata_points = 4000000000rows = 50columns = int(data_points / rows)matrix = np.arange(data_points).reshape(rows, columns)s = time.time()res = np.mean(matrix, axis=1)e = time.time()print(e - s)

该脚本首先创建一个包含 40 亿个整数数据点的数组，范围从 10 亿到 40 亿（注意，这需要相当多的内存；如果应用程序因内存不足错误而崩溃，请降低此数字）。然后，我们将数组“重塑”为一个 50 行的矩阵。最后，我们调用 NumPy 的 mean 函数，axis=1，计算每行的均值。总的来说，这段脚本在 8 核 2.4 吉赫兹的 CPU 上大约需要 25-30 秒才能运行。让我们稍微调整一下代码，使其使用线程。我们将为每行在独立线程中运行均值，并使用 asyncio.gather 等待所有行的均值。

列表 7.19 使用 NumPy 的多线程

import functoolsfrom concurrent.futures.thread import ThreadPoolExecutorimport numpy as npimport asynciofrom util import async_timeddef mean_for_row(arr, row):    return np.mean(arr[row])data_points = 4000000000rows = 50columns = int(data_points / rows)matrix = np.arange(data_points).reshape(rows, columns)@async_timed()async def main():    loop = asyncio.get_running_loop()    with ThreadPoolExecutor() as pool:        tasks = []        for i in range(rows):            mean = functools.partial(mean_for_row, matrix, i)            tasks.append(loop.run_in_executor(pool, mean))        results = asyncio.gather(*tasks)asyncio.run(main())

首先，我们创建一个 mean_for_row 函数，用于计算一行的均值。由于我们的计划是为每一行在独立线程中计算均值，我们不能再像之前那样使用带 axis 的 mean 函数。然后，我们创建一个主协程，使用线程池执行器并为每一行创建一个任务来计算均值，使用 gather 等待所有计算完成。

在相同的机器上，这段代码运行时间约为 9-10 秒，性能提升了近 3 倍！多线程可以在某些情况下帮助我们使用 NumPy，尽管在撰写本文时，哪些功能可以从线程中受益的文档尚不完善。如有疑问，如果多线程对计算密集型工作有益，最好的方法是测试并进行基准测试。

另外，请记住，在尝试使用多线程或 multiprocessing 提高性能之前，你的 NumPy 代码应尽可能向量化。这意味着要避免像 Python 循环或 NumPy 的 apply_along_axis 这样的函数，这些函数只是隐藏了循环。在 NumPy 中，通过将尽可能多的计算推送到库的低级实现中，你通常会看到更好的性能。

总结