fix(benches): 🐛 measure peak memory with the kernel high-water mark

benches/_memory.py

@@ -12,12 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""RSS-based memory-measurement primitives for the benchmark suite.
-
-What the per-test peak means: **peak-of-sum, not sum-of-peaks.** The peak is
-``max over time`` of the summed RSS. Summing each rank's independently-timed peak
-counts transients that never coexisted. Comparable wall-clock timestamps let
-:func:`merge_peak_of_sum` recover the true peak-of-sum.
+"""Memory-measurement primitives for the benchmark suite.
+
+Two metrics, deliberately kept apart:
+
+:class:`HighWaterMark`
+    The kernel's own peak RSS over a resettable window. Exact -- there is no sampling, so
+    no transient can be missed -- and it needs no background thread, so the GIL cannot
+    hide anything from it. This is the metric to quote, and the only one that is
+    like-for-like against a non-Python library.
+
+:class:`PssSampler`
+    A sampled ``(wall_clock, pss)`` timeline. Needed only under MPI, where the job's
+    footprint is the **peak-of-sum, not the sum-of-peaks**: summing each rank's
+    independently-timed peak counts transients that never coexisted, so
+    :func:`merge_peak_of_sum` replays the timelines instead. A scalar per-rank peak
+    cannot reconstruct that, which is the one thing sampling still buys.
 """
 
 from __future__ import annotations
@@ -35,10 +45,12 @@
     from types import TracebackType
     from typing import Self
 
-# RSS sampling cadence. monoprop's heavy work runs in C++ with the GIL released, so
-# the background sampler costs an idle core, not the timed thread.
+# Sampling cadence for the MPI timeline.
 SAMPLE_INTERVAL_S = 0.005
 
+# Reset VmHWM to the current RSS, starting a new measurement window
+_CLEAR_REFS_MM_HIWATER_RSS = "5\n"
+
 
 def proc_field(path: str, key: str) -> int:
     """Return a ``/proc/self`` size field (kB → bytes); 0 if unavailable."""
@@ -60,6 +72,40 @@ def rss_bytes() -> int:
     return proc_field("/proc/self/status", "VmRSS:")
 
 
+def pss_bytes() -> int:
+    """Return this process's proportional set size (PSS) in bytes.
+
+    PSS divides each shared page among the processes mapping it, so PSS summed over the
+    ranks on a node counts every page exactly once. RSS charges a shared page (the Python
+    interpreter, libstdc++, a shared graph) in full to every rank, which inflates an
+    MPI sum by roughly the shared footprint times the rank count.
+    """
+    return proc_field("/proc/self/smaps_rollup", "Pss:")
+
+
+def peak_rss_bytes() -> int:
+    """Return the kernel's high-water mark of this process's RSS, in bytes.
+
+    ``VmHWM`` is maintained by the kernel on every RSS increase, so it is exact.
+    """
+    return proc_field("/proc/self/status", "VmHWM:")
+
+
+def reset_peak_rss() -> bool:
+    """Reset ``VmHWM`` to the current RSS, starting a new measurement window.
+
+    Returns:
+        ``True`` if the reset took effect, ``False`` where ``/proc/self/clear_refs`` is
+        unavailable (non-Linux, kernel < 4.0, or a restricted sandbox), in which case
+        ``VmHWM`` keeps counting from process start and callers must fall back.
+    """
+    try:
+        Path("/proc/self/clear_refs").write_text(_CLEAR_REFS_MM_HIWATER_RSS)
+    except OSError:  # pragma: no cover - platform dependent
+        return False
+    return True
+
+
 def heap_trim() -> None:
     """Ask the C allocator to return unused heap pages to the OS."""
     with contextlib.suppress(Exception):  # unsupported platform / allocator
@@ -73,15 +119,75 @@ def resting_rss_bytes() -> int:
     return rss_bytes()
 
 
