fix(notebook): keep the trainer's stdout out of notebook cell outputs - #305
Open
guillaume-byte wants to merge 1 commit into
Open
fix(notebook): keep the trainer's stdout out of notebook cell outputs#305guillaume-byte wants to merge 1 commit into
guillaume-byte wants to merge 1 commit into
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Training logs showed up in notebook cell outputs — a cell that never asked for
them filling with the trainer's tqdm bar:
Why
The studio notebook runs an embedded ipykernel inside the trainer's own
process, and
IPKernelApp.initialize()swapssys.stdout/sys.stderrprocess-wide for an
OutStreamthat 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_stdoutis process-global too.
Fix
Route the streams per write instead of per process.
_ThreadRoutedStreamwraps ipykernel's OutStreams: the thread currentlyrunning 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_executehooks, 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-capturesexactly the writes just routed back to the terminal.
_LiveStream(legacy kernel) got the same thread check, falling back to thepre-redirect stream so those writes still reach the console.
Testing
background thread's output must not appear in a cell, while the cell's own
print still does.
rather than being dropped.
write,and by stashing the fix: exit 1).
tests/trainer/services: 348 passed, 69 skipped. Ruff clean under CI's ruleset.
Trade-offs
the cell.
notebook. For an in-process kernel sharing a terminal with the trainer,
that's the better trade.