Cache separate headers on subshell threads (#1414)

Author: ianthomas23

.github/workflows/downstream.yml

@@ -41,6 +41,7 @@ jobs:
           test_command: pytest -vv -raXxs -W default --durations 10 --color=yes
 
   jupyter_client:
+    if: false
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
@@ -55,6 +56,7 @@ jobs:
           package_name: jupyter_client
 
   ipyparallel:
+    if: false
     runs-on: ubuntu-latest
     timeout-minutes: 20
     steps:

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):
@@ -302,6 +305,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."""
         # Ensure only one control message is processed at a time
@@ -737,8 +748,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.
@@ -763,7 +778,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,
@@ -1424,7 +1441,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,
         )
@@ -1441,7 +1458,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/conftest.py

@@ -4,7 +4,7 @@
 from typing import no_type_check
 from unittest.mock import MagicMock
 
-import pytest
+import pytest_asyncio
 import zmq
 from jupyter_client.session import Session
 from tornado.ioloop import IOLoop
@@ -143,15 +143,15 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
 
-@pytest.fixture()
+@pytest_asyncio.fixture()
 def kernel():
     kernel = MockKernel()
     kernel.io_loop = IOLoop.current()
     yield kernel
     kernel.destroy()
 
 
-@pytest.fixture()
+@pytest_asyncio.fixture()
 def ipkernel():
     kernel = MockIPyKernel()
     kernel.io_loop = IOLoop.current()

tests/test_subshells.py

@@ -6,11 +6,12 @@
 
 import platform
 import time
+from collections import Counter
 
 import pytest
 from jupyter_client.blocking.client import BlockingKernelClient
 
-from .utils import TIMEOUT, assemble_output, get_replies, get_reply, new_kernel
+from .utils import TIMEOUT, assemble_output, get_replies, get_reply, new_kernel, wait_for_idle
 
 # Helpers
 
@@ -258,3 +259,57 @@ 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
+        ]
+
+        # Wait for all idle status messages to be received.
+        for _ in range(1 + sum(are_subshells)):
+            wait_for_idle(kc)
+
+        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)

tests/test_zmq_shell.py

@@ -245,7 +245,7 @@ def test_zmq_interactive_shell(kernel):
         shell.data_pub
     shell.kernel = kernel
     shell.set_next_input("hi")
-    assert shell.get_parent() is None
+    assert shell.get_parent() == {}
     if os.name == "posix":
         shell.system_piped("ls")
     else: