Cache separate headers on subshell threads

Author: ianthomas23

ipykernel/displayhook.py

@@ -7,10 +7,11 @@
 import builtins
 import sys
 import typing as t
+from contextvars import ContextVar
 
 from IPython.core.displayhook import DisplayHook
 from jupyter_client.session import Session, extract_header
-from traitlets import Any, Dict, Instance
+from traitlets import Any, Instance
 
 from ipykernel.jsonutil import encode_images, json_clean
 
@@ -25,7 +26,9 @@ def __init__(self, session, pub_socket):
         """Initialize the hook."""
         self.session = session
         self.pub_socket = pub_socket
-        self.parent_header = {}
+
+        self._parent_header: ContextVar[dict[str, Any]] = ContextVar("parent_header")
+        self._parent_header.set({})
 
     def get_execution_count(self):
         """This method is replaced in kernelapp"""
@@ -45,12 +48,20 @@ def __call__(self, obj):
             "metadata": {},
         }
         self.session.send(
-            self.pub_socket, "execute_result", contents, parent=self.parent_header, ident=self.topic
+            self.pub_socket,
+            "execute_result",
+            contents,
+            parent=self.parent_header,
+            ident=self.topic,
         )
 
+    @property
+    def parent_header(self):
+        return self._parent_header.get()
+
     def set_parent(self, parent):
         """Set the parent header."""
-        self.parent_header = extract_header(parent)
+        self._parent_header.set(extract_header(parent))
 
 
 class ZMQShellDisplayHook(DisplayHook):
@@ -62,12 +73,21 @@ class ZMQShellDisplayHook(DisplayHook):
 
     session = Instance(Session, allow_none=True)
     pub_socket = Any(allow_none=True)
-    parent_header = Dict({})
+    _parent_header: ContextVar[dict[str, Any]]
     msg: dict[str, t.Any] | None
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._parent_header = ContextVar("parent_header")
+        self._parent_header.set({})
+
+    @property
+    def parent_header(self):
+        return self._parent_header.get()
+
     def set_parent(self, parent):
         """Set the parent for outbound messages."""
-        self.parent_header = extract_header(parent)
+        self._parent_header.set(extract_header(parent))
 
     def start_displayhook(self):
         """Start the display hook."""

ipykernel/kernelbase.py

@@ -16,6 +16,7 @@
 import typing as t
 import uuid
 import warnings
+from contextvars import ContextVar
 from datetime import datetime
 from functools import partial
 from signal import SIGINT, SIGTERM, Signals, default_int_handler, signal
@@ -194,8 +195,10 @@ def _default_ident(self):
 
     # track associations with current request
     _allow_stdin = Bool(False)
-    _parents: Dict[str, t.Any] = Dict({"shell": {}, "control": {}})
-    _parent_ident = Dict({"shell": b"", "control": b""})
+    _control_parent: Dict[str, t.Any] = Dict({})
+    _control_parent_ident: bytes = b""
+    _shell_parent: ContextVar[dict[str, Any]]
+    _shell_parent_ident: ContextVar[bytes]
 
     @property
     def _parent_header(self):
@@ -299,6 +302,14 @@ def __init__(self, **kwargs):
             self.do_execute, ["cell_meta", "cell_id"]
         )
 
+        self._control_parent = {}
+        self._control_parent_ident = b""
+
+        self._shell_parent = ContextVar("shell_parent")
+        self._shell_parent.set({})
+        self._shell_parent_ident = ContextVar("shell_parent_ident")
+        self._shell_parent_ident.set(b"")
+
     async def dispatch_control(self, msg):
         """Dispatch a control request, ensuring only one message is processed at a time."""
         async with asyncio.Lock():
@@ -721,8 +732,12 @@ def set_parent(self, ident, parent, channel="shell"):
         The parent identity is used to route input_request messages
         on the stdin channel.
         """
-        self._parent_ident[channel] = ident
-        self._parents[channel] = parent
+        if channel == "control":
+            self._control_parent_ident = ident
+            self._control_parent = parent
+        else:
+            self._shell_parent_ident.set(ident)
+            self._shell_parent.set(parent)
 
     def get_parent(self, channel=None):
         """Get the parent request associated with a channel.
@@ -747,7 +762,9 @@ def get_parent(self, channel=None):
             else:
                 channel = "shell"
 
-        return self._parents.get(channel, {})
+        if channel == "control":
+            return self._control_parent
+        return self._shell_parent.get()
 
     def send_response(
         self,
@@ -1408,7 +1425,7 @@ def getpass(self, prompt="", stream=None):
             )
         return self._input_request(
             prompt,
-            self._parent_ident["shell"],
+            self._shell_parent_ident.get(),
             self.get_parent("shell"),
             password=True,
         )
@@ -1425,7 +1442,7 @@ def raw_input(self, prompt=""):
             raise StdinNotImplementedError(msg)
         return self._input_request(
             str(prompt),
-            self._parent_ident["shell"],
+            self._shell_parent_ident.get(),
             self.get_parent("shell"),
             password=False,
         )

ipykernel/zmqshell.py

@@ -14,6 +14,7 @@
 # Copyright (c) IPython Development Team.
 # Distributed under the terms of the Modified BSD License.
 
+import contextvars
 import os
 import sys
 import threading
