From 9cce3220e0fc6892774b7a7f1839d464ad486c9c Mon Sep 17 00:00:00 2001 From: lmoresi Date: Sun, 10 May 2026 20:39:40 +1000 Subject: [PATCH 1/3] Add memprobe diagnostic module for tracking memory growth (#176) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Long parallel runs occasionally OOM on HPC even when each step looks small. memprobe gives a way to "light up" memory tracking on demand, sample at regular intervals, and pin which subsystem is growing. Components: - src/underworld3/utilities/memprobe.py — snapshot/diff/probe/instrument API. Snapshots capture process RSS (resource.getrusage), KDTree live + total-constructed counts, and (with full=True) per-class Python instance counts via gc.get_objects. Diffs filter out unchanged keys and sort py_classes by absolute change so dominant suspects surface first. - ckdtree.pyx — module-level live-instance counter, accurate via Cython's deterministic __cinit__/__dealloc__. Exposed as uw.kdtree.live_count() and uw.kdtree.total_constructed(). - Stokes.solve() and NavierStokes.solve() decorated with @memprobe.instrument(...). The decorator fast-returns when ENABLED is False (one bool check, sub-microsecond), so it's safe on hot paths. Activation: - UW_MEMPROBE=1 env var → enables instrumentation hooks at import time. - memprobe.enable() / disable() for runtime toggling. - with memprobe.probe("label"): … for ad-hoc blocks. Skipped: parsing PETSc.Log object tables from Python. The runtime flags -log_view and -malloc_dump give the same information more reliably; documented in the guide and exposed via memprobe.dump_petsc_leaks_at_finalize(). Adds tests/test_0780_memprobe.py (9 smoke tests) and docs/developer/guides/memory-diagnostics.md with debugging recipes. Does not fix issue #176 (Ben's OOM on HPC) — provides the instrumentation he needs to bisect it. Underworld development team with AI support from Claude Code --- docs/developer/guides/memory-diagnostics.md | 145 +++++++++++ src/underworld3/ckdtree.pyx | 23 ++ src/underworld3/systems/solvers.py | 3 + src/underworld3/utilities/__init__.py | 1 + src/underworld3/utilities/memprobe.py | 265 ++++++++++++++++++++ tests/test_0780_memprobe.py | 140 +++++++++++ 6 files changed, 577 insertions(+) create mode 100644 docs/developer/guides/memory-diagnostics.md create mode 100644 src/underworld3/utilities/memprobe.py create mode 100644 tests/test_0780_memprobe.py diff --git a/docs/developer/guides/memory-diagnostics.md b/docs/developer/guides/memory-diagnostics.md new file mode 100644 index 000000000..6e658e755 --- /dev/null +++ b/docs/developer/guides/memory-diagnostics.md @@ -0,0 +1,145 @@ +# Memory diagnostics + +Long parallel runs occasionally OOM on HPC even when each step looks small. +The `uw.utilities.memprobe` module gives you a way to "light up" memory +tracking on demand, sample at regular intervals, and pin which subsystem is +growing. + +## Quick start + +```bash +UW_MEMPROBE=1 mpirun -n 16 python my_long_run.py +``` + +With this flag set, the `Stokes.solve()` and `NavierStokes.solve()` paths +emit a one-line growth report each time they're called: + +``` +[memprobe] Stokes.solve: + RSS +0.42 MiB + kdtree: live +1, total_constructed +1 +``` + +A clean run shows mostly zero deltas. A leak shows the same component +growing on every step. + +## What's tracked + +| Signal | Source | Cost | +|---|---|---| +| Process RSS (MiB) | `resource.getrusage` | free | +| KDTree live count | `uw.kdtree.live_count()` | free | +| KDTree total constructed | `uw.kdtree.total_constructed()` | free | +| Per-class Python instance counts | `gc.get_objects()` walk | slow — gated behind `full=True` | + +`KDTree` instances are tracked via Cython class counters in `__cinit__` and +`__dealloc__`. Cython's deterministic destruction makes the live count +accurate without weak-references. + +PETSc-side object and allocation tracking is **not** parsed from Python — +PETSc's own `-log_view` and `-malloc_dump` runtime flags give the same +information more reliably. To enable them from Python: + +```python +from underworld3.utilities import memprobe +memprobe.dump_petsc_leaks_at_finalize() # equivalent to -malloc_dump -objects_dump +``` + +## API + +### Snapshots and diffs + +```python +from underworld3.utilities import memprobe + +before = memprobe.snapshot() +do_work() +after = memprobe.snapshot() +print(memprobe.format_diff("after-work", memprobe.diff(before, after))) +``` + +Add `full=True` to also walk Python-class counts: + +```python +snap = memprobe.snapshot(full=True) +# snap["py_classes"] = {"underworld3.swarm.Swarm": 3, ...} +``` + +### Probe context manager + +```python +with memprobe.probe("step 42"): + advance_one_step() +# On exit, the diff is emitted via `print` (configurable). +``` + +The `emit` keyword takes any callable: `with probe(..., emit=logger.info):` +to route into the logging system, or a rank-aware writer for parallel runs: + +```python +import underworld3 as uw +emit = lambda s: uw.pprint(0, s) # rank-0 only +with memprobe.probe("step 42", emit=emit): + ... +``` + +### Decorator + +```python +@memprobe.instrument("my-hot-loop") +def step(): + ... +``` + +When `memprobe.ENABLED` is `False` (the default) the decorator's wrapper is +a single attribute lookup + branch — sub-microsecond — so it's safe to +leave on hot paths permanently. + +`Stokes.solve()` and `NavierStokes.solve()` are pre-decorated. Add more if +you want them. + +### Runtime toggles + +```python +memprobe.enable() +# ...probed region... +memprobe.disable() +``` + +`UW_MEMPROBE=1` flips it on at import time. + +## Debugging recipes + +### "RSS grows X MiB per step — which component is it?" + +1. Set `UW_MEMPROBE=1` and run for ~20 steps to confirm the per-solve + growth pattern. +2. Add `with memprobe.probe("step N", full=True):` around your step loop. + The `full=True` walks `gc.get_objects()` and lists Python class growth + sorted by absolute change — usually the dominant suspect is on top. +3. If RSS grows but no Python class does, the leak is in PETSc memory or + C extensions. Re-run with `-log_view -malloc_dump` and inspect the PETSc + reports written at finalize. + +### "Are kd-trees being released properly?" + +Check `uw.kdtree.live_count()` directly, or look for the `kdtree: live +N` +line in the diff. KDTrees should drop to zero when their owning object +(typically a `Mesh` or `Swarm`) is destroyed. + +### Parallel runs + +`memprobe` runs on each rank independently. For meaningful aggregate +output, pipe `emit` through a rank filter: + +```python +import underworld3 as uw +def root_only(s): + if uw.mpi.rank == 0: + print(s) + +with memprobe.probe("step", emit=root_only): + ... +``` + +Or compare per-rank snapshots manually for a load-imbalance view. diff --git a/src/underworld3/ckdtree.pyx b/src/underworld3/ckdtree.pyx index 99a13087f..1f7fdf699 100644 --- a/src/underworld3/ckdtree.pyx +++ b/src/underworld3/ckdtree.pyx @@ -15,6 +15,23 @@ cdef extern from "kdtree_interface.hpp" nogil: void find_closest_point( size_t num_coords, const double* coords, long unsigned int* indices, double* out_dist_sqr, bool* found ) size_t knnSearch(const double* query_point, const size_t num_closest, long unsigned int* indices, double* out_dist_sqr ) +# Module-level live-instance counter for memory introspection. +# Incremented in __cinit__ and decremented in __dealloc__ — Cython +# guarantees deterministic destruction so this stays accurate across +# normal use. Read via uw.utilities.memprobe.snapshot(), or directly +# via uw.kdtree.live_count(). +cdef long _live_instances = 0 +cdef long _total_constructed = 0 + +def live_count(): + """Number of KDTree instances currently alive on this rank.""" + return _live_instances + +def total_constructed(): + """Total KDTree instances ever constructed on this rank.""" + return _total_constructed + + cdef class KDTree: """ Unit-aware KD-Tree for spatial indexing and queries. @@ -84,10 +101,16 @@ cdef class KDTree: self.points = points self.index = new KDTree_Interface( &points[0][0], points.shape[0], points.shape[1]) + global _live_instances, _total_constructed + _live_instances += 1 + _total_constructed += 1 + super().__init__() def __dealloc__(self): del self.index + global _live_instances + _live_instances -= 1 @property def n(self): diff --git a/src/underworld3/systems/solvers.py b/src/underworld3/systems/solvers.py index 8b2af5d88..5e26d4cc2 100644 --- a/src/underworld3/systems/solvers.py +++ b/src/underworld3/systems/solvers.py @@ -62,6 +62,7 @@ from underworld3.cython.generic_solvers import SNES_MultiComponent from underworld3 import VarType import underworld3.timing as timing +from underworld3.utilities import memprobe from underworld3.utilities._api_tools import ( uw_object, SymbolicProperty, @@ -1232,6 +1233,7 @@ def _create_stress_history_ddt(self, order=2): self.Unknowns.DFDt.enable_source_snapshot() @timing.routine_timer_decorator + @memprobe.instrument("Stokes.solve") def solve( self, zero_init_guess: bool = True, @@ -3724,6 +3726,7 @@ def penalty(self, value): self._penalty.sym = value @timing.routine_timer_decorator + @memprobe.instrument("NavierStokes.solve") def solve( self, zero_init_guess: bool = True, diff --git a/src/underworld3/utilities/__init__.py b/src/underworld3/utilities/__init__.py index 495301e31..1731805b3 100644 --- a/src/underworld3/utilities/__init__.py +++ b/src/underworld3/utilities/__init__.py @@ -83,3 +83,4 @@ def _append_petsc_path(): ) from . import retention_curves +from . import memprobe diff --git a/src/underworld3/utilities/memprobe.py b/src/underworld3/utilities/memprobe.py new file mode 100644 index 000000000..b82256e84 --- /dev/null +++ b/src/underworld3/utilities/memprobe.py @@ -0,0 +1,265 @@ +""" +Lightweight memory-growth diagnostics. + +Designed to surface slow leaks in long parallel runs (HPC OOM-class bugs) +without slowing normal use. Three signal sources: + +* **Process RSS** via :mod:`resource` (peak + current). +* **Live ``uw.kdtree.KDTree`` count** — UW's nanoflann wrapper isn't visible + to PETSc's object tracker, so it's a common silent-leak suspect. +* **Per-class Python instance counts** for the ``underworld3`` package, gated + behind ``full=True`` because :func:`gc.get_objects` is slow at scale. + +For PETSc-side object/allocation tracking, use the runtime flags this module +documents (``-log_view``, ``-malloc_dump``) — parsing PETSc's viewer output +from Python is brittle and adds no signal beyond what those flags give. + +Activation +---------- +- ``UW_MEMPROBE=1`` env var → enables instrumentation hooks at import time. +- :func:`enable` / :func:`disable` for runtime toggling. +- :func:`probe` context manager for ad-hoc snapshot+diff blocks. +- :func:`instrument` decorator for per-method instrumentation; a no-op when + disabled (one attribute lookup + branch, sub-microsecond). + +Examples +-------- +>>> import underworld3 as uw +>>> from underworld3.utilities import memprobe +>>> +>>> with memprobe.probe("step 42"): +... do_one_step() # diff is logged on exit +>>> +>>> @memprobe.instrument("stokes-solve") +... def solve(...): +... ... + +For a long run, set ``UW_MEMPROBE=1`` and add a probe block per step. To pin +PETSc-side leaks, run with ``-log_view -malloc_dump`` (or call +:func:`dump_petsc_leaks_at_finalize`). +""" +from __future__ import annotations + +import functools +import gc +import os +import resource +import sys +from collections import Counter +from contextlib import contextmanager +from typing import Any + +__all__ = [ + "ENABLED", + "enable", + "disable", + "snapshot", + "diff", + "probe", + "instrument", + "dump_petsc_leaks_at_finalize", + "format_diff", +] + +ENABLED: bool = bool(int(os.environ.get("UW_MEMPROBE", "0"))) + + +def enable() -> None: + """Turn instrumentation on at runtime.""" + global ENABLED + ENABLED = True + + +def disable() -> None: + """Turn instrumentation off at runtime.""" + global ENABLED + ENABLED = False + + +def _rss_mb() -> float: + """Current resident-set size in MiB. + + On Linux ``ru_maxrss`` is in KiB; on macOS it is in bytes. We probe the + platform once to avoid getting it wrong silently. + """ + rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + if sys.platform == "darwin": + return rss / (1024 * 1024) + return rss / 1024 + + +def _kdtree_counts() -> dict[str, int]: + try: + from underworld3 import kdtree + + return { + "live": kdtree.live_count(), + "total_constructed": kdtree.total_constructed(), + } + except (ImportError, AttributeError): + return {"live": -1, "total_constructed": -1} + + +def _python_class_counts(prefix: str = "underworld3") -> Counter[str]: + """Count live Python objects whose class lives under ``prefix``. + + Walks ``gc.get_objects()`` once. Slow at large counts (~10–100 ms for + typical UW runs). Only called when ``full=True`` is requested. + """ + counts: Counter[str] = Counter() + for obj in gc.get_objects(): + cls = type(obj) + mod = getattr(cls, "__module__", "") or "" + # Some objects expose a non-string descriptor as ``__module__``; + # skip those rather than crashing the walk. + if isinstance(mod, str) and mod.startswith(prefix): + counts[f"{mod}.{cls.__name__}"] += 1 + return counts + + +def snapshot(full: bool = False) -> dict[str, Any]: + """Capture a memory snapshot. + + Parameters + ---------- + full + If ``True``, also walk ``gc.get_objects()`` to count live Python + instances by class (slow). Defaults to ``False``. + + Returns + ------- + dict + Keys: ``rss_mb``, ``kdtree`` (mapping with ``live`` and + ``total_constructed``), ``py_classes`` (mapping; only present when + ``full=True``). + """ + snap: dict[str, Any] = { + "rss_mb": _rss_mb(), + "kdtree": _kdtree_counts(), + } + if full: + snap["py_classes"] = dict(_python_class_counts()) + return snap + + +def diff(before: dict[str, Any], after: dict[str, Any]) -> dict[str, Any]: + """Compute a structured diff between two snapshots. + + Returns deltas only (excludes anything unchanged or absent). Python + class entries are sorted by absolute change so the dominant suspects + appear first. + """ + delta: dict[str, Any] = {} + + drss = after["rss_mb"] - before["rss_mb"] + if drss != 0: + delta["rss_mb"] = drss + + kd_delta: dict[str, int] = {} + for key in ("live", "total_constructed"): + d = after["kdtree"].get(key, 0) - before["kdtree"].get(key, 0) + if d != 0: + kd_delta[key] = d + if kd_delta: + delta["kdtree"] = kd_delta + + if "py_classes" in before and "py_classes" in after: + b = before["py_classes"] + a = after["py_classes"] + keys = set(b) | set(a) + py_delta = {k: a.get(k, 0) - b.get(k, 0) for k in keys} + py_delta = {k: v for k, v in py_delta.items() if v != 0} + if py_delta: + delta["py_classes"] = dict( + sorted(py_delta.items(), key=lambda kv: abs(kv[1]), reverse=True) + ) + + return delta + + +def format_diff(label: str, delta: dict[str, Any]) -> str: + """Render a diff for stdout / logs. One line for trivial deltas, more if needed.""" + if not delta: + return f"[memprobe] {label}: no change" + + parts = [f"[memprobe] {label}:"] + if "rss_mb" in delta: + parts.append(f" RSS {delta['rss_mb']:+.2f} MiB") + if "kdtree" in delta: + kd = delta["kdtree"] + chunk = ", ".join(f"{k} {v:+d}" for k, v in kd.items()) + parts.append(f" kdtree: {chunk}") + if "py_classes" in delta: + parts.append(" py_classes (top 10 by |Δ|):") + for cls, d in list(delta["py_classes"].items())[:10]: + parts.append(f" {cls}: {d:+d}") + return "\n".join(parts) + + +@contextmanager +def probe(label: str, full: bool = False, *, emit=print): + """Context manager: snapshot before, snapshot after, log the diff. + + Parameters + ---------- + label + Short identifier for the block, included in the emitted diff. + full + Walk ``gc.get_objects()`` for per-class deltas (slow). + emit + Callable receiving the formatted diff string. Defaults to + :func:`print`. Use a logger or rank-aware writer for parallel runs. + """ + if not ENABLED: + yield + return + before = snapshot(full=full) + try: + yield + finally: + after = snapshot(full=full) + emit(format_diff(label, diff(before, after))) + + +def instrument(label: str, full: bool = False, *, emit=print): + """Decorator wrapping a method/function with :func:`probe`. + + Fast-returns when :data:`ENABLED` is ``False`` (sub-microsecond overhead), + so it's safe to leave permanent decorations on hot paths. + + Examples + -------- + >>> @instrument("stokes-solve") + ... def solve(self, ...): + ... ... + """ + def decorator(func): + @functools.wraps(func) + def wrapper(*args, **kw): + if not ENABLED: + return func(*args, **kw) + with probe(label, full=full, emit=emit): + return func(*args, **kw) + return wrapper + return decorator + + +def dump_petsc_leaks_at_finalize(filename: str | None = None) -> None: + """Ask PETSc to dump unfreed objects at finalize. + + Sets the runtime options ``-malloc_dump`` and ``-objects_dump`` so PETSc + writes its leak report when the process tears down. ``filename`` is + advisory — if provided, also sets ``-malloc_view`` / ``-options_view`` + for a more verbose dump. + + Equivalent to running with ``-malloc_dump`` on the command line. Useful + when you can't change the launch command but can call this from Python + early in the run. + """ + from petsc4py import PETSc + + opts = PETSc.Options() + opts.setValue("-malloc_dump", "") + opts.setValue("-objects_dump", "") + if filename: + opts.setValue("-malloc_view", filename) diff --git a/tests/test_0780_memprobe.py b/tests/test_0780_memprobe.py new file mode 100644 index 000000000..f6265ddce --- /dev/null +++ b/tests/test_0780_memprobe.py @@ -0,0 +1,140 @@ +""" +Smoke tests for the memprobe diagnostic module. + +These tests check the *plumbing*: instrumentation enable/disable, snapshot +shape, KDTree live-count tracking, decorator no-op-when-disabled, and +context-manager diff emission. They do not validate that any real leak is +or isn't present — that's what users do with the tool, not what the tool +itself can verify. +""" +import gc +import os +import numpy as np +import pytest + +import underworld3 as uw +from underworld3.utilities import memprobe + + +@pytest.fixture(autouse=True) +def _restore_enabled_flag(): + """Each test starts with memprobe disabled regardless of env.""" + saved = memprobe.ENABLED + memprobe.disable() + yield + memprobe.ENABLED = saved + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_snapshot_shape(): + """Snapshot must contain rss_mb (float) and kdtree (dict with two keys).""" + snap = memprobe.snapshot() + + assert isinstance(snap["rss_mb"], float) + assert snap["rss_mb"] > 0 + assert set(snap["kdtree"].keys()) == {"live", "total_constructed"} + + # full=False should NOT walk gc + assert "py_classes" not in snap + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_snapshot_full_includes_py_classes(): + snap = memprobe.snapshot(full=True) + assert "py_classes" in snap + # Sanity: gc walked something + assert isinstance(snap["py_classes"], dict) + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_kdtree_live_count_tracks_construction_destruction(): + """KDTree __cinit__/__dealloc__ must update the live-instance counter.""" + pts = np.random.random((20, 2)) + + before = uw.kdtree.live_count() + total_before = uw.kdtree.total_constructed() + + tree = uw.kdtree.KDTree(pts) + assert uw.kdtree.live_count() == before + 1 + assert uw.kdtree.total_constructed() == total_before + 1 + + del tree + gc.collect() + assert uw.kdtree.live_count() == before + # total_constructed never decreases + assert uw.kdtree.total_constructed() == total_before + 1 + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_diff_reports_kdtree_growth(): + pts = np.random.random((20, 2)) + + before = memprobe.snapshot() + tree = uw.kdtree.KDTree(pts) # noqa: F841 — kept alive across the diff + after = memprobe.snapshot() + + delta = memprobe.diff(before, after) + assert delta["kdtree"]["live"] == 1 + assert delta["kdtree"]["total_constructed"] == 1 + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_diff_drops_unchanged_keys(): + snap = memprobe.snapshot() + delta = memprobe.diff(snap, snap) + # Identical snapshots produce no deltas + assert "kdtree" not in delta + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_probe_context_emits_when_enabled(capsys): + memprobe.enable() + with memprobe.probe("test-block"): + _ = uw.kdtree.KDTree(np.random.random((10, 2))) + + captured = capsys.readouterr() + assert "[memprobe] test-block" in captured.out + assert "kdtree" in captured.out + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_probe_context_silent_when_disabled(capsys): + assert memprobe.ENABLED is False + with memprobe.probe("silent"): + _ = uw.kdtree.KDTree(np.random.random((10, 2))) + + captured = capsys.readouterr() + assert captured.out == "" + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_instrument_decorator_no_op_when_disabled(capsys): + """Decorator must not emit (or noticeably slow) when ENABLED is False.""" + @memprobe.instrument("test-fn") + def f(x): + return x * 2 + + assert f(3) == 6 + captured = capsys.readouterr() + assert captured.out == "" + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_instrument_decorator_emits_when_enabled(capsys): + @memprobe.instrument("test-fn") + def f(x): + return x * 2 + + memprobe.enable() + assert f(3) == 6 + captured = capsys.readouterr() + assert "[memprobe] test-fn" in captured.out From 389b1c9712e00bad9203f069b9b2b2175e2fbeba Mon Sep 17 00:00:00 2001 From: lmoresi Date: Sun, 10 May 2026 20:59:58 +1000 Subject: [PATCH 2/3] Address Copilot review on PR #179 + fix flaky KDTree count test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - _rss_mb(): switch from resource.ru_maxrss (peak/high-water RSS) to psutil.Process().memory_info().rss (current RSS), with /proc/self/statm Linux fallback and ru_maxrss kept only as a last resort. Current RSS is preferred so freed memory shows as a negative delta. (Copilot #1, #6) - diff(): apply 0.01 MiB threshold before including rss_mb in the delta. Smaller noise prints as "+0.00 MiB" and defeats the "no change" fast path. (Copilot #2) - dump_petsc_leaks_at_finalize(): drop leading "-" from PETSc option keys (codebase convention is no dash; "-malloc_dump" was a no-op), and align docstring with what the function actually sets. (Copilot #3) - Remove unused `os` import from test module. (Copilot #4) - Soften ckdtree.pyx comment claiming Cython "guarantees deterministic destruction" — that's CPython refcounting, which can lag for objects trapped in reference cycles until cyclic GC runs. (Copilot #5) - Test: assert KDTree count deltas relative to immediate before/after, never to absolute baselines. Earlier tests in the same pytest session can leave KDTree refs alive that the cyclic GC may collect at any time, shifting the absolute count. CI surfaced this with "assert 332 == 338" — six trees from earlier tests collected during our gc.collect(). Underworld development team with AI support from Claude Code --- docs/developer/guides/memory-diagnostics.md | 14 ++++- src/underworld3/ckdtree.pyx | 10 ++-- src/underworld3/utilities/memprobe.py | 63 ++++++++++++++++----- tests/test_0780_memprobe.py | 31 +++++++--- 4 files changed, 87 insertions(+), 31 deletions(-) diff --git a/docs/developer/guides/memory-diagnostics.md b/docs/developer/guides/memory-diagnostics.md index 6e658e755..8dea38e70 100644 --- a/docs/developer/guides/memory-diagnostics.md +++ b/docs/developer/guides/memory-diagnostics.md @@ -27,14 +27,22 @@ growing on every step. | Signal | Source | Cost | |---|---|---| -| Process RSS (MiB) | `resource.getrusage` | free | +| Process RSS (MiB, current) | `psutil.Process().memory_info().rss`, fallback `/proc/self/statm` (Linux), last resort `resource.ru_maxrss` | free | | KDTree live count | `uw.kdtree.live_count()` | free | | KDTree total constructed | `uw.kdtree.total_constructed()` | free | | Per-class Python instance counts | `gc.get_objects()` walk | slow — gated behind `full=True` | +The RSS source is **current** RSS where possible — it should drop when memory +is freed, not just rise. The last-resort `resource.ru_maxrss` fallback returns +the **peak** (high-water-mark) RSS instead and never decreases, so on systems +where neither psutil nor `/proc/self/statm` is available you'll see growth but +not recovery; install `psutil` to fix that. + `KDTree` instances are tracked via Cython class counters in `__cinit__` and -`__dealloc__`. Cython's deterministic destruction makes the live count -accurate without weak-references. +`__dealloc__`. CPython refcounting calls `__dealloc__` promptly when the +refcount hits zero, so the count is accurate for typical use; it can lag if a +KDTree ends up in a reference cycle that only the cyclic garbage collector +can break. Call `gc.collect()` before reading if that matters. PETSc-side object and allocation tracking is **not** parsed from Python — PETSc's own `-log_view` and `-malloc_dump` runtime flags give the same diff --git a/src/underworld3/ckdtree.pyx b/src/underworld3/ckdtree.pyx index 1f7fdf699..19acaedbb 100644 --- a/src/underworld3/ckdtree.pyx +++ b/src/underworld3/ckdtree.pyx @@ -16,10 +16,12 @@ cdef extern from "kdtree_interface.hpp" nogil: size_t knnSearch(const double* query_point, const size_t num_closest, long unsigned int* indices, double* out_dist_sqr ) # Module-level live-instance counter for memory introspection. -# Incremented in __cinit__ and decremented in __dealloc__ — Cython -# guarantees deterministic destruction so this stays accurate across -# normal use. Read via uw.utilities.memprobe.snapshot(), or directly -# via uw.kdtree.live_count(). +# Incremented in __cinit__, decremented in __dealloc__. CPython refcounting +# calls __dealloc__ promptly when the refcount hits zero, so the count is +# accurate for typical use; it can lag if a KDTree ends up in a reference +# cycle that only the cyclic garbage collector can break — call +# gc.collect() before reading if that matters. Read via +# uw.utilities.memprobe.snapshot() or directly via uw.kdtree.live_count(). cdef long _live_instances = 0 cdef long _total_constructed = 0 diff --git a/src/underworld3/utilities/memprobe.py b/src/underworld3/utilities/memprobe.py index b82256e84..52e22860c 100644 --- a/src/underworld3/utilities/memprobe.py +++ b/src/underworld3/utilities/memprobe.py @@ -4,7 +4,10 @@ Designed to surface slow leaks in long parallel runs (HPC OOM-class bugs) without slowing normal use. Three signal sources: -* **Process RSS** via :mod:`resource` (peak + current). +* **Process RSS** — current resident-set size via :mod:`psutil` if available, + falling back to ``/proc/self/statm`` on Linux, then to :mod:`resource`'s + high-water-mark RSS as a last resort. Current RSS is preferred so that + freed memory shows up as a negative delta. * **Live ``uw.kdtree.KDTree`` count** — UW's nanoflann wrapper isn't visible to PETSc's object tracker, so it's a common silent-leak suspect. * **Per-class Python instance counts** for the ``underworld3`` package, gated @@ -43,7 +46,6 @@ import functools import gc import os -import resource import sys from collections import Counter from contextlib import contextmanager @@ -79,9 +81,33 @@ def disable() -> None: def _rss_mb() -> float: """Current resident-set size in MiB. - On Linux ``ru_maxrss`` is in KiB; on macOS it is in bytes. We probe the - platform once to avoid getting it wrong silently. + Prefers ``psutil.Process().memory_info().rss`` (true current RSS, drops + when memory is freed). Falls back to ``/proc/self/statm`` on Linux when + psutil isn't available, then to ``resource.ru_maxrss`` (peak/high-water + RSS — does not decrease) as a last resort. The fallback is good enough + for spotting growth, less good for spotting recovery. """ + try: + import psutil + + return psutil.Process().memory_info().rss / (1024 * 1024) + except ImportError: + pass + + if sys.platform.startswith("linux"): + try: + with open("/proc/self/statm", "r") as f: + # statm: size resident shared text lib data dt (in pages) + resident_pages = int(f.read().split()[1]) + page_size = os.sysconf("SC_PAGESIZE") + return (resident_pages * page_size) / (1024 * 1024) + except (OSError, ValueError): + pass + + # Last resort: peak RSS — does NOT decrease, so growth shows but + # recovery does not. ru_maxrss is KiB on Linux, bytes on macOS. + import resource + rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss if sys.platform == "darwin": return rss / (1024 * 1024) @@ -152,7 +178,10 @@ class entries are sorted by absolute change so the dominant suspects delta: dict[str, Any] = {} drss = after["rss_mb"] - before["rss_mb"] - if drss != 0: + # Threshold below the format precision (0.01 MiB ≈ 10 KiB): smaller + # noise — inevitable when the source is RSS in KiB pages — would print + # as "+0.00 MiB" and defeat the "no change" fast path. + if abs(drss) >= 0.01: delta["rss_mb"] = drss kd_delta: dict[str, int] = {} @@ -247,19 +276,23 @@ def wrapper(*args, **kw): def dump_petsc_leaks_at_finalize(filename: str | None = None) -> None: """Ask PETSc to dump unfreed objects at finalize. - Sets the runtime options ``-malloc_dump`` and ``-objects_dump`` so PETSc - writes its leak report when the process tears down. ``filename`` is - advisory — if provided, also sets ``-malloc_view`` / ``-options_view`` - for a more verbose dump. + Sets the runtime options ``malloc_dump`` and ``objects_dump`` so PETSc + writes its leak report when the process tears down. When ``filename`` + is given, also sets ``malloc_view`` to the same path for a more verbose + allocation dump. + + Note that ``petsc4py``'s :class:`Options` API expects keys *without* the + leading ``-`` that you'd type on the command line — same convention used + elsewhere in this codebase. - Equivalent to running with ``-malloc_dump`` on the command line. Useful - when you can't change the launch command but can call this from Python - early in the run. + Equivalent to running with ``-malloc_dump -objects_dump`` on the command + line. Useful when you can't change the launch command but can call this + from Python early in the run. """ from petsc4py import PETSc opts = PETSc.Options() - opts.setValue("-malloc_dump", "") - opts.setValue("-objects_dump", "") + opts.setValue("malloc_dump", "") + opts.setValue("objects_dump", "") if filename: - opts.setValue("-malloc_view", filename) + opts.setValue("malloc_view", filename) diff --git a/tests/test_0780_memprobe.py b/tests/test_0780_memprobe.py index f6265ddce..d1b4b5414 100644 --- a/tests/test_0780_memprobe.py +++ b/tests/test_0780_memprobe.py @@ -8,7 +8,6 @@ itself can verify. """ import gc -import os import numpy as np import pytest @@ -51,21 +50,35 @@ def test_snapshot_full_includes_py_classes(): @pytest.mark.level_1 @pytest.mark.tier_a def test_kdtree_live_count_tracks_construction_destruction(): - """KDTree __cinit__/__dealloc__ must update the live-instance counter.""" + """KDTree __cinit__/__dealloc__ must update the live-instance counter. + + The test asserts only on deltas around the operation under test, never + on absolute counts, because earlier tests in the same pytest session + can leave KDTree references alive that the cyclic GC may collect at + any time, shifting the baseline. + """ pts = np.random.random((20, 2)) - before = uw.kdtree.live_count() - total_before = uw.kdtree.total_constructed() + # Flush any pending cyclic-GC clean-ups so the baseline doesn't shift + # under us between snapshots. + gc.collect() + before_live = uw.kdtree.live_count() + before_total = uw.kdtree.total_constructed() tree = uw.kdtree.KDTree(pts) - assert uw.kdtree.live_count() == before + 1 - assert uw.kdtree.total_constructed() == total_before + 1 + assert uw.kdtree.live_count() - before_live == 1 + assert uw.kdtree.total_constructed() - before_total == 1 + + after_create_live = uw.kdtree.live_count() del tree gc.collect() - assert uw.kdtree.live_count() == before - # total_constructed never decreases - assert uw.kdtree.total_constructed() == total_before + 1 + + # The only thing we care about: this construction/destruction pair + # produced a +1/-1 swing relative to its immediate before/after, and + # total_constructed never went backwards. + assert uw.kdtree.live_count() - after_create_live == -1 + assert uw.kdtree.total_constructed() >= before_total + 1 @pytest.mark.level_1 From 5f64de2ab896a899a9e7c94fcd11d9b1cc86e7e8 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Sun, 10 May 2026 22:13:22 +1000 Subject: [PATCH 3/3] Fix dictionary-iteration race in UWexpression._ephemeral_expr_names MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI on PR #179 surfaced a flaky failure in test_bc_accepts_raw_numbers: any(k[0] == name for k in UWexpression._ephemeral_expr_names) RuntimeError: dictionary changed size during iteration The dict is mutated asynchronously by weakref finalizers running during cyclic GC, so iterating it directly races against those callbacks. The fix is to snapshot the keys with list(...) before iterating — at most a few hundred entries, negligible cost. Pre-existing bug — surfaces depending on test ordering, GC pressure, and timing. The memprobe PR's added tests happen to shift teardown state enough to trigger it consistently. Worth landing here so #179 isn't blocked. Underworld development team with AI support from Claude Code --- src/underworld3/function/expressions.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/underworld3/function/expressions.py b/src/underworld3/function/expressions.py index 2f24914b3..2a35b6312 100644 --- a/src/underworld3/function/expressions.py +++ b/src/underworld3/function/expressions.py @@ -659,8 +659,13 @@ def __new__( # Check both dicts for name collisions name_exists_persistent = name in UWexpression._expr_names - # Check ephemeral dict - need to look for any key starting with this name - name_exists_ephemeral = any(k[0] == name for k in UWexpression._ephemeral_expr_names) + # Check ephemeral dict - need to look for any key starting with this name. + # Snapshot keys via list(...) before iterating: the underlying dict can + # be mutated mid-iteration by weakref finalizers running asynchronously + # during cyclic GC, raising "dictionary changed size during iteration". + name_exists_ephemeral = any( + k[0] == name for k in list(UWexpression._ephemeral_expr_names) + ) # Determine unique ID for disambiguation # When _unique_name_generation=True, ALWAYS use instance_no as _uw_id