From ddd1e5cfe4fd8074f1ff6befa50946c84ee93fe4 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Sun, 10 May 2026 12:05:40 +1000 Subject: [PATCH 01/13] Fix swarm particle loss across rank boundaries during advection (#175) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Swarm.advection() ended with self.migrate(delete_lost_points=True, max_its=1). max_its=1 only consults the closest domain centroid for each unclaimed particle, so any boundary particle whose owner happened to be the 2nd-or-3rd-closest rank failed points_in_domain on its sole try and was silently deleted. The override looks like leftover debugging — there is no reason for the end-of-advection migrate to truncate the kdtree retry. Drop it so the default max_its=10 is used, restoring the boundary-claim retry path. Adds a parallel regression test (tests/parallel/test_0765) adapted from @bknight1's reproducer in the bug report. With the fix in place: 4 ranks, no particles lost. Without it: 30/726 lost on order=1 advection, 66/726 on order=2 rotation. Reported and diagnosed by @bknight1. Underworld development team with AI support from Claude Code --- src/underworld3/swarm.py | 8 +- .../test_0765_swarm_advection_no_loss.py | 96 +++++++++++++++++++ 2 files changed, 102 insertions(+), 2 deletions(-) create mode 100644 tests/parallel/test_0765_swarm_advection_no_loss.py diff --git a/src/underworld3/swarm.py b/src/underworld3/swarm.py index ff8dc593d..22662a95f 100644 --- a/src/underworld3/swarm.py +++ b/src/underworld3/swarm.py @@ -4537,10 +4537,14 @@ def advection( # # - # Remove points no longer in the domain + # Re-route particles to their owning ranks and remove any that + # have genuinely left the domain. Use the default max_its so that + # boundary particles whose owner is the 2nd/3rd-closest centroid + # get reclaimed via the kdtree retry — max_its=1 here was an + # accidental regression that deleted boundary particles (issue #175, + # reported by @bknight1). self.migrate( delete_lost_points=True, - max_its=1, ) return diff --git a/tests/parallel/test_0765_swarm_advection_no_loss.py b/tests/parallel/test_0765_swarm_advection_no_loss.py new file mode 100644 index 000000000..6a68d3f52 --- /dev/null +++ b/tests/parallel/test_0765_swarm_advection_no_loss.py @@ -0,0 +1,96 @@ +""" +Regression test for swarm particle loss during advection across processor boundaries. + +Bug: ``Swarm.advection()`` finished with ``self.migrate(delete_lost_points=True, +max_its=1)``. ``max_its=1`` only tries the *closest* domain centroid for each +unclaimed particle, so any particle whose owning rank happened to be the +2nd-or-3rd closest centroid (typical near a process boundary) failed +``points_in_domain`` on its sole try and got deleted. + +Reproducer adapted from @bknight1's report on issue #175. + +See: https://github.com/underworldcode/underworld3/issues/175 + +Run with: + mpirun -n 2 python -m pytest --with-mpi tests/parallel/test_0765_swarm_advection_no_loss.py + mpirun -n 4 python -m pytest --with-mpi tests/parallel/test_0765_swarm_advection_no_loss.py +""" + +import pytest +import numpy as np +import sympy +import underworld3 as uw +from mpi4py import MPI + +pytestmark = [pytest.mark.mpi(min_size=2), pytest.mark.timeout(60)] + + +@pytest.mark.mpi(min_size=2) +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_advection_preserves_particle_count_order1(): + """Order-1 advection across rank boundaries must not lose particles.""" + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=0.1, + ) + + swarm = uw.swarm.Swarm(mesh) + var = uw.swarm.SwarmVariable("test_var", swarm, 1) + swarm.populate(fill_param=1) + var.data[...] = uw.mpi.rank + + comm = uw.mpi.comm + initial_count = comm.allreduce(swarm.dm.getLocalSize(), op=MPI.SUM) + + # Constant rightward velocity; dt=0.6 sweeps particles ~6 cells across the + # domain in one step. With max_its=1 in the post-advection migrate this + # used to drop boundary particles. + v_fn = sympy.Matrix([1.0, 0.0]) + swarm.advection(v_fn, 0.6, order=1) + + final_count = comm.allreduce(swarm.dm.getLocalSize(), op=MPI.SUM) + + assert final_count == initial_count, ( + f"Lost {initial_count - final_count} particles " + f"(initial={initial_count}, final={final_count})" + ) + + +@pytest.mark.mpi(min_size=2) +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_advection_preserves_particle_count_order2(): + """Order-2 mid-point advection across rank boundaries must not lose particles. + + Uses solid-body rotation about the domain centre so particles stay inside + the unit box for the whole rotation. This isolates rank-boundary loss + from genuine domain-exit, while still forcing every particle across + multiple processor cuts over many substeps. + """ + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=0.1, + ) + + swarm = uw.swarm.Swarm(mesh) + swarm.populate(fill_param=1) + + comm = uw.mpi.comm + initial_count = comm.allreduce(swarm.dm.getLocalSize(), op=MPI.SUM) + + # Solid-body rotation about (0.5, 0.5) — radius <= sqrt(0.5) so nothing + # can escape the unit box. omega=2*pi makes one revolution in dt=1. + x, y = mesh.X + omega = 2 * sympy.pi + v_fn = sympy.Matrix([-omega * (y - 0.5), omega * (x - 0.5)]) + swarm.advection(v_fn, 0.5, order=2, step_limit=True) + + final_count = comm.allreduce(swarm.dm.getLocalSize(), op=MPI.SUM) + + assert final_count == initial_count, ( + f"Lost {initial_count - final_count} particles " + f"(initial={initial_count}, final={final_count})" + ) From 9cce3220e0fc6892774b7a7f1839d464ad486c9c Mon Sep 17 00:00:00 2001 From: lmoresi Date: Sun, 10 May 2026 20:39:40 +1000 Subject: [PATCH 02/13] 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 03/13] 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 04/13] 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 From 6bb1794f720da0dbbc52a1c4bcb5a7f2730152ae Mon Sep 17 00:00:00 2001 From: Ben Knight <55677727+bknight1@users.noreply.github.com> Date: Mon, 11 May 2026 13:34:37 +0800 Subject: [PATCH 05/13] Fix memory leaks in PtrContainer Cython class Implement __dealloc__ to ensure malloc'd function pointer arrays are freed when the object is garbage collected. Added re-allocation guards in allocate() to prevent leaks when a container is reused. This resolves heap growth observed when solvers are repeatedly instantiated and destroyed during parameter sweeps or time-looping. --- src/underworld3/cython/petsc_types.pyx | 29 +++++++++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/src/underworld3/cython/petsc_types.pyx b/src/underworld3/cython/petsc_types.pyx index 609dad2a6..247401832 100644 --- a/src/underworld3/cython/petsc_types.pyx +++ b/src/underworld3/cython/petsc_types.pyx @@ -1,9 +1,36 @@ -from libc.stdlib cimport malloc +from libc.stdlib cimport malloc, free cdef class PtrContainer: + def __cinit__(self): + self.fns_residual = NULL + self.fns_bcs = NULL + self.fns_jacobian = NULL + self.fns_bd_residual = NULL + self.fns_bd_jacobian = NULL + + def __dealloc__(self): + if self.fns_residual != NULL: + free(self.fns_residual) + if self.fns_bcs != NULL: + free(self.fns_bcs) + if self.fns_jacobian != NULL: + free(self.fns_jacobian) + if self.fns_bd_residual != NULL: + free(self.fns_bd_residual) + if self.fns_bd_jacobian != NULL: + free(self.fns_bd_jacobian) + cpdef allocate(self, int n_res, int n_bcs, int n_jac, int n_bd_res, int n_bd_jac): """Allocate function pointer arrays of the given sizes.""" + + # Free existing memory if already allocated + if self.fns_residual != NULL: free(self.fns_residual) + if self.fns_bcs != NULL: free(self.fns_bcs) + if self.fns_jacobian != NULL: free(self.fns_jacobian) + if self.fns_bd_residual != NULL: free(self.fns_bd_residual) + if self.fns_bd_jacobian != NULL: free(self.fns_bd_jacobian) + self.fns_residual = malloc(n_res * sizeof(PetscDSResidualFn)) self.fns_bcs = malloc(n_bcs * sizeof(PetscDSResidualFn)) self.fns_jacobian = malloc(n_jac * sizeof(PetscDSJacobianFn)) From 545eb8d2f8652d4109936cfee952ff5050e5ce71 Mon Sep 17 00:00:00 2001 From: Ben Knight <55677727+bknight1@users.noreply.github.com> Date: Mon, 11 May 2026 13:35:44 +0800 Subject: [PATCH 06/13] Harden array callbacks against object destruction and recursion Updated mesh and variable update callbacks to verify 'array.owner' before processing. This prevents AttributeErrors during complex object teardowns (e.g. at exit or during mesh replacement). Integrated explicit recursion guards and mesh_update_lock checks in MeshVariable callbacks to ensure PETSc synchronization doesn't occur during sensitive coordinate deformations, preventing state corruption. --- .../discretisation/discretisation_mesh.py | 26 +++++--- .../discretisation_mesh_variables.py | 62 +++++++++++-------- src/underworld3/swarm.py | 8 ++- 3 files changed, 60 insertions(+), 36 deletions(-) diff --git a/src/underworld3/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index bf3e40752..a77b9c928 100644 --- a/src/underworld3/discretisation/discretisation_mesh.py +++ b/src/underworld3/discretisation/discretisation_mesh.py @@ -594,14 +594,18 @@ class replacement_boundaries(Enum): # to handle that so we just wrap it here. def mesh_update_callback(array, change_context): + mesh = array.owner + if mesh is None: + return + print(f"Mesh update callback - mesh deform") - coords = array.reshape(-1, array.owner.cdim) - self._deform_mesh(coords, verbose=True) + coords = array.reshape(-1, mesh.cdim) + mesh._deform_mesh(coords, verbose=True) # Increment mesh version to notify registered swarms of coordinate changes - with self._mesh_update_lock: - self._mesh_version += 1 - print(f"Mesh version incremented to {self._mesh_version}") + with mesh._mesh_update_lock: + mesh._mesh_version += 1 + print(f"Mesh version incremented to {mesh._mesh_version}") return @@ -1337,10 +1341,14 @@ def _re_extract_from_parent(self, verbose=False): ) def mesh_update_callback(array, change_context): - coords = array.reshape(-1, array.owner.cdim) - self._deform_mesh(coords, verbose=False) - with self._mesh_update_lock: - self._mesh_version += 1 + mesh = array.owner + if mesh is None: + return + + coords = array.reshape(-1, mesh.cdim) + mesh._deform_mesh(coords, verbose=False) + with mesh._mesh_update_lock: + mesh._mesh_version += 1 return self._coords.add_callback(mesh_update_callback) diff --git a/src/underworld3/discretisation/discretisation_mesh_variables.py b/src/underworld3/discretisation/discretisation_mesh_variables.py index 510787459..db2e46ced 100644 --- a/src/underworld3/discretisation/discretisation_mesh_variables.py +++ b/src/underworld3/discretisation/discretisation_mesh_variables.py @@ -492,36 +492,40 @@ def _create_variable_array(self, initial_data=None): # Single callback function (following mesh_update_callback pattern) def variable_update_callback(array, change_context): """Callback to sync variable changes back to PETSc (like mesh.points)""" + var = array.owner + if var is None: + return + # Only act on data-changing operations (following mesh.points pattern) data_changed = change_context.get("data_has_changed", True) if not data_changed: return # Prevent recursion by checking if we're already in a callback - if hasattr(self, "_in_callback") and self._in_callback: + if hasattr(var, "_in_callback") and var._in_callback: return # Set recursion guard - self._in_callback = True + var._in_callback = True try: # Skip updates during mesh coordinate changes to prevent corruption # Check if mesh is currently being updated - if hasattr(self.mesh, "_mesh_update_lock"): + if hasattr(var.mesh, "_mesh_update_lock"): # Try to acquire lock without blocking - if we can't, skip update - if not self.mesh._mesh_update_lock.acquire(blocking=False): + if not var.mesh._mesh_update_lock.acquire(blocking=False): return try: # Persist changes to PETSc (like mesh callback updates coordinates) - self.pack_uw_data_to_petsc(array, sync=True) + var.pack_uw_data_to_petsc(array, sync=True) finally: - self.mesh._mesh_update_lock.release() + var.mesh._mesh_update_lock.release() else: # Fallback if no lock exists - self.pack_uw_data_to_petsc(array, sync=True) + var.pack_uw_data_to_petsc(array, sync=True) finally: # Clear recursion guard - self._in_callback = False + var._in_callback = False # Register the callback (following mesh.points pattern) array_obj.add_callback(variable_update_callback) @@ -556,34 +560,38 @@ def _create_flat_data_array(self, initial_data=None): # Callback for flat data format def flat_data_update_callback(array, change_context): """Callback to sync flat data changes back to PETSc""" + var = array.owner + if var is None: + return + # Only act on data-changing operations data_changed = change_context.get("data_has_changed", True) if not data_changed: return # Prevent recursion by checking if we're already in a callback - if hasattr(self, "_in_flat_callback") and self._in_flat_callback: + if hasattr(var, "_in_flat_callback") and var._in_flat_callback: return # Set recursion guard - self._in_flat_callback = True + var._in_flat_callback = True try: # Skip updates during mesh coordinate changes to prevent corruption - if hasattr(self.mesh, "_mesh_update_lock"): - if not self.mesh._mesh_update_lock.acquire(blocking=False): + if hasattr(var.mesh, "_mesh_update_lock"): + if not var.mesh._mesh_update_lock.acquire(blocking=False): return try: # Use pack_raw for flat data format - self.pack_raw_data_to_petsc(array, sync=True) + var.pack_raw_data_to_petsc(array, sync=True) finally: - self.mesh._mesh_update_lock.release() + var.mesh._mesh_update_lock.release() else: # Fallback if no lock exists - self.pack_raw_data_to_petsc(array, sync=True) + var.pack_raw_data_to_petsc(array, sync=True) finally: # Clear recursion guard - self._in_flat_callback = False + var._in_flat_callback = False # Register the callback array_obj.add_callback(flat_data_update_callback) @@ -2541,6 +2549,10 @@ def _create_canonical_data_array(self): # Single canonical callback for PETSc synchronization def canonical_data_callback(array, change_context): """ONLY callback that handles PETSc synchronization - prevents conflicts""" + var = array.owner + if var is None: + return + # Only act on data-changing operations data_changed = change_context.get("data_has_changed", True) if not data_changed: @@ -2556,26 +2568,26 @@ def canonical_data_callback(array, change_context): canonical_array = np.atleast_2d(array) - if canonical_array.shape != (canonical_array.shape[0], self.num_components): + if canonical_array.shape != (canonical_array.shape[0], var.num_components): # Only reshape if we actually need to - canonical_array = canonical_array.reshape(-1, self.num_components) + canonical_array = canonical_array.reshape(-1, var.num_components) # Skip updates during mesh coordinate changes to prevent corruption - if hasattr(self.mesh, "_mesh_update_lock"): - if not self.mesh._mesh_update_lock.acquire(blocking=False): + if hasattr(var.mesh, "_mesh_update_lock"): + if not var.mesh._mesh_update_lock.acquire(blocking=False): return try: # STEP 1: Sync to PETSc using established method with correct shape - self.pack_raw_data_to_petsc(canonical_array, sync=True) + var.pack_raw_data_to_petsc(canonical_array, sync=True) finally: - self.mesh._mesh_update_lock.release() + var.mesh._mesh_update_lock.release() else: # Fallback if no lock exists - self.pack_raw_data_to_petsc(canonical_array, sync=True) + var.pack_raw_data_to_petsc(canonical_array, sync=True) # STEP 2: Handle variable-specific updates (extensible like SwarmVariable) - if hasattr(self, "_on_data_changed"): - self._on_data_changed() + if hasattr(var, "_on_data_changed"): + var._on_data_changed() array_obj.add_callback(canonical_data_callback) return array_obj diff --git a/src/underworld3/swarm.py b/src/underworld3/swarm.py index ff8dc593d..04a2c43e1 100644 --- a/src/underworld3/swarm.py +++ b/src/underworld3/swarm.py @@ -393,17 +393,21 @@ def _create_variable_array(self, initial_data=None): # Single callback function (following swarm_update_callback pattern) def variable_update_callback(array, change_context): """Callback to sync variable changes back to PETSc (like swarm.points)""" + var = array.owner + if var is None: + return + # Only act on data-changing operations (following swarm.points pattern) data_changed = change_context.get("data_has_changed", True) if not data_changed: return # Skip updates during coordinate changes to prevent corruption - if hasattr(self.swarm, "_migration_disabled") and self.swarm._migration_disabled: + if hasattr(var.swarm, "_migration_disabled") and var.swarm._migration_disabled: return # Persist changes to PETSc (like swarm callback updates coordinates) - self.pack_uw_data_to_petsc(array, sync=True) + var.pack_uw_data_to_petsc(array, sync=True) # Register the callback (following swarm.points pattern) array_obj.add_callback(variable_update_callback) From 5579fb18ffbbc351c0f92ce0108fc625cee0a652 Mon Sep 17 00:00:00 2001 From: Ben Knight <55677727+bknight1@users.noreply.github.com> Date: Mon, 11 May 2026 13:34:37 +0800 Subject: [PATCH 07/13] Fix memory leaks in PtrContainer Cython class Implement __dealloc__ to ensure malloc'd function pointer arrays are freed when the object is garbage collected. Added re-allocation guards in allocate() to prevent leaks when a container is reused. This resolves heap growth observed when solvers are repeatedly instantiated and destroyed during parameter sweeps or time-looping. --- src/underworld3/cython/petsc_types.pyx | 29 +++++++++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/src/underworld3/cython/petsc_types.pyx b/src/underworld3/cython/petsc_types.pyx index 609dad2a6..247401832 100644 --- a/src/underworld3/cython/petsc_types.pyx +++ b/src/underworld3/cython/petsc_types.pyx @@ -1,9 +1,36 @@ -from libc.stdlib cimport malloc +from libc.stdlib cimport malloc, free cdef class PtrContainer: + def __cinit__(self): + self.fns_residual = NULL + self.fns_bcs = NULL + self.fns_jacobian = NULL + self.fns_bd_residual = NULL + self.fns_bd_jacobian = NULL + + def __dealloc__(self): + if self.fns_residual != NULL: + free(self.fns_residual) + if self.fns_bcs != NULL: + free(self.fns_bcs) + if self.fns_jacobian != NULL: + free(self.fns_jacobian) + if self.fns_bd_residual != NULL: + free(self.fns_bd_residual) + if self.fns_bd_jacobian != NULL: + free(self.fns_bd_jacobian) + cpdef allocate(self, int n_res, int n_bcs, int n_jac, int n_bd_res, int n_bd_jac): """Allocate function pointer arrays of the given sizes.""" + + # Free existing memory if already allocated + if self.fns_residual != NULL: free(self.fns_residual) + if self.fns_bcs != NULL: free(self.fns_bcs) + if self.fns_jacobian != NULL: free(self.fns_jacobian) + if self.fns_bd_residual != NULL: free(self.fns_bd_residual) + if self.fns_bd_jacobian != NULL: free(self.fns_bd_jacobian) + self.fns_residual = malloc(n_res * sizeof(PetscDSResidualFn)) self.fns_bcs = malloc(n_bcs * sizeof(PetscDSResidualFn)) self.fns_jacobian = malloc(n_jac * sizeof(PetscDSJacobianFn)) From 160805d095b330a9aa3558b472a5cabe7b26f315 Mon Sep 17 00:00:00 2001 From: Ben Knight <55677727+bknight1@users.noreply.github.com> Date: Mon, 11 May 2026 13:35:44 +0800 Subject: [PATCH 08/13] Harden array callbacks against object destruction and recursion Updated mesh and variable update callbacks to verify 'array.owner' before processing. This prevents AttributeErrors during complex object teardowns (e.g. at exit or during mesh replacement). Integrated explicit recursion guards and mesh_update_lock checks in MeshVariable callbacks to ensure PETSc synchronization doesn't occur during sensitive coordinate deformations, preventing state corruption. --- .../discretisation/discretisation_mesh.py | 26 +++++--- .../discretisation_mesh_variables.py | 62 +++++++++++-------- src/underworld3/swarm.py | 8 ++- 3 files changed, 60 insertions(+), 36 deletions(-) diff --git a/src/underworld3/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index bf3e40752..a77b9c928 100644 --- a/src/underworld3/discretisation/discretisation_mesh.py +++ b/src/underworld3/discretisation/discretisation_mesh.py @@ -594,14 +594,18 @@ class replacement_boundaries(Enum): # to handle that so we just wrap it here. def mesh_update_callback(array, change_context): + mesh = array.owner + if mesh is None: + return + print(f"Mesh update callback - mesh deform") - coords = array.reshape(-1, array.owner.cdim) - self._deform_mesh(coords, verbose=True) + coords = array.reshape(-1, mesh.cdim) + mesh._deform_mesh(coords, verbose=True) # Increment mesh version to notify registered swarms of coordinate changes - with self._mesh_update_lock: - self._mesh_version += 1 - print(f"Mesh version incremented to {self._mesh_version}") + with mesh._mesh_update_lock: + mesh._mesh_version += 1 + print(f"Mesh version incremented to {mesh._mesh_version}") return @@ -1337,10 +1341,14 @@ def _re_extract_from_parent(self, verbose=False): ) def mesh_update_callback(array, change_context): - coords = array.reshape(-1, array.owner.cdim) - self._deform_mesh(coords, verbose=False) - with self._mesh_update_lock: - self._mesh_version += 1 + mesh = array.owner + if mesh is None: + return + + coords = array.reshape(-1, mesh.cdim) + mesh._deform_mesh(coords, verbose=False) + with mesh._mesh_update_lock: + mesh._mesh_version += 1 return self._coords.add_callback(mesh_update_callback) diff --git a/src/underworld3/discretisation/discretisation_mesh_variables.py b/src/underworld3/discretisation/discretisation_mesh_variables.py index 510787459..db2e46ced 100644 --- a/src/underworld3/discretisation/discretisation_mesh_variables.py +++ b/src/underworld3/discretisation/discretisation_mesh_variables.py @@ -492,36 +492,40 @@ def _create_variable_array(self, initial_data=None): # Single callback function (following mesh_update_callback pattern) def variable_update_callback(array, change_context): """Callback to sync variable changes back to PETSc (like mesh.points)""" + var = array.owner + if var is None: + return + # Only act on data-changing operations (following mesh.points pattern) data_changed = change_context.get("data_has_changed", True) if not data_changed: return # Prevent recursion by checking if we're already in a callback - if hasattr(self, "_in_callback") and self._in_callback: + if hasattr(var, "_in_callback") and var._in_callback: return # Set recursion guard - self._in_callback = True + var._in_callback = True try: # Skip updates during mesh coordinate changes to prevent corruption # Check if mesh is currently being updated - if hasattr(self.mesh, "_mesh_update_lock"): + if hasattr(var.mesh, "_mesh_update_lock"): # Try to acquire lock without blocking - if we can't, skip update - if not self.mesh._mesh_update_lock.acquire(blocking=False): + if not var.mesh._mesh_update_lock.acquire(blocking=False): return try: # Persist changes to PETSc (like mesh callback updates coordinates) - self.pack_uw_data_to_petsc(array, sync=True) + var.pack_uw_data_to_petsc(array, sync=True) finally: - self.mesh._mesh_update_lock.release() + var.mesh._mesh_update_lock.release() else: # Fallback if no lock exists - self.pack_uw_data_to_petsc(array, sync=True) + var.pack_uw_data_to_petsc(array, sync=True) finally: # Clear recursion guard - self._in_callback = False + var._in_callback = False # Register the callback (following mesh.points pattern) array_obj.add_callback(variable_update_callback) @@ -556,34 +560,38 @@ def _create_flat_data_array(self, initial_data=None): # Callback for flat data format def flat_data_update_callback(array, change_context): """Callback to sync flat data changes back to PETSc""" + var = array.owner + if var is None: + return + # Only act on data-changing operations data_changed = change_context.get("data_has_changed", True) if not data_changed: return # Prevent recursion by checking if we're already in a callback - if hasattr(self, "_in_flat_callback") and self._in_flat_callback: + if hasattr(var, "_in_flat_callback") and var._in_flat_callback: return # Set recursion guard - self._in_flat_callback = True + var._in_flat_callback = True try: # Skip updates during mesh coordinate changes to prevent corruption - if hasattr(self.mesh, "_mesh_update_lock"): - if not self.mesh._mesh_update_lock.acquire(blocking=False): + if hasattr(var.mesh, "_mesh_update_lock"): + if not var.mesh._mesh_update_lock.acquire(blocking=False): return try: # Use pack_raw for flat data format - self.pack_raw_data_to_petsc(array, sync=True) + var.pack_raw_data_to_petsc(array, sync=True) finally: - self.mesh._mesh_update_lock.release() + var.mesh._mesh_update_lock.release() else: # Fallback if no lock exists - self.pack_raw_data_to_petsc(array, sync=True) + var.pack_raw_data_to_petsc(array, sync=True) finally: # Clear recursion guard - self._in_flat_callback = False + var._in_flat_callback = False # Register the callback array_obj.add_callback(flat_data_update_callback) @@ -2541,6 +2549,10 @@ def _create_canonical_data_array(self): # Single canonical callback for PETSc synchronization def canonical_data_callback(array, change_context): """ONLY callback that handles PETSc synchronization - prevents conflicts""" + var = array.owner + if var is None: + return + # Only act on data-changing operations data_changed = change_context.get("data_has_changed", True) if not data_changed: @@ -2556,26 +2568,26 @@ def canonical_data_callback(array, change_context): canonical_array = np.atleast_2d(array) - if canonical_array.shape != (canonical_array.shape[0], self.num_components): + if canonical_array.shape != (canonical_array.shape[0], var.num_components): # Only reshape if we actually need to - canonical_array = canonical_array.reshape(-1, self.num_components) + canonical_array = canonical_array.reshape(-1, var.num_components) # Skip updates during mesh coordinate changes to prevent corruption - if hasattr(self.mesh, "_mesh_update_lock"): - if not self.mesh._mesh_update_lock.acquire(blocking=False): + if hasattr(var.mesh, "_mesh_update_lock"): + if not var.mesh._mesh_update_lock.acquire(blocking=False): return try: # STEP 1: Sync to PETSc using established method with correct shape - self.pack_raw_data_to_petsc(canonical_array, sync=True) + var.pack_raw_data_to_petsc(canonical_array, sync=True) finally: - self.mesh._mesh_update_lock.release() + var.mesh._mesh_update_lock.release() else: # Fallback if no lock exists - self.pack_raw_data_to_petsc(canonical_array, sync=True) + var.pack_raw_data_to_petsc(canonical_array, sync=True) # STEP 2: Handle variable-specific updates (extensible like SwarmVariable) - if hasattr(self, "_on_data_changed"): - self._on_data_changed() + if hasattr(var, "_on_data_changed"): + var._on_data_changed() array_obj.add_callback(canonical_data_callback) return array_obj diff --git a/src/underworld3/swarm.py b/src/underworld3/swarm.py index 22662a95f..72f659179 100644 --- a/src/underworld3/swarm.py +++ b/src/underworld3/swarm.py @@ -393,17 +393,21 @@ def _create_variable_array(self, initial_data=None): # Single callback function (following swarm_update_callback pattern) def variable_update_callback(array, change_context): """Callback to sync variable changes back to PETSc (like swarm.points)""" + var = array.owner + if var is None: + return + # Only act on data-changing operations (following swarm.points pattern) data_changed = change_context.get("data_has_changed", True) if not data_changed: return # Skip updates during coordinate changes to prevent corruption - if hasattr(self.swarm, "_migration_disabled") and self.swarm._migration_disabled: + if hasattr(var.swarm, "_migration_disabled") and var.swarm._migration_disabled: return # Persist changes to PETSc (like swarm callback updates coordinates) - self.pack_uw_data_to_petsc(array, sync=True) + var.pack_uw_data_to_petsc(array, sync=True) # Register the callback (following swarm.points pattern) array_obj.add_callback(variable_update_callback) From 872a9871a41edd8b4825871421d983f7f51a72e3 Mon Sep 17 00:00:00 2001 From: Ben Knight <55677727+bknight1@users.noreply.github.com> Date: Tue, 12 May 2026 08:16:09 +0800 Subject: [PATCH 09/13] docs: clarify callback safety and coordinate locking - Added safety comments to NDArray callbacks in MeshVariable and SwarmVariable explaining owner-existence checks for object teardown. - Documented _mesh_update_lock in Mesh and ensured _deform_mesh is properly wrapped to prevent PETSc sync conflicts during coordinate changes. - Verified test stability for integrals and swarm statistics. --- .../discretisation/discretisation_mesh.py | 78 ++++++++++--------- .../discretisation_mesh_variables.py | 1 + src/underworld3/swarm.py | 3 + 3 files changed, 46 insertions(+), 36 deletions(-) diff --git a/src/underworld3/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index a77b9c928..fe5994ede 100644 --- a/src/underworld3/discretisation/discretisation_mesh.py +++ b/src/underworld3/discretisation/discretisation_mesh.py @@ -285,6 +285,11 @@ def __init__( self._registered_swarms = weakref.WeakSet() self._registered_surfaces = weakref.WeakSet() # Surfaces using this mesh self._registered_submeshes = weakref.WeakSet() # Submeshes from extract_region + + # _mesh_update_lock: Re-entrant lock to coordinate mesh deformation. + # Held by mesh_update_callback during _deform_mesh(). Checked by + # MeshVariable callbacks (blocking=False) to skip PETSc sync during + # sensitive coordinate changes. self._mesh_update_lock = threading.RLock() comm = PETSc.COMM_WORLD @@ -1785,42 +1790,43 @@ def _deform_mesh(self, new_coords: numpy.ndarray, verbose=False): The coord array that is passed in should match the shape of self.data """ - coord_vec = self.dm.getCoordinatesLocal() - coords = coord_vec.array.reshape(-1, self.cdim) - coords[...] = new_coords[...] - - self.dm.setCoordinatesLocal(coord_vec) - self.nuke_coords_and_rebuild() - - # Rebuild the _coords array view. nuke_coords_and_rebuild may - # replace the coordinate vector internally (createCoordinateSpace), - # leaving self._coords as a stale numpy view of the old buffer. - import underworld3.utilities - old_callbacks = getattr(self._coords, "_callbacks", []) - self._coords = underworld3.utilities.NDArray_With_Callback( - numpy.ndarray.view( - self.dm.getCoordinatesLocal().array.reshape(-1, self.cdim) - ), - owner=self, - ) - for cb in old_callbacks: - self._coords.add_callback(cb) - - # BUGFIX(#122): mark registered solvers for rebuild. Since PR #127 - # ("Trust JIT cache: skip DM rebuild on constant-only parameter - # changes") a solver with is_setup=True trusts its cached PETSc DM - # / SNES assembly and skips rebuild on the next solve(). After a - # coordinate change the cached DM still carries pre-deform - # coordinates, so F(v_prev) ≈ 0 and the solver converges in 0 - # iterations without updating the solution. mesh.adapt() already - # does this; _deform_mesh must match. - for solver in self._equation_systems_register: - if solver is not None and hasattr(solver, "is_setup"): - solver.is_setup = False - - # Propagate coordinate changes to registered submeshes - for submesh in self._registered_submeshes: - submesh.sync_coordinates_from_parent() + with self._mesh_update_lock: + coord_vec = self.dm.getCoordinatesLocal() + coords = coord_vec.array.reshape(-1, self.cdim) + coords[...] = new_coords[...] + + self.dm.setCoordinatesLocal(coord_vec) + self.nuke_coords_and_rebuild() + + # Rebuild the _coords array view. nuke_coords_and_rebuild may + # replace the coordinate vector internally (createCoordinateSpace), + # leaving self._coords as a stale numpy view of the old buffer. + import underworld3.utilities + old_callbacks = getattr(self._coords, "_callbacks", []) + self._coords = underworld3.utilities.NDArray_With_Callback( + numpy.ndarray.view( + self.dm.getCoordinatesLocal().array.reshape(-1, self.cdim) + ), + owner=self, + ) + for cb in old_callbacks: + self._coords.add_callback(cb) + + # BUGFIX(#122): mark registered solvers for rebuild. Since PR #127 + # ("Trust JIT cache: skip DM rebuild on constant-only parameter + # changes") a solver with is_setup=True trusts its cached PETSc DM + # / SNES assembly and skips rebuild on the next solve(). After a + # coordinate change the cached DM still carries pre-deform + # coordinates, so F(v_prev) ≈ 0 and the solver converges in 0 + # iterations without updating the solution. mesh.adapt() already + # does this; _deform_mesh must match. + for solver in self._equation_systems_register: + if solver is not None and hasattr(solver, "is_setup"): + solver.is_setup = False + + # Propagate coordinate changes to registered submeshes + for submesh in self._registered_submeshes: + submesh.sync_coordinates_from_parent() return diff --git a/src/underworld3/discretisation/discretisation_mesh_variables.py b/src/underworld3/discretisation/discretisation_mesh_variables.py index db2e46ced..47c0ed701 100644 --- a/src/underworld3/discretisation/discretisation_mesh_variables.py +++ b/src/underworld3/discretisation/discretisation_mesh_variables.py @@ -494,6 +494,7 @@ def variable_update_callback(array, change_context): """Callback to sync variable changes back to PETSc (like mesh.points)""" var = array.owner if var is None: + # Array is being accessed during object teardown (owner is gone) return # Only act on data-changing operations (following mesh.points pattern) diff --git a/src/underworld3/swarm.py b/src/underworld3/swarm.py index 72f659179..e705b9b80 100644 --- a/src/underworld3/swarm.py +++ b/src/underworld3/swarm.py @@ -395,6 +395,9 @@ def variable_update_callback(array, change_context): """Callback to sync variable changes back to PETSc (like swarm.points)""" var = array.owner if var is None: + # This guard handles cases where the array is modified during + # object teardown (e.g. at application exit), where the owning + # Python variable has already been garbage collected. return # Only act on data-changing operations (following swarm.points pattern) From 3aba9132efe7ad2e6639e38ecf3449bb48aeb24c Mon Sep 17 00:00:00 2001 From: Ben Knight <55677727+bknight1@users.noreply.github.com> Date: Tue, 12 May 2026 08:43:00 +0800 Subject: [PATCH 10/13] docs: improve safety guard comments in variable/mesh callbacks Documented the ownership check guards in NDArray callbacks to explain that they handle teardown race conditions during Python garbage collection (e.g. at application exit or mesh rebuilds). --- src/underworld3/discretisation/discretisation_mesh.py | 4 ++++ .../discretisation/discretisation_mesh_variables.py | 5 ++++- src/underworld3/swarm.py | 7 ++++--- 3 files changed, 12 insertions(+), 4 deletions(-) diff --git a/src/underworld3/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index fe5994ede..05f47c102 100644 --- a/src/underworld3/discretisation/discretisation_mesh.py +++ b/src/underworld3/discretisation/discretisation_mesh.py @@ -601,6 +601,10 @@ class replacement_boundaries(Enum): def mesh_update_callback(array, change_context): mesh = array.owner if mesh is None: + # This guard handles cases where the array is accessed during + # object teardown (e.g. at application exit or during mesh + # replacement), where the owning Python mesh object has already + # been garbage collected but the NDArray proxy still exists. return print(f"Mesh update callback - mesh deform") diff --git a/src/underworld3/discretisation/discretisation_mesh_variables.py b/src/underworld3/discretisation/discretisation_mesh_variables.py index 47c0ed701..38cc8905f 100644 --- a/src/underworld3/discretisation/discretisation_mesh_variables.py +++ b/src/underworld3/discretisation/discretisation_mesh_variables.py @@ -494,7 +494,10 @@ def variable_update_callback(array, change_context): """Callback to sync variable changes back to PETSc (like mesh.points)""" var = array.owner if var is None: - # Array is being accessed during object teardown (owner is gone) + # This guard handles cases where the array is accessed during + # object teardown (e.g. at application exit or mesh rebuilds), + # where the owning Python variable has already been garbage + # collected but the NDArray proxy still exists. return # Only act on data-changing operations (following mesh.points pattern) diff --git a/src/underworld3/swarm.py b/src/underworld3/swarm.py index e705b9b80..41da07950 100644 --- a/src/underworld3/swarm.py +++ b/src/underworld3/swarm.py @@ -395,9 +395,10 @@ def variable_update_callback(array, change_context): """Callback to sync variable changes back to PETSc (like swarm.points)""" var = array.owner if var is None: - # This guard handles cases where the array is modified during - # object teardown (e.g. at application exit), where the owning - # Python variable has already been garbage collected. + # This guard handles cases where the array is accessed during + # object teardown (e.g. at application exit or mesh rebuilds), + # where the owning Python variable has already been garbage + # collected but the NDArray proxy still exists. return # Only act on data-changing operations (following swarm.points pattern) From 9722c78504bd915acc37f4feb8bd1f124a244be5 Mon Sep 17 00:00:00 2001 From: Ben Knight <55677727+bknight1@users.noreply.github.com> Date: Tue, 12 May 2026 09:02:10 +0800 Subject: [PATCH 11/13] test: make swarm integration statistics test MPI-aware Use global_sum and global_size for arithmetic mean calculation to ensure rank-independent assertions when comparing with global integration results. --- tests/test_0852_swarm_integration_statistics.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/test_0852_swarm_integration_statistics.py b/tests/test_0852_swarm_integration_statistics.py index f2ef4aef1..2606c583f 100644 --- a/tests/test_0852_swarm_integration_statistics.py +++ b/tests/test_0852_swarm_integration_statistics.py @@ -122,7 +122,11 @@ def test_clustered_swarm_shows_difference(self): # Arithmetic mean: biased toward left half values # (75% particles at x ≈ 0.25) + (25% particles at x ≈ 0.75) # ≈ 0.75 * (1 + 0.25) + 0.25 * (1 + 0.75) = 1.375 - arithmetic_mean = s_var.array.mean() + # + # CRITICAL MPI FIX: Use global swarm statistics for arithmetic mean + # to ensure comparison with integration mean (which is always global) + # works correctly on all ranks. + arithmetic_mean = s_var.global_sum() / s_var.global_size() # Integration-based mean: weights by volume equally # ∫∫(1 + x) dA = [x + 0.5x²]₀¹ × 1 = (1 + 0.5) = 1.5 From 20f4d8cdd8faa6e3caca76b84407406d827984b7 Mon Sep 17 00:00:00 2001 From: Ben Knight <55677727+bknight1@users.noreply.github.com> Date: Tue, 12 May 2026 09:15:38 +0800 Subject: [PATCH 12/13] docs: silence mesh update callback prints and use rank-0 only pprint --- .../discretisation/discretisation_mesh.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/src/underworld3/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index 05f47c102..35e321906 100644 --- a/src/underworld3/discretisation/discretisation_mesh.py +++ b/src/underworld3/discretisation/discretisation_mesh.py @@ -607,14 +607,17 @@ def mesh_update_callback(array, change_context): # been garbage collected but the NDArray proxy still exists. return - print(f"Mesh update callback - mesh deform") + if verbose: + uw.pprint(0, f"Mesh update callback - mesh deform") + coords = array.reshape(-1, mesh.cdim) - mesh._deform_mesh(coords, verbose=True) + mesh._deform_mesh(coords, verbose=verbose) # Increment mesh version to notify registered swarms of coordinate changes with mesh._mesh_update_lock: mesh._mesh_version += 1 - print(f"Mesh version incremented to {mesh._mesh_version}") + if verbose: + uw.pprint(0, f"Mesh version incremented to {mesh._mesh_version}") return @@ -3795,12 +3798,15 @@ def adapt(self, metric_field, verbose=False): # Rebuild the callback for mesh deformation def mesh_update_callback(array, change_context): - print(f"Mesh update callback - mesh deform") + if verbose: + uw.pprint(0, f"Mesh update callback - mesh deform") + coords = array.reshape(-1, array.owner.cdim) - self._deform_mesh(coords, verbose=True) + self._deform_mesh(coords, verbose=verbose) with self._mesh_update_lock: self._mesh_version += 1 - print(f"Mesh version incremented to {self._mesh_version}") + if verbose: + uw.pprint(0, f"Mesh version incremented to {self._mesh_version}") return self._coords.add_callback(mesh_update_callback) From 6f617f46cbb4c0a50556301f6b313f4a419ecdf7 Mon Sep 17 00:00:00 2001 From: Ben Knight <55677727+bknight1@users.noreply.github.com> Date: Tue, 12 May 2026 10:03:26 +0800 Subject: [PATCH 13/13] fix: resolve swarm particle loss and mesh variable sync issues - Restore default max_its in Swarm.migrate to fix advection particle loss (issue #175) - Ensure canonical data array has owner for MeshVariable to enable PETSc sync - Confirmed with parallel advection and integral tests --- .../discretisation/discretisation_mesh_variables.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/underworld3/discretisation/discretisation_mesh_variables.py b/src/underworld3/discretisation/discretisation_mesh_variables.py index 38cc8905f..9cd5943d3 100644 --- a/src/underworld3/discretisation/discretisation_mesh_variables.py +++ b/src/underworld3/discretisation/discretisation_mesh_variables.py @@ -2547,10 +2547,10 @@ def _create_canonical_data_array(self): # Create NDArray_With_Callback with proper shape and data from underworld3.utilities import NDArray_With_Callback - - array_obj = NDArray_With_Callback(flat_petsc_data) + array_obj = NDArray_With_Callback(flat_petsc_data, owner=self) # Single canonical callback for PETSc synchronization + def canonical_data_callback(array, change_context): """ONLY callback that handles PETSc synchronization - prevents conflicts""" var = array.owner