Fix routing of background thread output when no parent is set explicitly

- removes guessing about whether threads should route output to their originating cell
- thread-local value is persisted, if set explicitly
- if no thread-local value is set, fallback on a global value (usually the latest cell)

There's lots of duplication around, but I couldn't see a way to resolve that without breaking things.

Author: minrk

ipykernel/displayhook.py

@@ -29,6 +29,7 @@ def __init__(self, session, pub_socket):
 
         self._parent_header: ContextVar[dict[str, Any]] = ContextVar("parent_header")
         self._parent_header.set({})
+        self._parent_header_global = {}
 
     def get_execution_count(self):
         """This method is replaced in kernelapp"""
@@ -57,11 +58,16 @@ def __call__(self, obj):
 
     @property
     def parent_header(self):
-        return self._parent_header.get()
+        try:
+            return self._parent_header.get()
+        except LookupError:
+            return self._parent_header_global
 
     def set_parent(self, parent):
         """Set the parent header."""
-        self._parent_header.set(extract_header(parent))
+        parent_header = extract_header(parent)
+        self._parent_header.set(parent_header)
+        self._parent_header_global = parent_header
 
 
 class ZMQShellDisplayHook(DisplayHook):
@@ -83,11 +89,16 @@ def __init__(self, *args, **kwargs):
 
     @property
     def parent_header(self):
-        return self._parent_header.get()
+        try:
+            return self._parent_header.get()
+        except LookupError:
+            return self._parent_header_global
 
     def set_parent(self, parent):
-        """Set the parent for outbound messages."""
-        self._parent_header.set(extract_header(parent))
+        """Set the parent header."""
+        parent_header = extract_header(parent)
+        self._parent_header.set(parent_header)
+        self._parent_header_global = parent_header
 
     def start_displayhook(self):
         """Start the display hook."""

ipykernel/iostream.py

@@ -456,8 +456,6 @@ def __init__(
             "parent_header"
         )
         self._parent_header.set({})
-        self._thread_to_parent = {}
-        self._thread_to_parent_header = {}
         self._parent_header_global = {}
         self._master_pid = os.getpid()
         self._flush_pending = False
