Skip to content

Fix output routing when threads inherit context #1559

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
emecii wants to merge 1 commit into
base: main
Choose a base branch
from
+52 −28
Open

Fix output routing when threads inherit context #1559

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 16 additions & 10 deletions ipykernel/displayhook.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,8 @@ def __init__(self, session, pub_socket):
self.session = session
self.pub_socket = pub_socket

self._parent_header: ContextVar[dict[str, Any]] = ContextVar("parent_header")
self._parent_header.set({})
self._parent_header: ContextVar[tuple[int, dict[str, Any]]] = ContextVar("parent_header")
self._parent_header.set((threading.get_ident(), {}))
self._parent_header_global = {}

def get_execution_count(self):
Expand Down Expand Up @@ -60,18 +60,21 @@ def __call__(self, obj):
@property
def parent_header(self):
try:
return self._parent_header.get()
thread_id, parent_header = self._parent_header.get()
except LookupError:
return self._parent_header_global
if thread_id != threading.get_ident():
return self._parent_header_global
return parent_header

@parent_header.setter
def parent_header(self, value):
self._parent_header.set(value)
self._parent_header.set((threading.get_ident(), value))
self._parent_header_global = value

def set_thread_parent(self, parent):
"""Set the parent header for the calling thread only. Returns a reset token that can be used with reset_thread_parent."""
return self._parent_header.set(extract_header(parent))
return self._parent_header.set((threading.get_ident(), extract_header(parent)))

def reset_thread_parent(self, token):
"""Reset the parent header to undo the set_thread_parent call that returned the token."""
Expand All @@ -91,14 +94,14 @@ class ZMQShellDisplayHook(DisplayHook):

session = Instance(Session, allow_none=True)
pub_socket = Any(allow_none=True)
_parent_header: ContextVar[dict[str, Any]]
_parent_header: ContextVar[tuple[int, dict[str, Any]]]
_thread_local = 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({})
self._parent_header.set((threading.get_ident(), {}))
self._parent_header_global = {}

@default("_thread_local")
Expand Down Expand Up @@ -131,18 +134,21 @@ def unregister_hook(self, hook):
@property
def parent_header(self):
try:
return self._parent_header.get()
thread_id, parent_header = self._parent_header.get()
except LookupError:
return self._parent_header_global
if thread_id != threading.get_ident():
return self._parent_header_global
return parent_header

@parent_header.setter
def parent_header(self, value):
self._parent_header.set(value)
self._parent_header.set((threading.get_ident(), value))
self._parent_header_global = value

def set_thread_parent(self, parent):
"""Set the parent header for the calling thread only. Returns a reset token that can be used with reset_thread_parent."""
return self._parent_header.set(extract_header(parent))
return self._parent_header.set((threading.get_ident(), extract_header(parent)))

