AsyncIO can cause deadlocks in the InteractiveShell - Impossible to wait on Comm events

I was looking for a way to delay module imports until some AsyncIO code has run.

Before anyone calls me mad, the reason I want to do this: is to check for an ipykernel.comm.Comm to have setup. I am waiting on an RPC call that verifies my Jupyter extension has loaded; this solves the problem that Jupyter extensions load asynchronously. When JavaScript comm target it registered it sends a success message back.

def create_future():
    future = asyncio.Future()
    comm = Comm('my_ext')

    @comm.on_msg
    def on_message(msg):
        future.set_result(comm)
        comm.on_msg(None)

    return future

async def wait_for_future():
    return await create_future()

# Something along the lines of this...
loop.run_until_complete(wait_for_future())

In the context of a Jupyter Notebook: I wanted to have the above code block my module import (with a short timeout, obviously) until on_message was invoked. I tried wrapping that code in an asyncio.Future, but ran into problems actually executing it because Jupyter has an event loop already running.

I tried to work around this with threads and loop.call threadsafe; then waiting for the thread to complete. I also tried something similar with a nested event loop (async_nested).

What I think the core problem is

The problem is that I can’t wait for the asyncio.Future without waiting on first run_cell_async, pausing all other execution on the event loop. The deadlock can not be released because it depends on the on_message future resolving, but Comm messages are sent by that same event loop.

The bigger problem

I think it’s impossible to wait on Comms.

You can’t even do this:

def create_comm_future():
    """
    Wait for the extension to send back
    it's heartbeat.
    """
    future = asyncio.Future()
    comm = Comm('my_ext')

    @comm.on_msg
    def on_message(msg):
        future.set_result(comm)
        # Unset the callback
        comm.on_msg(None)

    return future

Then in a Jupyter notebook:

[1]: from my_module import create_comm_future

# This will never complete - because on message will never be called. Infuriating
[2]: await create_comm_future()

[sonthonaxrk]

@Carreau I’ve found solution.

I’ve written a Kernel demonstrating it. Now I can await comm events in my notebooks! It’s just a little hack.

https://gist.github.com/sonthonaxrk/ecf2222a3e8f0035928467c36c443428

    def schedule_dispatch(self, priority, dispatch, *args):
        """
        Changes the schedule_dispatch dispatch method to
        always dispatch comm events.
        """
        # quick and dirty checks for message type

        if priority == SHELL_PRIORITY:
            stream, unparsed_msg = args
            indent, msg = self._parse_message(unparsed_msg)
            new_args = (stream, msg, indent)
            if msg['header']['msg_type'] in self.comm_msg_types:
                # Dispatch now
                return dispatch(*new_args)
        elif priority == CONTROL_PRIORITY:
            # One arg
            (unparsed_msg,) = args
            indent, msg = self._parse_message(unparsed_msg)
            new_args = (msg, indent)
        elif priority == ABORT_PRIORITY: 
            new_args = args

        # Anything else, keep things the same.
        idx = next(self._message_counter)

        self.msg_queue.put_nowait(
            (
                priority,
                idx,
                dispatch,
                new_args,
            )
        )
        # ensure the eventloop wakes up
        self.io_loop.add_callback(lambda: None)

---

However, I’m thinking about a more robust way of doing this. I really hope to be able to implement something more maintainable. It could really open up some cool things in IPython.

I think I can refactor kernelbase.py just a little bit without changing its public APIs so that I can subclass the kernel or add hooks easily.

[sonthonaxrk]

@Carreau https://github.com/sonthonaxrk/async_gui_ipython_kernel

I’ve written a kernel subclass that kinda fixes this issue.