Skip to content

Add pytest plugin (pytest --scalene) — fixes #70 - #1059

Open
emeryberger wants to merge 10 commits into
masterfrom
feature/pytest-plugin-issue-70
Open

Add pytest plugin (pytest --scalene) — fixes #70#1059
emeryberger wants to merge 10 commits into
masterfrom
feature/pytest-plugin-issue-70

Conversation

@emeryberger

Copy link
Copy Markdown
Member

Summary

Adds a pytest plugin so test suites can be profiled with Scalene directly, resolving the long-standing request in #70. No more writing a throwaway driver script — just add a flag to your usual pytest command:

pytest --scalene                  # in-process CPU profiling
pytest --scalene -k test_hot      # profile only the selected tests
pytest --scalene --scalene-memory # full CPU+memory profile (re-execs)

The plugin is registered via a pytest11 entry point, so a bare pytest --scalene works after install.

Behavior

  • In-process (default): --scalene sets up Scalene's signal-based CPU sampler inside the running pytest process — no re-exec, so it composes with IDE runners and -k selection. Writes scalene-profile.json, viewable with scalene view.
  • Re-exec for memory/GPU: --scalene-memory / --scalene-gpu re-launch pytest under scalene run -m pytest (allocation tracking needs libscalene preloaded). When already running under scalene run -m pytest, the plugin detects it and drives the existing profiler instead of re-execing again.
  • Granularity: whole session by default; @pytest.mark.scalene narrows profiling to only the marked tests.
  • Empty-profile fix: when running under scalene run -m pytest, the plugin repoints the program path to the pytest rootdir, so the test files are actually profiled instead of producing an empty profile (the exact pain point reporters worked around with --profile-all).

Options: --scalene-outfile PATH, --scalene-args='...' (forward extra args to scalene run).

