Using Trio inside Jupyter notebook?

Jupyter has built-in support for trio so you can run coroutines inside a cell (doesn’t have to be inside an async def).

This is pretty nifty, but its usefulness is limited by structured concurrency. For example, if I have a connection object that relies on a background task (like a Trio Websocket), then that object’s lifetime has to exist entirely inside of a nursery. Jupyter does not run cells concurrently, which means I can’t create a connection object in one cell and use it another. Here’s an example:

In the second cell, I want to create a connection. The connection needs a nursery to spawn a background task into, but where can I obtain such a nursery from? In ordinary code, I would create a new nursery:

But this nursery will not exit until all tasks finish, including the connection’s background task. This means I can’t use the connection in any other cell.

Of course, this is exactly how Trio is supposed to work! It’s structured concurrency after all. That’s why I’m posing this here instead of GitHub. Has anybody tried to do something like this and are there any tips for making it work? I’m willing to use a hacky solution, since this is just for experimenting and not for production code.

I tried some various sketchy ideas but nothing worked:

I think this probably requires some support in Jupyter itself, e.g. a global nursery that is created for you. (Jupyter’s trio integration uses on each cell, which is also a big blocker, I think…)

I looked into this a bit more, and I shouldn’t have been surprised to find that Nathaniel has already commented on this over on the IPython project:

I opened up a new issue to track this:

I use IPython in the terminal for most things, but I’d love to see it supported in Jupyter notebooks. I’d probably use notebooks a lot more with more Trio support.

As an alternative to a global nursery object, there could be magic to run notebook cells concurrently as tasks. You could open a nursery block which runs in the background as you edit other notebook cells:

In [1]: %task async with trio.open_nursery() as nursery:
            nursery.start_soon(ws_reverse_server, 8888)

I also saw a feature request in ipython to “Create a %with magic”, which would expose the context manager variable to the interactive shell. Extending that idea, I could see a cell with a magic background task, running a context manager with an interactive prompt:

In [2]: %task async with open_websocket_url("ws://localhost:8888") as ws:
        >>> await ws.send_message(b"Hello, world.")
        >>> await ws.get_message()
        b'.dlrow ,olleH'

Not really sure how viable these ideas are, but it’s interesting to consider the possibilities…

I tried some approaches to getting your example code to work with the current IPython terminal and came up with a few results (one of them actually works).

The first one was:

import atexit
from functools import partial

import IPython
import trio

def ipython_embed(nursery):
    class NurseryWrapper:
        def __init__(self, nursery):
            self._nursery = nursery

        def start_soon(self, fn, *args):
            trio.from_thread.run_sync(self._nursery.start_soon, fn, *args)

    nursery = NurseryWrapper(nursery)

    # Avoid ipython-history-sqlite3-threading error

async def main():
    async with trio.open_nursery() as nursery:
        await trio.to_thread.run_sync(ipython_embed, nursery)

This one works for some really basic stuff (i.e. run a background task that prints and sleeps in a loop), but the nursery can’t be used with %autoawait, and it’s probably broken in many other ways.

Next, I started looking at the ipython code and got confused. I decided to write my own REPL to get a feel for what kind of patterns I’d expect to recognize in the ipython code. In particular, I was curious about integration with Python Prompt Toolkit. I made a gist with two examples that sort of work (one for prompt-toolkit 2 and one for version 3). The script for version 2 seems more stable. I think they both end up losing the ability to print to stdout after a while, so not really usable…

Anyway, the most successful attempt was just forking prompt-toolkit and making it work with Trio natively. Version 3 of prompt-toolkit is asyncio native, so I just went in with brute force and put Trio code where I saw asyncio code. The Trio REPL example script can do this (with syntax hilighting!):

>>> import trio_websocket                                                                  
>>> async def tock(n=5):                                                                   
...     for i in range(n):                                                                 
...         print(i + 1)                                                                   
...         await trio.sleep(i + 1)                                                        
>>> await tock(3)                                                                          
>>> conn = await trio_websocket.connect_websocket_url(nursery, "ws://localhost:8888")      
>>> nursery.start_soon(tock, 42)                                                           
>>> await conn.send_message("Hello, world.")                                               
>>> await conn.get_message()                                                               
.dlrow ,olleH

I think I’ll continue working on this idea in the context of the Trio monitor, the maintenance of which is an open issue. I’m hoping that this work could eventually be a path towards Trio support in Jupyter notebooks.

Thanks for the inspiration to look into all this stuff! I’ve been curious, but didn’t take a serious look until you made this post.

@zthompson47 I know there’s still a lot of work to do, but that’s SUPER AWESOME. A trio-enabled ptt would be a tremendous step towards trio REPL, trio monitor, being able to ssh into your trio server to poke around and debug it…

Hey there,

I saw Nathaniel ping from twitter, so a couple of notes:

  1. I’d like to avoid magics when we can use Python constructs. There will also anyway be issues if you try to run cell concurrently as there are some assumptions cells run in order in the jupyter protocol.

  2. The notebook and CLI implementation of async differs; in particular the kernel when using a notebook has a persistent running eventloop (tornado), I think it might be easier to just swap that for trio at startup to get native trio features.

  3. Currently Terminal IPython start and stop the event loop between each user input, so no BG task can run when waiting for user input. I’ve started to work on this some time ago.

Yes the IPython codebase is quite complex, if you need pointers, feel free to open issues (I’m most watching ipython/ipython); on GitHub, I’ll do my best when I have time to try to explain it.

I would also strongly suggest for you to do all those experiments with Python 3.8 which gained the ability to compile top-level await, so no need to do all those ast and source code munging.

I’m not having much time to hack on it these days; so PR welcome I’ll do my best to review when I can.