如何在FastAPI路由中运行Jupyter内核并解决任务调度错误？

我尝试实现一个简单功能——在FastAPI路由中执行Jupyter Notebook内核代码，以下是我的代码实现：

from fastapi import FastAPI, HTTPException, Body
from jupyter_client import KernelManager
from jupyter_client.blocking.client import BlockingKernelClient

app = FastAPI()
km = KernelManager(kernel_name="python3")
client = km.client()

def execute_code():
    msg_id = client.execute("time.sleep(10)")
    print(msg_id)
    while True:
        try:
            reply = client.get_shell_msg(timeout=1)
            if reply['parent_header']['msg_id'] == msg_id:
                if reply['content']['status'] == 'error':
                    return {"error_type": "failed_to_execute", 'error': '\n'.join(reply['content']['traceback'])}
                break
        except Empty:
            pass

        try:
            iopub_msg = client.get_iopub_msg(timeout=1)
            print(iopub_msg)
            break
        except Empty:
            pass

    return {"result": "success"}
                
@app.post("/execute_cell")
async def execute_cell() -> dict:
    try:
        execute_code()
    except Exception as e:
        print(e)
        raise HTTPException(status_code=500, detail={"error_type": "failed_to_execute_internal", "error": str(e)})

运行后出现如下错误：

Traceback (most recent call last):
  File "/home/vedantroy/miniconda3/lib/python3.10/asyncio/events.py", line 80, in _run
    self._context.run(self._callback, *self._args)
RuntimeError: Cannot enter into task <Task pending name='Task-1' coro=<Server.serve() running at /home/vedantroy/miniconda3/lib/python3.10/site-packages/uvicorn/server.py:81> wait_for=<Future finished result=None> cb=[_run_until_complete_cb() at /home/vedantroy/miniconda3/lib/python3.10/asyncio/base_events.py:184]> while another task <Task pending name='Task-5' coro=<RequestResponseCycle.run_asgi() running at /home/vedantroy/miniconda3/lib/python3.10/site-packages/uvicorn/protocols/http/h11_impl.py:428> cb=[set.discard()]> is being executed.

该错误是因为阻塞的同步代码在FastAPI的异步事件循环中运行，导致事件循环被阻塞，引发任务调度冲突。修复方案如下：

修复步骤 #

• 启动Jupyter内核：初始化KernelManager后必须调用start_kernel()，否则客户端无法连接到内核。
• 用线程池隔离阻塞代码：将execute_code这类同步阻塞函数放到线程池中执行，避免阻塞异步事件循环。
• 补充缺失的异常导入：导入jupyter_client.exceptions.Empty以处理超时异常。

修改后的完整代码 #

from fastapi import FastAPI, HTTPException
from jupyter_client import KernelManager
from jupyter_client.blocking.client import BlockingKernelClient
from jupyter_client.exceptions import Empty
from concurrent.futures import ThreadPoolExecutor
import asyncio

app = FastAPI()
# 初始化线程池
executor = ThreadPoolExecutor(max_workers=2)

# 初始化并启动Jupyter内核
km = KernelManager(kernel_name="python3")
km.start_kernel()
client = km.client()
client.start_channels()

def execute_code():
    msg_id = client.execute("import time; time.sleep(10)")
    print(msg_id)
    while True:
        try:
            reply = client.get_shell_msg(timeout=1)
            if reply['parent_header']['msg_id'] == msg_id:
                if reply['content']['status'] == 'error':
                    return {"error_type": "failed_to_execute", 'error': '\n'.join(reply['content']['traceback'])}
                break
        except Empty:
            pass

        try:
            iopub_msg = client.get_iopub_msg(timeout=1)
            print(iopub_msg)
        except Empty:
            pass

    return {"result": "success"}
                
@app.post("/execute_cell")
async def execute_cell() -> dict:
    try:
        # 在线程池中执行阻塞函数
        result = await asyncio.get_event_loop().run_in_executor(executor, execute_code)
        return result
    except Exception as e:
        print(e)
        raise HTTPException(status_code=500, detail={"error_type": "failed_to_execute_internal", "error": str(e)})

关键说明 #

• ThreadPoolExecutor将同步任务放到独立线程运行，不会占用FastAPI的异步事件循环，避免任务调度冲突。
• 必须调用client.start_channels()来启动与内核的通信通道，否则无法接收内核的消息回复。
• 代码中补充了import time到执行的代码里，避免因未导入模块引发额外错误。

内容的提问来源于stack exchange，提问作者Foobar