From 4ea2d046ce345838c7501f043d6e2c0889b5ae5e Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Wed, 8 May 2024 11:34:51 +0100 Subject: [PATCH 01/37] Separate thread to handle incoming shell channel messages --- ipykernel/control.py | 30 ++++------------------- ipykernel/kernelapp.py | 9 +++++++ ipykernel/kernelbase.py | 50 ++++++++++++++++++++++++++++++--------- ipykernel/shellchannel.py | 42 ++++++++++++++++++++++++++++++++ ipykernel/thread.py | 37 +++++++++++++++++++++++++++++ 5 files changed, 131 insertions(+), 37 deletions(-) create mode 100644 ipykernel/shellchannel.py create mode 100644 ipykernel/thread.py diff --git a/ipykernel/control.py b/ipykernel/control.py index a70377c0..14a86a18 100644 --- a/ipykernel/control.py +++ b/ipykernel/control.py @@ -1,40 +1,18 @@ """A thread for a control channel.""" -from threading import Event, Thread -from anyio import create_task_group, run, to_thread +from .thread import BaseThread CONTROL_THREAD_NAME = "Control" -class ControlThread(Thread): +class ControlThread(BaseThread): """A thread for a control channel.""" def __init__(self, **kwargs): """Initialize the thread.""" - Thread.__init__(self, name=CONTROL_THREAD_NAME, **kwargs) - self.pydev_do_not_trace = True - self.is_pydev_daemon_thread = True - self.__stop = Event() - self._task = None - - def set_task(self, task): - self._task = task + super().__init__(name=CONTROL_THREAD_NAME, **kwargs) def run(self): """Run the thread.""" self.name = CONTROL_THREAD_NAME - run(self._main) - - async def _main(self): - async with create_task_group() as tg: - if self._task is not None: - tg.start_soon(self._task) - await to_thread.run_sync(self.__stop.wait) - tg.cancel_scope.cancel() - - def stop(self): - """Stop the thread. - - This method is threadsafe. - """ - self.__stop.set() + super().run() diff --git a/ipykernel/kernelapp.py b/ipykernel/kernelapp.py index c02c3cf3..642fb08b 100644 --- a/ipykernel/kernelapp.py +++ b/ipykernel/kernelapp.py @@ -53,8 +53,10 @@ from .iostream import IOPubThread from .ipkernel import IPythonKernel from .parentpoller import ParentPollerUnix, ParentPollerWindows +from .shellchannel import ShellChannelThread from .zmqshell import ZMQInteractiveShell + # ----------------------------------------------------------------------------- # Flags and Aliases # ----------------------------------------------------------------------------- @@ -143,6 +145,7 @@ class IPKernelApp(BaseIPythonApplication, InteractiveShellApp, ConnectionFileMix iopub_socket = Any() iopub_thread = Any() control_thread = Any() + shell_channel_thread = Any() _ports = Dict() @@ -367,6 +370,7 @@ def init_control(self, context): self.control_socket.router_handover = 1 self.control_thread = ControlThread(daemon=True) + self.shell_channel_thread = ShellChannelThread(context, daemon=True) def init_iopub(self, context): """Initialize the iopub channel.""" @@ -406,6 +410,10 @@ def close(self): self.log.debug("Closing control thread") self.control_thread.stop() self.control_thread.join() + if self.shell_channel_thread and self.shell_channel_thread.is_alive(): + self.log.debug("Closing shell channel thread") + self.shell_channel_thread.stop() + self.shell_channel_thread.join() if self.debugpy_socket and not self.debugpy_socket.closed: self.debugpy_socket.close() @@ -562,6 +570,7 @@ def init_kernel(self): debug_shell_socket=self.debug_shell_socket, shell_socket=self.shell_socket, control_thread=self.control_thread, + shell_channel_thread=self.shell_channel_thread, iopub_thread=self.iopub_thread, iopub_socket=self.iopub_socket, stdin_socket=self.stdin_socket, diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index 050f57be..0751fe0a 100644 --- a/ipykernel/kernelbase.py +++ b/ipykernel/kernelbase.py @@ -103,6 +103,7 @@ class Kernel(SingletonConfigurable): debug_shell_socket = Any() control_thread = Any() + shell_channel_thread = Any() iopub_socket = Any() iopub_thread = Any() stdin_socket = Any() @@ -259,7 +260,6 @@ async def process_control(self): while True: await self.process_control_message() except BaseException as e: - print("base exception") if self.control_stop.is_set(): return raise e @@ -356,28 +356,47 @@ async def advance_eventloop(): def _message_counter_default(self): return itertools.count() - async def shell_main(self): + async def shell_channel_thread_main(self): + assert self.shell_socket is not None + assert self.session is not None + assert self.shell_channel_thread + + try: + while True: + msg = await self.shell_socket.recv_multipart() + send_socket = self.shell_channel_thread.get_send_inproc_socket() + assert send_socket is not None + send_socket.send_multipart(msg, copy=False) + except BaseException as e: + if self.shell_stop.is_set(): + return + raise e + + async def shell_main(self, recv_socket): async with create_task_group() as tg: - tg.start_soon(self.process_shell) + tg.start_soon(self.process_shell, recv_socket) await to_thread.run_sync(self.shell_stop.wait) tg.cancel_scope.cancel() - async def process_shell(self): + async def process_shell(self, recv_socket=None): try: while True: - await self.process_shell_message() + await self.process_shell_message(recv_socket=recv_socket) except BaseException as e: if self.shell_stop.is_set(): return raise e - async def process_shell_message(self, msg=None): - assert self.shell_socket is not None + async def process_shell_message(self, msg=None, recv_socket=None): + if not recv_socket: + recv_socket = self.shell_socket + + assert recv_socket is not None assert self.session is not None - no_msg = msg is None if self._is_test else not await self.shell_socket.poll(0) + no_msg = msg is None if self._is_test else not await recv_socket.poll(0) + msg = msg or await recv_socket.recv_multipart(copy=False) - msg = msg or await self.shell_socket.recv_multipart() received_time = time.monotonic() copy = not isinstance(msg[0], zmq.Message) idents, msg = self.session.feed_identities(msg, copy=copy) @@ -474,8 +493,17 @@ async def start(self, *, task_status: TaskStatus = TASK_STATUS_IGNORED) -> None: self.shell_is_awaiting = False self.shell_is_blocking = False self.shell_stop = threading.Event() - if not self._is_test and self.shell_socket is not None: - tg.start_soon(self.shell_main) + + if self.shell_channel_thread: + # Each subshell thread listens on inproc socket for messages from shell channel thread + recv_socket = self.shell_channel_thread.create_inproc_sockets() + tg.start_soon(self.shell_main, recv_socket) + + self.shell_channel_thread.set_task(self.shell_channel_thread_main) + self.shell_channel_thread.start() + else: + if not self._is_test and self.shell_socket is not None: + tg.start_soon(self.shell_main, self.shell_socket) def stop(self): if not self._eventloop_set.is_set(): diff --git a/ipykernel/shellchannel.py b/ipykernel/shellchannel.py new file mode 100644 index 00000000..7da5f64f --- /dev/null +++ b/ipykernel/shellchannel.py @@ -0,0 +1,42 @@ +"""A thread for a shell channel.""" + +import zmq + +from .thread import BaseThread + +SHELL_CHANNEL_THREAD_NAME = "Shell channel" + + +class ShellChannelThread(BaseThread): + """A thread for a shell channel. + + Communicates with shell execute threads via pairs of ZMQ inproc sockets. + """ + + def __init__(self, context, **kwargs): + """Initialize the thread.""" + super().__init__(name=SHELL_CHANNEL_THREAD_NAME, **kwargs) + self._context = context + + def create_inproc_sockets(self): + # Socket used in another thread to receive messages from this thread. + self._recv_socket = self._context.socket(zmq.PAIR) + self._recv_socket.bind(f"inproc://shell") + + # Socket used in this thread to send messages. + self._send_socket = self._context.socket(zmq.PAIR) + self._send_socket.connect(f"inproc://shell") + + return self._recv_socket + + def get_send_inproc_socket(self): + return self._send_socket + + def run(self): + """Run the thread.""" + self.name = SHELL_CHANNEL_THREAD_NAME + super().run() + + for socket in (self._send_socket, self._recv_socket): + if socket and not socket.closed: + socket.close() diff --git a/ipykernel/thread.py b/ipykernel/thread.py new file mode 100644 index 00000000..16d6f321 --- /dev/null +++ b/ipykernel/thread.py @@ -0,0 +1,37 @@ +"""Base class for threads.""" +from threading import Event, Thread + +from anyio import create_task_group, run, to_thread + + +class BaseThread(Thread): + """Base class for threads.""" + + def __init__(self, name, **kwargs): + """Initialize the thread.""" + Thread.__init__(self, name=name, **kwargs) + self.pydev_do_not_trace = True + self.is_pydev_daemon_thread = True + self.__stop = Event() + self._task = None + + def set_task(self, task): + self._task = task + + def run(self): + """Run the thread.""" + run(self._main) + + async def _main(self): + async with create_task_group() as tg: + if self._task is not None: + tg.start_soon(self._task) + await to_thread.run_sync(self.__stop.wait) + tg.cancel_scope.cancel() + + def stop(self): + """Stop the thread. + + This method is threadsafe. + """ + self.__stop.set() From b801b79d376e286aff386f03f1ffbe611070e279 Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Wed, 8 May 2024 13:21:03 +0100 Subject: [PATCH 02/37] Add supported_features entry to kernel_info_reply --- ipykernel/kernelbase.py | 8 +++++++- tests/test_ipkernel_direct.py | 1 + tests/test_kernel_direct.py | 1 + tests/test_message_spec.py | 2 ++ 4 files changed, 11 insertions(+), 1 deletion(-) diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index 0751fe0a..22ebeeb7 100644 --- a/ipykernel/kernelbase.py +++ b/ipykernel/kernelbase.py @@ -834,7 +834,7 @@ async def connect_request(self, socket, ident, parent): @property def kernel_info(self): - return { + info = { "protocol_version": kernel_protocol_version, "implementation": self.implementation, "implementation_version": self.implementation_version, @@ -842,6 +842,9 @@ def kernel_info(self): "banner": self.banner, "help_links": self.help_links, } + if self._supports_kernel_subshells(): + info["supported_features"] = ["kernel subshells"] + return info async def kernel_info_request(self, socket, ident, parent): """Handle a kernel info request.""" @@ -1302,3 +1305,6 @@ async def _at_shutdown(self): ident=self._topic("shutdown"), ) self.log.debug("%s", self._shutdown_message) + + def _supports_kernel_subshells(self): + return self.shell_channel_thread is not None diff --git a/tests/test_ipkernel_direct.py b/tests/test_ipkernel_direct.py index cea2ec99..f650cc20 100644 --- a/tests/test_ipkernel_direct.py +++ b/tests/test_ipkernel_direct.py @@ -27,6 +27,7 @@ async def test_properties(ipkernel: IPythonKernel) -> None: async def test_direct_kernel_info_request(ipkernel): reply = await ipkernel.test_shell_message("kernel_info_request", {}) assert reply["header"]["msg_type"] == "kernel_info_reply" + assert "supported_features" not in reply["content"] or "kernel subshells" not in reply["content"]["supported_features"] async def test_direct_execute_request(ipkernel: MockIPyKernel) -> None: diff --git a/tests/test_kernel_direct.py b/tests/test_kernel_direct.py index ea3c6fe7..3b6af4c9 100644 --- a/tests/test_kernel_direct.py +++ b/tests/test_kernel_direct.py @@ -16,6 +16,7 @@ async def test_direct_kernel_info_request(kernel): reply = await kernel.test_shell_message("kernel_info_request", {}) assert reply["header"]["msg_type"] == "kernel_info_reply" + assert "supported_features" not in reply["content"] or "kernel subshells" not in reply["content"]["supported_features"] async def test_direct_execute_request(kernel): diff --git a/tests/test_message_spec.py b/tests/test_message_spec.py index d98503ee..1c2e4122 100644 --- a/tests/test_message_spec.py +++ b/tests/test_message_spec.py @@ -498,6 +498,8 @@ def test_kernel_info_request(): msg_id = KC.kernel_info() reply = get_reply(KC, msg_id, TIMEOUT) validate_message(reply, "kernel_info_reply", msg_id) + assert "supported_features" in reply["content"] + assert "kernel subshells" in reply["content"]["supported_features"] def test_connect_request(): From 0ddc1b7f2c56d13488a99d71cac9286d65ae80ea Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Wed, 8 May 2024 15:24:43 +0100 Subject: [PATCH 03/37] Dummy subshell control message implementations --- ipykernel/kernelbase.py | 62 ++++++++++++++++++++++++++++++++++++++ tests/test_message_spec.py | 40 ++++++++++++++++++++++++ tests/test_subshells.py | 59 ++++++++++++++++++++++++++++++++++++ 3 files changed, 161 insertions(+) create mode 100644 tests/test_subshells.py diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index 22ebeeb7..399989b0 100644 --- a/ipykernel/kernelbase.py +++ b/ipykernel/kernelbase.py @@ -206,6 +206,9 @@ def _parent_header(self): # execution count we store in the shell. execution_count = 0 + # Dictionary of subshell_id to subshell info. + _subshells = {} + msg_types = [ "execute_request", "complete_request", @@ -227,6 +230,9 @@ def _parent_header(self): "abort_request", "debug_request", "usage_request", + "create_subshell_request", + "delete_subshell_request", + "list_subshell_request", ] _eventloop_set: Event = Event() @@ -1015,6 +1021,62 @@ async def usage_request(self, socket, ident, parent): async def do_debug_request(self, msg): raise NotImplementedError + # --------------------------------------------------------------------------- + # Subshell control message handlers + # --------------------------------------------------------------------------- + + async def create_subshell_request(self, socket, ident, parent): + if not self._supports_kernel_subshells(): + self.log.error("KERNEL SUBSHELLS NOT SUPPORTED") + return + + subshell_id = str(uuid.uuid4()) + self._subshells[subshell_id] = {} + + content = { + "status": "ok", + "subshell_id": subshell_id, + } + self.session.send(socket, "create_subshell_reply", content, parent, ident) + + async def delete_subshell_request(self, socket, ident, parent): + if not self._supports_kernel_subshells(): + self.log.error("KERNEL SUBSHELLS NOT SUPPORTED") + return + + try: + content = parent["content"] + subshell_id = content["subshell_id"] + except Exception: + self.log.error("Got bad msg: ") + self.log.error("%s", parent) + return + + content: dict[str, t.Any] = {"status": "ok"} + try: + self._subshells.pop(subshell_id) + except KeyError as err: + import traceback + content = { + "status": "error", + "traceback": traceback.format_stack(), + "ename": str(type(err).__name__), + "evalue": str(err), + } + + self.session.send(socket, "delete_subshell_reply", content, parent, ident) + + async def list_subshell_request(self, socket, ident, parent): + if not self._supports_kernel_subshells(): + self.log.error("KERNEL SUBSHELLS NOT SUPPORTED") + return + + content = { + "status": "ok", + "subshell_id": list(self._subshells), + } + self.session.send(socket, "list_subshell_reply", content, parent, ident) + # --------------------------------------------------------------------------- # Engine methods (DEPRECATED) # --------------------------------------------------------------------------- diff --git a/tests/test_message_spec.py b/tests/test_message_spec.py index 1c2e4122..109a33a9 100644 --- a/tests/test_message_spec.py +++ b/tests/test_message_spec.py @@ -239,6 +239,20 @@ class HistoryReply(Reply): history = List(List()) +# Subshell control messages + +class CreateSubshellReply(Reply): + subshell_id = Unicode() + + +class DeleteSubshellReply(Reply): + pass + + +class ListSubshellReply(Reply): + subshell_id = List(Unicode()) + + references = { "execute_reply": ExecuteReply(), "inspect_reply": InspectReply(), @@ -255,6 +269,9 @@ class HistoryReply(Reply): "stream": Stream(), "display_data": DisplayData(), "header": RHeader(), + "create_subshell_reply": CreateSubshellReply(), + "delete_subshell_reply": DeleteSubshellReply(), + "list_subshell_reply": ListSubshellReply(), } # ----------------------------------------------------------------------------- @@ -511,6 +528,29 @@ def test_connect_request(): validate_message(reply, "connect_reply", msg_id) +def test_subshell(): + flush_channels() + + msg = KC.session.msg("create_subshell_request") + KC.control_channel.send(msg) + msg_id = msg["header"]["msg_id"] + reply = get_reply(KC, msg_id, TIMEOUT, channel="control") + validate_message(reply, "create_subshell_reply", msg_id) + subshell_id = reply["content"]["subshell_id"] + + msg = KC.session.msg("list_subshell_request") + KC.control_channel.send(msg) + msg_id = msg["header"]["msg_id"] + reply = get_reply(KC, msg_id, TIMEOUT, channel="control") + validate_message(reply, "list_subshell_reply", msg_id) + + msg = KC.session.msg("delete_subshell_request", {"subshell_id": subshell_id}) + KC.control_channel.send(msg) + msg_id = msg["header"]["msg_id"] + reply = get_reply(KC, msg_id, TIMEOUT, channel="control") + validate_message(reply, "delete_subshell_reply", msg_id) + + @pytest.mark.skipif( version_info < (5, 0), reason="earlier Jupyter Client don't have comm_info", diff --git a/tests/test_subshells.py b/tests/test_subshells.py new file mode 100644 index 00000000..6bfca3b0 --- /dev/null +++ b/tests/test_subshells.py @@ -0,0 +1,59 @@ +"""Test kernel subshells.""" + +# Copyright (c) IPython Development Team. +# Distributed under the terms of the Modified BSD License. + +from .utils import TIMEOUT, get_reply, kernel + + +# Helpers + +def create_subshell_helper(kc): + msg = kc.session.msg("create_subshell_request") + kc.control_channel.send(msg) + msg_id = msg["header"]["msg_id"] + reply = get_reply(kc, msg_id, TIMEOUT, channel="control") + return reply["content"] + + +def delete_subshell_helper(kc, subshell_id): + msg = kc.session.msg("delete_subshell_request", {"subshell_id": subshell_id}) + kc.control_channel.send(msg) + msg_id = msg["header"]["msg_id"] + reply = get_reply(kc, msg_id, TIMEOUT, channel="control") + return reply["content"] + + +def list_subshell_helper(kc): + msg = kc.session.msg("list_subshell_request") + kc.control_channel.send(msg) + msg_id = msg["header"]["msg_id"] + reply = get_reply(kc, msg_id, TIMEOUT, channel="control") + return reply["content"] + + +# Tests + +def test_supported(): + with kernel() as kc: + msg_id = kc.kernel_info() + reply = get_reply(kc, msg_id, TIMEOUT) + assert "supported_features" in reply["content"] + assert "kernel subshells" in reply["content"]["supported_features"] + + +def test_subshell_id_lifetime(): + with kernel() as kc: + assert list_subshell_helper(kc)["subshell_id"] == [] + subshell_id = create_subshell_helper(kc)["subshell_id"] + assert list_subshell_helper(kc)["subshell_id"] == [subshell_id] + delete_subshell_helper(kc, subshell_id) + assert list_subshell_helper(kc)["subshell_id"] == [] + + +def test_delete_non_existent(): + with kernel() as kc: + reply = delete_subshell_helper(kc, "unknown_subshell_id") + assert reply["status"] == "error" + for key in ("ename", "evalue", "traceback"): + assert key in reply From 17bd0f6a7115834757d5e2ff2cba4b290181e0b3 Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Fri, 10 May 2024 09:40:00 +0100 Subject: [PATCH 04/37] SubshellCache --- ipykernel/kernelbase.py | 17 +++---- ipykernel/shellchannel.py | 27 ++++------ ipykernel/subshell_cache.py | 99 +++++++++++++++++++++++++++++++++++++ 3 files changed, 116 insertions(+), 27 deletions(-) create mode 100644 ipykernel/subshell_cache.py diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index 399989b0..8ef901dd 100644 --- a/ipykernel/kernelbase.py +++ b/ipykernel/kernelbase.py @@ -206,9 +206,6 @@ def _parent_header(self): # execution count we store in the shell. execution_count = 0 - # Dictionary of subshell_id to subshell info. - _subshells = {} - msg_types = [ "execute_request", "complete_request", @@ -370,7 +367,8 @@ async def shell_channel_thread_main(self): try: while True: msg = await self.shell_socket.recv_multipart() - send_socket = self.shell_channel_thread.get_send_inproc_socket() + subshell_id = None # Would obtain this from msg + send_socket = self.shell_channel_thread.cache.get_send_socket(subshell_id) assert send_socket is not None send_socket.send_multipart(msg, copy=False) except BaseException as e: @@ -486,6 +484,7 @@ def post_handler_hook(self): async def start(self, *, task_status: TaskStatus = TASK_STATUS_IGNORED) -> None: """Process messages on shell and control channels""" + async with create_task_group() as tg: self.control_stop = threading.Event() if not self._is_test and self.control_socket is not None: @@ -502,8 +501,7 @@ async def start(self, *, task_status: TaskStatus = TASK_STATUS_IGNORED) -> None: if self.shell_channel_thread: # Each subshell thread listens on inproc socket for messages from shell channel thread - recv_socket = self.shell_channel_thread.create_inproc_sockets() - tg.start_soon(self.shell_main, recv_socket) + tg.start_soon(self.shell_main, self.shell_channel_thread.cache.get_recv_socket(None)) self.shell_channel_thread.set_task(self.shell_channel_thread_main) self.shell_channel_thread.start() @@ -1031,7 +1029,7 @@ async def create_subshell_request(self, socket, ident, parent): return subshell_id = str(uuid.uuid4()) - self._subshells[subshell_id] = {} + self.shell_channel_thread.cache.create(subshell_id) content = { "status": "ok", @@ -1054,7 +1052,8 @@ async def delete_subshell_request(self, socket, ident, parent): content: dict[str, t.Any] = {"status": "ok"} try: - self._subshells.pop(subshell_id) + # Should error here give traceback to the user? Probably not. + self.shell_channel_thread.cache.remove(subshell_id) except KeyError as err: import traceback content = { @@ -1073,7 +1072,7 @@ async def list_subshell_request(self, socket, ident, parent): content = { "status": "ok", - "subshell_id": list(self._subshells), + "subshell_id": self.shell_channel_thread.cache.list(), } self.session.send(socket, "list_subshell_reply", content, parent, ident) diff --git a/ipykernel/shellchannel.py b/ipykernel/shellchannel.py index 7da5f64f..9b13ac79 100644 --- a/ipykernel/shellchannel.py +++ b/ipykernel/shellchannel.py @@ -1,7 +1,6 @@ """A thread for a shell channel.""" -import zmq - +from .subshell_cache import SubshellCache from .thread import BaseThread SHELL_CHANNEL_THREAD_NAME = "Shell channel" @@ -16,27 +15,19 @@ class ShellChannelThread(BaseThread): def __init__(self, context, **kwargs): """Initialize the thread.""" super().__init__(name=SHELL_CHANNEL_THREAD_NAME, **kwargs) + self._cache: SubshellCache | None = None self._context = context - def create_inproc_sockets(self): - # Socket used in another thread to receive messages from this thread. - self._recv_socket = self._context.socket(zmq.PAIR) - self._recv_socket.bind(f"inproc://shell") - - # Socket used in this thread to send messages. - self._send_socket = self._context.socket(zmq.PAIR) - self._send_socket.connect(f"inproc://shell") - - return self._recv_socket - - def get_send_inproc_socket(self): - return self._send_socket + @property + def cache(self): + if self._cache is None: + self._cache = SubshellCache(self._context) + return self._cache def run(self): """Run the thread.""" self.name = SHELL_CHANNEL_THREAD_NAME super().run() - for socket in (self._send_socket, self._recv_socket): - if socket and not socket.closed: - socket.close() + if self._cache: + self._cache.close() diff --git a/ipykernel/subshell_cache.py b/ipykernel/subshell_cache.py new file mode 100644 index 00000000..bb44bb8c --- /dev/null +++ b/ipykernel/subshell_cache.py @@ -0,0 +1,99 @@ +"""A cache for subshell information.""" +# Copyright (c) IPython Development Team. +# Distributed under the terms of the Modified BSD License. + +from threading import Lock +import zmq + + +class SubshellCache: + """A cache for subshell information. + + Care needed with threadsafe access here. After the cache is created, all write + access is performed by the control thread so there is only ever one write access at + any one time. Reading of cache information is performed by a number of different + threads: + + 1) Receive inproc socket for a child subshell is passed to the subshell task (by + the control thread) when it is created. + 2) Send inproc socket is looked up for every message received by the shell channel + thread so that the message is sent to the correct subshell thread. + 3) Control thread reads shell_ids for list_subshell_request message. + + Python dictionary reads and writes are atomic and therefore threadsafe because of + the GIL in conventional CPython. But wise to use a mutex to support non-GIL + python. + + Read-only access to the parent subshell sockets never needs to be wrapped in a + mutex as there is only one pair over the whole lifetime of this object. + """ + def __init__(self, context: zmq.Context): + self._context: zmq.Context = context + self._cache = {} + self._lock: Lock = Lock() + + # Parent subshell sockets. + self._parent_send_socket, self._parent_recv_socket = \ + self._create_inproc_sockets(None) + + def close(self): + for socket in (self._parent_send_socket, self._parent_recv_socket): + if socket and not socket.closed: + socket.close() + + self._parent_recv_socket = None + self._parent_send_socket = None + assert not self._cache # Should not be anything left in cache. + + def create(self, subshell_id: str) -> None: + # check if subshell_id already exists... + # assume it doesn't + with self._lock: + assert subshell_id not in self._cache + send_socket, recv_socket = self._create_inproc_sockets(subshell_id) + self._cache[subshell_id] = { + send_socket: send_socket, + recv_socket: recv_socket, + } + + def get_recv_socket(self, subshell_id: str | None): + if subshell_id is None: + return self._parent_recv_socket + else: + with self._lock: + return self._cache[subshell_id]["recv_socket"] + + def get_send_socket(self, subshell_id: str | None): + if subshell_id is None: + return self._parent_send_socket + else: + with self._lock: + return self._cache[subshell_id]["send_socket"] + + def list(self) -> list[str]: + with self._lock: + return list(self._cache) + + def remove(self, subshell_id: str) -> None: + """Raises key error if subshell_id not in cache""" + with self._lock: + dict_ = self._cache.pop(subshell_id) + for socket in (dict_["send_socket"], dict_["recv_socket"]): + if socket and not socket.closed: + socket.close() + + def _create_inproc_sockets(self, subshell_id: str | None): + """Create a pair of inproc sockets to communicate with a subshell. + """ + name = f"shell-{subshell_id}" if subshell_id else "shell" + address = f"inproc://{name}" + + # Socket used in subshell thread to receive messages. + recv_socket = self._context.socket(zmq.PAIR) + recv_socket.bind(address) + + # Socket used in shell channel thread to send messages. + send_socket = self._context.socket(zmq.PAIR) + send_socket.connect(address) + + return send_socket, recv_socket From 62e5d2199e9a2e6716814adfa9e7fbfe583d66d8 Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Fri, 10 May 2024 14:58:42 +0100 Subject: [PATCH 05/37] Subshell threads with correct lifetime but don't yet execute code --- ipykernel/kernelbase.py | 8 +++++++- ipykernel/subshell.py | 12 ++++++++++++ ipykernel/subshell_cache.py | 16 +++++++++++++--- tests/test_subshells.py | 26 ++++++++++++++++++++++++++ 4 files changed, 58 insertions(+), 4 deletions(-) create mode 100644 ipykernel/subshell.py diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index 8ef901dd..c48c484b 100644 --- a/ipykernel/kernelbase.py +++ b/ipykernel/kernelbase.py @@ -19,6 +19,7 @@ from signal import SIGINT, SIGTERM, Signals from .control import CONTROL_THREAD_NAME +from .subshell import SubshellThread if sys.platform != "win32": from signal import SIGKILL @@ -1029,7 +1030,11 @@ async def create_subshell_request(self, socket, ident, parent): return subshell_id = str(uuid.uuid4()) - self.shell_channel_thread.cache.create(subshell_id) + thread = SubshellThread(subshell_id) + self.shell_channel_thread.cache.create(subshell_id, thread) + + #thread.set_task() + thread.start() content = { "status": "ok", @@ -1051,6 +1056,7 @@ async def delete_subshell_request(self, socket, ident, parent): return content: dict[str, t.Any] = {"status": "ok"} + try: # Should error here give traceback to the user? Probably not. self.shell_channel_thread.cache.remove(subshell_id) diff --git a/ipykernel/subshell.py b/ipykernel/subshell.py new file mode 100644 index 00000000..b1f6b3f1 --- /dev/null +++ b/ipykernel/subshell.py @@ -0,0 +1,12 @@ +"""A thread for a subshell.""" + +from .thread import BaseThread + + +class SubshellThread(BaseThread): + """A thread for a subshell. + """ + + def __init__(self, subshell_id: str, **kwargs): + """Initialize the thread.""" + super().__init__(name=f"subshell-{subshell_id}", **kwargs) diff --git a/ipykernel/subshell_cache.py b/ipykernel/subshell_cache.py index bb44bb8c..b481e27c 100644 --- a/ipykernel/subshell_cache.py +++ b/ipykernel/subshell_cache.py @@ -5,6 +5,8 @@ from threading import Lock import zmq +from .subshell import SubshellThread + class SubshellCache: """A cache for subshell information. @@ -45,15 +47,17 @@ def close(self): self._parent_send_socket = None assert not self._cache # Should not be anything left in cache. - def create(self, subshell_id: str) -> None: + def create(self, subshell_id: str, thread: SubshellThread) -> None: # check if subshell_id already exists... # assume it doesn't + with self._lock: assert subshell_id not in self._cache send_socket, recv_socket = self._create_inproc_sockets(subshell_id) self._cache[subshell_id] = { - send_socket: send_socket, - recv_socket: recv_socket, + "send_socket": send_socket, + "recv_socket": recv_socket, + "thread": thread, } def get_recv_socket(self, subshell_id: str | None): @@ -78,6 +82,12 @@ def remove(self, subshell_id: str) -> None: """Raises key error if subshell_id not in cache""" with self._lock: dict_ = self._cache.pop(subshell_id) + + thread = dict_["thread"] + if thread and thread.is_alive(): + thread.stop() + thread.join() + for socket in (dict_["send_socket"], dict_["recv_socket"]): if socket and not socket.closed: socket.close() diff --git a/tests/test_subshells.py b/tests/test_subshells.py index 6bfca3b0..41025a33 100644 --- a/tests/test_subshells.py +++ b/tests/test_subshells.py @@ -32,6 +32,19 @@ def list_subshell_helper(kc): return reply["content"] +def get_thread_count(kc): + code = "import threading as t; print(t.active_count())" + msg_id = kc.execute(code=code) + while True: + msg = kc.get_iopub_msg() + # Get the stream message corresponding to msg_id + if msg["msg_type"] == "stream" and msg["parent_header"]["msg_id"] == msg_id: + content = msg["content"] + #assert content["name"] == "stdout" + break + return int(content["text"].strip()) + + # Tests def test_supported(): @@ -57,3 +70,16 @@ def test_delete_non_existent(): assert reply["status"] == "error" for key in ("ename", "evalue", "traceback"): assert key in reply + + +def test_thread_counts(): + with kernel() as kc: + nthreads = get_thread_count(kc) + + subshell_id = create_subshell_helper(kc)["subshell_id"] + nthreads2 = get_thread_count(kc) + assert nthreads2 > nthreads + + delete_subshell_helper(kc, subshell_id) + nthreads3 = get_thread_count(kc) + assert nthreads3 == nthreads From 2f8273b30c3aa9c803d9f58368942bd88545183a Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Mon, 13 May 2024 08:07:29 +0100 Subject: [PATCH 06/37] Extract subshell_id, and handle message in correct subshell --- ipykernel/kernelbase.py | 30 +++++++++++++---- ipykernel/thread.py | 6 ++-- tests/test_subshells.py | 72 +++++++++++++++++++++++++++++++++++++---- 3 files changed, 92 insertions(+), 16 deletions(-) diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index c48c484b..e524c225 100644 --- a/ipykernel/kernelbase.py +++ b/ipykernel/kernelbase.py @@ -368,7 +368,15 @@ async def shell_channel_thread_main(self): try: while True: msg = await self.shell_socket.recv_multipart() - subshell_id = None # Would obtain this from msg + + # Deserialize whole message just 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.deserialize(msg2, content=False, copy=copy) + subshell_id = msg2["header"].get("subshell_id") + send_socket = self.shell_channel_thread.cache.get_send_socket(subshell_id) assert send_socket is not None send_socket.send_multipart(msg, copy=False) @@ -377,11 +385,18 @@ async def shell_channel_thread_main(self): return raise e - async def shell_main(self, recv_socket): + async def shell_main(self, subshell_id: str | None): + if self._supports_kernel_subshells(): + recv_socket = self.shell_channel_thread.cache.get_recv_socket(subshell_id) + else: + recv_socket = self.shell_socket + async with create_task_group() as tg: tg.start_soon(self.process_shell, recv_socket) - await to_thread.run_sync(self.shell_stop.wait) - tg.cancel_scope.cancel() + if subshell_id is None: + # Main subshell. + await to_thread.run_sync(self.shell_stop.wait) + tg.cancel_scope.cancel() async def process_shell(self, recv_socket=None): try: @@ -502,13 +517,13 @@ async def start(self, *, task_status: TaskStatus = TASK_STATUS_IGNORED) -> None: if self.shell_channel_thread: # Each subshell thread listens on inproc socket for messages from shell channel thread - tg.start_soon(self.shell_main, self.shell_channel_thread.cache.get_recv_socket(None)) + tg.start_soon(self.shell_main, None) self.shell_channel_thread.set_task(self.shell_channel_thread_main) self.shell_channel_thread.start() else: if not self._is_test and self.shell_socket is not None: - tg.start_soon(self.shell_main, self.shell_socket) + tg.start_soon(self.shell_main, None) def stop(self): if not self._eventloop_set.is_set(): @@ -1032,8 +1047,9 @@ async def create_subshell_request(self, socket, ident, parent): subshell_id = str(uuid.uuid4()) thread = SubshellThread(subshell_id) self.shell_channel_thread.cache.create(subshell_id, thread) + recv_socket = self.shell_channel_thread.cache.get_recv_socket(subshell_id) - #thread.set_task() + thread.set_task(self.shell_main, subshell_id) thread.start() content = { diff --git a/ipykernel/thread.py b/ipykernel/thread.py index 16d6f321..92448b5a 100644 --- a/ipykernel/thread.py +++ b/ipykernel/thread.py @@ -14,9 +14,11 @@ def __init__(self, name, **kwargs): self.is_pydev_daemon_thread = True self.__stop = Event() self._task = None + self._task_args = () - def set_task(self, task): + def set_task(self, task, *args): self._task = task + self._task_args = args def run(self): """Run the thread.""" @@ -25,7 +27,7 @@ def run(self): async def _main(self): async with create_task_group() as tg: if self._task is not None: - tg.start_soon(self._task) + tg.start_soon(self._task, *self._task_args) await to_thread.run_sync(self.__stop.wait) tg.cancel_scope.cancel() diff --git a/tests/test_subshells.py b/tests/test_subshells.py index 41025a33..294a2ab3 100644 --- a/tests/test_subshells.py +++ b/tests/test_subshells.py @@ -3,12 +3,16 @@ # Copyright (c) IPython Development Team. # Distributed under the terms of the Modified BSD License. +from datetime import datetime, timedelta + +from jupyter_client.blocking.client import BlockingKernelClient + from .utils import TIMEOUT, get_reply, kernel # Helpers -def create_subshell_helper(kc): +def create_subshell_helper(kc: BlockingKernelClient): msg = kc.session.msg("create_subshell_request") kc.control_channel.send(msg) msg_id = msg["header"]["msg_id"] @@ -16,7 +20,7 @@ def create_subshell_helper(kc): return reply["content"] -def delete_subshell_helper(kc, subshell_id): +def delete_subshell_helper(kc: BlockingKernelClient, subshell_id: str): msg = kc.session.msg("delete_subshell_request", {"subshell_id": subshell_id}) kc.control_channel.send(msg) msg_id = msg["header"]["msg_id"] @@ -24,7 +28,7 @@ def delete_subshell_helper(kc, subshell_id): return reply["content"] -def list_subshell_helper(kc): +def list_subshell_helper(kc: BlockingKernelClient): msg = kc.session.msg("list_subshell_request") kc.control_channel.send(msg) msg_id = msg["header"]["msg_id"] @@ -32,9 +36,11 @@ def list_subshell_helper(kc): return reply["content"] -def get_thread_count(kc): - code = "import threading as t; print(t.active_count())" - msg_id = kc.execute(code=code) +def execute_request_subshell_id(kc: BlockingKernelClient, code: str, subshell_id: str | None): + msg = kc.session.msg("execute_request", {"code": code}) + msg["header"]["subshell_id"] = subshell_id + msg_id = msg["msg_id"] + kc.shell_channel.send(msg) while True: msg = kc.get_iopub_msg() # Get the stream message corresponding to msg_id @@ -42,7 +48,17 @@ def get_thread_count(kc): content = msg["content"] #assert content["name"] == "stdout" break - return int(content["text"].strip()) + return content["text"].strip() + + +def get_thread_count(kc: BlockingKernelClient) -> int: + code = "import threading as t; print(t.active_count())" + return int(execute_request_subshell_id(kc, code, None)) + + +def get_thread_ids(kc: BlockingKernelClient, subshell_id: str | None = None) -> tuple[str, str]: + code = "import threading as t; print(t.get_ident(), t.main_thread().ident)" + return execute_request_subshell_id(kc, code, subshell_id).split() # Tests @@ -83,3 +99,45 @@ def test_thread_counts(): delete_subshell_helper(kc, subshell_id) nthreads3 = get_thread_count(kc) assert nthreads3 == nthreads + + +def test_thread_ids(): + with kernel() as kc: + subshell_id = create_subshell_helper(kc)["subshell_id"] + + thread_id, main_thread_id = get_thread_ids(kc) + assert thread_id == main_thread_id + + thread_id, main_thread_id = get_thread_ids(kc, subshell_id) + assert thread_id != main_thread_id + + delete_subshell_helper(kc, subshell_id) + + +def test_run_concurrently(): + with kernel() as kc: + subshell_id = create_subshell_helper(kc)["subshell_id"] + + # Prepare messages + times = (0.05, 0.1) # Sleep seconds + msgs = [] + for id, sleep in zip((None, subshell_id), times): + code = f"import time; time.sleep({sleep})" + msg = kc.session.msg("execute_request", {"code": code}) + msg["header"]["subshell_id"] = id + msgs.append(msg) + + # Send messages + start = datetime.now() + for msg in msgs: + kc.shell_channel.send(msg) + + # Wait for replies + _ = [get_reply(kc, msg["msg_id"]) for msg in msgs] + end = datetime.now() + + duration = end - start + assert duration >= timedelta(seconds=max(times)) + assert duration < timedelta(seconds=sum(times)) + + delete_subshell_helper(kc, subshell_id) From 1cc88923b8f11969b11594743d4c7c0a44d31806 Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Wed, 22 May 2024 12:34:31 +0100 Subject: [PATCH 07/37] Add subshell test_execution_count --- ipykernel/kernelbase.py | 1 - tests/test_subshells.py | 47 +++++++++++++++++++++++++++++++---------- tests/utils.py | 24 +++++++++++++++++++++ 3 files changed, 60 insertions(+), 12 deletions(-) diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index e524c225..68077688 100644 --- a/ipykernel/kernelbase.py +++ b/ipykernel/kernelbase.py @@ -1047,7 +1047,6 @@ async def create_subshell_request(self, socket, ident, parent): subshell_id = str(uuid.uuid4()) thread = SubshellThread(subshell_id) self.shell_channel_thread.cache.create(subshell_id, thread) - recv_socket = self.shell_channel_thread.cache.get_recv_socket(subshell_id) thread.set_task(self.shell_main, subshell_id) thread.start() diff --git a/tests/test_subshells.py b/tests/test_subshells.py index 294a2ab3..f674378b 100644 --- a/tests/test_subshells.py +++ b/tests/test_subshells.py @@ -7,7 +7,7 @@ from jupyter_client.blocking.client import BlockingKernelClient -from .utils import TIMEOUT, get_reply, kernel +from .utils import TIMEOUT, get_reply, get_replies, kernel # Helpers @@ -51,12 +51,12 @@ def execute_request_subshell_id(kc: BlockingKernelClient, code: str, subshell_id return content["text"].strip() -def get_thread_count(kc: BlockingKernelClient) -> int: +def execute_thread_count(kc: BlockingKernelClient) -> int: code = "import threading as t; print(t.active_count())" return int(execute_request_subshell_id(kc, code, None)) -def get_thread_ids(kc: BlockingKernelClient, subshell_id: str | None = None) -> tuple[str, str]: +def execute_thread_ids(kc: BlockingKernelClient, subshell_id: str | None = None) -> tuple[str, str]: code = "import threading as t; print(t.get_ident(), t.main_thread().ident)" return execute_request_subshell_id(kc, code, subshell_id).split() @@ -90,14 +90,14 @@ def test_delete_non_existent(): def test_thread_counts(): with kernel() as kc: - nthreads = get_thread_count(kc) + nthreads = execute_thread_count(kc) subshell_id = create_subshell_helper(kc)["subshell_id"] - nthreads2 = get_thread_count(kc) + nthreads2 = execute_thread_count(kc) assert nthreads2 > nthreads delete_subshell_helper(kc, subshell_id) - nthreads3 = get_thread_count(kc) + nthreads3 = execute_thread_count(kc) assert nthreads3 == nthreads @@ -105,10 +105,10 @@ def test_thread_ids(): with kernel() as kc: subshell_id = create_subshell_helper(kc)["subshell_id"] - thread_id, main_thread_id = get_thread_ids(kc) + thread_id, main_thread_id = execute_thread_ids(kc) assert thread_id == main_thread_id - thread_id, main_thread_id = get_thread_ids(kc, subshell_id) + thread_id, main_thread_id = execute_thread_ids(kc, subshell_id) assert thread_id != main_thread_id delete_subshell_helper(kc, subshell_id) @@ -119,7 +119,7 @@ def test_run_concurrently(): subshell_id = create_subshell_helper(kc)["subshell_id"] # Prepare messages - times = (0.05, 0.1) # Sleep seconds + times = (0.05, 0.1) # Sleep seconds. Test can fail with (0.05, 0.05) msgs = [] for id, sleep in zip((None, subshell_id), times): code = f"import time; time.sleep({sleep})" @@ -132,8 +132,7 @@ def test_run_concurrently(): for msg in msgs: kc.shell_channel.send(msg) - # Wait for replies - _ = [get_reply(kc, msg["msg_id"]) for msg in msgs] + _ = get_replies(kc, [msg["msg_id"] for msg in msgs]) end = datetime.now() duration = end - start @@ -141,3 +140,29 @@ def test_run_concurrently(): assert duration < timedelta(seconds=sum(times)) delete_subshell_helper(kc, subshell_id) + + +def test_execution_count(): + with kernel() as kc: + subshell_id = create_subshell_helper(kc)["subshell_id"] + + # Prepare messages + times = (0.1, 0.05, 0.2, 0.07) # Sleep seconds + msgs = [] + for id, sleep in zip((None, subshell_id, None, subshell_id), times): + code = f"import time; time.sleep({sleep})" + msg = kc.session.msg("execute_request", {"code": code}) + msg["header"]["subshell_id"] = id + msgs.append(msg) + + for msg in msgs: + kc.shell_channel.send(msg) + + # Wait for replies, may be in any order. + replies = get_replies(kc, [msg["msg_id"] for msg in msgs]) + + execution_counts = [r["content"]["execution_count"] for r in replies] + ec = execution_counts[0] + assert execution_counts == [ec, ec-1, ec+2, ec+1] + + delete_subshell_helper(kc, subshell_id) diff --git a/tests/utils.py b/tests/utils.py index b1b4119f..5fc18a82 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -68,6 +68,30 @@ def get_reply(kc, msg_id, timeout=TIMEOUT, channel="shell"): return reply +def get_replies(kc, msg_ids: list[str], timeout=TIMEOUT, channel="shell"): + # Get replies which may arrive in any order as they may be running on different subshells. + # Replies are returned in the same order as the msg_ids, not in the order of arrival. + t0 = time() + count = 0 + replies = [None]*len(msg_ids) + while True: + get_msg = getattr(kc, f"get_{channel}_msg") + reply = get_msg(timeout=timeout) + try: + msg_id = reply["parent_header"]["msg_id"] + replies[msg_ids.index(msg_id)] = reply + count += 1 + if count == len(msg_ids): + break + except ValueError: + # Allow debugging ignored replies + print(f"Ignoring reply not to any of {msg_ids}: {reply}") + t1 = time() + timeout -= t1 - t0 + t0 = t1 + return replies + + def execute(code="", kc=None, **kwargs): """wrapper for doing common steps for validating an execution request""" from .test_message_spec import validate_message From f5b664d11bf25b8d3bc5f2c0c51b16c8a5dae1b1 Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Wed, 22 May 2024 15:14:08 +0100 Subject: [PATCH 08/37] Use dataclass for Subshell, clean up correctly --- ipykernel/kernelbase.py | 2 +- ipykernel/subshell_cache.py | 60 +++++++++++++++++++++++-------------- 2 files changed, 38 insertions(+), 24 deletions(-) diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index 68077688..fc07d74c 100644 --- a/ipykernel/kernelbase.py +++ b/ipykernel/kernelbase.py @@ -1074,7 +1074,7 @@ async def delete_subshell_request(self, socket, ident, parent): try: # Should error here give traceback to the user? Probably not. - self.shell_channel_thread.cache.remove(subshell_id) + self.shell_channel_thread.cache.delete(subshell_id) except KeyError as err: import traceback content = { diff --git a/ipykernel/subshell_cache.py b/ipykernel/subshell_cache.py index b481e27c..12961d3d 100644 --- a/ipykernel/subshell_cache.py +++ b/ipykernel/subshell_cache.py @@ -2,12 +2,20 @@ # Copyright (c) IPython Development Team. # Distributed under the terms of the Modified BSD License. +from dataclasses import dataclass from threading import Lock import zmq from .subshell import SubshellThread +@dataclass +class Subshell: + thread: SubshellThread + send_socket: zmq.Socket + recv_socket: zmq.Socket + + class SubshellCache: """A cache for subshell information. @@ -31,7 +39,7 @@ class SubshellCache: """ def __init__(self, context: zmq.Context): self._context: zmq.Context = context - self._cache = {} + self._cache: dict[str, Subshell] = {} self._lock: Lock = Lock() # Parent subshell sockets. @@ -45,7 +53,14 @@ def close(self): self._parent_recv_socket = None self._parent_send_socket = None - assert not self._cache # Should not be anything left in cache. + + with self._lock: + while True: + try: + _, subshell = self._cache.popitem() + except KeyError: + break + self._stop_subshell(subshell) def create(self, subshell_id: str, thread: SubshellThread) -> None: # check if subshell_id already exists... @@ -54,44 +69,33 @@ def create(self, subshell_id: str, thread: SubshellThread) -> None: with self._lock: assert subshell_id not in self._cache send_socket, recv_socket = self._create_inproc_sockets(subshell_id) - self._cache[subshell_id] = { - "send_socket": send_socket, - "recv_socket": recv_socket, - "thread": thread, - } + self._cache[subshell_id] = Subshell(thread, send_socket, recv_socket) + + def delete(self, subshell_id: str) -> None: + """Raises key error if subshell_id not in cache""" + with self._lock: + subshell = self._cache.pop(subshell_id) + + self._stop_subshell(subshell) def get_recv_socket(self, subshell_id: str | None): if subshell_id is None: return self._parent_recv_socket else: with self._lock: - return self._cache[subshell_id]["recv_socket"] + return self._cache[subshell_id].recv_socket def get_send_socket(self, subshell_id: str | None): if subshell_id is None: return self._parent_send_socket else: with self._lock: - return self._cache[subshell_id]["send_socket"] + return self._cache[subshell_id].send_socket def list(self) -> list[str]: with self._lock: return list(self._cache) - def remove(self, subshell_id: str) -> None: - """Raises key error if subshell_id not in cache""" - with self._lock: - dict_ = self._cache.pop(subshell_id) - - thread = dict_["thread"] - if thread and thread.is_alive(): - thread.stop() - thread.join() - - for socket in (dict_["send_socket"], dict_["recv_socket"]): - if socket and not socket.closed: - socket.close() - def _create_inproc_sockets(self, subshell_id: str | None): """Create a pair of inproc sockets to communicate with a subshell. """ @@ -107,3 +111,13 @@ def _create_inproc_sockets(self, subshell_id: str | None): send_socket.connect(address) return send_socket, recv_socket + + def _stop_subshell(self, subshell: Subshell): + thread = subshell.thread + if thread and thread.is_alive(): + thread.stop() + thread.join() + + for socket in (subshell.send_socket, subshell.recv_socket): + if socket and not socket.closed: + socket.close() From c3fd89ebded5aae718f36d30c5b1939a2d80f803 Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Wed, 22 May 2024 16:39:09 +0100 Subject: [PATCH 09/37] Test can create subshell whilst main subshell is executing --- tests/test_subshells.py | 23 ++++++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/tests/test_subshells.py b/tests/test_subshells.py index f674378b..5312041d 100644 --- a/tests/test_subshells.py +++ b/tests/test_subshells.py @@ -160,9 +160,30 @@ def test_execution_count(): # Wait for replies, may be in any order. replies = get_replies(kc, [msg["msg_id"] for msg in msgs]) - + execution_counts = [r["content"]["execution_count"] for r in replies] ec = execution_counts[0] assert execution_counts == [ec, ec-1, ec+2, ec+1] delete_subshell_helper(kc, subshell_id) + + +def test_create_while_execute(): + with kernel() as kc: + # Send request to execute code on main subshell. + msg = kc.session.msg("execute_request", {"code": "import time; time.sleep(0.05)"}) + kc.shell_channel.send(msg) + + # Create subshell via control channel. + control_msg = kc.session.msg("create_subshell_request") + kc.control_channel.send(control_msg) + control_reply = get_reply(kc, control_msg["header"]["msg_id"], TIMEOUT, channel="control") + subshell_id = control_reply["content"]["subshell_id"] + control_date = control_reply["header"]["date"] + + # Get result message from main subshell. + shell_date = get_reply(kc, msg["msg_id"])["header"]["date"] + + delete_subshell_helper(kc, subshell_id) + + assert control_date < shell_date From 7fc1d6eb1c51f18afd806c893a3c095ddd74b615 Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Thu, 23 May 2024 13:09:16 +0100 Subject: [PATCH 10/37] Add %subshell magic --- ipykernel/subshell_cache.py | 16 +++++++++++++++- ipykernel/zmqshell.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/ipykernel/subshell_cache.py b/ipykernel/subshell_cache.py index 12961d3d..076cd749 100644 --- a/ipykernel/subshell_cache.py +++ b/ipykernel/subshell_cache.py @@ -3,7 +3,7 @@ # Distributed under the terms of the Modified BSD License. from dataclasses import dataclass -from threading import Lock +from threading import Lock, main_thread import zmq from .subshell import SubshellThread @@ -96,6 +96,20 @@ def list(self) -> list[str]: with self._lock: return list(self._cache) + def subshell_id_from_thread_id(self, thread_id) -> str | None: + """Return subshell_id of the specified thread_id. + + Raises RuntimeError if thread_id is not the main shell or a subshell. + """ + with self._lock: + if thread_id == main_thread().ident: + return None + for id, subshell in self._cache.items(): + if subshell.thread.ident == thread_id: + return id + msg = f"Thread id '{thread_id} does not correspond to a subshell of this kernel" + raise RuntimeError(msg) + def _create_inproc_sockets(self, subshell_id: str | None): """Create a pair of inproc sockets to communicate with a subshell. """ diff --git a/ipykernel/zmqshell.py b/ipykernel/zmqshell.py index bc99d000..51efda7b 100644 --- a/ipykernel/zmqshell.py +++ b/ipykernel/zmqshell.py @@ -439,6 +439,37 @@ def autosave(self, arg_s): else: print("Autosave disabled") + @line_magic + def subshell(self, arg_s): + from ipykernel.kernelapp import IPKernelApp + + if not IPKernelApp.initialized(): + msg = "Not in a running Kernel" + raise RuntimeError(msg) + + app = IPKernelApp.instance() + kernel = app.kernel + + if not (hasattr(kernel, "_supports_kernel_subshells") and + kernel._supports_kernel_subshells()): + print("Kernel does not support subshells") + return + + import threading + thread_id = threading.current_thread().ident + try: + subshell_id = kernel.shell_channel_thread.cache.subshell_id_from_thread_id(thread_id) + except RuntimeError: + subshell_id = "unknown" + subshell_id_list = kernel.shell_channel_thread.cache.list() + + print(f"subshell id: {subshell_id}") + print(f"thread id: {thread_id}") + print(f"main thread id: {threading.main_thread().ident}") + print(f"pid: {os.getpid()}") + print(f"thread count: {threading.active_count()}") + print(f"subshell list: {subshell_id_list}") + class ZMQInteractiveShell(InteractiveShell): """A subclass of InteractiveShell for ZMQ.""" From b8b7af5d620bc3115585058f73beb8192fe3effb Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Wed, 29 May 2024 16:41:14 +0100 Subject: [PATCH 11/37] Temporary mutex for execute_reply from subshells --- ipykernel/kernelbase.py | 16 +++++++++++----- tests/test_subshells.py | 9 ++++++--- 2 files changed, 17 insertions(+), 8 deletions(-) diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index fc07d74c..14b11094 100644 --- a/ipykernel/kernelbase.py +++ b/ipykernel/kernelbase.py @@ -732,15 +732,21 @@ async def execute_request(self, socket, ident, parent): reply_content = json_clean(reply_content) metadata = self.finish_metadata(parent, metadata, reply_content) - reply_msg = self.session.send( - socket, - "execute_reply", - reply_content, - parent, + kwargs = dict( + stream=socket, + msg_or_type="execute_reply", + content=reply_content, + parent=parent, metadata=metadata, ident=ident, ) + if self._supports_kernel_subshells(): + with self.LOCK: + reply_msg = self.session.send(**kwargs) + else: + reply_msg = self.session.send(**kwargs) + self.log.debug("%s", reply_msg) assert reply_msg is not None diff --git a/tests/test_subshells.py b/tests/test_subshells.py index 5312041d..87c5e23a 100644 --- a/tests/test_subshells.py +++ b/tests/test_subshells.py @@ -4,6 +4,7 @@ # Distributed under the terms of the Modified BSD License. from datetime import datetime, timedelta +import pytest from jupyter_client.blocking.client import BlockingKernelClient @@ -114,12 +115,14 @@ def test_thread_ids(): delete_subshell_helper(kc, subshell_id) -def test_run_concurrently(): +@pytest.mark.parametrize("times", [(0.1, 0.05), (0.05, 0.1), (0.05, 0.05)]) +def test_run_concurrently(times): with kernel() as kc: subshell_id = create_subshell_helper(kc)["subshell_id"] - # Prepare messages - times = (0.05, 0.1) # Sleep seconds. Test can fail with (0.05, 0.05) + # Prepare messages, times are sleep times in seconds. + # Identical times for both subshells is a harder test as preparing and sending + # the execute_reply messages may overlap. msgs = [] for id, sleep in zip((None, subshell_id), times): code = f"import time; time.sleep({sleep})" From 7dd34071e81dd47f04cf93de1119cc11b598e2f1 Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Thu, 30 May 2024 13:09:36 +0100 Subject: [PATCH 12/37] Process subshell control messages in SubshellCache --- ipykernel/kernelbase.py | 49 ++++++++++++----------------- ipykernel/subshell_cache.py | 63 +++++++++++++++++++++++++++++++++++-- ipykernel/thread.py | 12 +++---- tests/test_subshells.py | 3 +- 4 files changed, 88 insertions(+), 39 deletions(-) diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index 14b11094..f5bb1872 100644 --- a/ipykernel/kernelbase.py +++ b/ipykernel/kernelbase.py @@ -520,6 +520,8 @@ async def start(self, *, task_status: TaskStatus = TASK_STATUS_IGNORED) -> None: tg.start_soon(self.shell_main, None) self.shell_channel_thread.set_task(self.shell_channel_thread_main) + cache = self.shell_channel_thread.cache + self.shell_channel_thread.set_task(cache._TASK, self.shell_main) self.shell_channel_thread.start() else: if not self._is_test and self.shell_socket is not None: @@ -1050,24 +1052,21 @@ async def create_subshell_request(self, socket, ident, parent): self.log.error("KERNEL SUBSHELLS NOT SUPPORTED") return - subshell_id = str(uuid.uuid4()) - thread = SubshellThread(subshell_id) - self.shell_channel_thread.cache.create(subshell_id, thread) + # Check this is called in the control thread only (if it exists). - thread.set_task(self.shell_main, subshell_id) - thread.start() + s = self.shell_channel_thread.cache.control_recv_socket + await s.send_json( {"type": "create"}) + reply = await s.recv_json() - content = { - "status": "ok", - "subshell_id": subshell_id, - } - self.session.send(socket, "create_subshell_reply", content, parent, ident) + self.session.send(socket, "create_subshell_reply", reply, parent, ident) async def delete_subshell_request(self, socket, ident, parent): if not self._supports_kernel_subshells(): self.log.error("KERNEL SUBSHELLS NOT SUPPORTED") return + # Check this is called in the control thread only (if it exists). + try: content = parent["content"] subshell_id = content["subshell_id"] @@ -1076,32 +1075,24 @@ async def delete_subshell_request(self, socket, ident, parent): self.log.error("%s", parent) return - content: dict[str, t.Any] = {"status": "ok"} + s = self.shell_channel_thread.cache.control_recv_socket + await s.send_json({"type": "delete", "subshell_id": subshell_id}) + reply = await s.recv_json() - try: - # Should error here give traceback to the user? Probably not. - self.shell_channel_thread.cache.delete(subshell_id) - except KeyError as err: - import traceback - content = { - "status": "error", - "traceback": traceback.format_stack(), - "ename": str(type(err).__name__), - "evalue": str(err), - } - - self.session.send(socket, "delete_subshell_reply", content, parent, ident) + self.session.send(socket, "delete_subshell_reply", reply, parent, ident) async def list_subshell_request(self, socket, ident, parent): if not self._supports_kernel_subshells(): self.log.error("KERNEL SUBSHELLS NOT SUPPORTED") return - content = { - "status": "ok", - "subshell_id": self.shell_channel_thread.cache.list(), - } - self.session.send(socket, "list_subshell_reply", content, parent, ident) + # Check this is called in the control thread only (if it exists). + + s = self.shell_channel_thread.cache.control_recv_socket + await s.send_json({"type": "list"}) + reply = await s.recv_json() + + self.session.send(socket, "list_subshell_reply", reply, parent, ident) # --------------------------------------------------------------------------- # Engine methods (DEPRECATED) diff --git a/ipykernel/subshell_cache.py b/ipykernel/subshell_cache.py index 076cd749..88025194 100644 --- a/ipykernel/subshell_cache.py +++ b/ipykernel/subshell_cache.py @@ -2,8 +2,11 @@ # Copyright (c) IPython Development Team. # Distributed under the terms of the Modified BSD License. +from anyio import create_task_group from dataclasses import dataclass from threading import Lock, main_thread +import typing as t +import uuid import zmq from .subshell import SubshellThread @@ -46,13 +49,60 @@ def __init__(self, context: zmq.Context): self._parent_send_socket, self._parent_recv_socket = \ self._create_inproc_sockets(None) + + # Socket names are poor. Better to have this/cache end and other end. + self.control_send_socket, self.control_recv_socket = \ + self._create_inproc_sockets("control") + + + def _process_control_request(self, request, subshell_task): + try: + type = request["type"] + reply: dict[str, t.Any] = {"status": "ok"} + + if type == "create": + reply["subshell_id"] = self.create(subshell_task) + elif type == "delete": + subshell_id = request["subshell_id"] + self.delete(subshell_id) + elif type == "list": + reply["subshell_id"] = self.list() + else: + raise RuntimeError(f"Unrecognised message type {type}") + except BaseException as err: + # Not sure what information to return here. + reply = { + "status": "error", + "evalue": str(err), + } + return reply + + + async def _TASK(self, subshell_task): # messages from control channel via inproc pair sockets. + while True: + request = await self.control_send_socket.recv_json() + reply = self._process_control_request(request, subshell_task) + await self.control_send_socket.send_json(reply) + + + async def list_from_control(self): + async with create_task_group() as tg: + tg.start_soon(self._TASK) + + def close(self): for socket in (self._parent_send_socket, self._parent_recv_socket): if socket and not socket.closed: socket.close() + for socket in (self.control_send_socket, self.control_recv_socket): + if socket and not socket.closed: + socket.close() + self._parent_recv_socket = None self._parent_send_socket = None + self.control_send_socket = None + self.control_recv_socket = None with self._lock: while True: @@ -62,15 +112,24 @@ def close(self): break self._stop_subshell(subshell) - def create(self, subshell_id: str, thread: SubshellThread) -> None: + def create(self, subshell_task) -> str: + # Create new subshell thread and start it. + # check if subshell_id already exists... - # assume it doesn't + # assume it doesn't so far. + + subshell_id = str(uuid.uuid4()) + thread = SubshellThread(subshell_id) with self._lock: assert subshell_id not in self._cache send_socket, recv_socket = self._create_inproc_sockets(subshell_id) self._cache[subshell_id] = Subshell(thread, send_socket, recv_socket) + thread.set_task(subshell_task, subshell_id) + thread.start() + return subshell_id + def delete(self, subshell_id: str) -> None: """Raises key error if subshell_id not in cache""" with self._lock: diff --git a/ipykernel/thread.py b/ipykernel/thread.py index 92448b5a..d4f68745 100644 --- a/ipykernel/thread.py +++ b/ipykernel/thread.py @@ -13,12 +13,12 @@ def __init__(self, name, **kwargs): self.pydev_do_not_trace = True self.is_pydev_daemon_thread = True self.__stop = Event() - self._task = None - self._task_args = () + self._tasks = [] + self._task_args = [] def set_task(self, task, *args): - self._task = task - self._task_args = args + self._tasks.append(task) + self._task_args.append(args) def run(self): """Run the thread.""" @@ -26,8 +26,8 @@ def run(self): async def _main(self): async with create_task_group() as tg: - if self._task is not None: - tg.start_soon(self._task, *self._task_args) + for task, args in zip(self._tasks, self._task_args): + tg.start_soon(task, *args) await to_thread.run_sync(self.__stop.wait) tg.cancel_scope.cancel() diff --git a/tests/test_subshells.py b/tests/test_subshells.py index 87c5e23a..c1ad6eda 100644 --- a/tests/test_subshells.py +++ b/tests/test_subshells.py @@ -85,8 +85,7 @@ def test_delete_non_existent(): with kernel() as kc: reply = delete_subshell_helper(kc, "unknown_subshell_id") assert reply["status"] == "error" - for key in ("ename", "evalue", "traceback"): - assert key in reply + assert "evalue" in reply def test_thread_counts(): From 27ead0275ca4f7ff4d08d63e8dd0fcaabe23324b Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Wed, 5 Jun 2024 15:05:52 +0100 Subject: [PATCH 13/37] Send shell reply messages via shell channel thread --- ipykernel/kernelapp.py | 2 +- ipykernel/kernelbase.py | 25 ++++++++++++------------- ipykernel/shellchannel.py | 11 +++++++++-- ipykernel/subshell_cache.py | 33 ++++++++++++++++++++++++--------- ipykernel/thread.py | 3 ++- ipykernel/zmqshell.py | 2 +- tests/test_subshells.py | 26 +++++++++++++++++++++++--- 7 files changed, 72 insertions(+), 30 deletions(-) diff --git a/ipykernel/kernelapp.py b/ipykernel/kernelapp.py index 642fb08b..9ca2ef44 100644 --- a/ipykernel/kernelapp.py +++ b/ipykernel/kernelapp.py @@ -370,7 +370,7 @@ def init_control(self, context): self.control_socket.router_handover = 1 self.control_thread = ControlThread(daemon=True) - self.shell_channel_thread = ShellChannelThread(context, daemon=True) + self.shell_channel_thread = ShellChannelThread(context, self.shell_socket, daemon=True) def init_iopub(self, context): """Initialize the iopub channel.""" diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index f5bb1872..857f5ad4 100644 --- a/ipykernel/kernelbase.py +++ b/ipykernel/kernelbase.py @@ -408,6 +408,7 @@ async def process_shell(self, recv_socket=None): raise e async def process_shell_message(self, msg=None, recv_socket=None): + #### Do not like the name recv_socket if not recv_socket: recv_socket = self.shell_socket @@ -463,7 +464,7 @@ async def process_shell_message(self, msg=None, recv_socket=None): except Exception: self.log.debug("Unable to signal in pre_handler_hook:", exc_info=True) try: - result = handler(self.shell_socket, idents, msg) + result = handler(recv_socket, idents, msg) if inspect.isawaitable(result): await result except Exception: @@ -505,7 +506,7 @@ async def start(self, *, task_status: TaskStatus = TASK_STATUS_IGNORED) -> None: self.control_stop = threading.Event() if not self._is_test and self.control_socket is not None: if self.control_thread: - self.control_thread.set_task(self.control_main) + self.control_thread.add_task(self.control_main) self.control_thread.start() else: tg.start_soon(self.control_main) @@ -519,9 +520,13 @@ async def start(self, *, task_status: TaskStatus = TASK_STATUS_IGNORED) -> None: # Each subshell thread listens on inproc socket for messages from shell channel thread tg.start_soon(self.shell_main, None) - self.shell_channel_thread.set_task(self.shell_channel_thread_main) + # Needs tidying. cache = self.shell_channel_thread.cache - self.shell_channel_thread.set_task(cache._TASK, self.shell_main) + parent_send_socket = cache._parent_send_socket + + self.shell_channel_thread.add_task(self.shell_channel_thread_main) + self.shell_channel_thread.add_task(cache._TASK, self.shell_main) + self.shell_channel_thread.add_task(cache._from_subshell_task, parent_send_socket) self.shell_channel_thread.start() else: if not self._is_test and self.shell_socket is not None: @@ -734,21 +739,15 @@ async def execute_request(self, socket, ident, parent): reply_content = json_clean(reply_content) metadata = self.finish_metadata(parent, metadata, reply_content) - kwargs = dict( - stream=socket, - msg_or_type="execute_reply", + reply_msg = self.session.send( + socket, + "execute_reply", content=reply_content, parent=parent, metadata=metadata, ident=ident, ) - if self._supports_kernel_subshells(): - with self.LOCK: - reply_msg = self.session.send(**kwargs) - else: - reply_msg = self.session.send(**kwargs) - self.log.debug("%s", reply_msg) assert reply_msg is not None diff --git a/ipykernel/shellchannel.py b/ipykernel/shellchannel.py index 9b13ac79..fc6a2270 100644 --- a/ipykernel/shellchannel.py +++ b/ipykernel/shellchannel.py @@ -12,16 +12,23 @@ class ShellChannelThread(BaseThread): Communicates with shell execute threads via pairs of ZMQ inproc sockets. """ - def __init__(self, context, **kwargs): + def __init__(self, context, shell_socket, **kwargs): """Initialize the thread.""" super().__init__(name=SHELL_CHANNEL_THREAD_NAME, **kwargs) self._cache: SubshellCache | None = None self._context = context + self._shell_socket = shell_socket @property def cache(self): + + + # Would like to not have to have this lazy construction. + # Without it broke some tests. + + if self._cache is None: - self._cache = SubshellCache(self._context) + self._cache = SubshellCache(self._context, self._shell_socket) return self._cache def run(self): diff --git a/ipykernel/subshell_cache.py b/ipykernel/subshell_cache.py index 88025194..f3cf6242 100644 --- a/ipykernel/subshell_cache.py +++ b/ipykernel/subshell_cache.py @@ -19,6 +19,9 @@ class Subshell: recv_socket: zmq.Socket + +# SubshellManager perhaps??????? + class SubshellCache: """A cache for subshell information. @@ -40,8 +43,13 @@ class SubshellCache: Read-only access to the parent subshell sockets never needs to be wrapped in a mutex as there is only one pair over the whole lifetime of this object. """ - def __init__(self, context: zmq.Context): + def __init__(self, context: zmq.Context, shell_socket): + + + # assert this is only called from the main thread... + self._context: zmq.Context = context + self._shell_socket = shell_socket self._cache: dict[str, Subshell] = {} self._lock: Lock = Lock() @@ -49,7 +57,6 @@ def __init__(self, context: zmq.Context): self._parent_send_socket, self._parent_recv_socket = \ self._create_inproc_sockets(None) - # Socket names are poor. Better to have this/cache end and other end. self.control_send_socket, self.control_recv_socket = \ self._create_inproc_sockets("control") @@ -61,12 +68,12 @@ def _process_control_request(self, request, subshell_task): reply: dict[str, t.Any] = {"status": "ok"} if type == "create": - reply["subshell_id"] = self.create(subshell_task) + reply["subshell_id"] = self.create_subshell(subshell_task) elif type == "delete": subshell_id = request["subshell_id"] - self.delete(subshell_id) + self.delete_subshell(subshell_id) elif type == "list": - reply["subshell_id"] = self.list() + reply["subshell_id"] = self.list_subshell() else: raise RuntimeError(f"Unrecognised message type {type}") except BaseException as err: @@ -85,6 +92,12 @@ async def _TASK(self, subshell_task): # messages from control channel via inpr await self.control_send_socket.send_json(reply) + async def _from_subshell_task(self, send_socket): + while True: + msg = await send_socket.recv_multipart(copy=False) + with self._lock: + self._shell_socket.send_multipart(msg) + async def list_from_control(self): async with create_task_group() as tg: tg.start_soon(self._TASK) @@ -112,7 +125,7 @@ def close(self): break self._stop_subshell(subshell) - def create(self, subshell_task) -> str: + def create_subshell(self, subshell_task) -> str: # Create new subshell thread and start it. # check if subshell_id already exists... @@ -126,11 +139,13 @@ def create(self, subshell_task) -> str: send_socket, recv_socket = self._create_inproc_sockets(subshell_id) self._cache[subshell_id] = Subshell(thread, send_socket, recv_socket) - thread.set_task(subshell_task, subshell_id) + thread.add_task(subshell_task, subshell_id) + thread.add_task(self._from_subshell_task, send_socket) thread.start() + return subshell_id - def delete(self, subshell_id: str) -> None: + def delete_subshell(self, subshell_id: str) -> None: """Raises key error if subshell_id not in cache""" with self._lock: subshell = self._cache.pop(subshell_id) @@ -151,7 +166,7 @@ def get_send_socket(self, subshell_id: str | None): with self._lock: return self._cache[subshell_id].send_socket - def list(self) -> list[str]: + def list_subshell(self) -> list[str]: with self._lock: return list(self._cache) diff --git a/ipykernel/thread.py b/ipykernel/thread.py index d4f68745..642b1cd3 100644 --- a/ipykernel/thread.py +++ b/ipykernel/thread.py @@ -16,7 +16,8 @@ def __init__(self, name, **kwargs): self._tasks = [] self._task_args = [] - def set_task(self, task, *args): + def add_task(self, task, *args): + # May only add tasks before the thread is started. self._tasks.append(task) self._task_args.append(args) diff --git a/ipykernel/zmqshell.py b/ipykernel/zmqshell.py index 51efda7b..4de13708 100644 --- a/ipykernel/zmqshell.py +++ b/ipykernel/zmqshell.py @@ -461,7 +461,7 @@ def subshell(self, arg_s): subshell_id = kernel.shell_channel_thread.cache.subshell_id_from_thread_id(thread_id) except RuntimeError: subshell_id = "unknown" - subshell_id_list = kernel.shell_channel_thread.cache.list() + subshell_id_list = kernel.shell_channel_thread.cache.list_subshell() print(f"subshell id: {subshell_id}") print(f"thread id: {thread_id}") diff --git a/tests/test_subshells.py b/tests/test_subshells.py index c1ad6eda..394204bb 100644 --- a/tests/test_subshells.py +++ b/tests/test_subshells.py @@ -4,11 +4,13 @@ # Distributed under the terms of the Modified BSD License. from datetime import datetime, timedelta +import platform import pytest +import time from jupyter_client.blocking.client import BlockingKernelClient -from .utils import TIMEOUT, get_reply, get_replies, kernel +from .utils import TIMEOUT, get_reply, get_replies, kernel, new_kernel, wait_for_idle # Helpers @@ -114,11 +116,11 @@ def test_thread_ids(): delete_subshell_helper(kc, subshell_id) -@pytest.mark.parametrize("times", [(0.1, 0.05), (0.05, 0.1), (0.05, 0.05)]) -def test_run_concurrently(times): +def test_run_concurrently(): with kernel() as kc: subshell_id = create_subshell_helper(kc)["subshell_id"] + times = (0.05, 0.05) # Prepare messages, times are sleep times in seconds. # Identical times for both subshells is a harder test as preparing and sending # the execute_reply messages may overlap. @@ -189,3 +191,21 @@ def test_create_while_execute(): delete_subshell_helper(kc, subshell_id) assert control_date < shell_date + + +@pytest.mark.skipif( + platform.python_implementation() == "PyPy", + reason="does not work on PyPy", +) +def test_shutdown_with_subshell(): + with new_kernel() as kc: + km = kc.parent + subshell_id = create_subshell_helper(kc)["subshell_id"] + assert list_subshell_helper(kc)["subshell_id"] == [subshell_id] + kc.shutdown() + for _ in range(100): # 10 s timeout + if km.is_alive(): + time.sleep(0.1) + else: + break + assert not km.is_alive() From 1e79f4cc031acabdce80b7a60e6c0dd375e648cc Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Wed, 12 Jun 2024 16:32:24 +0100 Subject: [PATCH 14/37] Poll subshell inproc pair sockets in SubshellManager --- ipykernel/kernelbase.py | 107 +++++++++------- ipykernel/shellchannel.py | 26 ++-- ipykernel/subshell.py | 27 +++- ipykernel/subshell_cache.py | 211 ------------------------------ ipykernel/subshell_manager.py | 235 ++++++++++++++++++++++++++++++++++ ipykernel/zmqshell.py | 11 +- tests/test_ipkernel_direct.py | 5 +- tests/test_kernel_direct.py | 5 +- tests/test_message_spec.py | 1 + tests/test_subshells.py | 30 +++-- 10 files changed, 369 insertions(+), 289 deletions(-) delete mode 100644 ipykernel/subshell_cache.py create mode 100644 ipykernel/subshell_manager.py diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index 857f5ad4..03cfb904 100644 --- a/ipykernel/kernelbase.py +++ b/ipykernel/kernelbase.py @@ -19,7 +19,6 @@ from signal import SIGINT, SIGTERM, Signals from .control import CONTROL_THREAD_NAME -from .subshell import SubshellThread if sys.platform != "win32": from signal import SIGKILL @@ -272,6 +271,8 @@ async def process_control_message(self, msg=None): """dispatch control requests""" assert self.control_socket is not None assert self.session is not 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) @@ -361,9 +362,16 @@ def _message_counter_default(self): return itertools.count() async def shell_channel_thread_main(self): + """Main loop for shell channel thread. + + Listen for incoming messages on kernel shell_socket. For each message + received, extract the subshell_id from the message header and forward the + message to the correct subshell via ZMQ inproc pair socket. + """ assert self.shell_socket is not None assert self.session is not None - assert self.shell_channel_thread + assert self.shell_channel_thread is not None + assert threading.current_thread() == self.shell_channel_thread try: while True: @@ -377,46 +385,63 @@ async def shell_channel_thread_main(self): msg2 = self.session.deserialize(msg2, content=False, copy=copy) subshell_id = msg2["header"].get("subshell_id") - send_socket = self.shell_channel_thread.cache.get_send_socket(subshell_id) - assert send_socket is not None - send_socket.send_multipart(msg, copy=False) + # Find inproc pair socket to use to send message to correct subshell. + socket = self.shell_channel_thread.manager.get_shell_channel_socket(subshell_id) + assert socket is not None + socket.send_multipart(msg, copy=False) except BaseException as e: if self.shell_stop.is_set(): return raise e async def shell_main(self, subshell_id: str | None): + """Main loop for a single subshell.""" if self._supports_kernel_subshells(): - recv_socket = self.shell_channel_thread.cache.get_recv_socket(subshell_id) + if subshell_id is None: + assert threading.current_thread() == threading.main_thread() + else: + assert threading.current_thread() not in ( + self.shell_channel_thread, + threading.main_thread(), + ) + # Inproc pair socket that this subshell uses to talk to shell channel thread. + socket = self.shell_channel_thread.manager.get_other_socket(subshell_id) else: - recv_socket = self.shell_socket + assert subshell_id is None + assert threading.current_thread() == threading.main_thread() + socket = self.shell_socket async with create_task_group() as tg: - tg.start_soon(self.process_shell, recv_socket) + tg.start_soon(self.process_shell, socket) if subshell_id is None: # Main subshell. await to_thread.run_sync(self.shell_stop.wait) tg.cancel_scope.cancel() - async def process_shell(self, recv_socket=None): + async def process_shell(self, socket=None): try: while True: - await self.process_shell_message(recv_socket=recv_socket) + await self.process_shell_message(socket=socket) except BaseException as e: if self.shell_stop.is_set(): return raise e - async def process_shell_message(self, msg=None, recv_socket=None): - #### Do not like the name recv_socket - if not recv_socket: - recv_socket = self.shell_socket + async def process_shell_message(self, msg=None, socket=None): + if not socket: + socket = self.shell_socket - assert recv_socket is not None assert self.session is not None + if self._supports_kernel_subshells(): + assert threading.current_thread() not in ( + self.control_thread, + self.shell_channel_thread, + ) + else: + assert threading.current_thread() == threading.main_thread() - no_msg = msg is None if self._is_test else not await recv_socket.poll(0) - msg = msg or await recv_socket.recv_multipart(copy=False) + no_msg = msg is None if self._is_test else not await socket.poll(0) + msg = msg or await socket.recv_multipart(copy=False) received_time = time.monotonic() copy = not isinstance(msg[0], zmq.Message) @@ -464,7 +489,7 @@ async def process_shell_message(self, msg=None, recv_socket=None): except Exception: self.log.debug("Unable to signal in pre_handler_hook:", exc_info=True) try: - result = handler(recv_socket, idents, msg) + result = handler(socket, idents, msg) if inspect.isawaitable(result): await result except Exception: @@ -501,7 +526,6 @@ def post_handler_hook(self): async def start(self, *, task_status: TaskStatus = TASK_STATUS_IGNORED) -> None: """Process messages on shell and control channels""" - async with create_task_group() as tg: self.control_stop = threading.Event() if not self._is_test and self.control_socket is not None: @@ -517,16 +541,13 @@ async def start(self, *, task_status: TaskStatus = TASK_STATUS_IGNORED) -> None: self.shell_stop = threading.Event() if self.shell_channel_thread: - # Each subshell thread listens on inproc socket for messages from shell channel thread tg.start_soon(self.shell_main, None) - # Needs tidying. - cache = self.shell_channel_thread.cache - parent_send_socket = cache._parent_send_socket - + # Assign tasks to and start shell channel thread. + manager = self.shell_channel_thread.manager self.shell_channel_thread.add_task(self.shell_channel_thread_main) - self.shell_channel_thread.add_task(cache._TASK, self.shell_main) - self.shell_channel_thread.add_task(cache._from_subshell_task, parent_send_socket) + self.shell_channel_thread.add_task(manager.listen_from_control, self.shell_main) + self.shell_channel_thread.add_task(manager.listen_from_subshells) self.shell_channel_thread.start() else: if not self._is_test and self.shell_socket is not None: @@ -1048,14 +1069,14 @@ async def do_debug_request(self, msg): async def create_subshell_request(self, socket, ident, parent): if not self._supports_kernel_subshells(): - self.log.error("KERNEL SUBSHELLS NOT SUPPORTED") + self.log.error("Subshells are not supported by this kernel") return - # Check this is called in the control thread only (if it exists). - - s = self.shell_channel_thread.cache.control_recv_socket - await s.send_json( {"type": "create"}) - reply = await s.recv_json() + # This should only be called in the control thread if it exists. + # Request is passed to shell channel thread to process. + other_socket = self.shell_channel_thread.manager.get_control_other_socket() + await other_socket.send_json({"type": "create"}) + reply = await other_socket.recv_json() self.session.send(socket, "create_subshell_reply", reply, parent, ident) @@ -1064,8 +1085,6 @@ async def delete_subshell_request(self, socket, ident, parent): self.log.error("KERNEL SUBSHELLS NOT SUPPORTED") return - # Check this is called in the control thread only (if it exists). - try: content = parent["content"] subshell_id = content["subshell_id"] @@ -1074,22 +1093,24 @@ async def delete_subshell_request(self, socket, ident, parent): self.log.error("%s", parent) return - s = self.shell_channel_thread.cache.control_recv_socket - await s.send_json({"type": "delete", "subshell_id": subshell_id}) - reply = await s.recv_json() + # This should only be called in the control thread if it exists. + # Request is passed to shell channel thread to process. + other_socket = self.shell_channel_thread.manager.get_control_other_socket() + await other_socket.send_json({"type": "delete", "subshell_id": subshell_id}) + reply = await other_socket.recv_json() self.session.send(socket, "delete_subshell_reply", reply, parent, ident) async def list_subshell_request(self, socket, ident, parent): if not self._supports_kernel_subshells(): - self.log.error("KERNEL SUBSHELLS NOT SUPPORTED") + self.log.error("Subshells are not supported by this kernel") return - # Check this is called in the control thread only (if it exists). - - s = self.shell_channel_thread.cache.control_recv_socket - await s.send_json({"type": "list"}) - reply = await s.recv_json() + # This should only be called in the control thread if it exists. + # Request is passed to shell channel thread to process. + other_socket = self.shell_channel_thread.manager.get_control_other_socket() + await other_socket.send_json({"type": "list"}) + reply = await other_socket.recv_json() self.session.send(socket, "list_subshell_reply", reply, parent, ident) diff --git a/ipykernel/shellchannel.py b/ipykernel/shellchannel.py index fc6a2270..7341f9e4 100644 --- a/ipykernel/shellchannel.py +++ b/ipykernel/shellchannel.py @@ -1,6 +1,5 @@ """A thread for a shell channel.""" - -from .subshell_cache import SubshellCache +from .subshell_manager import SubshellManager from .thread import BaseThread SHELL_CHANNEL_THREAD_NAME = "Shell channel" @@ -9,32 +8,27 @@ class ShellChannelThread(BaseThread): """A thread for a shell channel. - Communicates with shell execute threads via pairs of ZMQ inproc sockets. + Communicates with shell/subshell threads via pairs of ZMQ inproc sockets. """ def __init__(self, context, shell_socket, **kwargs): """Initialize the thread.""" super().__init__(name=SHELL_CHANNEL_THREAD_NAME, **kwargs) - self._cache: SubshellCache | None = None + self._manager: SubshellManager | None = None self._context = context self._shell_socket = shell_socket @property - def cache(self): - - - # Would like to not have to have this lazy construction. - # Without it broke some tests. - - - if self._cache is None: - self._cache = SubshellCache(self._context, self._shell_socket) - return self._cache + def manager(self): + # Lazy initialisation. + if self._manager is None: + self._manager = SubshellManager(self._context, self._shell_socket) + return self._manager def run(self): """Run the thread.""" self.name = SHELL_CHANNEL_THREAD_NAME super().run() - if self._cache: - self._cache.close() + if self._manager: + self._manager.close() diff --git a/ipykernel/subshell.py b/ipykernel/subshell.py index b1f6b3f1..5a23dd08 100644 --- a/ipykernel/subshell.py +++ b/ipykernel/subshell.py @@ -1,12 +1,35 @@ """A thread for a subshell.""" +from threading import current_thread + +import zmq.asyncio + from .thread import BaseThread class SubshellThread(BaseThread): - """A thread for a subshell. - """ + """A thread for a subshell.""" def __init__(self, subshell_id: str, **kwargs): """Initialize the thread.""" super().__init__(name=f"subshell-{subshell_id}", **kwargs) + + # Inproc PAIR socket, for communication with shell channel thread. + self._pair_socket: zmq.asyncio.Socket | None = None + + async def create_pair_socket(self, context: zmq.asyncio.Context, address: str): + """Create inproc PAIR socket, for communication with shell channel thread. + + Should be called from this thread, so usually via add_task before the + thread is started. + """ + assert current_thread() == self + self._pair_socket = context.socket(zmq.PAIR) + self._pair_socket.connect(address) + + def run(self): + super().run() + + if self._pair_socket is not None: + self._pair_socket.close() + self._pair_socket = None diff --git a/ipykernel/subshell_cache.py b/ipykernel/subshell_cache.py deleted file mode 100644 index f3cf6242..00000000 --- a/ipykernel/subshell_cache.py +++ /dev/null @@ -1,211 +0,0 @@ -"""A cache for subshell information.""" -# Copyright (c) IPython Development Team. -# Distributed under the terms of the Modified BSD License. - -from anyio import create_task_group -from dataclasses import dataclass -from threading import Lock, main_thread -import typing as t -import uuid -import zmq - -from .subshell import SubshellThread - - -@dataclass -class Subshell: - thread: SubshellThread - send_socket: zmq.Socket - recv_socket: zmq.Socket - - - -# SubshellManager perhaps??????? - -class SubshellCache: - """A cache for subshell information. - - Care needed with threadsafe access here. After the cache is created, all write - access is performed by the control thread so there is only ever one write access at - any one time. Reading of cache information is performed by a number of different - threads: - - 1) Receive inproc socket for a child subshell is passed to the subshell task (by - the control thread) when it is created. - 2) Send inproc socket is looked up for every message received by the shell channel - thread so that the message is sent to the correct subshell thread. - 3) Control thread reads shell_ids for list_subshell_request message. - - Python dictionary reads and writes are atomic and therefore threadsafe because of - the GIL in conventional CPython. But wise to use a mutex to support non-GIL - python. - - Read-only access to the parent subshell sockets never needs to be wrapped in a - mutex as there is only one pair over the whole lifetime of this object. - """ - def __init__(self, context: zmq.Context, shell_socket): - - - # assert this is only called from the main thread... - - self._context: zmq.Context = context - self._shell_socket = shell_socket - self._cache: dict[str, Subshell] = {} - self._lock: Lock = Lock() - - # Parent subshell sockets. - self._parent_send_socket, self._parent_recv_socket = \ - self._create_inproc_sockets(None) - - # Socket names are poor. Better to have this/cache end and other end. - self.control_send_socket, self.control_recv_socket = \ - self._create_inproc_sockets("control") - - - def _process_control_request(self, request, subshell_task): - try: - type = request["type"] - reply: dict[str, t.Any] = {"status": "ok"} - - if type == "create": - reply["subshell_id"] = self.create_subshell(subshell_task) - elif type == "delete": - subshell_id = request["subshell_id"] - self.delete_subshell(subshell_id) - elif type == "list": - reply["subshell_id"] = self.list_subshell() - else: - raise RuntimeError(f"Unrecognised message type {type}") - except BaseException as err: - # Not sure what information to return here. - reply = { - "status": "error", - "evalue": str(err), - } - return reply - - - async def _TASK(self, subshell_task): # messages from control channel via inproc pair sockets. - while True: - request = await self.control_send_socket.recv_json() - reply = self._process_control_request(request, subshell_task) - await self.control_send_socket.send_json(reply) - - - async def _from_subshell_task(self, send_socket): - while True: - msg = await send_socket.recv_multipart(copy=False) - with self._lock: - self._shell_socket.send_multipart(msg) - - async def list_from_control(self): - async with create_task_group() as tg: - tg.start_soon(self._TASK) - - - def close(self): - for socket in (self._parent_send_socket, self._parent_recv_socket): - if socket and not socket.closed: - socket.close() - - for socket in (self.control_send_socket, self.control_recv_socket): - if socket and not socket.closed: - socket.close() - - self._parent_recv_socket = None - self._parent_send_socket = None - self.control_send_socket = None - self.control_recv_socket = None - - with self._lock: - while True: - try: - _, subshell = self._cache.popitem() - except KeyError: - break - self._stop_subshell(subshell) - - def create_subshell(self, subshell_task) -> str: - # Create new subshell thread and start it. - - # check if subshell_id already exists... - # assume it doesn't so far. - - subshell_id = str(uuid.uuid4()) - thread = SubshellThread(subshell_id) - - with self._lock: - assert subshell_id not in self._cache - send_socket, recv_socket = self._create_inproc_sockets(subshell_id) - self._cache[subshell_id] = Subshell(thread, send_socket, recv_socket) - - thread.add_task(subshell_task, subshell_id) - thread.add_task(self._from_subshell_task, send_socket) - thread.start() - - return subshell_id - - def delete_subshell(self, subshell_id: str) -> None: - """Raises key error if subshell_id not in cache""" - with self._lock: - subshell = self._cache.pop(subshell_id) - - self._stop_subshell(subshell) - - def get_recv_socket(self, subshell_id: str | None): - if subshell_id is None: - return self._parent_recv_socket - else: - with self._lock: - return self._cache[subshell_id].recv_socket - - def get_send_socket(self, subshell_id: str | None): - if subshell_id is None: - return self._parent_send_socket - else: - with self._lock: - return self._cache[subshell_id].send_socket - - def list_subshell(self) -> list[str]: - with self._lock: - return list(self._cache) - - def subshell_id_from_thread_id(self, thread_id) -> str | None: - """Return subshell_id of the specified thread_id. - - Raises RuntimeError if thread_id is not the main shell or a subshell. - """ - with self._lock: - if thread_id == main_thread().ident: - return None - for id, subshell in self._cache.items(): - if subshell.thread.ident == thread_id: - return id - msg = f"Thread id '{thread_id} does not correspond to a subshell of this kernel" - raise RuntimeError(msg) - - def _create_inproc_sockets(self, subshell_id: str | None): - """Create a pair of inproc sockets to communicate with a subshell. - """ - name = f"shell-{subshell_id}" if subshell_id else "shell" - address = f"inproc://{name}" - - # Socket used in subshell thread to receive messages. - recv_socket = self._context.socket(zmq.PAIR) - recv_socket.bind(address) - - # Socket used in shell channel thread to send messages. - send_socket = self._context.socket(zmq.PAIR) - send_socket.connect(address) - - return send_socket, recv_socket - - def _stop_subshell(self, subshell: Subshell): - thread = subshell.thread - if thread and thread.is_alive(): - thread.stop() - thread.join() - - for socket in (subshell.send_socket, subshell.recv_socket): - if socket and not socket.closed: - socket.close() diff --git a/ipykernel/subshell_manager.py b/ipykernel/subshell_manager.py new file mode 100644 index 00000000..a7e9629b --- /dev/null +++ b/ipykernel/subshell_manager.py @@ -0,0 +1,235 @@ +"""Manager of subshells in a kernel.""" + +# Copyright (c) IPython Development Team. +# Distributed under the terms of the Modified BSD License. + +import typing as t +import uuid +from dataclasses import dataclass +from threading import Lock, current_thread, main_thread + +import zmq +import zmq.asyncio +from anyio import sleep + +from .subshell import SubshellThread + + +@dataclass +class Subshell: + thread: SubshellThread + shell_channel_socket: zmq.asyncio.Socket + + +class SubshellManager: + """A manager of subshells. + + Controls the lifetimes of subshell threads and their associated ZMQ sockets. + Runs mostly in the shell channel thread. + + Care needed with threadsafe access here. All write access to the cache occurs in + the shell channel thread so there is only ever one write access at any one time. + Reading of cache information can be performed by other threads, so all reads are + protected by a lock so that they are atomic. + """ + + def __init__(self, context: zmq.asyncio.Context, shell_socket): + assert current_thread() == main_thread() + + self._context: zmq.asyncio.Context = context + self._shell_socket = shell_socket + self._cache: dict[str, Subshell] = {} + self._lock: Lock = Lock() + + # Inproc pair sockets for control channel and main shell (parent subshell). + # Each inproc pair has a "shell_channel" socket used in the shell channel + # thread, and an "other" socket used in the other thread. + self._control_shell_channel_socket = self._create_inproc_pair_socket("control", True) + self._control_other_socket = self._create_inproc_pair_socket("control", False) + self._parent_shell_channel_socket = self._create_inproc_pair_socket(None, True) + self._parent_other_socket = self._create_inproc_pair_socket(None, False) + + self._poller = zmq.Poller() + self._poller.register(self._parent_shell_channel_socket, zmq.POLLIN) + + def close(self): + """Stop all subshells and close all resources.""" + assert current_thread().name == "Shell channel" + + for socket in ( + self._control_shell_channel_socket, + self._control_other_socket, + self._parent_shell_channel_socket, + self._parent_other_socket, + ): + if socket is not None: + socket.close() + + with self._lock: + while True: + try: + _, subshell = self._cache.popitem() + except KeyError: + break + self._stop_subshell(subshell) + + def get_control_other_socket(self): + return self._control_other_socket + + def get_other_socket(self, subshell_id: str | None): + """Return the other inproc pair socket for a subshell. + + This socket is accessed from the subshell thread. + """ + if subshell_id is None: + return self._parent_other_socket + with self._lock: + return self._cache[subshell_id].thread._pair_socket + + def get_shell_channel_socket(self, subshell_id: str | None): + """Return the shell channel inproc pair socket for a subshell. + + This socket is accessed from the shell channel thread. + """ + if subshell_id is None: + return self._parent_shell_channel_socket + with self._lock: + return self._cache[subshell_id].shell_channel_socket + + def list_subshell(self) -> list[str]: + """Return list of current subshell ids. + + Can be called by any subshell using %subshell magic. + """ + with self._lock: + return list(self._cache) + + async def listen_from_control(self, subshell_task): + """Listen for messages on the control inproc socket, handle those messages and + return replies on the same socket. Runs in the shell channel thread. + """ + assert current_thread().name == "Shell channel" + + socket = self._control_shell_channel_socket + while True: + request = await socket.recv_json() + reply = await self._process_control_request(request, subshell_task) + await socket.send_json(reply) + + async def listen_from_subshells(self): + """Listen for reply messages on inproc sockets of all subshells and resend + those messages to the client via the shell_socket. + + Runs in the shell channel thread. + """ + assert current_thread().name == "Shell channel" + + while True: + for socket, _ in self._poller.poll(0): + msg = await socket.recv_multipart(copy=False) + self._shell_socket.send_multipart(msg) + + # Yield to other tasks. + await sleep(0) + + def subshell_id_from_thread_id(self, thread_id) -> str | None: + """Return subshell_id of the specified thread_id. + + Raises RuntimeError if thread_id is not the main shell or a subshell. + + Only used by %subshell magic so does not have to be fast/cached. + """ + with self._lock: + if thread_id == main_thread().ident: + return None + for id, subshell in self._cache.items(): + if subshell.thread.ident == thread_id: + return id + msg = f"Thread id '{thread_id} does not correspond to a subshell of this kernel" + raise RuntimeError(msg) + + def _create_inproc_pair_socket(self, name: str | None, shell_channel_end: bool): + """Create and return a single ZMQ inproc pair socket.""" + address = self._get_inproc_socket_address(name) + socket = self._context.socket(zmq.PAIR) + if shell_channel_end: + socket.bind(address) + else: + socket.connect(address) + return socket + + async def _create_subshell(self, subshell_task) -> str: + """Create and start a new subshell thread.""" + assert current_thread().name == "Shell channel" + + subshell_id = str(uuid.uuid4()) + thread = SubshellThread(subshell_id) + + with self._lock: + assert subshell_id not in self._cache + shell_channel_socket = self._create_inproc_pair_socket(subshell_id, True) + self._cache[subshell_id] = Subshell(thread, shell_channel_socket) + + address = self._get_inproc_socket_address(subshell_id) + thread.add_task(thread.create_pair_socket, self._context, address) + thread.add_task(subshell_task, subshell_id) + thread.start() + + self._poller.register(shell_channel_socket, zmq.POLLIN) + + return subshell_id + + def _delete_subshell(self, subshell_id: str) -> None: + """Delete subshell identified by subshell_id. + + Raises key error if subshell_id not in cache. + """ + assert current_thread().name == "Shell channel" + + with self._lock: + subshell = self._cache.pop(subshell_id) + + self._stop_subshell(subshell) + + def _get_inproc_socket_address(self, name: str | None): + full_name = f"subshell-{name}" if name else "subshell" + return f"inproc://{full_name}" + + async def _process_control_request(self, request, subshell_task): + """Process a control request message received on the control inproc + socket and return the reply. Runs in the shell channel thread. + """ + assert current_thread().name == "Shell channel" + + try: + type = request["type"] + reply: dict[str, t.Any] = {"status": "ok"} + + if type == "create": + reply["subshell_id"] = await self._create_subshell(subshell_task) + elif type == "delete": + subshell_id = request["subshell_id"] + self._delete_subshell(subshell_id) + elif type == "list": + reply["subshell_id"] = self.list_subshell() + else: + msg = f"Unrecognised message type {type}" + raise RuntimeError(msg) + except BaseException as err: + reply = { + "status": "error", + "evalue": str(err), + } + return reply + + def _stop_subshell(self, subshell: Subshell): + """Stop a subshell thread and close all of its resources.""" + assert current_thread().name == "Shell channel" + + thread = subshell.thread + if thread and thread.is_alive(): + thread.stop() + thread.join() + + self._poller.unregister(subshell.shell_channel_socket) + subshell.shell_channel_socket.close() diff --git a/ipykernel/zmqshell.py b/ipykernel/zmqshell.py index 4de13708..5ab7b188 100644 --- a/ipykernel/zmqshell.py +++ b/ipykernel/zmqshell.py @@ -450,18 +450,21 @@ def subshell(self, arg_s): app = IPKernelApp.instance() kernel = app.kernel - if not (hasattr(kernel, "_supports_kernel_subshells") and - kernel._supports_kernel_subshells()): + if not ( + hasattr(kernel, "_supports_kernel_subshells") and kernel._supports_kernel_subshells() + ): print("Kernel does not support subshells") return import threading + thread_id = threading.current_thread().ident + manager = kernel.shell_channel_thread.manager try: - subshell_id = kernel.shell_channel_thread.cache.subshell_id_from_thread_id(thread_id) + subshell_id = manager.subshell_id_from_thread_id(thread_id) except RuntimeError: subshell_id = "unknown" - subshell_id_list = kernel.shell_channel_thread.cache.list_subshell() + subshell_id_list = manager.list_subshell() print(f"subshell id: {subshell_id}") print(f"thread id: {thread_id}") diff --git a/tests/test_ipkernel_direct.py b/tests/test_ipkernel_direct.py index f650cc20..dfd0445c 100644 --- a/tests/test_ipkernel_direct.py +++ b/tests/test_ipkernel_direct.py @@ -27,7 +27,10 @@ async def test_properties(ipkernel: IPythonKernel) -> None: async def test_direct_kernel_info_request(ipkernel): reply = await ipkernel.test_shell_message("kernel_info_request", {}) assert reply["header"]["msg_type"] == "kernel_info_reply" - assert "supported_features" not in reply["content"] or "kernel subshells" not in reply["content"]["supported_features"] + assert ( + "supported_features" not in reply["content"] + or "kernel subshells" not in reply["content"]["supported_features"] + ) async def test_direct_execute_request(ipkernel: MockIPyKernel) -> None: diff --git a/tests/test_kernel_direct.py b/tests/test_kernel_direct.py index 3b6af4c9..50801b03 100644 --- a/tests/test_kernel_direct.py +++ b/tests/test_kernel_direct.py @@ -16,7 +16,10 @@ async def test_direct_kernel_info_request(kernel): reply = await kernel.test_shell_message("kernel_info_request", {}) assert reply["header"]["msg_type"] == "kernel_info_reply" - assert "supported_features" not in reply["content"] or "kernel subshells" not in reply["content"]["supported_features"] + assert ( + "supported_features" not in reply["content"] + or "kernel subshells" not in reply["content"]["supported_features"] + ) async def test_direct_execute_request(kernel): diff --git a/tests/test_message_spec.py b/tests/test_message_spec.py index 109a33a9..694de44b 100644 --- a/tests/test_message_spec.py +++ b/tests/test_message_spec.py @@ -241,6 +241,7 @@ class HistoryReply(Reply): # Subshell control messages + class CreateSubshellReply(Reply): subshell_id = Unicode() diff --git a/tests/test_subshells.py b/tests/test_subshells.py index 394204bb..fd81ec7b 100644 --- a/tests/test_subshells.py +++ b/tests/test_subshells.py @@ -3,18 +3,18 @@ # Copyright (c) IPython Development Team. # Distributed under the terms of the Modified BSD License. -from datetime import datetime, timedelta import platform -import pytest import time +from datetime import datetime, timedelta +import pytest from jupyter_client.blocking.client import BlockingKernelClient -from .utils import TIMEOUT, get_reply, get_replies, kernel, new_kernel, wait_for_idle - +from .utils import TIMEOUT, get_replies, get_reply, kernel, new_kernel # Helpers + def create_subshell_helper(kc: BlockingKernelClient): msg = kc.session.msg("create_subshell_request") kc.control_channel.send(msg) @@ -49,7 +49,7 @@ def execute_request_subshell_id(kc: BlockingKernelClient, code: str, subshell_id # Get the stream message corresponding to msg_id if msg["msg_type"] == "stream" and msg["parent_header"]["msg_id"] == msg_id: content = msg["content"] - #assert content["name"] == "stdout" + # assert content["name"] == "stdout" break return content["text"].strip() @@ -66,6 +66,7 @@ def execute_thread_ids(kc: BlockingKernelClient, subshell_id: str | None = None) # Tests + def test_supported(): with kernel() as kc: msg_id = kc.kernel_info() @@ -116,16 +117,20 @@ def test_thread_ids(): delete_subshell_helper(kc, subshell_id) -def test_run_concurrently(): +@pytest.mark.parametrize("include_main_shell", [True, False]) +def test_run_concurrently(include_main_shell): with kernel() as kc: - subshell_id = create_subshell_helper(kc)["subshell_id"] + subshell_ids = [ + None if include_main_shell else create_subshell_helper(kc)["subshell_id"], + create_subshell_helper(kc)["subshell_id"], + ] times = (0.05, 0.05) # Prepare messages, times are sleep times in seconds. # Identical times for both subshells is a harder test as preparing and sending # the execute_reply messages may overlap. msgs = [] - for id, sleep in zip((None, subshell_id), times): + for id, sleep in zip(subshell_ids, times): code = f"import time; time.sleep({sleep})" msg = kc.session.msg("execute_request", {"code": code}) msg["header"]["subshell_id"] = id @@ -143,7 +148,9 @@ def test_run_concurrently(): assert duration >= timedelta(seconds=max(times)) assert duration < timedelta(seconds=sum(times)) - delete_subshell_helper(kc, subshell_id) + for subshell_id in subshell_ids: + if subshell_id: + delete_subshell_helper(kc, subshell_id) def test_execution_count(): @@ -162,12 +169,12 @@ def test_execution_count(): for msg in msgs: kc.shell_channel.send(msg) - # Wait for replies, may be in any order. + # Wait for replies, may be in any order. replies = get_replies(kc, [msg["msg_id"] for msg in msgs]) execution_counts = [r["content"]["execution_count"] for r in replies] ec = execution_counts[0] - assert execution_counts == [ec, ec-1, ec+2, ec+1] + assert execution_counts == [ec, ec - 1, ec + 2, ec + 1] delete_subshell_helper(kc, subshell_id) @@ -198,6 +205,7 @@ def test_create_while_execute(): reason="does not work on PyPy", ) def test_shutdown_with_subshell(): + # Based on test_kernel.py::test_shutdown with new_kernel() as kc: km = kc.parent subshell_id = create_subshell_helper(kc)["subshell_id"] From 365d231575fe13a308a7fe83a418d377e6c7ebce Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 13 Jun 2024 10:02:55 +0000 Subject: [PATCH 15/37] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- ipykernel/kernelapp.py | 1 - tests/utils.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/ipykernel/kernelapp.py b/ipykernel/kernelapp.py index 9ca2ef44..2f462af4 100644 --- a/ipykernel/kernelapp.py +++ b/ipykernel/kernelapp.py @@ -56,7 +56,6 @@ from .shellchannel import ShellChannelThread from .zmqshell import ZMQInteractiveShell - # ----------------------------------------------------------------------------- # Flags and Aliases # ----------------------------------------------------------------------------- diff --git a/tests/utils.py b/tests/utils.py index 5fc18a82..fd165e85 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -73,7 +73,7 @@ def get_replies(kc, msg_ids: list[str], timeout=TIMEOUT, channel="shell"): # Replies are returned in the same order as the msg_ids, not in the order of arrival. t0 = time() count = 0 - replies = [None]*len(msg_ids) + replies = [None] * len(msg_ids) while True: get_msg = getattr(kc, f"get_{channel}_msg") reply = get_msg(timeout=timeout) From 90c3a20207fccd6ce5c74f979360a0deed6b0460 Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Thu, 13 Jun 2024 11:31:58 +0100 Subject: [PATCH 16/37] Better error checking --- ipykernel/kernelbase.py | 21 +++++++++++++++------ ipykernel/subshell_manager.py | 2 +- 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index 03cfb904..eb9de54a 100644 --- a/ipykernel/kernelbase.py +++ b/ipykernel/kernelbase.py @@ -382,13 +382,16 @@ async def shell_channel_thread_main(self): # 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.deserialize(msg2, content=False, copy=copy) - subshell_id = msg2["header"].get("subshell_id") + try: + msg3 = self.session.deserialize(msg2, content=False, copy=copy) + subshell_id = msg3["header"].get("subshell_id") - # Find inproc pair socket to use to send message to correct subshell. - socket = self.shell_channel_thread.manager.get_shell_channel_socket(subshell_id) - assert socket is not None - socket.send_multipart(msg, copy=False) + # Find inproc pair socket to use to send message to correct subshell. + socket = self.shell_channel_thread.manager.get_shell_channel_socket(subshell_id) + assert socket is not None + socket.send_multipart(msg, copy=False) + except Exception: + self.log.error("Invalid message", exc_info=True) # noqa: G201 except BaseException as e: if self.shell_stop.is_set(): return @@ -1068,6 +1071,8 @@ async def do_debug_request(self, msg): # --------------------------------------------------------------------------- async def create_subshell_request(self, socket, ident, parent): + if not self.session: + return if not self._supports_kernel_subshells(): self.log.error("Subshells are not supported by this kernel") return @@ -1081,6 +1086,8 @@ async def create_subshell_request(self, socket, ident, parent): self.session.send(socket, "create_subshell_reply", reply, parent, ident) async def delete_subshell_request(self, socket, ident, parent): + if not self.session: + return if not self._supports_kernel_subshells(): self.log.error("KERNEL SUBSHELLS NOT SUPPORTED") return @@ -1102,6 +1109,8 @@ async def delete_subshell_request(self, socket, ident, parent): self.session.send(socket, "delete_subshell_reply", reply, parent, ident) async def list_subshell_request(self, socket, ident, parent): + if not self.session: + return if not self._supports_kernel_subshells(): self.log.error("Subshells are not supported by this kernel") return diff --git a/ipykernel/subshell_manager.py b/ipykernel/subshell_manager.py index a7e9629b..d15c1d00 100644 --- a/ipykernel/subshell_manager.py +++ b/ipykernel/subshell_manager.py @@ -227,7 +227,7 @@ def _stop_subshell(self, subshell: Subshell): assert current_thread().name == "Shell channel" thread = subshell.thread - if thread and thread.is_alive(): + if thread.is_alive(): thread.stop() thread.join() From 2fa73650fca0f9695d5422b7b8fd6cb5fff1b50c Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Thu, 13 Jun 2024 14:48:17 +0100 Subject: [PATCH 17/37] Pin anyio <= 4.3.0 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index aeaeef46..328f9653 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,7 @@ dependencies = [ "pyzmq>=25.0", "psutil>=5.7", "packaging>=22", - "anyio>=4.2.0", + "anyio>=4.2.0,<=4.3.0", ] [project.urls] From ad32ea7960e7debc76d824a4de3124a5b878cf89 Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Mon, 1 Jul 2024 09:17:06 +0100 Subject: [PATCH 18/37] Remove anyio pin following fix in PR 1253 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 328f9653..aeaeef46 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,7 @@ dependencies = [ "pyzmq>=25.0", "psutil>=5.7", "packaging>=22", - "anyio>=4.2.0,<=4.3.0", + "anyio>=4.2.0", ] [project.urls] From f498045c06db95a31fb88f2dac7d60cd1f199ac7 Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Mon, 1 Jul 2024 09:52:58 +0100 Subject: [PATCH 19/37] Fix type annotations on old python --- ipykernel/subshell_manager.py | 1 + tests/test_subshells.py | 1 + tests/utils.py | 1 + 3 files changed, 3 insertions(+) diff --git a/ipykernel/subshell_manager.py b/ipykernel/subshell_manager.py index d15c1d00..1aa89e96 100644 --- a/ipykernel/subshell_manager.py +++ b/ipykernel/subshell_manager.py @@ -2,6 +2,7 @@ # Copyright (c) IPython Development Team. # Distributed under the terms of the Modified BSD License. +from __future__ import annotations import typing as t import uuid diff --git a/tests/test_subshells.py b/tests/test_subshells.py index fd81ec7b..42082276 100644 --- a/tests/test_subshells.py +++ b/tests/test_subshells.py @@ -2,6 +2,7 @@ # Copyright (c) IPython Development Team. # Distributed under the terms of the Modified BSD License. +from __future__ import annotations import platform import time diff --git a/tests/utils.py b/tests/utils.py index fd165e85..c7acbae9 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -2,6 +2,7 @@ # Copyright (c) IPython Development Team. # Distributed under the terms of the Modified BSD License. +from __future__ import annotations import atexit import os From 48f7630058a6545d1d885ee62d483dc35351922c Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Mon, 1 Jul 2024 15:01:40 +0100 Subject: [PATCH 20/37] Make tests more robust --- tests/test_subshells.py | 79 +++++++++++++++++++++++++++++++++-------- tests/utils.py | 4 +-- 2 files changed, 65 insertions(+), 18 deletions(-) diff --git a/tests/test_subshells.py b/tests/test_subshells.py index 42082276..f119a787 100644 --- a/tests/test_subshells.py +++ b/tests/test_subshells.py @@ -40,19 +40,26 @@ def list_subshell_helper(kc: BlockingKernelClient): return reply["content"] -def execute_request_subshell_id(kc: BlockingKernelClient, code: str, subshell_id: str | None): +def execute_request_subshell_id( + kc: BlockingKernelClient, code: str, subshell_id: str | None, terminator: str = "\n" +): msg = kc.session.msg("execute_request", {"code": code}) msg["header"]["subshell_id"] = subshell_id msg_id = msg["msg_id"] kc.shell_channel.send(msg) + stdout = "" while True: msg = kc.get_iopub_msg() - # Get the stream message corresponding to msg_id - if msg["msg_type"] == "stream" and msg["parent_header"]["msg_id"] == msg_id: - content = msg["content"] - # assert content["name"] == "stdout" - break - return content["text"].strip() + # Get the stream messages corresponding to msg_id + if ( + msg["msg_type"] == "stream" + and msg["parent_header"]["msg_id"] == msg_id + and msg["content"]["name"] == "stdout" + ): + stdout += msg["content"]["text"] + if stdout.endswith(terminator): + break + return stdout.strip() def execute_thread_count(kc: BlockingKernelClient) -> int: @@ -118,15 +125,53 @@ def test_thread_ids(): delete_subshell_helper(kc, subshell_id) +@pytest.mark.parametrize("are_subshells", [(False, True), (True, False), (True, True)]) +@pytest.mark.parametrize("overlap", [True, False]) +def test_run_concurrently_sequence(are_subshells, overlap): + with kernel() as kc: + subshell_ids = [ + create_subshell_helper(kc)["subshell_id"] if is_subshell else None + for is_subshell in are_subshells + ] + if overlap: + codes = [ + "import time; start0=True; end0=False; time.sleep(0.2); end0=True", + "assert start0; assert not end0; time.sleep(0.2); assert end0", + ] + else: + codes = [ + "import time; start0=True; end0=False; time.sleep(0.2); assert end1", + "assert start0; assert not end0; end1=True", + ] + + msgs = [] + for subshell_id, code in zip(subshell_ids, codes): + msg = kc.session.msg("execute_request", {"code": code}) + msg["header"]["subshell_id"] = subshell_id + kc.shell_channel.send(msg) + msgs.append(msg) + if len(msgs) == 1: + time.sleep(0.1) # Wait for first execute_request to start. + + replies = get_replies(kc, [msg["msg_id"] for msg in msgs]) + + for subshell_id in subshell_ids: + if subshell_id: + delete_subshell_helper(kc, subshell_id) + + for reply in replies: + assert reply["content"]["status"] == "ok" + + @pytest.mark.parametrize("include_main_shell", [True, False]) -def test_run_concurrently(include_main_shell): +def test_run_concurrently_timing(include_main_shell): with kernel() as kc: subshell_ids = [ None if include_main_shell else create_subshell_helper(kc)["subshell_id"], create_subshell_helper(kc)["subshell_id"], ] - times = (0.05, 0.05) + times = (0.2, 0.2) # Prepare messages, times are sleep times in seconds. # Identical times for both subshells is a harder test as preparing and sending # the execute_reply messages may overlap. @@ -145,14 +190,18 @@ def test_run_concurrently(include_main_shell): _ = get_replies(kc, [msg["msg_id"] for msg in msgs]) end = datetime.now() - duration = end - start - assert duration >= timedelta(seconds=max(times)) - assert duration < timedelta(seconds=sum(times)) - for subshell_id in subshell_ids: if subshell_id: delete_subshell_helper(kc, subshell_id) + duration = end - start + assert duration >= timedelta(seconds=max(times)) + # Care is needed with this test as runtime conditions such as gathering + # coverage can slow it down causing the following assert to fail. + # The sleep time of 0.2 is empirically determined to run OK in CI, but + # consider increasing it if the following fails. + assert duration < timedelta(seconds=sum(times)) + def test_execution_count(): with kernel() as kc: @@ -173,12 +222,12 @@ def test_execution_count(): # Wait for replies, may be in any order. replies = get_replies(kc, [msg["msg_id"] for msg in msgs]) + delete_subshell_helper(kc, subshell_id) + execution_counts = [r["content"]["execution_count"] for r in replies] ec = execution_counts[0] assert execution_counts == [ec, ec - 1, ec + 2, ec + 1] - delete_subshell_helper(kc, subshell_id) - def test_create_while_execute(): with kernel() as kc: diff --git a/tests/utils.py b/tests/utils.py index c7acbae9..b20e8fcb 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -75,15 +75,13 @@ def get_replies(kc, msg_ids: list[str], timeout=TIMEOUT, channel="shell"): t0 = time() count = 0 replies = [None] * len(msg_ids) - while True: + while count < len(msg_ids): get_msg = getattr(kc, f"get_{channel}_msg") reply = get_msg(timeout=timeout) try: msg_id = reply["parent_header"]["msg_id"] replies[msg_ids.index(msg_id)] = reply count += 1 - if count == len(msg_ids): - break except ValueError: # Allow debugging ignored replies print(f"Ignoring reply not to any of {msg_ids}: {reply}") From c79ba7911de033b6fc82a3f4748be1461e14d38c Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Wed, 7 Aug 2024 12:33:36 +0100 Subject: [PATCH 21/37] Isolate each subshell test --- tests/test_subshells.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/test_subshells.py b/tests/test_subshells.py index f119a787..b47126cb 100644 --- a/tests/test_subshells.py +++ b/tests/test_subshells.py @@ -11,7 +11,7 @@ import pytest from jupyter_client.blocking.client import BlockingKernelClient -from .utils import TIMEOUT, get_replies, get_reply, kernel, new_kernel +from .utils import TIMEOUT, get_replies, get_reply, new_kernel # Helpers @@ -76,7 +76,7 @@ def execute_thread_ids(kc: BlockingKernelClient, subshell_id: str | None = None) def test_supported(): - with kernel() as kc: + with new_kernel() as kc: msg_id = kc.kernel_info() reply = get_reply(kc, msg_id, TIMEOUT) assert "supported_features" in reply["content"] @@ -84,7 +84,7 @@ def test_supported(): def test_subshell_id_lifetime(): - with kernel() as kc: + with new_kernel() as kc: assert list_subshell_helper(kc)["subshell_id"] == [] subshell_id = create_subshell_helper(kc)["subshell_id"] assert list_subshell_helper(kc)["subshell_id"] == [subshell_id] @@ -93,14 +93,14 @@ def test_subshell_id_lifetime(): def test_delete_non_existent(): - with kernel() as kc: + with new_kernel() as kc: reply = delete_subshell_helper(kc, "unknown_subshell_id") assert reply["status"] == "error" assert "evalue" in reply def test_thread_counts(): - with kernel() as kc: + with new_kernel() as kc: nthreads = execute_thread_count(kc) subshell_id = create_subshell_helper(kc)["subshell_id"] @@ -113,7 +113,7 @@ def test_thread_counts(): def test_thread_ids(): - with kernel() as kc: + with new_kernel() as kc: subshell_id = create_subshell_helper(kc)["subshell_id"] thread_id, main_thread_id = execute_thread_ids(kc) @@ -128,7 +128,7 @@ def test_thread_ids(): @pytest.mark.parametrize("are_subshells", [(False, True), (True, False), (True, True)]) @pytest.mark.parametrize("overlap", [True, False]) def test_run_concurrently_sequence(are_subshells, overlap): - with kernel() as kc: + with new_kernel() as kc: subshell_ids = [ create_subshell_helper(kc)["subshell_id"] if is_subshell else None for is_subshell in are_subshells @@ -165,7 +165,7 @@ def test_run_concurrently_sequence(are_subshells, overlap): @pytest.mark.parametrize("include_main_shell", [True, False]) def test_run_concurrently_timing(include_main_shell): - with kernel() as kc: + with new_kernel() as kc: subshell_ids = [ None if include_main_shell else create_subshell_helper(kc)["subshell_id"], create_subshell_helper(kc)["subshell_id"], @@ -204,7 +204,7 @@ def test_run_concurrently_timing(include_main_shell): def test_execution_count(): - with kernel() as kc: + with new_kernel() as kc: subshell_id = create_subshell_helper(kc)["subshell_id"] # Prepare messages @@ -230,7 +230,7 @@ def test_execution_count(): def test_create_while_execute(): - with kernel() as kc: + with new_kernel() as kc: # Send request to execute code on main subshell. msg = kc.session.msg("execute_request", {"code": "import time; time.sleep(0.05)"}) kc.shell_channel.send(msg) From 97e3e91e2a0944a04df23bef40b65800abfcc7d8 Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Thu, 12 Sep 2024 12:33:09 +0100 Subject: [PATCH 22/37] Avoid high CPU load when waiting for subshell reply messages --- ipykernel/subshell_manager.py | 35 ++++++++++++++++++++++++++++------- 1 file changed, 28 insertions(+), 7 deletions(-) diff --git a/ipykernel/subshell_manager.py b/ipykernel/subshell_manager.py index 1aa89e96..e91ec636 100644 --- a/ipykernel/subshell_manager.py +++ b/ipykernel/subshell_manager.py @@ -11,7 +11,7 @@ import zmq import zmq.asyncio -from anyio import sleep +from anyio import Event, create_task_group from .subshell import SubshellThread @@ -50,8 +50,9 @@ def __init__(self, context: zmq.asyncio.Context, shell_socket): self._parent_shell_channel_socket = self._create_inproc_pair_socket(None, True) self._parent_other_socket = self._create_inproc_pair_socket(None, False) - self._poller = zmq.Poller() + self._poller = zmq.asyncio.Poller() self._poller.register(self._parent_shell_channel_socket, zmq.POLLIN) + self._subshell_change = Event() def close(self): """Stop all subshells and close all resources.""" @@ -126,12 +127,13 @@ async def listen_from_subshells(self): assert current_thread().name == "Shell channel" while True: - for socket, _ in self._poller.poll(0): - msg = await socket.recv_multipart(copy=False) - self._shell_socket.send_multipart(msg) + async with create_task_group() as tg: + tg.start_soon(self._listen_from_subshells) + await self._subshell_change.wait() + tg.cancel_scope.cancel() - # Yield to other tasks. - await sleep(0) + # anyio.Event is single use, so recreate to reuse + self._subshell_change = Event() def subshell_id_from_thread_id(self, thread_id) -> str | None: """Return subshell_id of the specified thread_id. @@ -177,6 +179,7 @@ async def _create_subshell(self, subshell_task) -> str: thread.start() self._poller.register(shell_channel_socket, zmq.POLLIN) + self._subshell_change.set() return subshell_id @@ -196,6 +199,23 @@ def _get_inproc_socket_address(self, name: str | None): full_name = f"subshell-{name}" if name else "subshell" return f"inproc://{full_name}" + async def _listen_from_subshells(self): + """Await next reply message from any subshell (parent or child) and resend + to the client via the shell_socket. + + If a subshell is created or deleted then the poller is updated and the task + executing this function is cancelled and then rescheduled with the updated + poller. + + Runs in the shell channel thread. + """ + assert current_thread().name == "Shell channel" + + while True: + for socket, _ in await self._poller.poll(): + msg = await socket.recv_multipart(copy=False) + self._shell_socket.send_multipart(msg) + async def _process_control_request(self, request, subshell_task): """Process a control request message received on the control inproc socket and return the reply. Runs in the shell channel thread. @@ -233,4 +253,5 @@ def _stop_subshell(self, subshell: Subshell): thread.join() self._poller.unregister(subshell.shell_channel_socket) + self._subshell_change.set() subshell.shell_channel_socket.close() From e19273cff732771f9b00774dd94319ca47b8f70e Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Fri, 13 Sep 2024 12:47:13 +0100 Subject: [PATCH 23/37] Replace use of zmq Poller with better anyio task handling --- ipykernel/subshell_manager.py | 66 +++++++++++++++++++++-------------- 1 file changed, 40 insertions(+), 26 deletions(-) diff --git a/ipykernel/subshell_manager.py b/ipykernel/subshell_manager.py index e91ec636..00809c7f 100644 --- a/ipykernel/subshell_manager.py +++ b/ipykernel/subshell_manager.py @@ -11,7 +11,7 @@ import zmq import zmq.asyncio -from anyio import Event, create_task_group +from anyio import create_memory_object_stream, create_task_group from .subshell import SubshellThread @@ -50,14 +50,17 @@ def __init__(self, context: zmq.asyncio.Context, shell_socket): self._parent_shell_channel_socket = self._create_inproc_pair_socket(None, True) self._parent_other_socket = self._create_inproc_pair_socket(None, False) - self._poller = zmq.asyncio.Poller() - self._poller.register(self._parent_shell_channel_socket, zmq.POLLIN) - self._subshell_change = Event() + # anyio memory object stream for async queue-like communication between tasks. + # Used by _create_subshell to tell listen_from_subshells to spawn a new task. + self._send_stream, self._receive_stream = create_memory_object_stream[str]() def close(self): """Stop all subshells and close all resources.""" assert current_thread().name == "Shell channel" + self._send_stream.close() + self._receive_stream.close() + for socket in ( self._control_shell_channel_socket, self._control_other_socket, @@ -126,14 +129,10 @@ async def listen_from_subshells(self): """ assert current_thread().name == "Shell channel" - while True: - async with create_task_group() as tg: - tg.start_soon(self._listen_from_subshells) - await self._subshell_change.wait() - tg.cancel_scope.cancel() - - # anyio.Event is single use, so recreate to reuse - self._subshell_change = Event() + async with create_task_group() as tg: + tg.start_soon(self._listen_for_subshell_reply, None) + async for subshell_id in self._receive_stream: + tg.start_soon(self._listen_for_subshell_reply, subshell_id) def subshell_id_from_thread_id(self, thread_id) -> str | None: """Return subshell_id of the specified thread_id. @@ -173,14 +172,15 @@ async def _create_subshell(self, subshell_task) -> str: shell_channel_socket = self._create_inproc_pair_socket(subshell_id, True) self._cache[subshell_id] = Subshell(thread, shell_channel_socket) + # Tell task running listen_from_subshells to create a new task to listen for + # reply messages from the new subshell to resend to the client. + await self._send_stream.send(subshell_id) + address = self._get_inproc_socket_address(subshell_id) thread.add_task(thread.create_pair_socket, self._context, address) thread.add_task(subshell_task, subshell_id) thread.start() - self._poller.register(shell_channel_socket, zmq.POLLIN) - self._subshell_change.set() - return subshell_id def _delete_subshell(self, subshell_id: str) -> None: @@ -199,22 +199,37 @@ def _get_inproc_socket_address(self, name: str | None): full_name = f"subshell-{name}" if name else "subshell" return f"inproc://{full_name}" - async def _listen_from_subshells(self): - """Await next reply message from any subshell (parent or child) and resend - to the client via the shell_socket. + def _get_shell_channel_socket(self, subshell_id: str | None) -> zmq.asyncio.Socket: + if subshell_id is None: + return self._parent_shell_channel_socket + with self._lock: + return self._cache[subshell_id].shell_channel_socket + + def _is_subshell(self, subshell_id: str | None) -> bool: + if subshell_id is None: + return True + with self._lock: + return subshell_id in self._cache - If a subshell is created or deleted then the poller is updated and the task - executing this function is cancelled and then rescheduled with the updated - poller. + async def _listen_for_subshell_reply(self, subshell_id: str | None): + """Listen for reply messages on specified subshell inproc socket and + resend to the client via the shell_socket. Runs in the shell channel thread. """ assert current_thread().name == "Shell channel" - while True: - for socket, _ in await self._poller.poll(): - msg = await socket.recv_multipart(copy=False) + shell_channel_socket = self._get_shell_channel_socket(subshell_id) + + try: + while True: + msg = await shell_channel_socket.recv_multipart(copy=False) self._shell_socket.send_multipart(msg) + except BaseException as e: + if not self._is_subshell(subshell_id): + # Subshell no longer exists so exit gracefully + return + raise e async def _process_control_request(self, request, subshell_task): """Process a control request message received on the control inproc @@ -252,6 +267,5 @@ def _stop_subshell(self, subshell: Subshell): thread.stop() thread.join() - self._poller.unregister(subshell.shell_channel_socket) - self._subshell_change.set() + # Closing the shell_channel_socket terminates the task that is listening on it. subshell.shell_channel_socket.close() From a050784fc261feb76b413c4a6b0f9d8cbd7afde2 Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Fri, 13 Sep 2024 13:13:31 +0100 Subject: [PATCH 24/37] Wrap sending via shell_socket with a lock --- ipykernel/subshell_manager.py | 27 ++++++++++++++++----------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/ipykernel/subshell_manager.py b/ipykernel/subshell_manager.py index 00809c7f..6574b8df 100644 --- a/ipykernel/subshell_manager.py +++ b/ipykernel/subshell_manager.py @@ -32,6 +32,9 @@ class SubshellManager: the shell channel thread so there is only ever one write access at any one time. Reading of cache information can be performed by other threads, so all reads are protected by a lock so that they are atomic. + + Sending reply messages via the shell_socket is wrapped by another lock to protect + against multiple subshells attempting to send at the same time. """ def __init__(self, context: zmq.asyncio.Context, shell_socket): @@ -40,7 +43,8 @@ def __init__(self, context: zmq.asyncio.Context, shell_socket): self._context: zmq.asyncio.Context = context self._shell_socket = shell_socket self._cache: dict[str, Subshell] = {} - self._lock: Lock = Lock() + self._lock_cache = Lock() + self._lock_shell_socket = Lock() # Inproc pair sockets for control channel and main shell (parent subshell). # Each inproc pair has a "shell_channel" socket used in the shell channel @@ -70,7 +74,7 @@ def close(self): if socket is not None: socket.close() - with self._lock: + with self._lock_cache: while True: try: _, subshell = self._cache.popitem() @@ -88,7 +92,7 @@ def get_other_socket(self, subshell_id: str | None): """ if subshell_id is None: return self._parent_other_socket - with self._lock: + with self._lock_cache: return self._cache[subshell_id].thread._pair_socket def get_shell_channel_socket(self, subshell_id: str | None): @@ -98,7 +102,7 @@ def get_shell_channel_socket(self, subshell_id: str | None): """ if subshell_id is None: return self._parent_shell_channel_socket - with self._lock: + with self._lock_cache: return self._cache[subshell_id].shell_channel_socket def list_subshell(self) -> list[str]: @@ -106,7 +110,7 @@ def list_subshell(self) -> list[str]: Can be called by any subshell using %subshell magic. """ - with self._lock: + with self._lock_cache: return list(self._cache) async def listen_from_control(self, subshell_task): @@ -141,7 +145,7 @@ def subshell_id_from_thread_id(self, thread_id) -> str | None: Only used by %subshell magic so does not have to be fast/cached. """ - with self._lock: + with self._lock_cache: if thread_id == main_thread().ident: return None for id, subshell in self._cache.items(): @@ -167,7 +171,7 @@ async def _create_subshell(self, subshell_task) -> str: subshell_id = str(uuid.uuid4()) thread = SubshellThread(subshell_id) - with self._lock: + with self._lock_cache: assert subshell_id not in self._cache shell_channel_socket = self._create_inproc_pair_socket(subshell_id, True) self._cache[subshell_id] = Subshell(thread, shell_channel_socket) @@ -190,7 +194,7 @@ def _delete_subshell(self, subshell_id: str) -> None: """ assert current_thread().name == "Shell channel" - with self._lock: + with self._lock_cache: subshell = self._cache.pop(subshell_id) self._stop_subshell(subshell) @@ -202,13 +206,13 @@ def _get_inproc_socket_address(self, name: str | None): def _get_shell_channel_socket(self, subshell_id: str | None) -> zmq.asyncio.Socket: if subshell_id is None: return self._parent_shell_channel_socket - with self._lock: + with self._lock_cache: return self._cache[subshell_id].shell_channel_socket def _is_subshell(self, subshell_id: str | None) -> bool: if subshell_id is None: return True - with self._lock: + with self._lock_cache: return subshell_id in self._cache async def _listen_for_subshell_reply(self, subshell_id: str | None): @@ -224,7 +228,8 @@ async def _listen_for_subshell_reply(self, subshell_id: str | None): try: while True: msg = await shell_channel_socket.recv_multipart(copy=False) - self._shell_socket.send_multipart(msg) + with self._lock_shell_socket: + self._shell_socket.send_multipart(msg) except BaseException as e: if not self._is_subshell(subshell_id): # Subshell no longer exists so exit gracefully From 9669fcd96c40fef35dc4e1097831ab12380f3c5f Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Fri, 13 Sep 2024 14:12:38 +0100 Subject: [PATCH 25/37] More type annotations --- ipykernel/subshell_manager.py | 34 +++++++++++++++++++--------------- 1 file changed, 19 insertions(+), 15 deletions(-) diff --git a/ipykernel/subshell_manager.py b/ipykernel/subshell_manager.py index 6574b8df..4fd29061 100644 --- a/ipykernel/subshell_manager.py +++ b/ipykernel/subshell_manager.py @@ -37,7 +37,7 @@ class SubshellManager: against multiple subshells attempting to send at the same time. """ - def __init__(self, context: zmq.asyncio.Context, shell_socket): + def __init__(self, context: zmq.asyncio.Context, shell_socket: zmq.asyncio.Socket): assert current_thread() == main_thread() self._context: zmq.asyncio.Context = context @@ -58,7 +58,7 @@ def __init__(self, context: zmq.asyncio.Context, shell_socket): # Used by _create_subshell to tell listen_from_subshells to spawn a new task. self._send_stream, self._receive_stream = create_memory_object_stream[str]() - def close(self): + def close(self) -> None: """Stop all subshells and close all resources.""" assert current_thread().name == "Shell channel" @@ -82,10 +82,10 @@ def close(self): break self._stop_subshell(subshell) - def get_control_other_socket(self): + def get_control_other_socket(self) -> zmq.asyncio.Socket: return self._control_other_socket - def get_other_socket(self, subshell_id: str | None): + def get_other_socket(self, subshell_id: str | None) -> zmq.asyncio.Socket: """Return the other inproc pair socket for a subshell. This socket is accessed from the subshell thread. @@ -93,9 +93,11 @@ def get_other_socket(self, subshell_id: str | None): if subshell_id is None: return self._parent_other_socket with self._lock_cache: - return self._cache[subshell_id].thread._pair_socket + socket = self._cache[subshell_id].thread._pair_socket + assert socket is not None + return socket - def get_shell_channel_socket(self, subshell_id: str | None): + def get_shell_channel_socket(self, subshell_id: str | None) -> zmq.asyncio.Socket: """Return the shell channel inproc pair socket for a subshell. This socket is accessed from the shell channel thread. @@ -113,7 +115,7 @@ def list_subshell(self) -> list[str]: with self._lock_cache: return list(self._cache) - async def listen_from_control(self, subshell_task): + async def listen_from_control(self, subshell_task) -> None: """Listen for messages on the control inproc socket, handle those messages and return replies on the same socket. Runs in the shell channel thread. """ @@ -121,11 +123,11 @@ async def listen_from_control(self, subshell_task): socket = self._control_shell_channel_socket while True: - request = await socket.recv_json() + request = await socket.recv_json() # type: ignore[misc] reply = await self._process_control_request(request, subshell_task) - await socket.send_json(reply) + await socket.send_json(reply) # type: ignore[func-returns-value] - async def listen_from_subshells(self): + async def listen_from_subshells(self) -> None: """Listen for reply messages on inproc sockets of all subshells and resend those messages to the client via the shell_socket. @@ -154,7 +156,9 @@ def subshell_id_from_thread_id(self, thread_id) -> str | None: msg = f"Thread id '{thread_id} does not correspond to a subshell of this kernel" raise RuntimeError(msg) - def _create_inproc_pair_socket(self, name: str | None, shell_channel_end: bool): + def _create_inproc_pair_socket( + self, name: str | None, shell_channel_end: bool + ) -> zmq.asyncio.Socket: """Create and return a single ZMQ inproc pair socket.""" address = self._get_inproc_socket_address(name) socket = self._context.socket(zmq.PAIR) @@ -199,7 +203,7 @@ def _delete_subshell(self, subshell_id: str) -> None: self._stop_subshell(subshell) - def _get_inproc_socket_address(self, name: str | None): + def _get_inproc_socket_address(self, name: str | None) -> str: full_name = f"subshell-{name}" if name else "subshell" return f"inproc://{full_name}" @@ -215,7 +219,7 @@ def _is_subshell(self, subshell_id: str | None) -> bool: with self._lock_cache: return subshell_id in self._cache - async def _listen_for_subshell_reply(self, subshell_id: str | None): + async def _listen_for_subshell_reply(self, subshell_id: str | None) -> None: """Listen for reply messages on specified subshell inproc socket and resend to the client via the shell_socket. @@ -236,7 +240,7 @@ async def _listen_for_subshell_reply(self, subshell_id: str | None): return raise e - async def _process_control_request(self, request, subshell_task): + async def _process_control_request(self, request, subshell_task) -> dict[str, t.Any]: """Process a control request message received on the control inproc socket and return the reply. Runs in the shell channel thread. """ @@ -263,7 +267,7 @@ async def _process_control_request(self, request, subshell_task): } return reply - def _stop_subshell(self, subshell: Subshell): + def _stop_subshell(self, subshell: Subshell) -> None: """Stop a subshell thread and close all of its resources.""" assert current_thread().name == "Shell channel" From b72d1da3ff66a0580c5cd75daf011b544e466730 Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Tue, 17 Sep 2024 17:25:42 +0100 Subject: [PATCH 26/37] Ensure abort reply messages are sent via the subshell manager --- ipykernel/kernelbase.py | 4 ++-- ipykernel/subshell_manager.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index eb9de54a..e03c89bc 100644 --- a/ipykernel/kernelbase.py +++ b/ipykernel/kernelbase.py @@ -469,7 +469,7 @@ async def process_shell_message(self, msg=None, socket=None): elif received_time - self._aborted_time > self.stop_on_error_timeout: self._aborting = False if self._aborting: - await self._send_abort_reply(self.shell_socket, msg, idents) + await self._send_abort_reply(socket, msg, idents) self._publish_status("idle", "shell") return @@ -479,7 +479,7 @@ async def process_shell_message(self, msg=None, socket=None): self.log.debug("\n*** MESSAGE TYPE:%s***", msg_type) self.log.debug(" Content: %s\n --->\n ", msg["content"]) - if not await self.should_handle(self.shell_socket, msg, idents): + if not await self.should_handle(socket, msg, idents): return handler = self.shell_handlers.get(msg_type) diff --git a/ipykernel/subshell_manager.py b/ipykernel/subshell_manager.py index 4fd29061..1f8a6199 100644 --- a/ipykernel/subshell_manager.py +++ b/ipykernel/subshell_manager.py @@ -233,7 +233,7 @@ async def _listen_for_subshell_reply(self, subshell_id: str | None) -> None: while True: msg = await shell_channel_socket.recv_multipart(copy=False) with self._lock_shell_socket: - self._shell_socket.send_multipart(msg) + await self._shell_socket.send_multipart(msg) except BaseException as e: if not self._is_subshell(subshell_id): # Subshell no longer exists so exit gracefully From e4fa75da473e6aeb74c5bb223b3e9879b7db6406 Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Tue, 17 Sep 2024 17:33:06 +0100 Subject: [PATCH 27/37] Add modules to docs API --- docs/api/ipykernel.rst | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/docs/api/ipykernel.rst b/docs/api/ipykernel.rst index 2e1cf20d..dd46d084 100644 --- a/docs/api/ipykernel.rst +++ b/docs/api/ipykernel.rst @@ -110,6 +110,30 @@ Submodules :show-inheritance: +.. automodule:: ipykernel.shellchannel + :members: + :undoc-members: + :show-inheritance: + + +.. automodule:: ipykernel.subshell + :members: + :undoc-members: + :show-inheritance: + + +.. automodule:: ipykernel.subshell_manager + :members: + :undoc-members: + :show-inheritance: + + +.. automodule:: ipykernel.thread + :members: + :undoc-members: + :show-inheritance: + + .. automodule:: ipykernel.trio_runner :members: :undoc-members: From a708e0ee8c31a5e91175b484e76f8af48be93235 Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Tue, 17 Sep 2024 17:34:10 +0100 Subject: [PATCH 28/37] Increase timings of test_subshells::test_execution_counts --- tests/test_subshells.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_subshells.py b/tests/test_subshells.py index b47126cb..f1328dda 100644 --- a/tests/test_subshells.py +++ b/tests/test_subshells.py @@ -208,7 +208,7 @@ def test_execution_count(): subshell_id = create_subshell_helper(kc)["subshell_id"] # Prepare messages - times = (0.1, 0.05, 0.2, 0.07) # Sleep seconds + times = (0.2, 0.1, 0.4, 0.15) # Sleep seconds msgs = [] for id, sleep in zip((None, subshell_id, None, subshell_id), times): code = f"import time; time.sleep({sleep})" From 2a91bcdb08ebada58059cae0aaabe9d30353c6dc Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Wed, 2 Oct 2024 13:00:53 +0100 Subject: [PATCH 29/37] Use _tasks_and_args in BaseThread --- ipykernel/thread.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/ipykernel/thread.py b/ipykernel/thread.py index 642b1cd3..80211ce7 100644 --- a/ipykernel/thread.py +++ b/ipykernel/thread.py @@ -1,4 +1,5 @@ """Base class for threads.""" +import typing as t from threading import Event, Thread from anyio import create_task_group, run, to_thread @@ -13,13 +14,11 @@ def __init__(self, name, **kwargs): self.pydev_do_not_trace = True self.is_pydev_daemon_thread = True self.__stop = Event() - self._tasks = [] - self._task_args = [] + self._tasks_and_args: t.List[t.Tuple[t.Callable, t.Tuple]] = [] - def add_task(self, task, *args): + def add_task(self, task: t.Callable, *args: t.Tuple): # May only add tasks before the thread is started. - self._tasks.append(task) - self._task_args.append(args) + self._tasks_and_args.append((task, args)) def run(self): """Run the thread.""" @@ -27,7 +26,7 @@ def run(self): async def _main(self): async with create_task_group() as tg: - for task, args in zip(self._tasks, self._task_args): + for task, args in self._tasks_and_args: tg.start_soon(task, *args) await to_thread.run_sync(self.__stop.wait) tg.cancel_scope.cancel() From 4bcd59e79ad014be6fb4acd6919594114ad52793 Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Wed, 2 Oct 2024 13:05:55 +0100 Subject: [PATCH 30/37] Use super when calling base class constructors --- ipykernel/heartbeat.py | 2 +- ipykernel/iostream.py | 2 +- ipykernel/thread.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/ipykernel/heartbeat.py b/ipykernel/heartbeat.py index d2890f67..9816959d 100644 --- a/ipykernel/heartbeat.py +++ b/ipykernel/heartbeat.py @@ -32,7 +32,7 @@ def __init__(self, context, addr=None): """Initialize the heartbeat thread.""" if addr is None: addr = ("tcp", localhost(), 0) - Thread.__init__(self, name="Heartbeat") + super().__init__(name="Heartbeat") self.context = context self.transport, self.ip, self.port = addr self.original_port = self.port diff --git a/ipykernel/iostream.py b/ipykernel/iostream.py index 6280905c..b83f8120 100644 --- a/ipykernel/iostream.py +++ b/ipykernel/iostream.py @@ -40,7 +40,7 @@ class _IOPubThread(Thread): def __init__(self, tasks, **kwargs): """Initialize the thread.""" - Thread.__init__(self, name="IOPub", **kwargs) + super().__init__(name="IOPub", **kwargs) self._tasks = tasks self.pydev_do_not_trace = True self.is_pydev_daemon_thread = True diff --git a/ipykernel/thread.py b/ipykernel/thread.py index 80211ce7..e66c141c 100644 --- a/ipykernel/thread.py +++ b/ipykernel/thread.py @@ -10,7 +10,7 @@ class BaseThread(Thread): def __init__(self, name, **kwargs): """Initialize the thread.""" - Thread.__init__(self, name=name, **kwargs) + super().__init__(name=name, **kwargs) self.pydev_do_not_trace = True self.is_pydev_daemon_thread = True self.__stop = Event() From b644f9e7e280d3057e1ced78292c2524bcec2f20 Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Wed, 2 Oct 2024 13:15:54 +0100 Subject: [PATCH 31/37] Improve thread classes --- ipykernel/control.py | 5 ----- ipykernel/shellchannel.py | 10 +++++----- ipykernel/subshell.py | 11 ++++++----- ipykernel/thread.py | 10 +++++----- 4 files changed, 16 insertions(+), 20 deletions(-) diff --git a/ipykernel/control.py b/ipykernel/control.py index 14a86a18..8fc04b43 100644 --- a/ipykernel/control.py +++ b/ipykernel/control.py @@ -11,8 +11,3 @@ class ControlThread(BaseThread): def __init__(self, **kwargs): """Initialize the thread.""" super().__init__(name=CONTROL_THREAD_NAME, **kwargs) - - def run(self): - """Run the thread.""" - self.name = CONTROL_THREAD_NAME - super().run() diff --git a/ipykernel/shellchannel.py b/ipykernel/shellchannel.py index 7341f9e4..eeaead7b 100644 --- a/ipykernel/shellchannel.py +++ b/ipykernel/shellchannel.py @@ -27,8 +27,8 @@ def manager(self): def run(self): """Run the thread.""" - self.name = SHELL_CHANNEL_THREAD_NAME - super().run() - - if self._manager: - self._manager.close() + try: + super().run() + finally: + if self._manager: + self._manager.close() diff --git a/ipykernel/subshell.py b/ipykernel/subshell.py index 5a23dd08..19aa66c4 100644 --- a/ipykernel/subshell.py +++ b/ipykernel/subshell.py @@ -28,8 +28,9 @@ async def create_pair_socket(self, context: zmq.asyncio.Context, address: str): self._pair_socket.connect(address) def run(self): - super().run() - - if self._pair_socket is not None: - self._pair_socket.close() - self._pair_socket = None + try: + super().run() + finally: + if self._pair_socket is not None: + self._pair_socket.close() + self._pair_socket = None diff --git a/ipykernel/thread.py b/ipykernel/thread.py index e66c141c..c41db727 100644 --- a/ipykernel/thread.py +++ b/ipykernel/thread.py @@ -8,9 +8,9 @@ class BaseThread(Thread): """Base class for threads.""" - def __init__(self, name, **kwargs): + def __init__(self, **kwargs): """Initialize the thread.""" - super().__init__(name=name, **kwargs) + super().__init__(**kwargs) self.pydev_do_not_trace = True self.is_pydev_daemon_thread = True self.__stop = Event() @@ -20,11 +20,11 @@ def add_task(self, task: t.Callable, *args: t.Tuple): # May only add tasks before the thread is started. self._tasks_and_args.append((task, args)) - def run(self): + def run(self) -> t.Any: """Run the thread.""" - run(self._main) + return run(self._main) - async def _main(self): + async def _main(self) -> None: async with create_task_group() as tg: for task, args in self._tasks_and_args: tg.start_soon(task, *args) From 924015ae8cc82acdf94cdef924d5c77a0e118b91 Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Wed, 2 Oct 2024 13:18:32 +0100 Subject: [PATCH 32/37] Always return supported_features in kernel_info --- ipykernel/kernelbase.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index e03c89bc..5d9d760a 100644 --- a/ipykernel/kernelbase.py +++ b/ipykernel/kernelbase.py @@ -892,9 +892,10 @@ def kernel_info(self): "language_info": self.language_info, "banner": self.banner, "help_links": self.help_links, + "supported_features": [], } if self._supports_kernel_subshells(): - info["supported_features"] = ["kernel subshells"] + info["supported_features"].append("kernel subshells") return info async def kernel_info_request(self, socket, ident, parent): From a1fb1de82975f9c148e1c75ecd4db942d60c2e3d Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Wed, 2 Oct 2024 13:27:20 +0100 Subject: [PATCH 33/37] Make _supports_kernel_subshells a property --- ipykernel/kernelbase.py | 13 +++++++------ ipykernel/zmqshell.py | 4 +--- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index 5d9d760a..f7a057ae 100644 --- a/ipykernel/kernelbase.py +++ b/ipykernel/kernelbase.py @@ -399,7 +399,7 @@ async def shell_channel_thread_main(self): async def shell_main(self, subshell_id: str | None): """Main loop for a single subshell.""" - if self._supports_kernel_subshells(): + if self._supports_kernel_subshells: if subshell_id is None: assert threading.current_thread() == threading.main_thread() else: @@ -435,7 +435,7 @@ async def process_shell_message(self, msg=None, socket=None): socket = self.shell_socket assert self.session is not None - if self._supports_kernel_subshells(): + if self._supports_kernel_subshells: assert threading.current_thread() not in ( self.control_thread, self.shell_channel_thread, @@ -894,7 +894,7 @@ def kernel_info(self): "help_links": self.help_links, "supported_features": [], } - if self._supports_kernel_subshells(): + if self._supports_kernel_subshells: info["supported_features"].append("kernel subshells") return info @@ -1074,7 +1074,7 @@ async def do_debug_request(self, msg): async def create_subshell_request(self, socket, ident, parent): if not self.session: return - if not self._supports_kernel_subshells(): + if not self._supports_kernel_subshells: self.log.error("Subshells are not supported by this kernel") return @@ -1089,7 +1089,7 @@ async def create_subshell_request(self, socket, ident, parent): async def delete_subshell_request(self, socket, ident, parent): if not self.session: return - if not self._supports_kernel_subshells(): + if not self._supports_kernel_subshells: self.log.error("KERNEL SUBSHELLS NOT SUPPORTED") return @@ -1112,7 +1112,7 @@ async def delete_subshell_request(self, socket, ident, parent): async def list_subshell_request(self, socket, ident, parent): if not self.session: return - if not self._supports_kernel_subshells(): + if not self._supports_kernel_subshells: self.log.error("Subshells are not supported by this kernel") return @@ -1415,5 +1415,6 @@ async def _at_shutdown(self): ) self.log.debug("%s", self._shutdown_message) + @property def _supports_kernel_subshells(self): return self.shell_channel_thread is not None diff --git a/ipykernel/zmqshell.py b/ipykernel/zmqshell.py index 5ab7b188..6c1273cd 100644 --- a/ipykernel/zmqshell.py +++ b/ipykernel/zmqshell.py @@ -450,9 +450,7 @@ def subshell(self, arg_s): app = IPKernelApp.instance() kernel = app.kernel - if not ( - hasattr(kernel, "_supports_kernel_subshells") and kernel._supports_kernel_subshells() - ): + if not getattr(kernel, "_supports_kernel_subshells", False): print("Kernel does not support subshells") return From be5c83a2cbc8552a0da403b9b4216f98976971c7 Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Wed, 2 Oct 2024 13:47:55 +0100 Subject: [PATCH 34/37] Clarify socket=None in process_shell_message --- ipykernel/kernelbase.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index f7a057ae..157e2dc2 100644 --- a/ipykernel/kernelbase.py +++ b/ipykernel/kernelbase.py @@ -422,6 +422,7 @@ async def shell_main(self, subshell_id: str | None): tg.cancel_scope.cancel() async def process_shell(self, socket=None): + # socket=None is valid if kernel subshells are not supported. try: while True: await self.process_shell_message(socket=socket) @@ -431,17 +432,20 @@ async def process_shell(self, socket=None): raise e async def process_shell_message(self, msg=None, socket=None): - if not socket: - socket = self.shell_socket - + # If socket is None kernel subshells are not supported so use socket=shell_socket. + # If msg is set, process that message. + # If msg is None, await the next message to arrive on the socket. assert self.session is not None if self._supports_kernel_subshells: assert threading.current_thread() not in ( self.control_thread, self.shell_channel_thread, ) + assert socket is not None else: assert threading.current_thread() == threading.main_thread() + assert socket is None + socket = self.shell_socket no_msg = msg is None if self._is_test else not await socket.poll(0) msg = msg or await socket.recv_multipart(copy=False) From 9109cc8154b9e9ea800c065b4048680b8b43bae4 Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Wed, 2 Oct 2024 13:56:31 +0100 Subject: [PATCH 35/37] Improved use of thread names --- ipykernel/control.py | 4 +--- ipykernel/shellchannel.py | 4 +--- ipykernel/subshell_manager.py | 17 +++++++++-------- ipykernel/thread.py | 3 +++ 4 files changed, 14 insertions(+), 14 deletions(-) diff --git a/ipykernel/control.py b/ipykernel/control.py index 8fc04b43..21d6d996 100644 --- a/ipykernel/control.py +++ b/ipykernel/control.py @@ -1,8 +1,6 @@ """A thread for a control channel.""" -from .thread import BaseThread - -CONTROL_THREAD_NAME = "Control" +from .thread import CONTROL_THREAD_NAME, BaseThread class ControlThread(BaseThread): diff --git a/ipykernel/shellchannel.py b/ipykernel/shellchannel.py index eeaead7b..1935f062 100644 --- a/ipykernel/shellchannel.py +++ b/ipykernel/shellchannel.py @@ -1,8 +1,6 @@ """A thread for a shell channel.""" from .subshell_manager import SubshellManager -from .thread import BaseThread - -SHELL_CHANNEL_THREAD_NAME = "Shell channel" +from .thread import SHELL_CHANNEL_THREAD_NAME, BaseThread class ShellChannelThread(BaseThread): diff --git a/ipykernel/subshell_manager.py b/ipykernel/subshell_manager.py index 1f8a6199..7c64190e 100644 --- a/ipykernel/subshell_manager.py +++ b/ipykernel/subshell_manager.py @@ -14,6 +14,7 @@ from anyio import create_memory_object_stream, create_task_group from .subshell import SubshellThread +from .thread import SHELL_CHANNEL_THREAD_NAME @dataclass @@ -60,7 +61,7 @@ def __init__(self, context: zmq.asyncio.Context, shell_socket: zmq.asyncio.Socke def close(self) -> None: """Stop all subshells and close all resources.""" - assert current_thread().name == "Shell channel" + assert current_thread().name == SHELL_CHANNEL_THREAD_NAME self._send_stream.close() self._receive_stream.close() @@ -119,7 +120,7 @@ async def listen_from_control(self, subshell_task) -> None: """Listen for messages on the control inproc socket, handle those messages and return replies on the same socket. Runs in the shell channel thread. """ - assert current_thread().name == "Shell channel" + assert current_thread().name == SHELL_CHANNEL_THREAD_NAME socket = self._control_shell_channel_socket while True: @@ -133,7 +134,7 @@ async def listen_from_subshells(self) -> None: Runs in the shell channel thread. """ - assert current_thread().name == "Shell channel" + assert current_thread().name == SHELL_CHANNEL_THREAD_NAME async with create_task_group() as tg: tg.start_soon(self._listen_for_subshell_reply, None) @@ -170,7 +171,7 @@ def _create_inproc_pair_socket( async def _create_subshell(self, subshell_task) -> str: """Create and start a new subshell thread.""" - assert current_thread().name == "Shell channel" + assert current_thread().name == SHELL_CHANNEL_THREAD_NAME subshell_id = str(uuid.uuid4()) thread = SubshellThread(subshell_id) @@ -196,7 +197,7 @@ def _delete_subshell(self, subshell_id: str) -> None: Raises key error if subshell_id not in cache. """ - assert current_thread().name == "Shell channel" + assert current_thread().name == SHELL_CHANNEL_THREAD_NAME with self._lock_cache: subshell = self._cache.pop(subshell_id) @@ -225,7 +226,7 @@ async def _listen_for_subshell_reply(self, subshell_id: str | None) -> None: Runs in the shell channel thread. """ - assert current_thread().name == "Shell channel" + assert current_thread().name == SHELL_CHANNEL_THREAD_NAME shell_channel_socket = self._get_shell_channel_socket(subshell_id) @@ -244,7 +245,7 @@ async def _process_control_request(self, request, subshell_task) -> dict[str, t. """Process a control request message received on the control inproc socket and return the reply. Runs in the shell channel thread. """ - assert current_thread().name == "Shell channel" + assert current_thread().name == SHELL_CHANNEL_THREAD_NAME try: type = request["type"] @@ -269,7 +270,7 @@ async def _process_control_request(self, request, subshell_task) -> dict[str, t. def _stop_subshell(self, subshell: Subshell) -> None: """Stop a subshell thread and close all of its resources.""" - assert current_thread().name == "Shell channel" + assert current_thread().name == SHELL_CHANNEL_THREAD_NAME thread = subshell.thread if thread.is_alive(): diff --git a/ipykernel/thread.py b/ipykernel/thread.py index c41db727..7efcabce 100644 --- a/ipykernel/thread.py +++ b/ipykernel/thread.py @@ -4,6 +4,9 @@ from anyio import create_task_group, run, to_thread +CONTROL_THREAD_NAME = "Control" +SHELL_CHANNEL_THREAD_NAME = "Shell channel" + class BaseThread(Thread): """Base class for threads.""" From d17c77f453f69ef3de963ed4783c8165ebcba49a Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Wed, 2 Oct 2024 14:05:59 +0100 Subject: [PATCH 36/37] Small changes --- ipykernel/iostream.py | 8 ++++---- ipykernel/kernelbase.py | 18 ++++++++---------- ipykernel/subshell_manager.py | 8 ++++---- ipykernel/zmqshell.py | 9 +++++---- 4 files changed, 21 insertions(+), 22 deletions(-) diff --git a/ipykernel/iostream.py b/ipykernel/iostream.py index b83f8120..beca44b1 100644 --- a/ipykernel/iostream.py +++ b/ipykernel/iostream.py @@ -170,10 +170,10 @@ async def _handle_event(self): for _ in range(n_events): event_f = self._events.popleft() event_f() - except Exception as e: + except Exception: if self.thread.__stop.is_set(): return - raise e + raise def _setup_pipe_in(self): """setup listening pipe for IOPub from forked subprocesses""" @@ -202,10 +202,10 @@ async def _handle_pipe_msgs(self): try: while True: await self._handle_pipe_msg() - except Exception as e: + except Exception: if self.thread.__stop.is_set(): return - raise e + raise async def _handle_pipe_msg(self, msg=None): """handle a pipe message from a subprocess""" diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index 157e2dc2..84db8660 100644 --- a/ipykernel/kernelbase.py +++ b/ipykernel/kernelbase.py @@ -262,10 +262,10 @@ async def process_control(self): try: while True: await self.process_control_message() - except BaseException as e: + except BaseException: if self.control_stop.is_set(): return - raise e + raise async def process_control_message(self, msg=None): """dispatch control requests""" @@ -392,10 +392,10 @@ async def shell_channel_thread_main(self): socket.send_multipart(msg, copy=False) except Exception: self.log.error("Invalid message", exc_info=True) # noqa: G201 - except BaseException as e: + except BaseException: if self.shell_stop.is_set(): return - raise e + raise async def shell_main(self, subshell_id: str | None): """Main loop for a single subshell.""" @@ -426,10 +426,10 @@ async def process_shell(self, socket=None): try: while True: await self.process_shell_message(socket=socket) - except BaseException as e: + except BaseException: if self.shell_stop.is_set(): return - raise e + raise async def process_shell_message(self, msg=None, socket=None): # If socket is None kernel subshells are not supported so use socket=shell_socket. @@ -718,8 +718,7 @@ async def execute_request(self, socket, ident, parent): cell_meta = parent.get("metadata", {}) cell_id = cell_meta.get("cellId") except Exception: - self.log.error("Got bad msg: ") - self.log.error("%s", parent) + self.log.error("Got bad msg from parent: %s", parent) return stop_on_error = content.get("stop_on_error", True) @@ -1101,8 +1100,7 @@ async def delete_subshell_request(self, socket, ident, parent): content = parent["content"] subshell_id = content["subshell_id"] except Exception: - self.log.error("Got bad msg: ") - self.log.error("%s", parent) + self.log.error("Got bad msg from parent: %s", parent) return # This should only be called in the control thread if it exists. diff --git a/ipykernel/subshell_manager.py b/ipykernel/subshell_manager.py index 7c64190e..3e4a9fd9 100644 --- a/ipykernel/subshell_manager.py +++ b/ipykernel/subshell_manager.py @@ -154,7 +154,7 @@ def subshell_id_from_thread_id(self, thread_id) -> str | None: for id, subshell in self._cache.items(): if subshell.thread.ident == thread_id: return id - msg = f"Thread id '{thread_id} does not correspond to a subshell of this kernel" + msg = f"Thread id {thread_id!r} does not correspond to a subshell of this kernel" raise RuntimeError(msg) def _create_inproc_pair_socket( @@ -235,11 +235,11 @@ async def _listen_for_subshell_reply(self, subshell_id: str | None) -> None: msg = await shell_channel_socket.recv_multipart(copy=False) with self._lock_shell_socket: await self._shell_socket.send_multipart(msg) - except BaseException as e: + except BaseException: if not self._is_subshell(subshell_id): # Subshell no longer exists so exit gracefully return - raise e + raise async def _process_control_request(self, request, subshell_task) -> dict[str, t.Any]: """Process a control request message received on the control inproc @@ -259,7 +259,7 @@ async def _process_control_request(self, request, subshell_task) -> dict[str, t. elif type == "list": reply["subshell_id"] = self.list_subshell() else: - msg = f"Unrecognised message type {type}" + msg = f"Unrecognised message type {type!r}" raise RuntimeError(msg) except BaseException as err: reply = { diff --git a/ipykernel/zmqshell.py b/ipykernel/zmqshell.py index 6c1273cd..3f97e817 100644 --- a/ipykernel/zmqshell.py +++ b/ipykernel/zmqshell.py @@ -16,9 +16,9 @@ import os import sys +import threading import warnings from pathlib import Path -from threading import local from IPython.core import page, payloadpage from IPython.core.autocall import ZMQExitAutocall @@ -69,7 +69,7 @@ def _flush_streams(self): @default("_thread_local") def _default_thread_local(self): """Initialize our thread local storage""" - return local() + return threading.local() @property def _hooks(self): @@ -441,6 +441,9 @@ def autosave(self, arg_s): @line_magic def subshell(self, arg_s): + """ + List all current subshells + """ from ipykernel.kernelapp import IPKernelApp if not IPKernelApp.initialized(): @@ -454,8 +457,6 @@ def subshell(self, arg_s): print("Kernel does not support subshells") return - import threading - thread_id = threading.current_thread().ident manager = kernel.shell_channel_thread.manager try: From 5715dd0993816d8aff463e706ae7a9266ac25a69 Mon Sep 17 00:00:00 2001 From: Ian Thomas Date: Wed, 2 Oct 2024 14:15:39 +0100 Subject: [PATCH 37/37] More type annotations --- ipykernel/kernelbase.py | 10 +++++----- ipykernel/shellchannel.py | 8 +++++--- ipykernel/subshell.py | 4 ++-- ipykernel/subshell_manager.py | 10 ++++++---- ipykernel/thread.py | 6 +++--- 5 files changed, 21 insertions(+), 17 deletions(-) diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index 84db8660..99358f9b 100644 --- a/ipykernel/kernelbase.py +++ b/ipykernel/kernelbase.py @@ -18,7 +18,7 @@ from datetime import datetime from signal import SIGINT, SIGTERM, Signals -from .control import CONTROL_THREAD_NAME +from .thread import CONTROL_THREAD_NAME if sys.platform != "win32": from signal import SIGKILL @@ -898,7 +898,7 @@ def kernel_info(self): "supported_features": [], } if self._supports_kernel_subshells: - info["supported_features"].append("kernel subshells") + info["supported_features"] = ["kernel subshells"] return info async def kernel_info_request(self, socket, ident, parent): @@ -1074,7 +1074,7 @@ async def do_debug_request(self, msg): # Subshell control message handlers # --------------------------------------------------------------------------- - async def create_subshell_request(self, socket, ident, parent): + async def create_subshell_request(self, socket, ident, parent) -> None: if not self.session: return if not self._supports_kernel_subshells: @@ -1089,7 +1089,7 @@ async def create_subshell_request(self, socket, ident, parent): self.session.send(socket, "create_subshell_reply", reply, parent, ident) - async def delete_subshell_request(self, socket, ident, parent): + async def delete_subshell_request(self, socket, ident, parent) -> None: if not self.session: return if not self._supports_kernel_subshells: @@ -1111,7 +1111,7 @@ async def delete_subshell_request(self, socket, ident, parent): self.session.send(socket, "delete_subshell_reply", reply, parent, ident) - async def list_subshell_request(self, socket, ident, parent): + async def list_subshell_request(self, socket, ident, parent) -> None: if not self.session: return if not self._supports_kernel_subshells: diff --git a/ipykernel/shellchannel.py b/ipykernel/shellchannel.py index 1935f062..bc0459c4 100644 --- a/ipykernel/shellchannel.py +++ b/ipykernel/shellchannel.py @@ -1,4 +1,6 @@ """A thread for a shell channel.""" +import zmq.asyncio + from .subshell_manager import SubshellManager from .thread import SHELL_CHANNEL_THREAD_NAME, BaseThread @@ -9,7 +11,7 @@ class ShellChannelThread(BaseThread): Communicates with shell/subshell threads via pairs of ZMQ inproc sockets. """ - def __init__(self, context, shell_socket, **kwargs): + def __init__(self, context: zmq.asyncio.Context, shell_socket: zmq.asyncio.Socket, **kwargs): """Initialize the thread.""" super().__init__(name=SHELL_CHANNEL_THREAD_NAME, **kwargs) self._manager: SubshellManager | None = None @@ -17,13 +19,13 @@ def __init__(self, context, shell_socket, **kwargs): self._shell_socket = shell_socket @property - def manager(self): + def manager(self) -> SubshellManager: # Lazy initialisation. if self._manager is None: self._manager = SubshellManager(self._context, self._shell_socket) return self._manager - def run(self): + def run(self) -> None: """Run the thread.""" try: super().run() diff --git a/ipykernel/subshell.py b/ipykernel/subshell.py index 19aa66c4..18e15ab3 100644 --- a/ipykernel/subshell.py +++ b/ipykernel/subshell.py @@ -17,7 +17,7 @@ def __init__(self, subshell_id: str, **kwargs): # Inproc PAIR socket, for communication with shell channel thread. self._pair_socket: zmq.asyncio.Socket | None = None - async def create_pair_socket(self, context: zmq.asyncio.Context, address: str): + async def create_pair_socket(self, context: zmq.asyncio.Context, address: str) -> None: """Create inproc PAIR socket, for communication with shell channel thread. Should be called from this thread, so usually via add_task before the @@ -27,7 +27,7 @@ async def create_pair_socket(self, context: zmq.asyncio.Context, address: str): self._pair_socket = context.socket(zmq.PAIR) self._pair_socket.connect(address) - def run(self): + def run(self) -> None: try: super().run() finally: diff --git a/ipykernel/subshell_manager.py b/ipykernel/subshell_manager.py index 3e4a9fd9..805d6f81 100644 --- a/ipykernel/subshell_manager.py +++ b/ipykernel/subshell_manager.py @@ -116,7 +116,7 @@ def list_subshell(self) -> list[str]: with self._lock_cache: return list(self._cache) - async def listen_from_control(self, subshell_task) -> None: + async def listen_from_control(self, subshell_task: t.Any) -> None: """Listen for messages on the control inproc socket, handle those messages and return replies on the same socket. Runs in the shell channel thread. """ @@ -141,7 +141,7 @@ async def listen_from_subshells(self) -> None: async for subshell_id in self._receive_stream: tg.start_soon(self._listen_for_subshell_reply, subshell_id) - def subshell_id_from_thread_id(self, thread_id) -> str | None: + def subshell_id_from_thread_id(self, thread_id: int) -> str | None: """Return subshell_id of the specified thread_id. Raises RuntimeError if thread_id is not the main shell or a subshell. @@ -169,7 +169,7 @@ def _create_inproc_pair_socket( socket.connect(address) return socket - async def _create_subshell(self, subshell_task) -> str: + async def _create_subshell(self, subshell_task: t.Any) -> str: """Create and start a new subshell thread.""" assert current_thread().name == SHELL_CHANNEL_THREAD_NAME @@ -241,7 +241,9 @@ async def _listen_for_subshell_reply(self, subshell_id: str | None) -> None: return raise - async def _process_control_request(self, request, subshell_task) -> dict[str, t.Any]: + async def _process_control_request( + self, request: dict[str, t.Any], subshell_task: t.Any + ) -> dict[str, t.Any]: """Process a control request message received on the control inproc socket and return the reply. Runs in the shell channel thread. """ diff --git a/ipykernel/thread.py b/ipykernel/thread.py index 7efcabce..a63011de 100644 --- a/ipykernel/thread.py +++ b/ipykernel/thread.py @@ -17,9 +17,9 @@ def __init__(self, **kwargs): self.pydev_do_not_trace = True self.is_pydev_daemon_thread = True self.__stop = Event() - self._tasks_and_args: t.List[t.Tuple[t.Callable, t.Tuple]] = [] + self._tasks_and_args: t.List[t.Tuple[t.Any, t.Any]] = [] - def add_task(self, task: t.Callable, *args: t.Tuple): + def add_task(self, task: t.Any, *args: t.Any) -> None: # May only add tasks before the thread is started. self._tasks_and_args.append((task, args)) @@ -34,7 +34,7 @@ async def _main(self) -> None: await to_thread.run_sync(self.__stop.wait) tg.cancel_scope.cancel() - def stop(self): + def stop(self) -> None: """Stop the thread. This method is threadsafe.