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.