diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index 7460cbce0..b46cc6600 100644 --- a/ipykernel/kernelbase.py +++ b/ipykernel/kernelbase.py @@ -59,6 +59,7 @@ from ._version import kernel_protocol_version from .iostream import OutStream +from .subshell_manager import UnknownSubshellError from .utils import LazyDict, _async_in_context psutil: t.Any | None = None @@ -594,14 +595,18 @@ async def shell_channel_thread_main(self, msg): # deserialize only the header to get subshell_id # Keep original message to send to subshell_id unmodified. - _, msg2 = self.session.feed_identities(msg, copy=False) + idents, msg2 = self.session.feed_identities(msg, copy=False) try: msg3 = self.session.deserialize(msg2, content=False, copy=False) subshell_id = msg3["header"].get("subshell_id") # Find inproc pair socket to use to send message to correct subshell. subshell_manager = self.shell_channel_thread.manager - socket = subshell_manager.get_shell_channel_to_subshell_socket(subshell_id) + try: + socket = subshell_manager.get_shell_channel_to_subshell_socket(subshell_id) + except UnknownSubshellError as err: + self._send_unknown_subshell_reply(idents, msg3, err) + return assert socket is not None socket.send_multipart(msg, copy=False) except Exception: @@ -1378,6 +1383,41 @@ def _send_abort_reply(self, stream, msg, idents): ident=idents, ) + def _send_unknown_subshell_reply(self, idents, msg, err: UnknownSubshellError) -> None: + """Send an error reply to a request addressed to a subshell that is not there. + + Runs in the shell channel thread, so it writes to the shell socket + directly instead of going through a subshell. + + The busy and idle status messages matter as much as the reply here. + A client tracks the completion of a request by the idle status that + carries it as parent, and for message types that have no reply, such as + the comm messages, that status is all it has to go on. + """ + if not self.session: + return + msg_type = msg["header"]["msg_type"] + self.log.warning("Cannot handle %s %s: %s", msg_type, msg["header"]["msg_id"], err) + self._publish_status("busy", "shell", parent=msg) + content = { + "status": "error", + "ename": type(err).__name__, + "evalue": str(err), + "traceback": [], + } + md = self.init_metadata(msg) + md = self.finish_metadata(msg, md, content) + md.update({"status": "error"}) + self.session.send( + self.shell_stream, + msg_type.rsplit("_", 1)[0] + "_reply", + metadata=md, + content=content, + parent=msg, + ident=idents, + ) + self._publish_status("idle", "shell", parent=msg) + def _no_raw_input(self): """Raise StdinNotImplementedError if active frontend doesn't support stdin.""" diff --git a/ipykernel/subshell_manager.py b/ipykernel/subshell_manager.py index 1f23085ae..ce4acccea 100644 --- a/ipykernel/subshell_manager.py +++ b/ipykernel/subshell_manager.py @@ -19,6 +19,20 @@ from .utils import _async_in_context +class UnknownSubshellError(KeyError): + """A subshell_id that does not name an existing subshell. + + .. versionadded:: 7.4 + """ + + def __init__(self, subshell_id: str) -> None: + super().__init__(subshell_id) + self.subshell_id = subshell_id + + def __str__(self) -> str: + return f"Unknown subshell_id {self.subshell_id!r}" + + class SubshellManager: """A manager of subshells. @@ -87,13 +101,24 @@ def close(self) -> None: self._main_to_shell_channel.close() self._shell_channel_to_main.close() + def _get_subshell(self, subshell_id: str) -> SubshellThread: + """Return the thread of the specified subshell. + + The caller must hold ``_lock_cache``. Raises ``UnknownSubshellError`` if + there is no such subshell. + """ + try: + return self._cache[subshell_id] + except KeyError: + raise UnknownSubshellError(subshell_id) from None + def get_shell_channel_to_subshell_pair(self, subshell_id: str | None) -> SocketPair: """Return the inproc socket pair used to send messages from the shell channel to a particular subshell or main shell.""" if subshell_id is None: return self._shell_channel_to_main with self._lock_cache: - return self._cache[subshell_id].shell_channel_to_subshell + return self._get_subshell(subshell_id).shell_channel_to_subshell def get_subshell_to_shell_channel_socket(self, subshell_id: str | None) -> zmq.Socket[t.Any]: """Return the socket used by a particular subshell or main shell to send @@ -102,7 +127,7 @@ def get_subshell_to_shell_channel_socket(self, subshell_id: str | None) -> zmq.S if subshell_id is None: return self._main_to_shell_channel.from_socket with self._lock_cache: - return self._cache[subshell_id].subshell_to_shell_channel.from_socket + return self._get_subshell(subshell_id).subshell_to_shell_channel.from_socket def get_shell_channel_to_subshell_socket(self, subshell_id: str | None) -> zmq.Socket[t.Any]: """Return the socket used by the shell channel to send messages to a particular @@ -113,12 +138,12 @@ def get_shell_channel_to_subshell_socket(self, subshell_id: str | None) -> zmq.S def get_subshell_aborting(self, subshell_id: str) -> bool: """Get the boolean aborting flag of the specified subshell.""" with self._lock_cache: - return self._cache[subshell_id].aborting + return self._get_subshell(subshell_id).aborting def get_subshell_asyncio_lock(self, subshell_id: str) -> asyncio.Lock: """Return the asyncio lock belonging to the specified subshell.""" with self._lock_cache: - return self._cache[subshell_id].asyncio_lock + return self._get_subshell(subshell_id).asyncio_lock def list_subshell(self) -> list[str]: """Return list of current subshell ids. @@ -141,7 +166,7 @@ def set_on_recv_callback(self, on_recv_callback): def set_subshell_aborting(self, subshell_id: str, aborting: bool) -> None: """Set the aborting flag of the specified subshell.""" with self._lock_cache: - self._cache[subshell_id].aborting = aborting + self._get_subshell(subshell_id).aborting = aborting def subshell_id_from_thread_id(self, thread_id: int) -> str | None: """Return subshell_id of the specified thread_id. @@ -185,14 +210,15 @@ def _create_subshell(self) -> str: def _delete_subshell(self, subshell_id: str) -> None: """Delete subshell identified by subshell_id. - Raises key error if subshell_id not in cache. + Raises ``UnknownSubshellError`` if subshell_id not in cache. """ assert current_thread().name == SHELL_CHANNEL_THREAD_NAME with self._lock_cache: - subshell_threwad = self._cache.pop(subshell_id) + subshell_thread = self._get_subshell(subshell_id) + del self._cache[subshell_id] - self._stop_subshell(subshell_threwad) + self._stop_subshell(subshell_thread) def _process_control_request( self, diff --git a/tests/test_subshells.py b/tests/test_subshells.py index 419c54dbf..a8e1a1abf 100644 --- a/tests/test_subshells.py +++ b/tests/test_subshells.py @@ -391,3 +391,51 @@ def test_silent_flag_in_subshells(): # Ensure subshell is always deleted if subshell_id: delete_subshell_helper(kc, subshell_id) + + +def test_unknown_subshell_id(): + # A request for a subshell that does not exist is answered with an error rather + # than dropped, so that a client waiting for it does not wait forever. + with new_kernel() as kc: + subshell_id = create_subshell_helper(kc)["subshell_id"] + delete_subshell_helper(kc, subshell_id) + + # Deleting it again names the missing subshell in the same way. + content = delete_subshell_helper(kc, subshell_id) + assert content["status"] == "error" + assert content["evalue"] == f"Unknown subshell_id {subshell_id!r}" + flush_channels(kc) + + msg = execute_request(kc, "a = 1", subshell_id) + msg_id = msg["header"]["msg_id"] + + reply = get_reply(kc, msg_id, TIMEOUT) + assert reply["content"]["status"] == "error" + assert reply["content"]["ename"] == "UnknownSubshellError" + assert reply["content"]["evalue"] == f"Unknown subshell_id {subshell_id!r}" + + states = [] + while True: + iopub_msg = kc.get_iopub_msg(timeout=TIMEOUT) + if iopub_msg["parent_header"].get("msg_id") != msg_id: + continue + assert iopub_msg["msg_type"] == "status" + states.append(iopub_msg["content"]["execution_state"]) + if states[-1] == "idle": + break + assert states == ["busy", "idle"] + + +def test_comm_close_on_deleted_subshell(): + # A comm message has no reply of its own, so the idle status is the only thing + # that tells the client the kernel is done with it. + with new_kernel() as kc: + subshell_id = create_subshell_helper(kc)["subshell_id"] + delete_subshell_helper(kc, subshell_id) + flush_channels(kc) + + msg = kc.session.msg("comm_close", {"comm_id": "comm-1", "data": {}}) + msg["header"]["subshell_id"] = subshell_id + kc.shell_channel.send(msg) + + wait_for_idle(kc, msg["header"]["msg_id"])