dask - 异步计算 dask 数组 block (Dask + FastAPI)

Question

我正在构建一个 FastAPI 应用程序，它将为 Dask 数组的 block 提供服务。我想利用 FastAPI's asynchronous functionality旁边Dask-distributed's ability to operate asynchronously .下面是一个 mcve，演示了我在应用程序的服务器端和客户端上尝试做的事情:

服务器端:

import time

import dask.array as da
import numpy as np
import uvicorn
from dask.distributed import Client
from fastapi import FastAPI

app = FastAPI()
# create a dask array that we can serve
data = da.from_array(np.arange(0, 1e6, dtype=np.int), chunks=100)


async def _get_block(block_id):
    """return one block of the dask array as a list"""
    block_data = data.blocks[block_id].compute()
    return block_data.tolist()


@app.get("/")
async def get_root():
    time.sleep(1)
    return {"Hello": "World"}


@app.get("/{block_id}")
async def get_block(block_id: int):
    time.sleep(1)  # so we can test concurrency
    my_list = await _get_block(block_id)
    return {"block": my_list}


if __name__ == "__main__":
    client = Client(n_workers=2)
    print(client)
    print(client.cluster.dashboard_link)
    uvicorn.run(app, host="0.0.0.0", port=9000, log_level="debug")

客户端

import dask
import requests
from dask.distributed import Client

client = Client()

responses = [
    dask.delayed(requests.get, pure=False)(f"http://127.0.0.1:9000/{i}") for i in range(10)
]
dask.compute(responses)

在此设置中，_get_block 中的compute() 调用是“阻塞”的，一次只计算一个 block 。我尝试了 Client(asynchronous=True) 和 client.compute(dask.compute(responses)) 的各种组合，但没有任何改进。是否可以等待 dask 数组的计算？

Answer 1

这一行

block_data = data.blocks[block_id].compute()

是阻塞调用。如果你改为执行 client.compute(data.blocks[block_id])，你将得到一个可以与你的 IOLoop 结合使用的可等待的 future ，只要 Dask 使用相同的循环。

请注意，Intake 服务器非常希望以这种方式工作(它也渴望按 block 流式传输数组和其他数据类型的数据)。