Added asyncio support to cell magic

[wolph]

The current version of the pyinstrument magic does not support async code at all. It doesn't matter whether you have async_mode set to enabled or disabled, it will break because it will try to run in a new event loop.

This seems to be a known upstream issue but I don't think anyone is working on it: ipython/ipython#11314

This modification is certainly very hacky, but it's working quite well for me and I don't foresee any issues with it.
If it does cause issues, it doesn't do anything when using --async_mode=disabled :)

[wolph]

To test the code, just add this to a notebook:

# Cell 0
%load_ext pyinstrument
# Cell 1
%%pyinstrument
import time
import asyncio

time.sleep(0.1)
await asyncio.sleep(0.1)

[joerick]

Thanks @wolph. This is cool. Do you think it would be possible to run this in the same thread, rather in a separate one? Perhaps by starting a asyncio run loop before the execution? I'm just wondering if things like threading.locals and contextvars could introduce bugs to users' programs.

[wolph]

That's not easily possible due to the upstream bug unfortunately. IPython needs to run its own event event loop and by default asyncio won't allow you to run a second event loop in the same thread.

One option would be to monkey-patch asyncio to allow for this, but I'm not sure if that's a better solution: https://github.com/erdewit/nest_asyncio

[joerick]

Gotcha. Well, maybe we can make it a different magic, perhaps %%pyinstrument_async? Or maybe we could detect the presence of await/async keywords in the code and switch this behaviour on that? I'd rather not have this hack affect the majority of non-async users, as I suspect that it might change behaviour in some cases.

[wolph]

This hack is already conditional of having asyncio enabled so we could default to having asyncio disabled.

I don't see a scenario where that would break existing behaviour for people, unless they already manually set the asyncio parameter

[wolph]

@joerick I've disabled async mode by default now.

To run the cell magic normally people can use:

%%pyinstrument

For asyncio they can use:

%%pyinstrument --async_mode=enabled

[joerick]

Great! Yeah this'll work. Last thing - if users don't know about this and try to do an await, they'll get the asyncio exception. Do you think it would be possible to catch the asyncio RuntimeError, check that 'event loop is already running' in str(exc) and if that's the case, print a message to users to let them know about the --async_mode=enabled option? I don't know if that would be in the ast transformer or just where the ip.run_cell function is called.

[wolph]

I can't prevent the exception output, but I can display a warning :)
I've updated it to this now:

[joerick]

I tweaked the detection logic to make a bit stricter - hopefully this works.

[wolph]

It seems to work great for me :)