def reset_thread_parent(self, token):
"""Reset the parent header to undo the set_thread_parent call that returned the token."""
Expand Down
17 changes: 11 additions & 6 deletions ipykernel/iostream.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -539,10 +539,10 @@ def __init__(
self.pub_thread = pub_thread
self.name = name
self.topic = b"stream." + name.encode()
self._parent_header: contextvars.ContextVar[dict[str, Any]] = contextvars.ContextVar(
"parent_header"
self._parent_header: contextvars.ContextVar[tuple[int, dict[str, Any]]] = (
contextvars.ContextVar("parent_header")
)
self._parent_header.set({})
self._parent_header.set((threading.get_ident(), {}))
self._parent_header_global = {}
self._master_pid = os.getpid()
self._flush_pending = False
Expand Down Expand Up @@ -598,14 +598,19 @@ def __init__(
def parent_header(self):
try:
# asyncio or thread-specific
return self._parent_header.get()
thread_id, parent_header = self._parent_header.get()
except LookupError:
# global (fallback)
return self._parent_header_global
if thread_id != threading.get_ident():
# Contexts can be inherited by a new thread, but the parent is
# thread-specific unless explicitly set in that thread.
return self._parent_header_global
return parent_header

@parent_header.setter
def parent_header(self, value):
self._parent_header.set(value)
self._parent_header.set((threading.get_ident(), value))
self._parent_header_global = value

def isatty(self):
Expand Down Expand Up @@ -634,7 +639,7 @@ def _is_master_process(self):

def set_thread_parent(self, parent):
"""Set the parent header for the calling thread only. Returns a reset token that can be used with reset_thread_parent."""
return self._parent_header.set(extract_header(parent))
return self._parent_header.set((threading.get_ident(), extract_header(parent)))

def reset_thread_parent(self, token):
"""Reset the parent header to undo the set_thread_parent call that returned the token."""
Expand Down
33 changes: 22 additions & 11 deletions ipykernel/zmqshell.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ class ZMQDisplayPublisher(DisplayPublisher):

session = Instance(Session, allow_none=True)
pub_socket = Any(allow_none=True)
_parent_header: contextvars.ContextVar[dict[str, Any]]
_parent_header: contextvars.ContextVar[tuple[int, dict[str, Any]]]
topic = CBytes(b"display_data")

store_display_history = Bool(
Expand All @@ -73,24 +73,27 @@ class ZMQDisplayPublisher(DisplayPublisher):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._parent_header = contextvars.ContextVar("parent_header")
self._parent_header.set({})
self._parent_header.set((threading.get_ident(), {}))
self._parent_header_global = {}

@property
def parent_header(self):
try:
return self._parent_header.get()
thread_id, parent_header = self._parent_header.get()
except LookupError:
return self._parent_header_global
if thread_id != threading.get_ident():
return self._parent_header_global
return parent_header

@parent_header.setter
def parent_header(self, value):
self._parent_header.set(value)
self._parent_header.set((threading.get_ident(), value))
self._parent_header_global = value

def set_thread_parent(self, parent):
"""Set the parent header for the calling thread only. Returns a reset token that can be used with reset_thread_parent."""
return self._parent_header.set(extract_header(parent))
return self._parent_header.set((threading.get_ident(), extract_header(parent)))

def reset_thread_parent(self, token):
"""Reset the parent header to undo the set_thread_parent call that returned the token."""
Expand Down Expand Up @@ -546,10 +549,10 @@ 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[dict[str, typing.Any]] = contextvars.ContextVar(
"parent_header"
self._parent_header: contextvars.ContextVar[tuple[int, dict[str, typing.Any]]] = (
contextvars.ContextVar("parent_header")
)
self._parent_header.set({})
self._parent_header.set((threading.get_ident(), {}))
self._parent_header_global = {}

displayhook_class = Type(ZMQShellDisplayHook)
Expand Down Expand Up @@ -732,13 +735,16 @@ def set_next_input(self, text, replace=False):
@property
def parent_header(self):
try:
return self._parent_header.get()
thread_id, parent_header = self._parent_header.get()
except LookupError:
return self._parent_header_global
if thread_id != threading.get_ident():
return self._parent_header_global
return parent_header

@parent_header.setter
def parent_header(self, value):
self._parent_header.set(value)
self._parent_header.set((threading.get_ident(), value))
self._parent_header_global = value

def set_parent(self, parent):
Expand All @@ -764,7 +770,12 @@ def get_parent(self):

def set_thread_parent(self, parent):
"""Set the parent header for only the current thread associating output with its triggering input"""
tokens = [(self._parent_header.reset, self._parent_header.set(parent))]
tokens = [
(
self._parent_header.reset,
self._parent_header.set((threading.get_ident(), parent)),
)
]
objs = [self.displayhook, self.display_pub, sys.stdout, sys.stderr]
if hasattr(self, "_data_pub"):
objs.append(self.data_pub)
Expand Down
4 changes: 3 additions & 1 deletion tests/test_kernel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,7 @@ def test_print_to_correct_cell_from_thread(explicit_parent: str):
get_ipython().set_thread_parent sets the thread-local parent for only the thread.
"""
code = f"""\
from contextvars import copy_context
from threading import Event, Thread
from time import sleep
from IPython.display import display
Expand Down Expand Up @@ -162,7 +163,8 @@ def thread_target():
print("after", flush=True)
display(3)

thread = Thread(target=thread_target)
# Free-threaded Python inherits the caller's context by default.
thread = Thread(target=copy_context().run, args=(thread_target,))
thread.start()
"""
outputs = {}
Expand Down