@@ -34,7 +35,7 @@
 from IPython.utils.process import arg_split, system  # type:ignore[attr-defined]
 from jupyter_client.session import Session, extract_header
 from jupyter_core.paths import jupyter_runtime_dir
-from traitlets import Any, CBool, CBytes, Dict, Instance, Type, default, observe
+from traitlets import Any, CBool, CBytes, Instance, Type, default, observe
 
 from ipykernel import connect_qtconsole, get_connection_file, get_connection_info
 from ipykernel.displayhook import ZMQShellDisplayHook
@@ -50,17 +51,26 @@ class ZMQDisplayPublisher(DisplayPublisher):
 
     session = Instance(Session, allow_none=True)
     pub_socket = Any(allow_none=True)
-    parent_header = Dict({})
+    _parent_header: contextvars.ContextVar[dict[str, Any]]
     topic = CBytes(b"display_data")
 
     # thread_local:
     # An attribute used to ensure the correct output message
     # is processed. See ipykernel Issue 113 for a discussion.
     _thread_local = Any()
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._parent_header = contextvars.ContextVar("parent_header")
+        self._parent_header.set({})
+
+    @property
+    def parent_header(self):
+        return self._parent_header.get()
+
     def set_parent(self, parent):
         """Set the parent for outbound messages."""
-        self.parent_header = extract_header(parent)
+        self._parent_header.set(extract_header(parent))
 
     def _flush_streams(self):
         """flush IO Streams prior to display"""
@@ -485,11 +495,14 @@ def __init__(self, *args, **kwargs):
         if "IPKernelApp" not in self.config:
             self.config.IPKernelApp.tqdm = "dummy value for https://github.com/tqdm/tqdm/pull/1628"
 
+        self._parent_header = contextvars.ContextVar("parent_header")
+        self._parent_header.set({})
+
     displayhook_class = Type(ZMQShellDisplayHook)
     display_pub_class = Type(ZMQDisplayPublisher)
     data_pub_class = Any()
     kernel = Any()
-    parent_header = Any()
+    _parent_header: contextvars.ContextVar[dict[str, Any]]
 
     @default("banner1")
     def _default_banner1(self):
@@ -658,9 +671,13 @@ def set_next_input(self, text, replace=False):
         )
         self.payload_manager.write_payload(payload)  # type:ignore[union-attr]
 
+    @property
+    def parent_header(self):
+        return self._parent_header.get()
+
     def set_parent(self, parent):
         """Set the parent header for associating output with its triggering input"""
-        self.parent_header = parent
+        self._parent_header.set(parent)
         self.displayhook.set_parent(parent)  # type:ignore[attr-defined]
         self.display_pub.set_parent(parent)  # type:ignore[attr-defined]
         if hasattr(self, "_data_pub"):

tests/test_subshells.py

@@ -6,6 +6,7 @@
 
 import platform
 import time
+from collections import Counter
 
 import pytest
 from jupyter_client.blocking.client import BlockingKernelClient
@@ -258,3 +259,53 @@ def test_execute_stop_on_error(are_subshells):
         for subshell_id in subshell_ids:
             if subshell_id:
                 delete_subshell_helper(kc, subshell_id)
+
+
+@pytest.mark.parametrize("are_subshells", [(False, True), (True, False), (True, True)])
+def test_idle_message_parent_headers(are_subshells):
+    with new_kernel() as kc:
+        # import time module on main shell.
+        msg = kc.session.msg("execute_request", {"code": "import time"})
+        kc.shell_channel.send(msg)
+
+        subshell_ids = [
+            create_subshell_helper(kc)["subshell_id"] if is_subshell else None
+            for is_subshell in are_subshells
+        ]
+
+        msg_ids = []
+        for subshell_id in subshell_ids:
+            msg = execute_request(kc, "time.sleep(0.5)", subshell_id)
+            msg_ids.append(msg["msg_id"])
+
+        # Expect 4 status messages (2 busy, 2 idle) on iopub channel for the two execute_requests
+        statuses = []
+        timeout = TIMEOUT  # Combined timeout to receive all the status messages
+        t0 = time.time()
+        while True:
+            status = kc.get_iopub_msg(timeout=timeout)
+            if status["msg_type"] != "status" or status["parent_header"]["msg_id"] not in msg_ids:
+                continue
+            statuses.append(status)
+            if len(statuses) == 4:
+                break
+            t1 = time.time()
+            timeout -= t1 - t0
+            t0 = t1
+
+        execution_states = Counter(msg["content"]["execution_state"] for msg in statuses)
+        assert execution_states["busy"] == 2
+        assert execution_states["idle"] == 2
+
+        parent_msg_ids = Counter(msg["parent_header"]["msg_id"] for msg in statuses)
+        assert parent_msg_ids[msg_ids[0]] == 2
+        assert parent_msg_ids[msg_ids[1]] == 2
+
+        parent_subshell_ids = Counter(msg["parent_header"].get("subshell_id") for msg in statuses)
+        assert parent_subshell_ids[subshell_ids[0]] == 2
+        assert parent_subshell_ids[subshell_ids[1]] == 2
+
+        # Cleanup
+        for subshell_id in subshell_ids:
+            if subshell_id:
+                delete_subshell_helper(kc, subshell_id)