使 FastAPI WebSockets 的 CPU 绑定任务异步答案

【问题标题】：Make an CPU-bound task asynchronous for FastAPI WebSockets使 FastAPI WebSockets 的 CPU 绑定任务异步
【发布时间】：2021-08-30 11:19:43
【问题描述】：

所以我有一个受 CPU 限制的长时间运行算法，我们称之为任务。假设它看起来像这样：

def task(parameters):
  result = 0
  for _ in range(10):
    for _ in range(10):
      for _ in range(10):
        result += do_things()
  return result

@app.get('/')
def results(parameters: BodyModel):
    return task(parameters)

如果我将其封装在def 路径操作函数中，一切正常，因为它是在不同的线程中启动的。我可以访问多个路径等。并发通过将我的 CPU 绑定任务推送到单独的线程来完成它的工作。但我现在想切换到 WebSockets，以传达中间结果。为此，我必须将我的整个事情标记为异步并将 WebSocket 传递给我的任务。所以它看起来像这样：

async def task(parameters):
  result = 0
  for _ in range(10):
    for _ in range(10):
      for _ in range(10):
        intermediate_result = do_things()
        await parameters.websocket.send_text(intermediate_result)
        result += intermediate_result
  return result

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    while True:
        parameters = await websocket.receive_text()
        parameters.websocket = websocket
        result = await task(parameters)
        await websocket.send_text(result)

发送中间结果就像一个魅力。但是现在我的算法阻止了 FastAPI，因为它本身并不是真正的异步。一旦我向“/ws”发布消息，FastAPI 就会被阻止，并且在我的任务完成之前不会响应任何其他请求。

所以我需要一些关于如何做的建议

a) 要么从同步 CPU 绑定任务中发送 WebSocket 消息（我没有找到同步的 send_text 替代方案），因此我可以使用 def 或
b) 如何让我的 CPU 绑定真正异步，以便在我使用 async def 时它不会再阻塞任何东西。

我尝试按照here 的描述使用 ProcessPoolExecuter，但无法腌制协程，据我所知，我必须将我的任务设为协程（使用异步）才能在其中使用 websocket.send_text()。

另外，我考虑将我的中间结果存储在某个地方，创建一个 HTTP POST 来启动我的任务，然后使用另一个 WebSocket 连接来读取和发送中间结果。但是我也可以类似地启动一个后台任务并实现一个常规的 HTTP 轮询机制。但我也不想要，主要是因为我打算使用谷歌云运行，它会在所有连接关闭时限制 CPU。而且我认为最好教我的任务如何直接通过 WebSocket 进行通信。

我希望我的问题很清楚。这是我第一个使用 FastAPI 和异步性的大型项目，之前没有真正使用过 AsyncIO。所以我可能只是错过了一些东西。谢谢你的建议。

【问题讨论】：

试试asyncio.to_thread(<blocking func>)。
没有尝试过，因为在文档中它说“由于 GIL，asyncio.to_thread() 通常只能用于使 IO 绑定函数非阻塞。”所以我现在试了一下。但这也带来了与多处理相同的问题。要么它不是异步的，要么我不能在其中使用 WebSockets。
您需要将执行计算任务的部分代码和通过await交互操作IO的部分分开，并通过队列将它们链接起来。例如，使用aioprocessing 库。

标签： python asynchronous websocket fastapi

【解决方案1】：

如果有人遇到这个问题，我现在会添加适合我的解决方案。

我在关注这个：https://docs.python.org/3/library/asyncio-eventloop.html#asyncio.loop.run_in_executor

关键是让它不阻塞。因此，例如，而不是：

 # 1. Run in the default loop's executor:
result = await loop.run_in_executor(None, blocking_io)
print('default thread pool', result)

我移动 await 并将代码更改为：

# 1. Run in the default loop's executor:
thread = loop.run_in_executor(None, blocking_io)
print('default thread pool', result)
while True:
    asyncio.sleep(1)
    websocket.send_text('status updates...'
    if internal_logger.blocking_stuff_finished:
        break
result = await thread
websocket.send_text('result:', result)
websocket.close()

这样，我将 cpu_bound 的东西放在一个单独的线程中，我不等待，一切正常。

创建自定义线程池也可以，但我们需要删除上下文管理器以使其成为非阻塞。

# 2. Run in a custom thread pool:
with concurrent.futures.ThreadPoolExecutor() as pool:
    result = await loop.run_in_executor(pool, blocking_io)

然后会变成：

pool = concurrent.futures.ThreadPoolExecutor()
thread = loop.run_in_executor(pool, blocking_io)

理论上，ProcessPoolExecutor 也是如此，但还需要做更多的工作，因为没有共享内存，并且我的终止条件不会像上面描述的那样工作。

是的，我知道 cpu_bound 的东西最好在不同的进程中完成，但是在我的情况下将它移动到一个单独的线程确实有效，而且我确实喜欢共享内存 atm。

【讨论】：