From 48e39cf97000785c0cc180f74359f15ab630dd2e Mon Sep 17 00:00:00 2001 From: Min RK Date: Fri, 25 Oct 2024 14:49:01 +0200 Subject: [PATCH] restore zero-copy recv on shell messages (#1280) --- ipykernel/kernelbase.py | 20 +++++++------------- 1 file changed, 7 insertions(+), 13 deletions(-) diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index c6b18eb7..1c623c08 100644 --- a/ipykernel/kernelbase.py +++ b/ipykernel/kernelbase.py @@ -274,10 +274,9 @@ async def process_control_message(self, msg=None): assert self.control_thread is None or threading.current_thread() == self.control_thread msg = msg or await self.control_socket.recv_multipart() - copy = not isinstance(msg[0], zmq.Message) - idents, msg = self.session.feed_identities(msg, copy=copy) + idents, msg = self.session.feed_identities(msg) try: - msg = self.session.deserialize(msg, content=True, copy=copy) + msg = self.session.deserialize(msg, content=True) except Exception: self.log.error("Invalid Control Message", exc_info=True) # noqa: G201 return @@ -375,15 +374,12 @@ async def shell_channel_thread_main(self): try: while True: - msg = await self.shell_socket.recv_multipart() - - # Deserialize whole message just to get subshell_id. + msg = await self.shell_socket.recv_multipart(copy=False) + # deserialize only the header to get subshell_id # Keep original message to send to subshell_id unmodified. - # Ideally only want to deserialize message once. - copy = not isinstance(msg[0], zmq.Message) - _, msg2 = self.session.feed_identities(msg, copy=copy) + _, msg2 = self.session.feed_identities(msg, copy=False) try: - msg3 = self.session.deserialize(msg2, content=False, copy=copy) + msg3 = self.session.deserialize(msg2, content=False, copy=False) subshell_id = msg3["header"].get("subshell_id") # Find inproc pair socket to use to send message to correct subshell. @@ -1210,9 +1206,7 @@ def do_clear(self): def _topic(self, topic): """prefixed topic for IOPub messages""" - base = "kernel.%s" % self.ident - - return (f"{base}.{topic}").encode() + return (f"kernel.{self.ident}.{topic}").encode() _aborting = Bool(False)