Skip to content

Add notebook profiler with live Rich tree display#887

Open
1kaiser wants to merge 2 commits into
nteract:mainfrom
1kaiser:feature/notebook-profiler
Open

Add notebook profiler with live Rich tree display#887
1kaiser wants to merge 2 commits into
nteract:mainfrom
1kaiser:feature/notebook-profiler

Conversation

@1kaiser
Copy link
Copy Markdown

@1kaiser 1kaiser commented May 22, 2026

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, and duration in every
cell's metadata.papermill block. This PR surfaces that data in two ways:

  1. papermill-profile <executed.ipynb> — reads any executed notebook and
    emits a structured JSON report showing which sections and cells took the
    most time, what output types were produced, and where the bottleneck is.

  2. papermill notebook.ipynb out.ipynb --live-tree — shows a live Rich
    tree during execution instead of the tqdm bar:

notebook.ipynb  6/10 cells  60%
├── ✓  Section 1                           0.1s
│   ├── ✓  [1] import numpy as np          0.1s
│   └── ✓  [2] data = np.zeros((n, m))     0.0s
├── ⟳  Section 2                           4.2s…
│   └── ⟳  [4] model.fit(X, y)             4.2s…
└── ·  Section 3                           pending

Section labels are auto-generated sequential numbers (Section 1,
Sub-section 1.2, …) derived from markdown heading depth — no raw notebook
source code appears in the display.

Changes

File Change
papermill/profile.py New — build_sections(), build_profile(), profile_notebook()
papermill/live_tree.py New — LiveTreeDisplay (optional Rich dependency)
papermill/tests/test_profile.py New — 20 unit tests
papermill/execute.py live_tree=False parameter
papermill/engines.py live_display= on NotebookExecutionManager; hooks into cell_start/complete/exception
papermill/cli.py --live-tree flag; papermill profile subcommand
papermill/__init__.py Export profile_notebook, build_profile, build_sections
pyproject.toml optional-dependencies.rich = ["rich>=13.0"]

New entry points

# Live tree during execution (requires pip install 'papermill[rich]')
papermill notebook.ipynb out.ipynb --live-tree

# Profile an already-executed notebook
papermill-profile executed.ipynb

# Python API
from papermill import profile_notebook
profile = profile_notebook("out.ipynb")
print(profile["bottleneck"])
# {'cell_index': 5, 'section': 'Sub-section 2.1', 'duration_s': 18.0, 'pct_of_total': 42.8}

Test plan

  • pytest papermill/tests/test_profile.py — 20 tests, all passing
  • papermill notebook.ipynb out.ipynb --live-tree --kernel python3 — tree renders live
  • papermill-profile out.ipynb — summary printed, .profile.json written
  • rich not installed → graceful fallback to tqdm with warning

🤖 Generated with Claude Code

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>
@1kaiser 1kaiser force-pushed the feature/notebook-profiler branch from e444129 to bf5f6b8 Compare May 22, 2026 21:36
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