使 FastAPI WebSockets 的 CPU 绑定任务异步

Question

所以我有一个受 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路径操作 function中，则一切正常，因为它是在不同的线程中启动的。 我可以访问多个路径等。并发通过将我的 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时它不再阻塞任何东西。

我尝试使用此处描述的 ProcessPoolExecuter，但不可能腌制协程，据我所知，我必须使我的任务成为协程（使用异步）才能在其中使用websocket.send_text() 。

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

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

Answer 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)

我移动等待并将代码更改为：

# 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 也是如此，但需要做更多的工作，因为没有共享的 memory 并且我的终止条件不会像上面描述的那样工作。

是的，我知道 cpu_bound 的东西最好在不同的进程中完成，但是在我的情况下将它移动到单独的线程不起作用，我确实喜欢共享的 memory atm。