Computing dask array chunks asynchronously (Dask + FastAPI) - dask

Question

I am building a FastAPI application that will serve chunks of a Dask Array. I would like to leverage FastAPI's asynchronous functionality alongside Dask-distributed's ability to operate asynchronously. Below is a mcve that demonstrates what I'm trying to do on both the server and client sides of the application:
Server-side:
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")
Client-side
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)
In this setup, the compute() call in _get_block is "blocking" and only one chunk is computed at a time. I've tried various combinations of Client(asynchronous=True) and client.compute(dask.compute(responses)) without any improvement. Is it possible to await the computation of the dask array?

Answer 1

This line
block_data = data.blocks[block_id].compute()
is a blocking call. If you instead did client.compute(data.blocks[block_id]), you would get an awaitable future that can be used in conjunction with you IOLoop, so long as Dask is using the same loop.
Note that the Intake server would very much like to work in this manner (it too aspires to stream data by chunks for arrays and other data types).