How can I execute code on a kernel inside a fastAPI route?

Kernels
1

I have the following code:

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

This code is supposed to run a simple time.sleep(10) instead of the jupyter kernel. But, in practice, I get the exception:

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.

How can I avoid this happening?

2

First: note that exposing an unauthenticated execute endpoint is… a really poor idea. It’s really not possible to ensure the input is “safe,” and usually some other tool such as kernel gateway or the JupyterHub architecture is used to ensure the requests are isolated with defense in depth, usually involving containers, VMs, or ideally, both.

All that being said: as FastAPI is async-first, perhaps using the async variants of the kernel classes to ensure the correct event loop is being used.

2 Likes