store and use a reference to the kernel launching thread instead of assuming it's always the main thread

Author: dfalbel

ipykernel/ipkernel.py

@@ -452,7 +452,7 @@ async def run_cell(*args, **kwargs):
 
                 cm = (
                     self._cancel_on_sigint
-                    if threading.current_thread() == threading.main_thread()
+                    if threading.current_thread() == self.shell_channel_thread.parent_thread
                     else self._dummy_context_manager
                 )
                 with cm(coro_future):

ipykernel/kernelbase.py

@@ -677,19 +677,19 @@ async def shell_main(self, subshell_id: str | None, msg):
         """Handler of shell messages for a single subshell"""
         if self._supports_kernel_subshells:
             if subshell_id is None:
-                assert threading.current_thread() == threading.main_thread()
+                assert threading.current_thread() == self.shell_channel_thread.parent_thread
                 asyncio_lock = self._main_asyncio_lock
             else:
                 assert threading.current_thread() not in (
                     self.shell_channel_thread,
-                    threading.main_thread(),
+                    self.shell_channel_thread.parent_thread,
                 )
                 asyncio_lock = self.shell_channel_thread.manager.get_subshell_asyncio_lock(
                     subshell_id
                 )
         else:
             assert subshell_id is None
-            assert threading.current_thread() == threading.main_thread()
+            assert threading.current_thread() == self.shell_channel_thread.parent_thread
             asyncio_lock = self._main_asyncio_lock
 
         # Whilst executing a shell message, do not accept any other shell messages on the

ipykernel/shellchannel.py

@@ -4,6 +4,7 @@
 
 import asyncio
 from typing import Any
+from threading import current_thread
 
 import zmq
 
@@ -28,13 +29,16 @@ def __init__(
         self._manager: SubshellManager | None = None
         self._zmq_context = context  # Avoid use of self._context
         self._shell_socket = shell_socket
+        # Record the parent thread - the thread that started the app (usually the main thread)
+        self.parent_thread = current_thread()
 
         self.asyncio_lock = asyncio.Lock()
 
     @property
     def manager(self) -> SubshellManager:
         # Lazy initialisation.
         if self._manager is None:
+            assert current_thread() == self.parent_thread
             self._manager = SubshellManager(
                 self._zmq_context,
                 self.io_loop,

ipykernel/subshell_manager.py

@@ -7,7 +7,7 @@
 import typing as t
 import uuid
 from functools import partial
-from threading import Lock, current_thread, main_thread
+from threading import Lock, current_thread
 
 import zmq
 from tornado.ioloop import IOLoop
@@ -41,7 +41,7 @@ def __init__(
         shell_socket: zmq.Socket[t.Any],
     ):
         """Initialize the subshell manager."""
-        assert current_thread() == main_thread()
+        self._parent_thread = current_thread()
 
         self._context: zmq.Context[t.Any] = context
         self._shell_channel_io_loop = shell_channel_io_loop
@@ -127,7 +127,7 @@ def set_on_recv_callback(self, on_recv_callback):
         """Set the callback used by the main shell and all subshells to receive
         messages sent from the shell channel thread.
         """
-        assert current_thread() == main_thread()
+        assert current_thread() == self._parent_thread
         self._on_recv_callback = on_recv_callback
         self._shell_channel_to_main.on_recv(IOLoop.current(), partial(self._on_recv_callback, None))
 
@@ -144,7 +144,7 @@ def subshell_id_from_thread_id(self, thread_id: int) -> str | None:
         Only used by %subshell magic so does not have to be fast/cached.
         """
         with self._lock_cache:
-            if thread_id == main_thread().ident:
+            if thread_id == self._parent_thread.ident:
                 return None
             for id, subshell in self._cache.items():
                 if subshell.ident == thread_id: