python - np.unique 使用 asyncio.to_thread 阻塞 CPU

Question

我已经设置了以下测试程序(Python 3.9.5，numpy 1.20.2):

import asyncio
from datetime import datetime
import numpy as np

async def calculate():
    print("=== unique")
    await asyncio.to_thread(lambda: np.unique(np.ones((2000, 50000)), axis=0))
    print("=== sort")
    await asyncio.to_thread(lambda: np.sort(np.ones((2000, 50000)), axis=0))
    print("=== cumsum")
    await asyncio.to_thread(lambda: np.cumsum(np.ones((2000, 100000)), axis=0))

async def ping():
    while True:
        print("async", datetime.utcnow())
        await asyncio.sleep(0.2)

async def main():
    p1 = asyncio.create_task(ping())
    c = asyncio.create_task(calculate())
    await asyncio.wait([p1, c], return_when=asyncio.FIRST_COMPLETED)
    p1.cancel()

asyncio.run(main())

输出如下:

async 2021-05-21 13:20:16.308948
=== unique
async 2021-05-21 13:20:16.531135
async 2021-05-21 13:20:40.142323
=== sort
async 2021-05-21 13:20:40.343306
async 2021-05-21 13:20:40.543658
async 2021-05-21 13:20:40.743989
async 2021-05-21 13:20:40.944312
async 2021-05-21 13:20:41.144664
async 2021-05-21 13:20:41.345007
=== cumsum
async 2021-05-21 13:20:41.545523
async 2021-05-21 13:20:41.745901
async 2021-05-21 13:20:41.946271
async 2021-05-21 13:20:42.146651
async 2021-05-21 13:20:42.347021
async 2021-05-21 13:20:42.547396

很明显 np.unique 需要大约 23 秒，并且不会像 np.cumsum 和 np.sort 那样被打断。

如果我对 asyncio.to_thread 和 GIL 的理解是正确的，那么任何在线程中运行的东西都应该被周期性地中断，以至少在一定程度上支持线程程序的多任务处理。 np.sort 和 np.cumsum 的行为支持这一点。 np.unique 中发生了什么以防止该线程被中断？

Answer 1

这是一个棘手的问题 ;-)

问题是 GIL 并没有在 np.unique 调用中实际释放。原因是 axis=0 参数(您可以验证没有它调用 np.unique 释放 GIL 并与 ping 交错调用)。

TLDR； axis 参数的语义对于 np.sort/cumsum 和 np.unique 调用是不同的:而对于 np.sort/cumsum 操作在该轴“中”执行矢量化(即，独立排序多个数组)，np.unique 在“沿”该轴的切片上执行，这些切片是非琐碎的数据类型，因此它们需要 Python 方法。

对于 axis=0，numpy 所做的是它在第一个轴上“切片”数组，创建一个 ndarray 形状为 (2000, 1)，每个元素都是一个“值的 n 元组”(它的数据类型是单个元素数据类型的数组)；这发生在 https://github.com/numpy/numpy/blob/7de0fa959e476900725d8a654775e0a38745de08/numpy/lib/arraysetops.py#L282-L294 .

然后在 https://github.com/numpy/numpy/blob/7de0fa959e476900725d8a654775e0a38745de08/numpy/lib/arraysetops.py#L333 处调用 ndarray.sort 方法.最后调用https://github.com/numpy/numpy/blob/7de0fa959e476900725d8a654775e0a38745de08/numpy/core/src/multiarray/item_selection.c#L1236 ，它试图在 https://github.com/numpy/numpy/blob/7de0fa959e476900725d8a654775e0a38745de08/numpy/core/src/multiarray/item_selection.c#L979 行释放 GIL , 其定义在 https://github.com/numpy/numpy/blob/7de0fa959e476900725d8a654775e0a38745de08/numpy/core/include/numpy/ndarraytypes.h#L1004-L1006 -- 因此仅当类型未声明 NPY_NEEDS_PYAPI 时才会发布 GIL。但是，考虑到此时各个数组元素都是非平凡类型，我假设它们声明了 NPY_NEEDS_PYAPI(例如，我希望通过 Python 进行比较)，并且 GIL 未发布。

干杯。