np.unique 块 CPU 和 asyncio.to_thread

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 所做的是它在第一个轴上“切片”数组，创建一个形状为 (2000, 1) 的 ndarray，每个元素都是一个“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. That in the end calls https://github.com/numpy/numpy/blob/7de0fa959e476900725d8a654775e0a38745de08/numpy/core/src/multiarray/item_selection.c#L1236, which tries to release GIL at line https://github.com/numpy/numpy/blob/7de0fa959e476900725d8a654775e0a38745de08/numpy/core/src/multiarray/item_selection.c#L979 , whose definition is at https://github.com/numpy/numpy/blob/7de0fa959e476900725d8a654775e0a38745de08/numpy/core/include/numpy/ndarraytypes.h#L1004-L1006 处调用 ndarray.sort 方法——因此仅当类型未声明 [=23=] 时才释放 GIL。但是，鉴于此时各个数组元素都是非平凡类型，我假设它们状态为 NPY_NEEDS_PYAPI（我希望例如比较通过 Python），并且 GIL 未发布。

干杯。

np.unique 块 CPU 和 asyncio.to_thread

np.unique blocks CPU with asyncio.to_thread

python

numpy

python-asyncio