diff --git a/benches/_memory.py b/benches/_memory.py index 9f135ce7..ffa22140 100644 --- a/benches/_memory.py +++ b/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 diff --git a/benches/conftest.py b/benches/conftest.py index b03f1224..194ab191 100644 --- a/benches/conftest.py +++ b/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"]) diff --git a/benches/report.py b/benches/report.py index 52049a19..d054a56f 100644 --- a/benches/report.py +++ b/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( diff --git a/tests/test_bench_memory.py b/tests/test_bench_memory.py index 89c7fc88..694dfdc9 100644 --- a/tests/test_bench_memory.py +++ b/tests/test_bench_memory.py @@ -14,14 +14,24 @@ """Unit tests for the benchmark memory primitives (``benches/_memory.py``). -The per-test peak under MPI is the *peak-of-sum* (the largest footprint that actually coexisted -across ranks), not the sum of per-rank lifetime peaks, which overcounts disjoint transients. +Two behaviours are pinned. The per-test peak under MPI is the *peak-of-sum* (the largest +footprint that actually coexisted across ranks), not the sum of per-rank lifetime peaks, +which overcounts disjoint transients. And :class:`HighWaterMark` sees transients that the +sampler cannot, which is the whole reason it exists. """ from __future__ import annotations import pytest -from _memory import RssSampler, merge_peak_of_sum, rss_bytes +from _memory import ( + HighWaterMark, + PssSampler, + merge_peak_of_sum, + peak_rss_bytes, + pss_bytes, + reset_peak_rss, + rss_bytes, +) MIB = 2**20 @@ -58,20 +68,61 @@ def test_merge_handles_empty_series() -> None: assert merge_peak_of_sum([[(0.0, 100)], []]) == 100 -def test_sampler_records_timeline_and_sees_a_transient() -> None: - if rss_bytes() == 0: - pytest.skip("/proc/self/status VmRSS unavailable (non-Linux)") - with RssSampler(interval=0.002) as sampler: - baseline = sampler.samples[0][1] +def test_sampler_records_an_ordered_timeline() -> None: + if pss_bytes() == 0: + pytest.skip("/proc/self/smaps_rollup Pss unavailable (non-Linux)") + with PssSampler(interval=0.002) as sampler: blob = bytearray(80 * MIB) for i in range(0, len(blob), 4096): # touch pages so they become resident blob[i] = 1 - # Hold across several sampling intervals so the transient is observed. - for _ in range(2_000_000): - pass del blob samples = sampler.samples assert len(samples) >= 2 # baseline on enter, final on exit assert all(isinstance(t, float) and isinstance(p, int) for t, p in samples) - assert max(p for _t, p in samples) > baseline + assert all(p > 0 for _t, p in samples) + assert [t for t, _p in samples] == sorted(t for t, _p in samples) + + +def test_pss_is_at_most_rss() -> None: + if pss_bytes() == 0: + pytest.skip("/proc/self/smaps_rollup Pss unavailable (non-Linux)") + # PSS splits each shared page across its mappers, so it can only be <= RSS. This is the + # property that makes summing PSS across ranks on a node meaningful. + assert 0 < pss_bytes() <= rss_bytes() + + +def test_high_water_mark_catches_a_freed_transient() -> None: + if not reset_peak_rss(): + pytest.skip("/proc/self/clear_refs unavailable (non-Linux or kernel < 4.0)") + with HighWaterMark() as window: + blob = bytearray(80 * MIB) + for i in range(0, len(blob), 4096): + blob[i] = 1 + del blob # gone before the window closes: only VmHWM still knows it existed + + assert window.exact + assert window.delta_bytes >= 70 * MIB + assert window.peak_bytes >= window.baseline_bytes + + +def test_high_water_mark_window_is_reset_per_block() -> None: + if not reset_peak_rss(): + pytest.skip("/proc/self/clear_refs unavailable (non-Linux or kernel < 4.0)") + with HighWaterMark() as first: + blob = bytearray(80 * MIB) + for i in range(0, len(blob), 4096): + blob[i] = 1 + del blob + with HighWaterMark() as second: + pass + + # The second window must not inherit the first's spike + assert first.delta_bytes >= 70 * MIB + assert second.delta_bytes < 10 * MIB + + +def test_peak_rss_never_below_current_rss() -> None: + if peak_rss_bytes() == 0: + pytest.skip("/proc/self/status VmHWM unavailable (non-Linux)") + assert peak_rss_bytes() >= rss_bytes() diff --git a/tests/test_bench_report.py b/tests/test_bench_report.py index 47612323..f80f563d 100644 --- a/tests/test_bench_report.py +++ b/tests/test_bench_report.py @@ -176,14 +176,14 @@ def test_build_report_includes_memory(tmp_path: Path) -> None: _write_timings(tmp_path) _write_results( tmp_path, - mem={ + memhwm={ "bench_random.py::test_random_energy[heisenberg]": 52428800, "bench_random.py::test_random_energy[schrodinger]": 104857600, }, ) md = _collapse(report.build_report(tmp_path)) - assert "Memory (RSS)" in md + assert "Memory (peak RSS)" in md assert "50.00 MiB" in md assert "100.00 MiB" in md