Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Replace nest_asyncio with greenlets #683

Merged

Conversation

douglas-raillard-arm
Copy link
Collaborator

@douglas-raillard-arm douglas-raillard-arm commented Apr 26, 2024

Considering #682, this PR updates devlib to use greenlets in cases where nest_asyncio cannot be used (e.g. using uvloop event loop).

This alternative path works by using greenlet to provide a way for nested coroutines separated from the top-level coroutine by blocking calls to yield their action.

TODO:

  • Handle exceptions raised in __await__, so that they are propagated to the correct coroutine rather than the top-level one
  • Write unit test to check everything works on various event loops (e.g. uvloop)
  • Add some rational in asyn.py to explain the moving pieces

Fixes #682

@douglas-raillard-arm douglas-raillard-arm changed the title [Draft] Fix nest asyncio [Draft] Replace nest_asyncio May 8, 2024
@douglas-raillard-arm
Copy link
Collaborator Author

douglas-raillard-arm commented May 9, 2024

PR updated with:

  1. implementation directly based on greenlet rather than greenback. This avoids some ctypes-based hacks that are not necessary. That implementation is compatible with any asyncio event loop implementation.
  2. Still use nest_asyncio when possible, as the greenlet version might execute the coroutine on a separate thread in some circumstances, whereas nest_asyncio never needs to. This important as the jupyter lab use case unfortunately falls in the "needing the thread" category, since the all code (including the imports) executes inside a pre-existing coroutine.
  3. A commit that allows recycling connection instances when a thread is done with them, rather than just destroying the connection. This means that a client frequently spawning threads (e.g. creating short lived thread pools) will not reconnect to the target all the time. The coroutine execution in a separate thread can also share the connection. This is implemented by returning the connection to the pool after every Target method that directly uses self.conn (see call_conn decorator). This has a user-visible effect that target.conn might give a different object across calls to Target methods.

@douglas-raillard-arm douglas-raillard-arm force-pushed the fix_nest_asyncio branch 2 times, most recently from 33f437b to 23db634 Compare May 10, 2024 12:48
@douglas-raillard-arm douglas-raillard-arm changed the title [Draft] Replace nest_asyncio [Draft] nest_asyncio fallback May 10, 2024
@douglas-raillard-arm
Copy link
Collaborator Author

PR updated with exception handling in coroutine and exception injection with .throw().
The state of this PR will stay as [Draft] until I add some unit tests to ensure correct operations (and fix all the items in the TODO list in this PR cover letter) as discussed with Vincent Coubard.

@douglas-raillard-arm
Copy link
Collaborator Author

PR updated with:

  • unit tests added for devlib.utils.asyn.run(), simulating various event loop setups including non-stdlib event loops (uvloop). A fairly comprehensive set of nested run() calls is checked, as well as exception and value injection (coro.send() and coro.throw())
  • nest_asyncio properly removed, as it relies on asyncio.get_event_loop() will not create an event loop automatically if one is not setup when called from a non-main thread. asyncio.get_event_loop() is deprecated and cannot be relied upon to create an event loop, so nest_asyncio usage is basically broken. Since we want devlib to be importable in the non-main thread (no reason for it to explode), this is not acceptable and we can just fully switch to our new implementation

@douglas-raillard-arm
Copy link
Collaborator Author

PR updated with extra comments. I now consider it ready so I'll remove the [Draft] status

@douglas-raillard-arm douglas-raillard-arm changed the title [Draft] nest_asyncio fallback Replace nest_asyncio with greenlets May 21, 2024
@douglas-raillard-arm
Copy link
Collaborator Author

douglas-raillard-arm commented May 22, 2024

Found an issue, please do not merge as-is.

EDIT: the issue is the following: devlib.utils.asyn.asynccontextmanager() provides a wrapper over contextlib.asynccontextmanager() to allow treating the async context manager as a regular blocking context manager. This is achieved by implementing the following:

    def __enter__(self, *args, **kwargs):
        return run(self.cm.__aenter__(*args, **kwargs))

    def __exit__(self, *args, **kwargs):
        return run(self.cm.__aexit__(*args, **kwargs))

This is fine as long as run() simply re-enters an existing event loop. If these run() calls are top-level calls, they will each create a new event loop and try to iterate over the async generator. This is a problem in two ways:

  1. An event loop closes all async generators upon exit. This means the 2nd run() call from __exit__() is ends up seeing a closed async gen, confusing the stdlib implem that then raises RuntimeError("generator didn't stop after athrow()")
  2. We end up "migrating" the async gen from one event loop to another, which is the same as migrating a coroutine across event loops.

Issue 1. can be worked around by hijacking the mechanism the event loop uses to be aware of new async gen.

