Add memprobe diagnostic module for memory-growth tracking (#176)

docs/developer/guides/memory-diagnostics.md

@@ -0,0 +1,153 @@
+# 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, 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__`. 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
+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.

src/underworld3/ckdtree.pyx

@@ -15,6 +15,25 @@ 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__, 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
+
+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 +103,16 @@ cdef class KDTree:
         self.points = points
         self.index = new KDTree_Interface(<const double *> &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):

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

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,

src/underworld3/utilities/__init__.py

@@ -83,3 +83,4 @@ def _append_petsc_path():
 )
 
 from . import retention_curves
+from . import memprobe