FEAT: support AnyIO

[davidbrochart]

What pyzmq version?

25.0.0

What libzmq version?

4.3.4

Python version (and how it was installed)

python 3.11.0, via conda-forge

OS

ubuntu 22.04

What happened?

I would like to use pyzmq with an async context inside an AnyIO event loop, using a trio backend.
From what I can see, pyzmq only supports asyncio. How do you feel about supporting AnyIO?

Code to reproduce bug

No response

Traceback, if applicable

No response

More info

No response

[minrk]

I don't object to it, but I'm not likely to implement it any time soon. The abstraction in zmq._future should make an anyio implementation pretty small, I would guess.

[davidbrochart]

Good to know, would you be open to a PR?

[minrk]

Yup!

[minrk]

I looked at this for a bit, and it appears anyio has decided against supporting FDs, which is required for zmq integration to be possible. So if zmq were to support anyio, it seems it would only work on asyncio, and require tornado's selector thread implementation.

tornado is the only event loop that appears to have bothered to support polling FDs on Windows, and pyzmq relies on tornado's implementation when working with asyncio on Windows.

It's possible that trio's workaround for pipes could be used for zmq FDs, but I don't know.

It might also be possible to port tornado's SelectorThread to anyio.

But the real answer is for these event loops to implement APIs to wait for an FD to be readable, as is possible in select, epoll, etc. It shouldn't be the responsibility of libraries to provide this basic functionality to the event loop.

[davidbrochart]

FWIW, I implemented something similar in Jupyverse (part of jupyter-server/jupyverse#388).

[minrk]

Yeah, I think that approach would work. Just need to be extremely careful about races on the edge-triggered FD when working with threads, because accessing zmq.EVENTS can prevent the select from waking. Looks like what you have would mean one select call and thread per zmq socket, which is not very scalable. The tornado SelectThread merges all readers into a single select call, which is nice for an application like Jupyter, which may work with quite a few zmq sockets.

I guess anyio brings the bad ProactorEventLoop situation to all platforms since it has to be least-common-denominator, which is consistent, I guess, if kind of a worst-case scenario.

I've managed to test that trio's lowlevel.wait_readable works for the zmq FDs. So pyzmq could work with anyio on at least trio or asyncio (the only real options), by picking based on sniffio.current_async_library:

if asyncio:
    use existing selector_thread implementation
elif trio:
    use trio.wait_readable
else:
    raise unsupported

[davidbrochart]

But I think you're right that it should be part of AnyIO, if Trio has this feature.

[minrk]

I did some tests, and I'm pretty sure this can work. Unfortunate that we can't use ~any anyio APIs to accomplish this, but I think we can make it compatible with trio, at least.

[davidbrochart]

Let's continue the discussion from ipython/ipykernel#1079 (comment) here, maybe @agronholm has ideas.

[agronholm]

@minrk Could you first help me understand why anyio.wait_socket_readable() is not good enough here?

[minrk]

I think maybe it can, I will have a look. I've never known what kind of socket it is, but I can probably figure it out. Thanks for the pointer.

Since socket.fromfd duplicates the FD, doing that in order to pass the duplicate FD to APIs where the existing one works is less than ideal when FD exhaustion is already an issue in applications like Jupyter. Especially when the underlying implementations support the integer FDs already.

from the docs:

This does NOT work on Windows when using the asyncio backend with a proactor event loop (default on py3.8+).

This would be a blocker for using it, though, since this is ~all Windows users. It would be wonderful if anyio ported tornado's SelectorThread to restore this functionality to asyncio for feature parity across implementations. CPython devs even agree it should be in asyncio itself, though nobody has the time to bring it in (copying trio's AFD_POLL would be even better, if harder). As it is now, since we have to detect asyncio and call lower-level APIs to make asyncio work, we might as well do the same for trio, which avoids the duplication of every FD.

[agronholm]

Ok, so the lack of proactor loop support for this is the crux of the problem. I'm open to allowing a hack like SelectorThread to be added to AnyIO for this, so long as I can agree with the implementation.

