Use zmq-anyio (#1291)

Co-authored-by: M Bussonnier <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

docs/api/ipykernel.rst

@@ -134,12 +134,6 @@ Submodules
    :show-inheritance:
 
 
-.. automodule:: ipykernel.trio_runner
-   :members:
-   :undoc-members:
-   :show-inheritance:
-
-
 .. automodule:: ipykernel.zmqshell
    :members:
    :undoc-members:

ipykernel/debugger.py

@@ -242,7 +242,7 @@ async def _send_request(self, msg):
         self.log.debug("DEBUGPYCLIENT:")
         self.log.debug(self.routing_id)
         self.log.debug(buf)
-        await self.debugpy_socket.send_multipart((self.routing_id, buf))
+        await self.debugpy_socket.asend_multipart((self.routing_id, buf)).wait()
 
     async def _wait_for_response(self):
         # Since events are never pushed to the message_queue
@@ -438,7 +438,7 @@ async def start(self):
                 (self.shell_socket.getsockopt(ROUTING_ID)),
             )
 
-            msg = await self.shell_socket.recv_multipart()
+            msg = await self.shell_socket.arecv_multipart().wait()
             ident, msg = self.session.feed_identities(msg, copy=True)
             try:
                 msg = self.session.deserialize(msg, content=True, copy=True)

ipykernel/heartbeat.py

@@ -92,13 +92,17 @@ def _bind_socket(self):
     def run(self):
         """Run the heartbeat thread."""
         self.name = "Heartbeat"
-        self.socket = self.context.socket(zmq.ROUTER)
-        self.socket.linger = 1000
+
         try:
+            self.socket = self.context.socket(zmq.ROUTER)
+            self.socket.linger = 1000
             self._bind_socket()
         except Exception:
-            self.socket.close()
-            raise
+            try:
+                self.socket.close()
+            except Exception:
+                pass
+            return
 
         while True:
             try:

ipykernel/inprocess/ipkernel.py

@@ -54,7 +54,7 @@ class InProcessKernel(IPythonKernel):
     _underlying_iopub_socket = Instance(DummySocket, (False,))
     iopub_thread: IOPubThread = Instance(IOPubThread)  # type:ignore[assignment]
 
-    shell_socket = Instance(DummySocket, (True,))  # type:ignore[arg-type]
+    shell_socket = Instance(DummySocket, (True,))
 
     @default("iopub_thread")
     def _default_iopub_thread(self):
@@ -207,7 +207,7 @@ def enable_pylab(self, gui=None, import_all=True, welcome_message=False):
         """Activate pylab support at runtime."""
         if not gui:
             gui = self.kernel.gui
-        return super().enable_pylab(gui, import_all, welcome_message)
+        return super().enable_pylab(gui, import_all, welcome_message)  # type: ignore[call-arg]
 
 
 InteractiveShellABC.register(InProcessInteractiveShell)

ipykernel/inprocess/session.py

