Add notebook profiler with live Rich tree display#887
Open
1kaiser wants to merge 2 commits into
Open
Conversation
Adds per-cell timing, output analysis, and section-grouped profiling to
papermill. Sections are auto-numbered from markdown headings so the display
shows "Section 1 / Sub-section 1.2" rather than raw notebook source text.
New modules:
- papermill/profile.py — build_sections(), build_profile(), profile_notebook()
- papermill/live_tree.py — LiveTreeDisplay (replaces tqdm with a Rich live tree)
New entry point:
- papermill-profile <executed.ipynb> — profile any executed notebook
Changes to existing modules:
- execute.py — live_tree=False parameter wired to execute_notebook()
- engines.py — NotebookExecutionManager accepts live_display=, calls its
on_cell_start/complete/exception hooks instead of tqdm
- cli.py — --live-tree / --no-live-tree flag; 'papermill profile' command
- __init__.py — export profile_notebook, build_profile, build_sections
- pyproject.toml — optional-dependencies.rich = ["rich>=13.0"]
Tests: papermill/tests/test_profile.py (20 tests, all passing)
Usage:
papermill notebook.ipynb out.ipynb --live-tree
papermill-profile executed.ipynb
python -c "from papermill import profile_notebook; print(profile_notebook('out.ipynb'))"
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
e444129 to
bf5f6b8
Compare
for more information, see https://pre-commit.ci
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.
What does this PR do?
Adds a notebook profiling subsystem to papermill — per-cell timing grouped by
markdown section headings, with an optional live Rich tree display that replaces
the tqdm progress bar during execution.
Motivation
Papermill already records
start_time,end_time, anddurationin everycell's
metadata.papermillblock. This PR surfaces that data in two ways:papermill-profile <executed.ipynb>— reads any executed notebook andemits a structured JSON report showing which sections and cells took the
most time, what output types were produced, and where the bottleneck is.
papermill notebook.ipynb out.ipynb --live-tree— shows a live Richtree during execution instead of the tqdm bar:
Section labels are auto-generated sequential numbers (
Section 1,Sub-section 1.2, …) derived from markdown heading depth — no raw notebooksource code appears in the display.
Changes
papermill/profile.pybuild_sections(),build_profile(),profile_notebook()papermill/live_tree.pyLiveTreeDisplay(optional Rich dependency)papermill/tests/test_profile.pypapermill/execute.pylive_tree=Falseparameterpapermill/engines.pylive_display=onNotebookExecutionManager; hooks intocell_start/complete/exceptionpapermill/cli.py--live-treeflag;papermill profilesubcommandpapermill/__init__.pyprofile_notebook,build_profile,build_sectionspyproject.tomloptional-dependencies.rich = ["rich>=13.0"]New entry points
Test plan
pytest papermill/tests/test_profile.py— 20 tests, all passingpapermill notebook.ipynb out.ipynb --live-tree --kernel python3— tree renders livepapermill-profile out.ipynb— summary printed,.profile.jsonwrittenrichnot installed → graceful fallback to tqdm with warning🤖 Generated with Claude Code