[agronholm]

That said, do you know why pyzmq needs this functionality? IIRC, AnyIO only uses it internally to implement UNIX domain sockets support.

[minrk]

Yes, very much so. I don't think I have the necessary expertise to port the selector thread into something that fits appropriately in the structured concurrency pattern, since where it is now is scoped only to the whole event loop, which I'm guessing is not the right fit, but I'm not sure.

[agronholm]

A potential solution would be to keep up that selector thread so long as there's any running calls to wait_socket_(readable|writable).

[minrk]

Yeah, that's one option. I think tornado's solution of the selector existing for the lifetime of the asyncio loop makes the most sense, but I can see how structured concurrency might not like that that. I think it is the best approach, though, given that this is really a missing feature in one implementation of the asyncio.EventLoop class, so fixing that seems to make the most sense at the EventLoop instance level, which would be consistent with all other asyncio.EventLoop implementations.

[davidbrochart]

@minrk Would you be open to changing the API of async sockets so that e.g. recv() or send() methods become proper async functions, instead of returning futures?
Having non-async functions doesn't play well with structured concurrency where creating a task group is already async. We could keep futures internally but I think public methods should be async.

[minrk]

I would need to do some benchmarks. It used to be that way, but there was a massive performance benefit (orders of magnitude) to using Futures that I don't want to lose. That may not be true anymore, though.

I'm not aware of problems solved by switching from futures to coroutines, can you give an example?

[agronholm]

Futures are an asyncio-only thing, and don't exist on Trio. If someone wants to use an async API with Trio, it should not deal directly with futures.

[minrk]

Coroutines don't solve that on their own, though, do they? It still needs the underlying event loop integration to wake when there's a socket event, whatever that looks like. So making this change on its own still keeps everything asyncio-specific, plus breaking backward compatibility of eager evaluation, and a potential performance hit, right? Can you share an example that would be improved by this? I don't have much experience with anyio, so I'm genuinely asking what the problems are, not trying to argue against it.

It may well be that a coroutines approach will work just as well without any performance cost using a wait_readable approach (I suspect it can if done right). It just needs testing, and careful implementation to deal with multiple waiters, which can't actually be allowed at a low level due to zmq's edge-triggered FD mess.

[agronholm]

The trouble with wait_readable() is that it doesn't work with IOCP (ProactorEventLoop), as Windows does not offer a public API for waiting until a socket is readable. There is an undocumented API call for that, used by Trio, but asyncio doesn't support it.

Thinking about what benefits could be had from this, I think AnyIO'fying just the socket stuff won't help immediately, but would be a gradual step towards abstracting away asyncio. I think that may be what @davidbrochart is aiming for?

I haven't worked with pyzmq in many years so I'm a bit out of touch with its codebase, and don't know what parts of asyncio it uses besides the socket APIs.

[davidbrochart]

My goal is to have ipykernel work with asyncio or Trio, since it's now based on AnyIO, and pyzmq is not allowing that.
I'll send a WIP PR to show how things look like.

[davidbrochart]

See #2045.

[minrk]

Yeah, I'm fully on board with anyio support and expect that to be coroutine at the top level. It's the partial support I don't quite understand the benefit of, since pyzmq already works fine in anyio as long as asyncio is underlying, as far as I understand it (methods return awaitables, need no special handling, etc.).

Thanks @davidbrochart for trying it out!

[davidbrochart]

It's the partial support I don't quite understand the benefit of, since pyzmq already works fine in anyio as long as asyncio is underlying

The goal is to get complete support of AnyIO, right? That would allow ipykernel user code to run on Trio as well.

[minrk]

Yes, absolutely. I think I misunderstood your original comment about coroutines - I thought you were proposing just changing the existing asyncio api as an incremental step (which I didn't understand, and maybe because it's not what you meant!), not whether changing to coroutines would be okay as part of full anyio support, which I think is totally sensible.