fix(notebook): keep the trainer's stdout out of notebook cell outputs

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):

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.