Implementation

  • scalene/pytest_scalene.py — the plugin.
  • scalene_profiler.py — two additive accessors for programmatic entry points: Scalene.get_initialized() (detect we're under scalene run) and set_program_path() (point tracing at the test rootdir).
  • pyproject.tomlpytest11 entry point.
  • mypy.ini — relax disallow_any_unimported for the plugin module (pytest is a test-only dep, absent in the lint env).
  • README / CLAUDE.md — usage and module docs.

Tests & CI

  • tests/test_pytest_plugin.py — unit tests for argv handling plus pytester-based integration tests (whole-session profiling, marker narrowing, no-op without the flag, marker registration). Integration tests rely on the auto-loaded entry point and skip cleanly when it isn't available.
  • .github/workflows/tests.yml — the regular pytest run picks up the new tests; an added "pytest plugin smoke test" step exercises the installed entry point end-to-end (bare pytest --scalene from a clean dir must auto-load, profile, and write a profile containing the test file).

Verified against a clean pip install -e . venv mirroring CI: all 8 tests pass, the full existing suite still passes with the plugin auto-loaded (confirming it's a no-op without --scalene), and all three execution paths (in-process CPU, selective marker, memory re-exec) produce correct attribution.

emeryberger and others added 10 commits June 5, 2026 12:06
Adds a pytest11 plugin so test suites can be profiled directly:

  pytest --scalene                  # in-process CPU profiling
  pytest --scalene -k test_hot      # profile only selected tests
  pytest --scalene --scalene-memory # full CPU+memory (re-execs)

By default the whole session is profiled and a single scalene-profile.json
is written; marking tests with @pytest.mark.scalene narrows profiling to
just those tests. Memory/GPU profiling re-execs pytest under `scalene run`
so libscalene is preloaded. When already running under `scalene run -m
pytest`, the plugin repoints the program path to the test rootdir so the
test files are actually profiled instead of producing an empty profile.

Supporting changes:
  - scalene_profiler.py: add Scalene.get_initialized() and
    set_program_path() accessors for programmatic entry points.
  - pyproject.toml: register the pytest11 entry point.
  - mypy.ini: relax disallow_any_unimported for the plugin module
    (pytest is a test-only dep, absent in the lint environment).
  - README/CLAUDE.md: document usage and module.
  - tests/test_pytest_plugin.py: unit tests for argv handling plus
    pytester-based integration tests (skip if native ext/plugin not
    importable from a subprocess).
The integration tests now rely on the plugin's pytest11 entry point
auto-loading (the real user path, and what `pip install -e .` provides in
CI) instead of passing `-p scalene.pytest_scalene`. When the package is
properly installed, the entry point already loads the module, so an
explicit `-p` would re-register it and raise 'Plugin already registered
under a different name' — caught by running against a clean install.

The readiness guard now checks that a subprocess `pytest --help` actually
advertises `--scalene` (entry point registered + native ext present),
skipping cleanly otherwise.

Adds a 'pytest plugin smoke test' CI step: a bare `pytest --scalene` from a
clean working dir must auto-load the plugin, profile, and write a profile
JSON containing the test file.
Installing the pytest plugin registers a pytest11 entry point, so every
`pytest` run now imports scalene.pytest_scalene -> the scalene package ->
scalene_profiler, whose module top-level sets `builtins.profile =
Scalene._profile`. A test module that uses the @Profile decorator (e.g.
test/profile_annotation_test.py) would then invoke Scalene._profile during
collection, which called pywhere.register_files_to_profile() before
libscalene was preloaded and crashed with 'Unable to find p_whereInPython',
breaking collection for the whole suite.

Guard the pywhere registration on Scalene.__initialized: during a real
`scalene run`, set_initialized() and populate_struct()/libscalene preload
happen before any user @Profile decorator executes, so behavior is
unchanged. When scalene is merely imported (no live profiling session),
@Profile is now a safe pass-through. Verified that @Profile still targets
only decorated functions under `scalene run`.
In the full CI suite, an earlier test constructs a Scalene instance, which
runs redirect_python and permanently rewrites sys.executable to a wrapper
that re-launches Python under `scalene run` (and prepends that wrapper's
dir to PATH). pytester.runpytest_subprocess reads sys.executable at call
time, so the integration tests' inner pytest either ran as
`scalene run -m pytest` (no terminal summary -> assert_outcomes raised
'Pytest terminal summary report not found') or, after a naive fix, picked a
different interpreter lacking the installed plugin (so --markers omitted the
scalene marker).

Capture the interpreter and PATH at module import time (before any
contamination) and restore them around each runpytest_subprocess call.
Reproduced locally by running a Scalene-constructing test before the plugin
tests; all pass with the fix.
CI surfaced two test bugs (461 passed, 1 failed on macos-3.10):

1. test_marker_narrows_profiling keyed its CPU lookup by bare function name,
   but Scalene's per-function record keys 'line' by the def line's *source
   text* (e.g. 'def _burn_marked():'), so both lookups returned the 0.0
   default and the assertion was effectively 0.0 > 0.0 -- which fails the
   moment any sample lands. Replaced with _cpu_by_helper(), which sums CPU by
   substring-matching the helper name across both function and line records.

2. Both integration tests used sub-second burns, so a busy CI runner could
   collect zero CPU samples (Scalene then writes no profile -- the documented
   'did not run long enough' flakiness). Burn 1-2s, retry up to 4x, and
   skip (not fail) if the profiler still collected nothing.

Verified directly that selectivity attribution is marked=99.4% /
unmarked=0.0%, confirming @pytest.mark.scalene narrows profiling as intended.
Also re-verified the full mixed ordering (Scalene-constructing test ->
@Profile test -> plugin tests) all pass together.
Brings in the non-ASCII path fix (#1087), the get_ipython() Optional
guard (#1088), and the free-threaded CI setup fix (#1089).
`pytest --scalene -n 2` succeeded and wrote a profile that was
essentially empty: the tests execute in xdist worker subprocesses, and
the sampler installed by the plugin lives in the controller, which does
almost nothing. Measured on a two-test workload: ~1% attributed to a
single line under `-n 2`, versus ~95% correctly spread across the three
hot lines when run serially. A silently near-empty profile is a worse
outcome than no profile, since nothing tells the user their numbers are
meaningless.

pytest_configure now raises pytest.UsageError (exit 4) with an
actionable message when --dist is anything other than "no". `-n 0` is
deliberately still allowed: xdist keeps every test in this process and
leaves dist == "no", and profiling is accurate there -- verified, and
covered by its own regression test so the guard can't grow into an
over-block.

Both new tests importorskip("xdist") since CI doesn't install it. The
rejection test fails against the unguarded plugin.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The pytest plugin was documented only in a FAQ entry near the bottom,
which is where you look after you already suspect the feature exists.
Jupyter support gets a top-level "Scalene with Jupyter" section; pytest
now gets the parallel treatment, placed right beside it, with the full
option table, the marker, and how to view the resulting profile.

Also documents the pytest-xdist restriction added in af3648d -- users
who run their suite with -n will hit it, and the error should not be the
first they hear of it -- and adds a bullet to "Other features". The FAQ
entry keeps its short answer and now links to the full section instead
of drifting out of sync with it.

Every claim checked against the running plugin: option names and help
text match `pytest --help` verbatim, and the marker's "unmarked tests
still run, with sampling suspended" wording matches the code.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Replaces the previous outright rejection of distributed runs with actual
support. It turned out not to need a new subsystem -- Scalene already has
everything required, it just wasn't being reached.

Under `scalene run`, Scalene installs a python alias on PATH and as
sys.executable. xdist's execnet launches its workers through it, so each
worker comes up as `scalene run --pid=<controller>`, profiles itself,
and dumps its stats; the controller's output path already calls
ScaleneStatistics.merge_stats to fold children in. That merge is
concurrency-correct for this case: per-line samples are summed while
elapsed_time is a max, which is what parallel workers need.

So a distributed run now joins memory/GPU in _needs_preload: it takes
the re-exec path even for plain --scalene, because in-process sampling
in the controller sees essentially none of the work. `-n 0` is not
distributed (xdist leaves dist == "no") and stays in-process. The
re-exec passes --cpu-only unless memory/GPU was actually requested, so
`--scalene` doesn't quietly become a slower memory profile.

Verified: `pytest --scalene -n 2` attributes both workers' hot lines
(test_a.py lines 3-4, test_b.py line 2) in one merged profile, where the
same command previously produced ~1% on a single line.

Memory is the exception, and it is a real limitation rather than an
oversight: allocations in xdist workers are not tracked at all --
max_footprint_mb == 0, empty memory timeline, 0 MB on every line, for a
workload that reports 43 MB serially. Rather than hand back a memory
profile that reads as "your tests allocate nothing", --scalene-memory
with -n is rejected with an actionable UsageError, checked before the
re-exec so it fails immediately. A second guard covers the case where a
distributed run somehow isn't under `scalene run` by the time the plugin
registers.

The xdist support test keeps its deterministic core -- the session must
pass and write a profile on every attempt, which is what catches a
regression in the routing -- while the stronger "both workers sampled"
assertion retries and then skips, since a short test can finish between
samples. Same tradeoff as tests/_scalene_subprocess.py.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The previous commit rejected --scalene-memory with -n because worker
allocations came back as literally nothing: max_footprint_mb == 0, an
empty memory timeline, 0 MB on every line, for a workload reporting
43 MB serially. That turned out to be a fixable bug, not an inherent
limit, so this replaces the rejection with support.

Bisecting the path with probes showed the interposer was in place and
simply never recorded anything: in a worker, libscalene was loaded
(dlsym finds its symbols), DYLD_INSERT_LIBRARIES was set, args.memory
was True, the SIGXCPU handler was installed and the mapfiles existed --
but the interposer's write cursor was 0 bytes, against 8302 in the
controller. So nothing was being dropped on the Python side; nothing was
ever written.

The cause is pywhere's native TraceConfig. sampleheap.hpp only writes an
allocation record when where(...) returns non-zero, and whereInPython
returns 0 immediately when TraceConfig is null. Scalene registers that
config while preparing to run a program file; an interpreter launched as
`python -c ...` -- which is how xdist's execnet starts its workers --
never reaches that code, so the config stayed null and every allocation
went unrecorded.

Fix it where the asymmetry actually is: Scalene.set_program_path() now
re-registers the files to profile when memory profiling is on. That is
the natural home for it, since pywhere holds its own copy of the program
path, and it fixes any programmatic entry point that repoints the path
rather than just this plugin. Confirmed by hand first -- calling
_register_files_to_profile() in a worker took it from 0 bytes/0 samples
to 101752 bytes/100 samples.

`pytest --scalene-memory -n 2` now reports 43.0 MB on the allocating
line, matching the serial run exactly. All six combinations of
{--scalene, --scalene-memory} x {serial, -n 0, -n 2} verified; CPU
percentages are lower under -n 2 only because elapsed time is wall-clock
across parallel workers, which is the correct denominator.

Full suite passes on 3.12 (438) and 3.14 (450).

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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