However, issue 2. is more tricky and is probably a real issue if e.g. the async generator tries to take an asyncio.Lock() across yield calls. The lock future would be handled by the first event loop, then cancelled, and the generator would probably fail to run the 2nd iteration on another event loop.

I'll experiment to see what possibilities exist to fix this problem. This is fortunately the only place that relies on migrating a coroutine between multiple event loops. Trio encountered a similar problem: python-trio/trio#2081

@douglas-raillard-arm
Copy link
Collaborator Author

PR updated with:

  • extra unit tests that reproduced the issue
  • _AsyncPolymorphicCM now spins an event loop if necessary in __enter__ and closes it in __exit__, so that both the __aenter__() and __aexit__() coroutines are executed on the same event loop.

@douglas-raillard-arm
Copy link
Collaborator Author

@marcbonnici, Vincent Coubard, Branislav Rankov, this should be ready for the last review/testing round. I consider it to be ready.

@douglas-raillard-arm
Copy link
Collaborator Author

PR rebased. As of today we started dogfooding that PR in LISA's vendored devlib tree, so it should get some more real-world exposure in the coming weeks.

@douglas-raillard-arm
Copy link
Collaborator Author

PR updated with extra tests to check the blocking API works when invoked from asyncio.to_thread(). Everything does work.

@douglas-raillard-arm
Copy link
Collaborator Author

PR updated with a fix to use run() on the return values of anext() for async generators, with the matching unit test.

@douglas-raillard-arm
Copy link
Collaborator Author

PR rebased

@douglas-raillard-arm douglas-raillard-arm force-pushed the fix_nest_asyncio branch 2 times, most recently from d071e9d to 7645db7 Compare July 16, 2024 13:30
@douglas-raillard-arm
Copy link
Collaborator Author

Update with a check to only call loop.shutdown_default_executor() if it exists, since it was added in Python 3.9

https://docs.python.org/3/library/asyncio-eventloop.html#asyncio.loop.shutdown_default_executor

@douglas-raillard-arm
Copy link
Collaborator Author

An issue was found, I'm currently working on fixing it, please do not merge until it's fixed.

@douglas-raillard-arm
Copy link
Collaborator Author

PR updated with:

  • Some refactoring of private functions involved in run() implementation.
  • Ensure that we really do use the same loop for both __aenter__ and __aexit__ for async context managers.
  • Ensure that each asyncio.Task gets its own top-level instrumented coroutine. This is critical as a Task would not be able to yield on behalf of a coroutine in another task. This would deadlock, since the event loop would be polling task B while a nested coroutine inside task B would be trying to yield through task A. Task A would be ignored by the event loop until task B yields, leading to a deadlock.

@douglas-raillard-arm
Copy link
Collaborator Author

One thing I realized and might want to change is that if an event loop is already running (e.g. in a jupyterlab notebook), we will dispatch the coroutine on a loop setup in a separate thread. It's all good except we have a single such thread.

This means that code making parallel devlib invocation of devlib in threads with a pre-setup event loop will end up being serialized. It shouldn't be very hard to change, I'll see if I can do that tomorrow

Set the "_name" attribute rather than trying to set the "name" read-only
property.
Once a thread exits, the connection instance it was using can be
returned to the pool so it can be reused by another thread.

Since there is no per-thread equivalent to atexit, this is achieved by
returning the connection to the pool after every top-level method call
that uses it directly, so the object the user can get by accessing
Target.conn can change after each call to Target method.
Prepare for providing our own implementation of asyncio.run() to work
without nest_asyncio package.
Provide an implementation of re-entrant asyncio.run() that is less
brittle than what greenback provides (e.g. no use of ctypes to poke
extension types).

The general idea of the implementation consists in treating the executed
coroutine as a generator, then turning that generator into a generator
implemented using greenlet. This allows a nested function to make the
top-level parent yield values on its behalf, as if every call was
annotated with "yield from".
@douglas-raillard-arm
Copy link
Collaborator Author

PR updated with:

  • significant simplification in the way the different cases are handled. There is now a _CoroRunner subclass for each case.
  • simplification of how async context managers are handled, thanks to _CoroRunner. All we now need is to get a runner that will be usable to execute multiple coroutines on the same event loop, even if that event loop is sitting in a separate thread.
  • Added contextvars test and fixed the code. contextvars are now propagated down through devlib.utils.asyn.run() and any update done inside the coroutine will be propagated back in the caller of run(), so it is completely transparent.
  • The threaded coroutine runner used as last resort when an event loop is already running in the current thread can now use an arbitrary number of threads. This prevents serialization of otherwise parallel code and ensures no deadlock can occur.

Considered it is a substantial change, I'd be more comfortable dogfooding it in LISA for a little while before we consider merging it.

@douglas-raillard-arm
Copy link
Collaborator Author

1.5 month later there has been no reported issue, so I think it's good to go

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Remove nest_asyncio dependency
2 participants