Skip to content

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

Open
guillaume-byte wants to merge 1 commit into
devfrom
fix/notebook-thread-output-leak
Open

fix(notebook): keep the trainer's stdout out of notebook cell outputs#305
guillaume-byte wants to merge 1 commit into
devfrom
fix/notebook-thread-output-leak

Conversation

@guillaume-byte

Copy link
Copy Markdown
Member

Problem

Training logs showed up in notebook cell outputs — a cell that never asked for
them filling with the trainer's tqdm bar:

Training: 193497 steps [2:26:49, 21.08it/s] | train_loss=1.4612 | test_loss=1.4871 | test_acc=97.2%

Why

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,
attributed to whichever cell was last executed. The training loop's tqdm bar —
a different thread, writing continuously — went straight into that stream.

The legacy in-process kernel leaked the same way: contextlib.redirect_stdout
is process-global too.

Fix

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 assumes which thread ipykernel picked for the shell channel.
  • capture_fd_output = False: otherwise ipykernel's fd 1/2 pipe re-captures
    exactly the writes just routed back to the terminal.
  • _LiveStream (legacy kernel) got the same thread check, falling back to the
    pre-redirect stream so those writes still reach the console.

Testing

  • One test 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.
  • One legacy-only test proving the other thread's writes reach the console
    rather than being dropped.
  • Both fail against the pre-fix code (verified by restoring the old write,
    and by stashing the fix: exit 1).
  • tests/trainer/services: 348 passed, 69 skipped. Ruff clean under CI's rule
    set.

Trade-offs

  • Output from a thread a cell itself spawns now goes to the terminal, not
    the cell.
  • Output written straight to fd 1/2 by C extensions no longer reaches the
    notebook. For an in-process kernel sharing a terminal with the trainer,
    that's the better trade.

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) <noreply@anthropic.com>
@guillaume-byte guillaume-byte self-assigned this Sep 9, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant