Add pytest plugin (pytest --scalene) — fixes #70 - #1059
Open
emeryberger wants to merge 10 commits into
Open
Conversation
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.
`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>
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.
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:
The plugin is registered via a
pytest11entry point, so a barepytest --scaleneworks after install.Behavior
--scalenesets up Scalene's signal-based CPU sampler inside the running pytest process — no re-exec, so it composes with IDE runners and-kselection. Writesscalene-profile.json, viewable withscalene view.--scalene-memory/--scalene-gpure-launch pytest underscalene run -m pytest(allocation tracking needs libscalene preloaded). When already running underscalene run -m pytest, the plugin detects it and drives the existing profiler instead of re-execing again.@pytest.mark.scalenenarrows profiling to only the marked tests.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 toscalene run).Implementation
scalene/pytest_scalene.py— the plugin.scalene_profiler.py— two additive accessors for programmatic entry points:Scalene.get_initialized()(detect we're underscalene run) andset_program_path()(point tracing at the test rootdir).pyproject.toml—pytest11entry point.mypy.ini— relaxdisallow_any_unimportedfor the plugin module (pytest is a test-only dep, absent in the lint env).Tests & CI
tests/test_pytest_plugin.py— unit tests for argv handling pluspytester-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 regularpytestrun picks up the new tests; an added "pytest plugin smoke test" step exercises the installed entry point end-to-end (barepytest --scalenefrom 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.