@@ -12,7 +12,7 @@ async def recv(  # type: ignore[override]
         mode, content, copy have no effect, but are present for superclass compatibility
 
         """
-        return await socket.recv_multipart()
+        return await socket.arecv_multipart().wait()
 
     def send(
         self,

ipykernel/inprocess/socket.py

@@ -65,4 +65,8 @@ async def poll(self, timeout=0):
         return statistics.current_buffer_used != 0
 
     def close(self):
-        pass
+        if self.is_shell:
+            self.in_send_stream.close()
+            self.in_receive_stream.close()
+            self.out_send_stream.close()
+            self.out_receive_stream.close()

ipykernel/iostream.py

@@ -20,6 +20,7 @@
 from typing import Any, Callable
 
 import zmq
+import zmq_anyio
 from anyio import sleep
 from jupyter_client.session import extract_header
 
@@ -48,7 +49,7 @@ class IOPubThread:
     whose IO is always run in a thread.
     """
 
-    def __init__(self, socket, pipe=False):
+    def __init__(self, socket: zmq_anyio.Socket, pipe: bool = False):
         """Create IOPub thread
 
         Parameters
@@ -61,10 +62,7 @@ def __init__(self, socket, pipe=False):
         """
         # ensure all of our sockets as sync zmq.Sockets
         # don't create async wrappers until we are within the appropriate coroutines
-        self.socket: zmq.Socket[bytes] | None = zmq.Socket(socket)
-        if self.socket.context is None:
-            # bug in pyzmq, shadow socket doesn't always inherit context attribute
-            self.socket.context = socket.context  # type:ignore[unreachable]
+        self.socket: zmq_anyio.Socket = socket
         self._context = socket.context
 
         self.background_socket = BackgroundSocket(self)
@@ -78,7 +76,7 @@ def __init__(self, socket, pipe=False):
         self._event_pipe_gc_lock: threading.Lock = threading.Lock()
         self._event_pipe_gc_seconds: float = 10
         self._setup_event_pipe()
-        tasks = [self._handle_event, self._run_event_pipe_gc]
+        tasks = [self._handle_event, self._run_event_pipe_gc, self.socket.start]
         if pipe:
             tasks.append(self._handle_pipe_msgs)
         self.thread = BaseThread(name="IOPub", daemon=True)
@@ -87,7 +85,7 @@ def __init__(self, socket, pipe=False):
 
     def _setup_event_pipe(self):
         """Create the PULL socket listening for events that should fire in this thread."""
-        self._pipe_in0 = self._context.socket(zmq.PULL, socket_class=zmq.Socket)
+        self._pipe_in0 = self._context.socket(zmq.PULL)
         self._pipe_in0.linger = 0
 
         _uuid = b2a_hex(os.urandom(16)).decode("ascii")
@@ -99,11 +97,11 @@ async def _run_event_pipe_gc(self):
         while True:
             await sleep(self._event_pipe_gc_seconds)
             try:
-                await self._event_pipe_gc()
+                self._event_pipe_gc()
             except Exception as e:
                 print(f"Exception in IOPubThread._event_pipe_gc: {e}", file=sys.__stderr__)
 
-    async def _event_pipe_gc(self):
+    def _event_pipe_gc(self):
         """run a single garbage collection on event pipes"""
         if not self._event_pipes:
             # don't acquire the lock if there's nothing to do
@@ -122,7 +120,7 @@ def _event_pipe(self):
         except AttributeError:
             # new thread, new event pipe
             # create sync base socket
-            event_pipe = self._context.socket(zmq.PUSH, socket_class=zmq.Socket)
+            event_pipe = self._context.socket(zmq.PUSH)
             event_pipe.linger = 0
             event_pipe.connect(self._event_interface)
             self._local.event_pipe = event_pipe
@@ -141,30 +139,28 @@ async def _handle_event(self):
         Whenever *an* event arrives on the event stream,
         *all* waiting events are processed in order.
         """
-        # create async wrapper within coroutine
-        pipe_in = zmq.asyncio.Socket(self._pipe_in0)
-        try:
-            while True:
-                await pipe_in.recv()
-                # freeze event count so new writes don't extend the queue
-                # while we are processing
-                n_events = len(self._events)
-                for _ in range(n_events):
-                    event_f = self._events.popleft()
-                    event_f()
-        except Exception:
-            if self.thread.stopped.is_set():
-                return
-            raise
+        pipe_in = zmq_anyio.Socket(self._pipe_in0)
+        async with pipe_in:
+            try:
+                while True:
+                    await pipe_in.arecv().wait()
+                    # freeze event count so new writes don't extend the queue
+                    # while we are processing
+                    n_events = len(self._events)
+                    for _ in range(n_events):
+                        event_f = self._events.popleft()
+                        event_f()
+            except Exception:
+                if self.thread.stopped.is_set():
+                    return
+                raise
 
     def _setup_pipe_in(self):
         """setup listening pipe for IOPub from forked subprocesses"""
-        ctx = self._context
-
         # use UUID to authenticate pipe messages
         self._pipe_uuid = os.urandom(16)
 
-        self._pipe_in1 = ctx.socket(zmq.PULL, socket_class=zmq.Socket)
+        self._pipe_in1 = zmq_anyio.Socket(self._context.socket(zmq.PULL))
         self._pipe_in1.linger = 0
 
         try:
@@ -181,19 +177,18 @@ def _setup_pipe_in(self):
 
     async def _handle_pipe_msgs(self):
         """handle pipe messages from a subprocess"""
-        # create async wrapper within coroutine
-        self._async_pipe_in1 = zmq.asyncio.Socket(self._pipe_in1)
-        try:
-            while True:
-                await self._handle_pipe_msg()
-        except Exception:
-            if self.thread.stopped.is_set():
-                return
-            raise
+        async with self._pipe_in1:
+            try:
+                while True:
+                    await self._handle_pipe_msg()
+            except Exception:
+                if self.thread.stopped.is_set():
+                    return
+                raise
 
     async def _handle_pipe_msg(self, msg=None):
         """handle a pipe message from a subprocess"""
-        msg = msg or await self._async_pipe_in1.recv_multipart()
+        msg = msg or await self._pipe_in1.arecv_multipart().wait()
         if not self._pipe_flag or not self._is_main_process():
             return
         if msg[0] != self._pipe_uuid:
@@ -246,7 +241,10 @@ def close(self):
         """Close the IOPub thread."""
         if self.closed:
             return
-        self._pipe_in0.close()
+        try:
+            self._pipe_in0.close()
+        except Exception:
+            pass
         if self._pipe_flag:
             self._pipe_in1.close()
         if self.socket is not None:

ipykernel/ipkernel.py

@@ -12,7 +12,7 @@
 from dataclasses import dataclass
 
 import comm
-import zmq.asyncio
+import zmq_anyio
 from anyio import TASK_STATUS_IGNORED, create_task_group, to_thread
 from anyio.abc import TaskStatus
 from IPython.core import release
@@ -93,7 +93,7 @@ class IPythonKernel(KernelBase):
         help="Set this flag to False to deactivate the use of experimental IPython completion APIs.",
     ).tag(config=True)
 
-    debugpy_socket = Instance(zmq.asyncio.Socket, allow_none=True)
+    debugpy_socket = Instance(zmq_anyio.Socket, allow_none=True)
 
     user_module = Any()
 
@@ -229,7 +229,8 @@ def __init__(self, **kwargs):
     }
 
     async def process_debugpy(self):
-        async with create_task_group() as tg:
+        assert self.debugpy_socket is not None
+        async with self.debug_shell_socket, self.debugpy_socket, create_task_group() as tg:
             tg.start_soon(self.receive_debugpy_messages)
             tg.start_soon(self.poll_stopped_queue)
             await to_thread.run_sync(self.debugpy_stop.wait)
@@ -252,7 +253,7 @@ async def receive_debugpy_message(self, msg=None):
 
         if msg is None:
             assert self.debugpy_socket is not None
-            msg = await self.debugpy_socket.recv_multipart()
+            msg = await self.debugpy_socket.arecv_multipart().wait()
         # The first frame is the socket id, we can drop it
         frame = msg[1].decode("utf-8")
         self.log.debug("Debugpy received: %s", frame)