@@ -512,21 +510,11 @@ def __init__(
     @property
     def parent_header(self):
         try:
-            # asyncio-specific
+            # asyncio or thread-specific
             return self._parent_header.get()
         except LookupError:
-            try:
-                # thread-specific
-                identity = threading.current_thread().ident
-                # retrieve the outermost (oldest ancestor,
-                # discounting the kernel thread) thread identity
-                while identity in self._thread_to_parent:
-                    identity = self._thread_to_parent[identity]
-                # use the header of the oldest ancestor
-                return self._thread_to_parent_header[identity]
-            except KeyError:
-                # global (fallback)
-                return self._parent_header_global
+            # global (fallback)
+            return self._parent_header_global
 
     @parent_header.setter
     def parent_header(self, value):

ipykernel/ipkernel.py

@@ -4,7 +4,6 @@
 
 import asyncio
 import builtins
-import gc
 import getpass
 import os
 import signal
@@ -17,15 +16,13 @@
 import comm
 from IPython.core import release
 from IPython.utils.tokenutil import line_at_cursor, token_at_cursor
-from jupyter_client.session import extract_header
 from traitlets import Any, Bool, HasTraits, Instance, List, Type, default, observe, observe_compat
 from zmq.eventloop.zmqstream import ZMQStream
 
 from .comm.comm import BaseComm
 from .comm.manager import CommManager
 from .compiler import XCachingCompiler
 from .eventloops import _use_appnope
-from .iostream import OutStream
 from .kernelbase import Kernel as KernelBase
 from .kernelbase import _accepts_parameters
 from .zmqshell import ZMQInteractiveShell
@@ -167,14 +164,6 @@ def __init__(self, **kwargs):
 
             appnope.nope()
 
-        self._new_threads_parent_header = {}
-        self._initialize_thread_hooks()
-
-        if hasattr(gc, "callbacks"):
-            # while `gc.callbacks` exists since Python 3.3, pypy does not
-            # implement it even as of 3.9.
-            gc.callbacks.append(self._clean_thread_parent_frames)
-
     help_links = List(
         [
             {
@@ -374,8 +363,6 @@ def _dummy_context_manager(self, *args):
 
     async def execute_request(self, stream, ident, parent):
         """Override for cell output - cell reconciliation."""
-        parent_header = extract_header(parent)
-        self._associate_new_top_level_threads_with(parent_header)
         await super().execute_request(stream, ident, parent)
 
     async def do_execute(
@@ -750,83 +737,6 @@ def do_clear(self):
             self.shell.reset(False)
         return dict(status="ok")
 
-    def _associate_new_top_level_threads_with(self, parent_header):
-        """Store the parent header to associate it with new top-level threads"""
-        self._new_threads_parent_header = parent_header
-
-    def _initialize_thread_hooks(self):
-        """Store thread hierarchy and thread-parent_header associations."""
-        stdout = self._stdout
-        stderr = self._stderr
-        kernel_thread_ident = threading.get_ident()
-        kernel = self
-        _threading_Thread_run = threading.Thread.run
-        _threading_Thread__init__ = threading.Thread.__init__
-
-        def run_closure(self: threading.Thread):
-            """Wrap the `threading.Thread.start` to intercept thread identity.
-
-            This is needed because there is no "start" hook yet, but there
-            might be one in the future: https://bugs.python.org/issue14073
-
-            This is a no-op if the `self._stdout` and `self._stderr` are not
-            sub-classes of `OutStream`.
-            """
-
-            try:
-                parent = self._ipykernel_parent_thread_ident  # type:ignore[attr-defined]
-            except AttributeError:
-                return
-            for stream in [stdout, stderr]:
-                if isinstance(stream, OutStream):
-                    if parent == kernel_thread_ident:
-                        stream._thread_to_parent_header[self.ident] = (
-                            kernel._new_threads_parent_header
-                        )
-                    else:
-                        stream._thread_to_parent[self.ident] = parent
-            _threading_Thread_run(self)
-
-        def init_closure(self: threading.Thread, *args, **kwargs):
-            _threading_Thread__init__(self, *args, **kwargs)
-            self._ipykernel_parent_thread_ident = threading.get_ident()  # type:ignore[attr-defined]
-
-        threading.Thread.__init__ = init_closure  # type:ignore[method-assign]
-        threading.Thread.run = run_closure  # type:ignore[method-assign]
-
-    def _clean_thread_parent_frames(
-        self, phase: t.Literal["start", "stop"], info: dict[str, t.Any]
-    ):
-        """Clean parent frames of threads which are no longer running.
-        This is meant to be invoked by garbage collector callback hook.
-
-        The implementation enumerates the threads because there is no "exit" hook yet,
-        but there might be one in the future: https://bugs.python.org/issue14073
-
-        This is a no-op if the `self._stdout` and `self._stderr` are not
-        sub-classes of `OutStream`.
-        """
-        # Only run before the garbage collector starts
-        if phase != "start":
-            return
-        active_threads = {thread.ident for thread in threading.enumerate()}
-        for stream in [self._stdout, self._stderr]:
-            if isinstance(stream, OutStream):
-                thread_to_parent_header = stream._thread_to_parent_header
-                for identity in list(thread_to_parent_header.keys()):
-                    if identity not in active_threads:
-                        try:
-                            del thread_to_parent_header[identity]
-                        except KeyError:
-                            pass
-                thread_to_parent = stream._thread_to_parent
-                for identity in list(thread_to_parent.keys()):
-                    if identity not in active_threads:
-                        try:
-                            del thread_to_parent[identity]
-                        except KeyError:
-                            pass
-
 
 # This exists only for backwards compatibility - use IPythonKernel instead
 

ipykernel/kernelbase.py

@@ -17,7 +17,7 @@
 import uuid
 import warnings
 from collections.abc import Mapping
-from contextvars import ContextVar
+from contextvars import Context, ContextVar, copy_context
 from datetime import datetime
 from functools import partial
 from signal import SIGINT, SIGTERM, Signals, default_int_handler, signal
@@ -201,6 +201,7 @@ def _default_ident(self):
     _control_parent_ident: bytes = b""
     _shell_parent: ContextVar[dict[str, Any]]
     _shell_parent_ident: ContextVar[bytes]
+    _shell_context: Context
     # Kept for backward-compatibility, accesses _control_parent_ident and _shell_parent_ident,
     # see https://github.com/jupyterlab/jupyterlab/issues/17785
     _parent_ident: Mapping[str, bytes]
@@ -320,13 +321,14 @@ def __init__(self, **kwargs):
         self._shell_parent.set({})
         self._shell_parent_ident = ContextVar("shell_parent_ident")
         self._shell_parent_ident.set(b"")
+        self._shell_context = copy_context()
 
         # For backward compatibility so that _parent_ident["shell"] and _parent_ident["control"]
         # work as they used to for ipykernel >= 7
         self._parent_ident = LazyDict(
             {
                 "control": lambda: self._control_parent_ident,
-                "shell": lambda: self._shell_parent_ident.get(),
+                "shell": lambda: self._get_shell_context_var(self._shell_parent_ident),
             }
         )
 
@@ -768,6 +770,8 @@ def set_parent(self, ident, parent, channel="shell"):
         else:
             self._shell_parent_ident.set(ident)
             self._shell_parent.set(parent)
+            # preserve the last call to set_parent
+            self._shell_context = copy_context()
 
     def get_parent(self, channel=None):
         """Get the parent request associated with a channel.
@@ -794,7 +798,20 @@ def get_parent(self, channel=None):
 
         if channel == "control":
             return self._control_parent
-        return self._shell_parent.get()
+
+        return self._get_shell_context_var(self._shell_parent)
+
+    def _get_shell_context_var(self, var: ContextVar):
+        """Lookup a ContextVar, falling back on the shell context
+
+        Allows for user-launched Threads to still resolve to the shell's mai context
+
+        necessary for e.g. display from threads.
+        """
+        try:
+            return var.get()
+        except LookupError:
+            return self._shell_context[var]
 
     def send_response(
         self,
@@ -1455,7 +1472,7 @@ def getpass(self, prompt="", stream=None):
             )
         return self._input_request(
             prompt,
-            self._shell_parent_ident.get(),
+            self._get_shell_context_var(self._shell_parent_ident),
             self.get_parent("shell"),
             password=True,
         )
@@ -1472,7 +1489,7 @@ def raw_input(self, prompt=""):
             raise StdinNotImplementedError(msg)
         return self._input_request(
             str(prompt),
-            self._shell_parent_ident.get(),
+            self._get_shell_context_var(self._shell_parent_ident),
             self.get_parent("shell"),
             password=False,
         )

ipykernel/zmqshell.py

@@ -63,14 +63,20 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self._parent_header = contextvars.ContextVar("parent_header")
         self._parent_header.set({})
+        self._parent_header_global = {}
 
     @property
     def parent_header(self):
-        return self._parent_header.get()
+        try:
+            return self._parent_header.get()
+        except LookupError:
+            return self._parent_header_global
 
     def set_parent(self, parent):
         """Set the parent for outbound messages."""
-        self._parent_header.set(extract_header(parent))
+        parent_header = extract_header(parent)
+        self._parent_header.set(parent_header)
+        self._parent_header_global = parent_header
 
     def _flush_streams(self):
         """flush IO Streams prior to display"""
@@ -673,11 +679,23 @@ def set_next_input(self, text, replace=False):
 
     @property
     def parent_header(self):
-        return self._parent_header.get()
+        try:
+            return self._parent_header.get()
+        except LookupError:
+            return self._parent_header_global
+
+    @parent_header.setter
+    def parent_header(self, value):
+        self._parent_header_global = value
+        self._parent_header.set(value)
 
     def set_parent(self, parent):
-        """Set the parent header for associating output with its triggering input"""
-        self._parent_header.set(parent)
+        """Set the parent header for associating output with its triggering input
+
+        When called from a thread, sets the thread-local value, which persists
+        until the next call from this thread.
+        """
+        self.parent_header = 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"):
@@ -688,7 +706,12 @@ def set_parent(self, parent):
             sys.stderr.set_parent(parent)
 
     def get_parent(self):
-        """Get the parent header."""
+        """Get the parent header.
+
+        If set_parent has never been called from the current thread,
+        the value from the last call to set_parent from _any_ thread will be used
+        (typically the currently running cell).
+        """
         return self.parent_header
 
     def init_magics(self):