diff --git a/pixi.toml b/pixi.toml index 2bc8f4c60..e2436c07e 100644 --- a/pixi.toml +++ b/pixi.toml @@ -89,6 +89,7 @@ pytest-mpi = ">=0.6,<0.7" pytest-timeout = ">=2.3,<3" pytest-forked = ">=1.6,<2" pytest-xdist = ">=3.8,<4" +psutil = ">=5.9,<8" # leak-loop regression test in tests/test_0111_swarm_lifecycle.py # Mesh generation pykdtree = ">=1.4,<2" diff --git a/src/underworld3/function/_function.pyx b/src/underworld3/function/_function.pyx index 3d550ef94..9dd71e15d 100644 --- a/src/underworld3/function/_function.pyx +++ b/src/underworld3/function/_function.pyx @@ -506,10 +506,7 @@ def global_evaluate_nd( expr, # Invalidate cached data after bare-bones dm.migrate — # particle count and values changed but Swarm.migrate() was bypassed. - evaluation_swarm._particle_coordinates._canonical_data = None - for var in evaluation_swarm._vars.values(): - if hasattr(var, "_canonical_data"): - var._canonical_data = None + evaluation_swarm._invalidate_canonical_data() # Pre-allocate with NaN so the shape is always correct. If any points # are lost during the migration round-trip, they remain NaN rather than diff --git a/src/underworld3/model.py b/src/underworld3/model.py index 2a398e57d..8e1801d2c 100644 --- a/src/underworld3/model.py +++ b/src/underworld3/model.py @@ -112,7 +112,17 @@ class Model(PintNativeModelMixin, BaseModel): # Declare private attributes for Pydantic v2 _meshes: Any = PrivateAttr(default_factory=dict) _primary_mesh_id: Optional[int] = PrivateAttr(default=None) - _swarms: Any = PrivateAttr(default_factory=dict) + # WeakValueDictionary so dropping the user's last reference to a swarm + # actually lets it be garbage-collected. Without this, transient swarms + # used inside functions (global expression evaluation, checkpoint reads, + # mesh-adapt transfers) accumulate forever via the registry's strong + # reference and ``Swarm.__del__`` never fires. Length and membership + # checks are dictionary-like; iteration is *not* stable in the way a + # plain dict's is, because entries can disappear asynchronously as + # their swarm value is collected. Call sites that need a stable + # traversal (summaries, snapshots) should iterate a copy such as + # ``list(self._swarms.items())``. + _swarms: Any = PrivateAttr(default_factory=weakref.WeakValueDictionary) _variables: Dict[str, Any] = PrivateAttr(default_factory=dict) _solvers: Dict[str, Any] = PrivateAttr(default_factory=dict) @@ -210,6 +220,31 @@ def _register_swarm(self, swarm): swarm_id = id(swarm) self._swarms[swarm_id] = swarm + def _unregister_swarm(self, swarm): + """ + Internal method to drop a swarm and all its variables from the registry. + + Without this, transient swarms (used for global expression evaluation, + checkpoint reads, mesh adaptation transfers) accumulate in the registry + forever, leaking the underlying PETSc DMs and field storage. Called + from :meth:`Swarm.__del__`. + """ + swarm_id = id(swarm) + self._swarms.pop(swarm_id, None) + # Drop any variables that belong to this swarm. Building the list of + # names first avoids mutating the dict during iteration. + try: + from .swarm import SwarmVariable + except ImportError: + return + names_to_drop = [ + name for name, var in self._variables.items() + if isinstance(var, SwarmVariable) and getattr(var, "_swarm_ref", None) is not None + and var._swarm_ref() is swarm + ] + for name in names_to_drop: + self._variables.pop(name, None) + def _register_variable(self, name, variable): """ Internal method to register a variable with this model. @@ -225,13 +260,15 @@ def _register_variable(self, name, variable): variable : MeshVariable or SwarmVariable Variable instance to register """ - # For SwarmVariables, ensure we keep strong reference to swarm to prevent garbage collection + # For SwarmVariables, make sure the parent swarm has been registered. + # Registration is via WeakValueDictionary, so it does not pin the + # swarm — when the user drops their last reference the swarm dies + # and ``Swarm.__del__`` cleans up the variable side of the registry. from .swarm import SwarmVariable if isinstance(variable, SwarmVariable): swarm = variable.swarm # This will raise if swarm already garbage collected swarm_id = id(swarm) - # Ensure swarm is registered with strong reference if swarm_id not in self._swarms: self._register_swarm(swarm) @@ -4245,7 +4282,7 @@ def view(self, verbose: int = 0, show_materials: bool = True, show_petsc: bool = swarm_count = len(self._swarms) lines.append(f"**Swarms:** {swarm_count} registered") if swarm_count > 0 and verbose >= 1: - for swarm_id, swarm in self._swarms.items(): + for swarm_id, swarm in list(self._swarms.items()): try: if hasattr(swarm, "data") and swarm.data is not None: particle_count = len(swarm.data) diff --git a/src/underworld3/swarm.py b/src/underworld3/swarm.py index ee2dbe4ad..83fadb386 100644 --- a/src/underworld3/swarm.py +++ b/src/underworld3/swarm.py @@ -2499,17 +2499,63 @@ def __init__(self, mesh, recycle_rate=0, verbose=False, clip_to_mesh=True): super().__init__() - # Register with default model for orchestration - uw.get_default_model()._register_swarm(self) + # Register with the same model already captured in self._model_ref + # above (not a fresh ``get_default_model()`` call) so that + # ``__del__`` deregisters from the same registry it registered with, + # even if the default model is swapped mid-session. + model._register_swarm(self) def __del__(self): - """Cleanup swarm by unregistering from mesh to prevent memory leaks""" + """Cleanup swarm: unregister from mesh and model, destroy the DM. + + Three steps: drop the mesh-side weak-set entry, drop the model-side + registry entry (which also forgets any SwarmVariables that belonged + to this swarm), and call ``self.dm.destroy()`` to release the PETSc + DMSwarm and its registered fields. + + Historically the model registry kept a strong reference to every + swarm and ``__del__`` did not call ``dm.destroy()`` — both leaks + accumulated quickly inside time-stepping loops that build transient + swarms (global expression evaluation, checkpoint reads, mesh + adaptation transfers). The registry is now a ``WeakValueDictionary`` + and ``__del__`` cleans up its own resources. + """ try: if hasattr(self, "mesh") and self.mesh is not None: self.mesh.unregister_swarm(self) except (AttributeError, ReferenceError, RuntimeError): # Mesh/Model may have already been garbage collected, which is fine pass + try: + model = self._model_ref() if hasattr(self, "_model_ref") else None + if model is not None: + model._unregister_swarm(self) + except (AttributeError, ReferenceError, RuntimeError): + pass + try: + if hasattr(self, "dm") and self.dm is not None: + self.dm.destroy() + except Exception: + # DM may already have been destroyed (e.g. by Mesh.adapt) or + # PETSc may be finalising. Either way, nothing more we can do. + pass + + def _invalidate_canonical_data(self): + """Drop cached array views on the coordinates and every registered variable. + + Required after any operation that bypasses :meth:`Swarm.migrate` but + still mutates particle layout — notably the bare ``self.dm.migrate(...)`` + used by round-trip evaluation patterns. Calling :meth:`Swarm.migrate` + already does this internally; call this directly only when you used the + underlying DMSwarm migrate. + """ + if hasattr(self, "_particle_coordinates") and self._particle_coordinates is not None: + self._particle_coordinates._canonical_data = None + # ``self._vars`` may be a WeakValueDictionary whose values disappear + # asynchronously during GC; iterate a snapshot. + for var in list(self._vars.values()): + if hasattr(var, "_canonical_data"): + var._canonical_data = None @property def mesh(self): @@ -3268,10 +3314,7 @@ def migrate( # Invalidate all cached data after migration. # Any particle movement (send, receive, or balanced swap) makes # cached arrays stale — both size and values may have changed. - self._particle_coordinates._canonical_data = None - for var in self._vars.values(): - if hasattr(var, "_canonical_data"): - var._canonical_data = None + self._invalidate_canonical_data() return diff --git a/tests/test_0111_swarm_lifecycle.py b/tests/test_0111_swarm_lifecycle.py new file mode 100644 index 000000000..5fbaa2319 --- /dev/null +++ b/tests/test_0111_swarm_lifecycle.py @@ -0,0 +1,224 @@ +"""Lifecycle regression tests for Swarm cleanup. + +The prep work for the global-point-routing redesign restores the swarm +lifecycle so that transient swarms (used inside ``global_evaluate_nd``, +checkpoint reads, mesh-adapt transfers) are actually freed on garbage +collection instead of accumulating in the model registry forever. + + - ``Swarm.__del__`` now calls ``self.dm.destroy()`` so the PETSc DMSwarm and + every registered field are released when the swarm is collected. + - ``Swarm._invalidate_canonical_data()`` consolidates the cache invalidation + that previously appeared inline in two places. + - ``Model._swarms`` is a ``WeakValueDictionary`` so that registration no + longer pins swarms beyond the user's last strong reference. + - ``Model._unregister_swarm`` drops the swarm's registered variables when + ``__del__`` fires. +""" + +import gc + +import pytest + +import underworld3 as uw +from underworld3.meshing import UnstructuredSimplexBox + +pytestmark = pytest.mark.level_1 + + +@pytest.fixture +def mesh(): + return UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=1.0 / 16.0, + ) + + +def test_swarm_del_calls_dm_destroy(): + """``Swarm.__del__`` actually contains a call to ``self.dm.destroy()``. + + petsc4py exposes DM methods as read-only Cython attributes, so we cannot + monkeypatch ``dm.destroy`` to spy on it. AST-based source check: parse + the function body and look for an actual ``self.dm.destroy(...)`` call + node. A bare substring search would also match the destroy reference + inside the docstring and pass even if the executable line were removed. + """ + import ast + import inspect + import textwrap + + src = textwrap.dedent(inspect.getsource(uw.swarm.Swarm.__del__)) + tree = ast.parse(src) + + def is_self_dm_destroy_call(node): + return ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "destroy" + and isinstance(node.func.value, ast.Attribute) + and node.func.value.attr == "dm" + and isinstance(node.func.value.value, ast.Name) + and node.func.value.value.id == "self" + ) + + assert any(is_self_dm_destroy_call(node) for node in ast.walk(tree)), ( + "Swarm.__del__ no longer contains a call to self.dm.destroy() — " + "lifecycle regression" + ) + + +def test_swarm_invalidate_canonical_data_method(mesh): + """The §1b extraction is callable and clears the expected caches.""" + swarm = uw.swarm.Swarm(mesh) + var = uw.swarm.SwarmVariable( + "test_var", swarm, + vtype=uw.VarType.SCALAR, dtype=float, _proxy=False, + ) + swarm.populate(fill_param=2) + # Touch caches so they exist + _ = swarm._particle_coordinates.array + _ = var.array + + # Set caches manually, then invalidate + swarm._particle_coordinates._canonical_data = "fake_cache" + var._canonical_data = "fake_cache" + swarm._invalidate_canonical_data() + assert swarm._particle_coordinates._canonical_data is None + assert var._canonical_data is None + + # Cleanup + uw.get_default_model()._unregister_swarm(swarm) + + +def test_model_unregister_swarm_drops_swarm_and_variables(mesh): + """``Model._unregister_swarm`` releases the swarm registry slot and any + registered SwarmVariables that belong to it. + + This is the building block consumer recipes use to ensure transient + swarms (used in checkpoint reads, mesh-adapt transfers, etc.) actually + get freed instead of accumulating in ``Model._swarms``. + """ + model = uw.get_default_model() + n_swarms_before = len(model._swarms) + n_vars_before = len(model._variables) + + swarm = uw.swarm.Swarm(mesh) + uw.swarm.SwarmVariable( + "_test_lifecycle_a", swarm, + vtype=uw.VarType.MATRIX, size=(1, 3), + dtype=float, _proxy=False, + ) + uw.swarm.SwarmVariable( + "_test_lifecycle_b", swarm, + vtype=uw.VarType.SCALAR, dtype=float, _proxy=False, + ) + + # Both registries grew + assert len(model._swarms) == n_swarms_before + 1 + assert len(model._variables) >= n_vars_before + 2 + + model._unregister_swarm(swarm) + + # Both registries shrank back + assert len(model._swarms) == n_swarms_before, ( + "swarm not removed from Model._swarms" + ) + # The two test variables are gone (other internal vars from the swarm — + # coord var, _remeshed if any — are also gone). + assert "_test_lifecycle_a" not in model._variables + assert "_test_lifecycle_b" not in model._variables + + +def test_model_unregister_idempotent(mesh): + """Calling ``_unregister_swarm`` twice is safe (no KeyError).""" + model = uw.get_default_model() + swarm = uw.swarm.Swarm(mesh) + model._unregister_swarm(swarm) + model._unregister_swarm(swarm) # must not raise + + +def test_swarm_del_fires_on_drop(mesh): + """A swarm with no remaining strong references is collected automatically. + + Before ``Model._swarms`` became a ``WeakValueDictionary`` the registry + pinned every swarm forever, so ``__del__`` never fired. The proof that + the cycle is broken: registry size returns to baseline after ``del``. + """ + model = uw.get_default_model() + n_before = len(model._swarms) + + swarm = uw.swarm.Swarm(mesh) + uw.swarm.SwarmVariable( + "_test_drop", swarm, + vtype=uw.VarType.SCALAR, dtype=float, _proxy=False, + ) + swarm.populate(fill_param=2) + assert len(model._swarms) == n_before + 1 + + del swarm + gc.collect() + + assert len(model._swarms) == n_before, ( + "Swarm not collected after del — the registry is still pinning it" + ) + assert "_test_drop" not in model._variables + + +def test_swarm_lifecycle_does_not_leak(mesh): + """End-to-end: many transient swarms do not grow current RSS unboundedly. + + Uses ``psutil.Process().memory_info().rss`` (current resident memory) so + OS-level peak tracking does not mask freed memory. The threshold is + deliberately generous — a real leak grows by tens of MB per 100 swarms. + """ + psutil = pytest.importorskip("psutil") + import os + p = psutil.Process(os.getpid()) + + def rss_mb(): + return p.memory_info().rss / (1024 * 1024) + + n_iters = 500 + sample_every = 100 + + # Warm-up: first iterations always cost RSS for one-off allocations + # (caches, JIT tables, lazy initialisation). + for _ in range(50): + s = uw.swarm.Swarm(mesh) + uw.swarm.SwarmVariable( + "v", s, + vtype=uw.VarType.MATRIX, size=(1, 3), + dtype=float, _proxy=False, + ) + s.populate(fill_param=2) + del s + gc.collect() + rss_samples = [rss_mb()] + + for i in range(n_iters): + s = uw.swarm.Swarm(mesh) + for j in range(3): + uw.swarm.SwarmVariable( + f"v_{j}", s, + vtype=uw.VarType.MATRIX, size=(1, 3), + dtype=float, _proxy=False, + ) + s.populate(fill_param=2) + for var in s._vars.values(): + _ = var.array + del s + + if (i + 1) % sample_every == 0: + gc.collect() + rss_samples.append(rss_mb()) + + growth_mb = rss_samples[-1] - rss_samples[0] + growth_per_100 = growth_mb / (n_iters / sample_every) + print(f"\nRSS samples (MB): {[round(x, 1) for x in rss_samples]}") + print(f"Growth per 100 swarms: {growth_per_100:.2f} MB") + # Without the WeakValueDictionary fix this comes back as ~25–30 MB / 100 + # (full DMSwarm + registered fields preserved in the registry). + assert growth_per_100 < 5.0, ( + f"Suspected swarm leak: {growth_per_100:.2f} MB per 100 swarms " + f"(samples: {rss_samples})" + )