From 028152cfbe118fd60ddd66195f41be708424d266 Mon Sep 17 00:00:00 2001 From: GuillaumePELLUET Date: Wed, 9 Sep 2026 17:22:49 +0200 Subject: [PATCH] fix(notebook): keep the trainer's stdout out of notebook cell outputs The studio notebook runs an embedded ipykernel inside the trainer's own process, and IPKernelApp.initialize() swaps sys.stdout/sys.stderr process-wide for an OutStream that publishes every write to iopub. So the training loop's tqdm bar -- a different thread, writing continuously -- surfaced in whatever cell was last executed ("Training: 193497 steps ... train_loss=1.4612" in a cell that never asked for it). The legacy in-process kernel leaked the same way: contextlib.redirect_stdout is process-global too. Route the streams per write instead of per process: * _ThreadRoutedStream wraps ipykernel's OutStreams -- the thread currently running a cell reaches the kernel stream, every other thread (at any time, including while the kernel is idle) gets the console stream the process had before the kernel existed. Ownership comes from the pre_execute/post_execute hooks, which run on the real execution thread, so nothing here assumes which thread ipykernel picked for the shell channel. * capture_fd_output=False, or ipykernel's fd 1/2 pipe would re-capture exactly the writes that were just routed back to the terminal. The cost is that output written straight to the fds by C extensions no longer reaches the notebook -- for an in-process kernel sharing a terminal with the trainer, that is the better trade. * _LiveStream (legacy kernel) got the same thread check, falling back to the pre-redirect stream so those writes still reach the console. Two tests: one in the shared contract class, so both kernels are held to it (a background thread's output must not appear in a cell, while the cell's own print still does), and a legacy-only one proving the other thread's writes reach the console rather than being dropped. Known consequence: output from a thread a cell itself spawns now goes to the terminal, not the cell. Co-Authored-By: Claude Opus 5 (1M context) --- .../services/test_notebook_service_unit.py | 62 ++++++++ .../trainer/services/notebook_service.py | 149 +++++++++++++++++- 2 files changed, 204 insertions(+), 7 deletions(-) diff --git a/tests/trainer/services/test_notebook_service_unit.py b/tests/trainer/services/test_notebook_service_unit.py index b02974f6..b8a6adfa 100644 --- a/tests/trainer/services/test_notebook_service_unit.py +++ b/tests/trainer/services/test_notebook_service_unit.py @@ -5,6 +5,8 @@ style of test_agent_service_unit.py. """ +import io +import sys import json import time import tempfile @@ -100,6 +102,37 @@ def test_stdout_streams_live_not_buffered_until_cell_finishes(self): f"the end ({total:.3f}s) -- looks buffered, not streamed live", ) + def test_other_threads_output_does_not_land_in_the_cell(self): + # Regression: the kernel shares its process with the trainer, and + # both kernels swap sys.stdout process-wide (redirect_stdout() in the + # legacy one, ipykernel's OutStream in the embedded one) -- so a + # training loop's tqdm bar, its own thread writing the whole time, + # surfaced in the output of whatever cell happened to be running. + stop = threading.Event() + + def _trainer(): + # Resolved per write, exactly as print() does, so this thread sees + # the kernel's swapped stream while a cell is running. + while not stop.is_set(): + print("Training: 193497 steps | train_loss=1.4612", file=sys.stdout) + time.sleep(0.01) + + noise = threading.Thread(target=_trainer, daemon=True) + noise.start() + try: + chunks = _run( + self.service, + "import time\nprint('cell-own-output')\ntime.sleep(0.3)", + ) + finally: + stop.set() + noise.join(timeout=2) + + outs = "".join(c.stdout for c in chunks if c.WhichOneof("payload") == "stdout") + self.assertIn("cell-own-output", outs, "the cell's own print was lost") + self.assertNotIn("Training:", outs, + "another thread's output leaked into the cell") + def test_interrupt_reports_false_when_nothing_running(self): resp = self.service.InterruptNotebookCell(pb2.InterruptNotebookCellRequest(), None) self.assertFalse(resp.ok) @@ -310,6 +343,35 @@ class TestNotebookKernelLegacy(_NotebookKernelContractTests, unittest.TestCase): def _make_service(self): return NotebookService(_fake_data_service(), root_log_dir=str(self.root)) + def test_other_threads_output_still_reaches_the_console(self): + # The flip side of the contract test above: writes the cell may not + # publish are handed to the stream that was in place before + # redirect_stdout(), not dropped on the floor -- the trainer's logs + # must keep showing up in the terminal while a cell runs. + console = io.StringIO() + stop = threading.Event() + + def _trainer(): + while not stop.is_set(): + print("Training: 193497 steps", file=sys.stdout) + time.sleep(0.01) + + real_stdout = sys.stdout + sys.stdout = console + noise = threading.Thread(target=_trainer, daemon=True) + noise.start() + try: + chunks = _run(self.service, "import time\ntime.sleep(0.3)") + finally: + stop.set() + noise.join(timeout=2) + sys.stdout = real_stdout + + outs = "".join(c.stdout for c in chunks if c.WhichOneof("payload") == "stdout") + self.assertNotIn("Training:", outs) + self.assertIn("Training:", console.getvalue(), + "the other thread's output never reached the console") + @unittest.skipUnless(_IPYKERNEL_AVAILABLE, "ipykernel/jupyter_client not installed") class TestNotebookKernelEmbedded(_NotebookKernelContractTests, unittest.TestCase): diff --git a/weightslab/trainer/services/notebook_service.py b/weightslab/trainer/services/notebook_service.py index de0d2e54..5679d741 100644 --- a/weightslab/trainer/services/notebook_service.py +++ b/weightslab/trainer/services/notebook_service.py @@ -28,6 +28,7 @@ import os import re import ast +import sys import json import time import ctypes @@ -409,6 +410,87 @@ def get_embedded_kernel_connection_file(wait_timeout: float = 8.0): return _EMBED_STATE["connection_file"] +# Ident of the thread currently running a notebook cell, or None while the +# kernel is idle. Set/cleared by the pre_execute/post_execute hooks, which run +# on the kernel's own execution thread -- so this needs no assumption about +# which thread ipykernel picked for the shell channel. +_CELL_THREAD = {"ident": None} + + +class _ThreadRoutedStream: + """stdout/stderr proxy that lets only the *cell's own* thread reach the + notebook; every other thread keeps writing to the real console. + + The embedded kernel shares its process with the trainer, and ipykernel's + ``init_io()`` swaps ``sys.stdout``/``sys.stderr`` process-wide for an + OutStream that ships everything to iopub. So a training loop's tqdm bar -- + another thread entirely, writing continuously -- landed in whichever cell + was last executed ("Training: 193497 steps ... train_loss=..." showing up + in a cell that never asked for it). Routing happens per write instead: + while a cell runs, its own thread reaches the kernel stream; anything else, + at any time, goes to the stream the process would have had without a + kernel. + """ + + def __init__(self, kernel_stream, console_stream): + self._kernel = kernel_stream + self._console = console_stream + + def _target(self): + ident = _CELL_THREAD["ident"] + if ident is not None and ident == threading.get_ident(): + return self._kernel + return self._console if self._console is not None else self._kernel + + def write(self, s): + return self._target().write(s) + + def writelines(self, lines): + target = self._target() + for line in lines: + target.write(line) + + def flush(self): + for stream in (self._kernel, self._console): + if stream is None: + continue + try: + stream.flush() + except Exception: # noqa: BLE001 -- a closed console must not break a cell + pass + + # Routed too, not delegated: tqdm asks isatty() once, when it is built, and + # a bar built on the training thread must get the console's answer ( + # refreshes) rather than the kernel OutStream's flat False. + def isatty(self): + try: + return bool(self._target().isatty()) + except Exception: # noqa: BLE001 + return False + + def fileno(self): + return self._target().fileno() + + def writable(self): + return True + + def __getattr__(self, name): + # encoding, errors, buffer, _original_stdstream_copy, ... -- whatever + # ipykernel or a library reaches for beyond the file protocol above. + return getattr(self._kernel, name) + + +def _install_thread_routed_streams(console_stdout, console_stderr) -> None: + """Wrap ipykernel's OutStreams so only cell threads publish to the notebook. + + Call after ``IPKernelApp.initialize()`` (which installs the OutStreams) and + before ``app.start()``. + """ + import sys as _sys + _sys.stdout = _ThreadRoutedStream(_sys.stdout, console_stdout) + _sys.stderr = _ThreadRoutedStream(_sys.stderr, console_stderr) + + def _run_embedded_kernel(connection_file: Path) -> None: import asyncio from ipykernel.kernelapp import IPKernelApp @@ -424,7 +506,19 @@ def _run_embedded_kernel(connection_file: Path) -> None: ns = build_notebook_namespace( _ACTIVE_BINDING["data_service"], _ACTIVE_BINDING["root_log_dir"]) + # The real console streams, grabbed before initialize() swaps them for + # ipykernel's OutStream -- _ThreadRoutedStream hands every non-cell thread + # back to these. + console_stdout, console_stderr = sys.stdout, sys.stderr + app = IPKernelApp.instance(connection_file=str(connection_file), matplotlib="inline") + # Without this, ipykernel replaces fd 1/2 with a pipe it forwards to iopub, + # which would swallow the console writes _ThreadRoutedStream routes back to + # the terminal (and re-publish the trainer's output into a cell anyway). + # The cost is that output written straight to the fds by C extensions no + # longer reaches the notebook -- for an in-process kernel sharing a + # terminal with the trainer, that is the better trade. + app.capture_fd_output = False # IPKernelApp.initialize() installs a SIGINT handler, and signal handlers # can only be installed on the main thread -- which an embedded kernel is # never on. ipykernel catches the resulting ValueError but logs it as @@ -482,6 +576,8 @@ def _run_embedded_kernel(connection_file: Path) -> None: if hasattr(_stream, "flush_interval"): _stream.flush_interval = 0.05 _install_kernel_hooks(app.shell) + # After initialize() (OutStreams exist), before start() (cells run). + _install_thread_routed_streams(console_stdout, console_stderr) logger.info("Embedded Jupyter kernel connection file: %s", connection_file) app.start() # blocks this thread forever (event loop) except Exception: @@ -496,6 +592,9 @@ def _install_kernel_hooks(shell) -> None: box = {"guard_cm": None} def _pre_execute(): + # This hook runs on the thread that executes the cell -- the one + # _ThreadRoutedStream lets through to the notebook. + _CELL_THREAD["ident"] = threading.get_ident() try: shell.user_ns["df"] = get_df(_ACTIVE_BINDING["data_service"]) except Exception: @@ -505,6 +604,7 @@ def _pre_execute(): box["guard_cm"] = cm def _post_execute(): + _CELL_THREAD["ident"] = None cm = box.pop("guard_cm", None) if cm is not None: cm.__exit__(None, None, None) @@ -639,19 +739,46 @@ class _LiveStream: """Write-only file-like object that forwards each write directly to ``emit(kind, text)`` instead of buffering into a StringIO -- lets stdout/ stderr reach the gRPC client as the cell actually prints, rather than only - after the whole cell finishes.""" + after the whole cell finishes. + + Only writes from the kernel worker thread are forwarded. redirect_stdout() + swaps ``sys.stdout`` for the whole process, and this kernel shares its + process with the trainer -- so without the thread check a training loop's + tqdm bar ends up in the output of whatever cell happens to be running. + Other threads keep writing to ``console``, the stream that was in place + before the redirect. + """ - def __init__(self, kind: str, emit): + def __init__(self, kind: str, emit, console=None, owner=None): self._kind = kind self._emit = emit + self._console = console + self._owner = owner if owner is not None else threading.get_ident() def write(self, s): - if s: - self._emit(self._kind, _capped(self._kind, s)) + if not s: + return 0 + if threading.get_ident() != self._owner: + if self._console is not None: + return self._console.write(s) + return len(s) + self._emit(self._kind, _capped(self._kind, s)) return len(s) def flush(self): - pass + if self._console is not None: + try: + self._console.flush() + except Exception: # noqa: BLE001 -- a closed console must not break a cell + pass + + def isatty(self): + if threading.get_ident() != self._owner and self._console is not None: + try: + return bool(self._console.isatty()) + except Exception: # noqa: BLE001 + return False + return False # --------------------------------------------------------------------------- @@ -771,10 +898,18 @@ def _run_on_kernel_thread(self, code: str, emit): except Exception: pass + # Captured before the redirect so _LiveStream can hand other + # threads' writes (the trainer's, typically) back to the console + # instead of publishing them into this cell's output. + console_stdout, console_stderr = sys.stdout, sys.stderr + owner = threading.get_ident() + try: with _WriteGuard.enforce(self._root_log_dir): - with contextlib.redirect_stdout(_LiveStream("stdout", emit)), \ - contextlib.redirect_stderr(_LiveStream("stderr", emit)): + with contextlib.redirect_stdout( + _LiveStream("stdout", emit, console_stdout, owner)), \ + contextlib.redirect_stderr( + _LiveStream("stderr", emit, console_stderr, owner)): result_repr = self._exec_with_last_expr(code) except BaseException: # noqa: BLE001 -- surface any user error (incl. an # interrupt() -injected KeyboardInterrupt) as a cell error, not a crash.