-class RssSampler:
-    """Background thread sampling this process's live RSS over time.
+class HighWaterMark:
+    """Exact peak RSS over the enclosed block, straight from the kernel.
+
+    Settles the process on entry (``gc.collect()`` + ``malloc_trim``) and resets ``VmHWM``
+    to that floor, so the peak reported is this block's own and not an earlier block's
+    retained garbage. The settling is what makes the number comparable across languages:
+    without it the figure tracks the allocator's or GC's willingness to return pages more
+    than it tracks what the code needed.
+
+    Use ``peak_bytes`` for the footprint (peak including everything already resident) and
+    ``delta_bytes`` for what this block added on top of its floor. Report the floor too --
+    an interpreter plus its imports is a 100+ MiB constant that has nothing to do with the
+    code under test.
+
+    ``exact`` is ``False`` when the kernel would not reset the window (see
+    :func:`reset_peak_rss`); the peak then degrades to the RSS observed on exit, which is
+    a lower bound. Callers that publish numbers should check it.
+    """
+
+    def __init__(self, *, settle: bool = True) -> None:
+        self._settle = settle
+        self.baseline_bytes = 0
+        self.peak_bytes = 0
+        self.exact = False
+
+    def __enter__(self) -> Self:
+        self.baseline_bytes = resting_rss_bytes() if self._settle else rss_bytes()
+        self.exact = reset_peak_rss()
+        self.peak_bytes = self.baseline_bytes
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        observed = peak_rss_bytes() if self.exact else rss_bytes()
+        self.peak_bytes = max(self.baseline_bytes, observed)
+
+    @property
+    def delta_bytes(self) -> int:
+        """Return the peak measured above the block's own starting floor."""
+        return self.peak_bytes - self.baseline_bytes
+
+    @property
+    def peak_mb(self) -> float:
+        """Return :attr:`peak_bytes` in MiB."""
+        return self.peak_bytes / 1024**2
+
+    @property
+    def baseline_mb(self) -> float:
+        """Return :attr:`baseline_bytes` in MiB."""
+        return self.baseline_bytes / 1024**2
+
+
+class PssSampler:
+    """Background thread sampling this process's live PSS over time.
+
+    Records ``(wall_clock, pss_bytes)`` pairs while active. Uses ``time.time`` (not
+    ``time.monotonic``) so timestamps are comparable across ranks sharing a node's clock,
+    which :func:`merge_peak_of_sum` needs to correlate readings.
 
-    Records ``(wall_clock, rss_bytes)`` pairs while active. Uses ``time.time``
-    (not ``time.monotonic``) so timestamps are comparable across ranks sharing a
-    node's clock, which :func:`merge_peak_of_sum` needs to correlate readings.
+    Use as a context manager around the operation; it samples on entry and exit, so even a
+    sub-interval operation yields a usable timeline.
 
-    Use as a context manager around the operation; it samples on entry and exit,
-    so even a sub-interval operation yields a usable timeline.
+    The timeline is only as dense as the timed thread's GIL releases allow, so treat it as
+    a lower bound on the true peak-of-sum and :class:`HighWaterMark` as the exact per-rank
+    figure. It is kept because no per-rank scalar can recover which peaks coexisted.
     """
 
     def __init__(self, interval: float = SAMPLE_INTERVAL_S) -> None:
@@ -92,11 +198,11 @@ def __init__(self, interval: float = SAMPLE_INTERVAL_S) -> None:
 
     def _run(self) -> None:
         while not self._stop.is_set():
-            self._samples.append((time.time(), rss_bytes()))
+            self._samples.append((time.time(), pss_bytes()))
             self._stop.wait(self._interval)
 
     def __enter__(self) -> Self:
-        self._samples.append((time.time(), rss_bytes()))
+        self._samples.append((time.time(), pss_bytes()))
         self._thread.start()
         return self
 
@@ -108,16 +214,16 @@ def __exit__(
     ) -> None:
         self._stop.set()
         self._thread.join()
-        self._samples.append((time.time(), rss_bytes()))
+        self._samples.append((time.time(), pss_bytes()))
 
     @property
     def samples(self) -> list[tuple[float, int]]:
-        """Return the recorded ``(wall_clock, rss_bytes)`` samples."""
+        """Return the recorded ``(wall_clock, pss_bytes)`` samples."""
         return self._samples
 
 
 def merge_peak_of_sum(per_rank: list[list[tuple[float, int]]]) -> int:
-    """Return the peak of the summed live RSS across ranks, in bytes.
+    """Return the peak of the summed live PSS across ranks, in bytes.
 
     ``per_rank[i]`` is rank ``i``'s samples. Walks all samples in time order,
     step-holding each rank's most recent reading, and tracks the maximum of the

benches/conftest.py

@@ -45,7 +45,12 @@
     build_random_propagator,
     make_random_problem,
 )
-from _memory import RssSampler, merge_peak_of_sum, resting_rss_bytes
+from _memory import (
+    HighWaterMark,
+    PssSampler,
+    merge_peak_of_sum,
+    resting_rss_bytes,
+)
 
 import monoprop
 
