Merge pull request #153 from underworldcode/feature/swarm-routed-point-eval

Fix Swarm lifecycle: stop leaking transient swarms via the model registry

Author: lmoresi

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"

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

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)

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