Make memory profiling work in pytest-xdist workers

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>

Author: emeryberger

CLAUDE.md

@@ -105,10 +105,18 @@ python3 -m pytest tests/test_coverup_83.py -v
   samples summed, `elapsed_time` maxed, which is correct for parallel workers.
   `-n 0` is *not* distributed (xdist leaves `dist == "no"`) and stays
   in-process. `--cpu-only` is passed on this re-exec unless the user asked for
-  memory/GPU, so `--scalene` keeps meaning CPU. Memory profiling does **not**
-  reach xdist workers (verified: `max_footprint_mb == 0` and 0 MB per line for
-  a workload reporting 43 MB serially), so `--scalene-memory` plus `-n` is
-  rejected with a `UsageError` rather than emitting an all-zero profile.
+  memory/GPU, so `--scalene` keeps meaning CPU.
+- **Memory profiling in `-c`-launched processes:** `Scalene.set_program_path()`
+  re-registers pywhere's native `TraceConfig` (via `_register_files_to_profile`)
+  when memory profiling is on. This matters far more than it looks:
+  `whereInPython` returns 0 outright when that config is missing, and
+  `sampleheap.hpp` only writes an allocation record when `where(...)` returns
+  non-zero — so a missing TraceConfig means libscalene records **nothing**, not
+  merely mis-attributed lines. Scalene's normal startup registers it while
+  preparing to run a program file; an interpreter started as `python -c ...`
+  (how xdist's execnet launches workers) never takes that path. Symptom when
+  broken: `max_footprint_mb == 0`, empty memory timeline, and 0 MB on every
+  line, for a workload that reports tens of MB serially.
 
 ### Replacement Modules (`replacement_*.py`)
 

README.md

@@ -543,20 +543,21 @@ interpreter starts. Equivalently, you can drive it the other way around:
 scalene run -m pytest your_test.py
 ```
 
-**Using `pytest-xdist`:** CPU profiling works with distributed runs
-(`-n 4`, `--dist`). Because the tests execute in worker subprocesses, Scalene
+**Using `pytest-xdist`:** distributed runs (`-n 4`, `--dist`) are supported for
+both CPU and memory. Because the tests execute in worker subprocesses, Scalene
 re-runs the session under `scalene run` so that each worker is profiled as a
 child process; their statistics are then merged into the single
 `scalene-profile.json`:
 
 ```console
-pytest --scalene -n 4                # profiles every worker, merged
+pytest --scalene -n 4                # CPU, every worker, merged
+pytest --scalene-memory -n 4         # CPU + memory, every worker, merged
 ```
 
-Memory profiling is the exception: allocations in xdist workers are not
-tracked, so `--scalene-memory` combined with `-n` reports an error rather than
-handing back a memory profile that reads as all zeros. Profile memory without
-`-n` (or with `-n 0`, which keeps the tests in the current process).
+Note that reported CPU percentages are relative to wall-clock time, which for
+a parallel run is the elapsed time of the whole session — so per-line
+percentages are lower than in a serial run even though the ranking is the
+same.
 
 </details>
 
@@ -664,10 +665,9 @@ You can also drive it the other way around, which is equivalent to
 scalene run -m pytest your_test.py
 ```
 
-CPU profiling works with distributed `pytest-xdist` runs (`pytest --scalene -n
-4`): Scalene profiles each worker and merges the results. Memory profiling
-does not — `--scalene-memory` together with `-n` reports an error instead of
-producing an all-zero memory profile.
+Distributed `pytest-xdist` runs work too (`pytest --scalene -n 4`, with or
+without `--scalene-memory`): Scalene profiles each worker and merges the
+results into one profile.
 
 See [Scalene with pytest](#scalene-with-pytest) for the full set of options.
 

scalene/pytest_scalene.py

@@ -105,24 +105,6 @@ def pytest_configure(config: pytest.Config) -> None:
     if not _scalene_requested(config):
         return
 
-    # CPU sampling aggregates across xdist workers, but memory tracking does
-    # not: a distributed `--scalene-memory` run comes back with
-    # max_footprint_mb == 0, an empty memory timeline and 0 MB on every line,
-    # even for a workload that reports 43 MB when run serially. Rather than
-    # hand back a memory profile that is uniformly zero, say what's wrong and
-    # what to do about it. Checked before the re-exec so we fail immediately
-    # instead of relaunching the whole session first.
-    if config.getoption("--scalene-memory") and _is_distributed(config):
-        import pytest
-
-        raise pytest.UsageError(
-            "--scalene-memory cannot be combined with a distributed "
-            "pytest-xdist run: allocations in the workers are not tracked, so "
-            "the memory profile would be all zeros. Use --scalene (CPU "
-            "profiling does aggregate across workers), or drop -n/--dist "
-            "(or use -n 0) to profile memory."
-        )
-
     # Memory/GPU profiling needs libscalene preloaded, and a distributed
     # xdist run needs Scalene's python alias in place so the workers are
     # profiled (see _needs_preload). Both are only possible by launching the

scalene/scalene_profiler.py

@@ -1768,6 +1768,16 @@ def set_program_path(path: str) -> None:
         Scalene.__program_path = Filename(os.path.abspath(os.path.expanduser(path)))
         if hasattr(Scalene, "_Scalene__tracing"):
             Scalene.__tracing.set_program_path(Scalene.__program_path)
+        # Refresh the native side too. pywhere keeps its own TraceConfig, and
+        # `whereInPython` returns 0 outright when that config is missing --
+        # which makes libscalene record *no* allocations at all, not merely
+        # mis-attributed ones. Scalene's usual startup registers it while
+        # setting up to run a program file; an interpreter launched as
+        # `python -c ...` (how pytest-xdist's execnet starts its workers)
+        # never goes through that path, so without this the workers report
+        # zero bytes. Same guard as the @profile decorator path above.
+        if Scalene.__initialized and getattr(Scalene.__args, "memory", False):
+            Scalene._register_files_to_profile()
 
     @staticmethod
     def main() -> None:

tests/test_pytest_plugin.py

@@ -367,27 +367,46 @@ def test_burn_two():
 
 
 @requires_native
-def test_memory_with_xdist_is_rejected(pytester: pytest.Pytester) -> None:
-    """`--scalene-memory -n 2` must refuse rather than report zero bytes.
-
-    CPU aggregates across workers, but allocation tracking does not: the run
-    comes back with max_footprint_mb == 0 and 0 MB on every line for a
-    workload that reports tens of MB serially. A memory profile that is
-    uniformly zero looks like "your tests allocate nothing," so this
-    combination is rejected outright.
+def test_memory_profiling_under_xdist(pytester: pytest.Pytester) -> None:
+    """`--scalene-memory -n 2` attributes allocations made inside workers.
+
+    This needs pywhere's native TraceConfig to exist in the worker: without
+    it `whereInPython` returns 0 and libscalene records *no* allocations at
+    all -- the whole run came back with max_footprint_mb == 0 and 0 MB on
+    every line. Scalene's usual startup registers that config while preparing
+    to run a program file, which an xdist worker (`python -c ...`) never does;
+    Scalene.set_program_path() now re-registers it.
+
+    The workload allocates ~40 MB in one test, so a threshold well above
+    incidental interpreter traffic still isn't sampling-sensitive.
     """
     pytest.importorskip("xdist")
     pytester.makepyfile(
         test_alloc="""
         def test_alloc():
             data = [[j for j in range(3000)] for _ in range(400)]
             assert len(data) == 400
+
+        def test_quiet():
+            assert 1 + 1 == 2
         """
     )
-    result = _run(pytester, "--scalene-memory", "-n", "2", "test_alloc.py")
-    assert result.ret != 0, "--scalene-memory with -n should not succeed"
-    result.stderr.fnmatch_lines(["*cannot be combined with a distributed*"])
-    assert not (pytester.path / "scalene-profile.json").exists()
+    result = _run(pytester, "--scalene-memory", "-n", "2", "-q")
+    result.assert_outcomes(passed=2)
+
+    profile = pytester.path / "scalene-profile.json"
+    assert profile.exists(), "distributed memory run wrote no profile"
+    data = json.loads(profile.read_text())
+    total_mb = sum(
+        line.get("n_malloc_mb", 0)
+        for fdata in data["files"].values()
+        for line in fdata["lines"]
+    )
+    assert total_mb > 5.0, (
+        f"only {total_mb:.1f} MB attributed across xdist workers; the "
+        "worker's allocations were probably not tracked at all "
+        f"(max_footprint_mb={data.get('max_footprint_mb')})"
+    )
 
 
 @requires_native