@@ -95,7 +100,8 @@ def _reduce_sum(comm: Any, value: int) -> int:
 _RESULTS: dict[str, Any] = {
     "meta": {},  # run configuration (ranks, threads, host, ...)
     "params": {},  # resolved random-problem hyperparameters
-    "mem": {},  # node id -> peak RSS bytes (per operation)
+    "mem": {},  # node id -> peak-of-sum PSS bytes (per operation; MPI lower bound)
+    "memhwm": {},  # node id -> summed per-rank peak RSS bytes (per operation; may fall back to exit RSS if VmHWM can't be reset)  # noqa: E501
     "opsize": {},  # picture / model -> {"terms": n}
     "memrest": {},  # picture / model -> resting RSS bytes
     "membase": {},  # fixed model -> resting RSS bytes before the model is built
@@ -120,9 +126,9 @@ def _results_path() -> Path | None:
 
 
 def _peak_of_sum(comm: Any, samples: list[tuple[float, int]]) -> int:
-    """Reduce per-rank RSS timelines to the job's peak summed RSS (bytes).
+    """Reduce per-rank PSS timelines to the job's peak summed footprint (bytes).
 
-    Gathers every rank's ``(wall_clock, rss)`` samples to rank 0 and merges them
+    Gathers every rank's ``(wall_clock, pss)`` samples to rank 0 and merges them
     via :func:`_memory.merge_peak_of_sum`. Collective; returns ``0`` off root.
     """
     if comm is None or comm.Get_size() == 1:
@@ -285,19 +291,32 @@ def _do(model: str, propagator: Any, baseline_rss: int) -> None:
 
 @pytest.fixture(autouse=True)
 def record_memory(request: pytest.FixtureRequest, bench_comm: Any) -> Iterator[None]:
-    """Record each benchmark's peak physical-memory footprint (RSS) for the report.
-
-    A background :class:`RssSampler` samples this rank's live RSS while the test
-    runs. It is a footprint: it includes structures already resident when the
-    operation starts (e.g. the shared :func:`built_graph`). Under MPI the per-rank
-    timelines are merged into the peak-of-sum (see :func:`_peak_of_sum`); the
-    gather is collective, but only rank 0 records.
+    """Record each benchmark's peak physical-memory footprint for the report.
+
+    Two numbers, because neither alone is both exact and MPI-aware:
+
+    ``memhwm``
+        The kernel's exact peak RSS (:class:`HighWaterMark`) summed over ranks. Exact per
+        rank, and an upper bound on the job total, since ranks that peak at different
+        moments are added as though they had peaked together.
+    ``mem``
+        The peak-of-sum of the sampled per-rank PSS timelines, which is the quantity that
+        actually bounds a node's RAM. A lower bound: the sampler only advances when a rank
+        drops the GIL, and monoprop's ``propagate()`` holds it for the whole call.
+
+    Both are footprints, not deltas: they include structures already resident when the
+    operation starts (e.g. the shared :func:`built_graph`). The gather is collective, but
+    only rank 0 records.
     """
-    with RssSampler() as sampler:
+    with HighWaterMark() as window, PssSampler() as sampler:
         yield
+    key = request.node.nodeid.split("/")[-1]
     mem = _peak_of_sum(bench_comm, sampler.samples)
     if mem:  # 0 => non-root rank or /proc unavailable: nothing to record
-        _record("mem", request.node.nodeid.split("/")[-1], mem)
+        _record("mem", key, mem)
+    hwm = _reduce_sum(bench_comm, window.peak_bytes)
+    if hwm:
+        _record("memhwm", key, hwm)
 
 
 @pytest.fixture(scope="session", params=["heisenberg", "schrodinger"])

benches/report.py

@@ -260,7 +260,7 @@ def sec(name: str) -> dict[str, dict]:
         sec("params"),
         sec("opsize"),
         sec("memrest"),
-        sec("mem"),
+        sec("memhwm"),
     )
 
     all_ops = sorted(
@@ -294,7 +294,7 @@ def ops_section(name: str, picture: str) -> list[str]:
                 level=3,
             ),
             *_section(
-                "Memory (RSS)",
+                "Memory (peak RSS)",
                 "",
                 "Operation",
                 ops,
@@ -308,9 +308,10 @@ def ops_section(name: str, picture: str) -> list[str]:
         "# monoprop benchmark report",
         "",
         f"Run labels: **{', '.join(labels)}**. Times are the mean over rounds; "
-        "memory is the peak resident footprint (RSS) during each operation. Under "
-        "MPI it is the peak of the RSS summed across ranks (shared pages counted "
-        "per rank, so an upper bound), not the sum of per-rank peaks.",
+        "memory is the kernel's exact peak resident footprint (`VmHWM`) during each "
+        "operation, measured from a window reset and settled per operation. Under MPI "
+        "it is summed across ranks, so ranks peaking at different moments are counted "
+        "together: an upper bound on the job total.",
         "",
         *_config_table(labels, results),
         *_section(