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

[lmoresi]

Summary

Fixes a long-standing leak in the Swarm lifecycle: transient swarms (used inside global_evaluate_nd, the upcoming read_timestep rewrite, mesh-adapt variable transfers, etc.) were pinned forever by Model._swarms, so Swarm.__del__ never fired and the underlying PETSc DMSwarm, registered fields, and cell-DM coupling all accumulated. Long-lived user swarms behave exactly as before.

This is the prep step for a wider redesign that routes global point evaluation through swarms instead of every-rank-reads-everywhere KDTree paths. Landing it ahead of the consumer rewrites so they have a clean lifecycle to build on.

What changed

Two small commits, intended to be reviewed individually:

96cc12d — Swarm lifecycle primitives

• Swarm.__del__ now calls self.dm.destroy() so the PETSc DMSwarm and registered fields are released when the swarm is collected. Previously only the mesh-side weak-set unregistration ran; PETSc objects leaked on every garbage collection.
• New Swarm._invalidate_canonical_data() consolidates the four-line cache-invalidation idiom that was duplicated in Swarm.migrate() and function/_function.pyx::global_evaluate_nd. Both call sites now use it.
• New Model._unregister_swarm(swarm) drops a swarm and its registered SwarmVariables from the global registry. Called from Swarm.__del__.

bc2068b — Make Model._swarms a WeakValueDictionary

• Closes the lifecycle cycle: dropping the user's last reference to a swarm now actually lets it be collected, __del__ fires, the cleanup hook drops the variables, and dm.destroy() releases the PETSc objects.
• Iteration, length, and membership checks behave the same as a regular dict.

Verification

• tests/test_0111_swarm_lifecycle.py (new) covers each building block individually plus an end-to-end leak-loop using current RSS via psutil (ru_maxrss is peak-only on macOS and would mask any fix).
• Leak-loop measurements over 500 create-and-discard iterations:
  • Before WVD fix: ~50 MB / 100 swarms growth
  • After: ~1.5 MB / 100 swarms (residual is allocator fragmentation)
  • Test threshold: 5 MB / 100 swarms
• 412 level_1 / early-numbered tests pass; full swarm + global_evaluate suites green (62 tests).

Why now

The next step is converting MeshVariable.read_timestep (currently 3.92 TB resident at 1152 ranks for a 1/128 spherical case — every rank reads the entire HDF5 file and builds a rank-local KDTree over it) to use a transient swarm to route saved (coord, value) pairs to their owner ranks. That conversion creates one swarm per call. Without this PR each call would leak a full DMSwarm worth of memory; with it the transient pattern is finally safe.

Test plan

• pytest tests/test_0111_swarm_lifecycle.py (6 tests, lifecycle + leak-loop)
• pytest tests/test_0110_basic_swarm.py tests/test_0003_swarm_variable_constraints.py tests/test_0510_enhanced_swarm_array.py tests/test_0005_IndexSwarmVariable.py tests/test_0505_rbf_swarm_mesh.py tests/test_0852_swarm_integration_statistics.py (36 swarm tests)
• pytest tests/test_0503_evaluate.py tests/test_0503_evaluate2.py tests/test_0730_evaluate_numpy_arrays.py tests/test_0755_evaluate_single_coordinate.py (26 evaluate tests — exercises the modified _invalidate_canonical_data call in _function.pyx)
• pytest -k "level_1 or test_00 or test_01" (412 passed, 10 skipped, 1 xfail — no regressions)
• CI parallel suite

Underworld development team with AI support from Claude Code

[Copilot]

Pull request overview

Fixes Swarm lifecycle leaks by ensuring transient swarms are no longer pinned indefinitely by the model registry and that PETSc DMSwarm resources are released when swarms are garbage-collected. This supports upcoming work that creates/destroys many short-lived swarms (global evaluation, checkpoint reads, mesh-adapt transfers) without unbounded memory growth.

Changes:

• Add explicit Swarm cleanup: unregister from model + call dm.destroy() in Swarm.__del__.
• Introduce Swarm._invalidate_canonical_data() and reuse it from Swarm.migrate() and global_evaluate_nd.
• Convert Model._swarms to a WeakValueDictionary and add Model._unregister_swarm(); add lifecycle regression tests (incl. optional RSS leak-loop).

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 6 comments.

File	Description
tests/test_0111_swarm_lifecycle.py	New lifecycle + leak regression tests for swarm GC/unregistration behavior.
src/underworld3/swarm.py	Adds DM destruction in __del__ and centralizes cache invalidation.
src/underworld3/model.py	Makes swarm registry weak and adds _unregister_swarm to drop swarm variables.
src/underworld3/function/_function.pyx	Uses the new centralized cache invalidation after bare dm.migrate.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

self._vars is a WeakValueDictionary, so iterating directly over self._vars.values() can raise RuntimeError: dictionary changed size during iteration if weakrefs are cleared during iteration (e.g. via GC). To make cache invalidation robust, iterate over a snapshot like list(self._vars.values()).

Suggested change

	forvarinself._vars.values():
	forvarinlist(self._vars.values()):

[Copilot]

Because _swarms is a WeakValueDictionary, code that iterates over _swarms.items() can see the dict shrink asynchronously as swarms are GC’d, which may raise RuntimeError: dictionary changed size during iteration. The comment here says iteration behaves the same as a regular dict; consider qualifying that statement or ensuring call sites iterate over a snapshot (e.g. list(self._swarms.items())) when producing summaries.

Suggested change

	# reference and ``Swarm.__del__`` never fires. Iteration, length, and
	# membership checks behave the same as a regular dict; the only
	# difference is that entries auto-disappear once their swarm value has
	# no other strong references.
	# reference and ``Swarm.__del__`` never fires. Length and membership
	# checks are dictionary-like, but iteration is not stable in the same
	# way as a regular dict because entries can disappear as swarm values
	# are garbage-collected. Call sites that need a stable traversal (for
	# example, producing summaries) should iterate over a snapshot such as
	# ``list(self._swarms.items())``.

[Copilot]

The Swarm.__del__ docstring says “the model registry holds a strong reference to the swarm”, but in this PR Model._swarms becomes a WeakValueDictionary (so it no longer pins swarms). Consider rephrasing this docstring to describe the historical behavior (“previously held a strong reference”) or to focus on the remaining leak cause (dm.destroy() missing) to avoid misleading future readers.

Suggested change

	leakoneachgarbagecollection: themodelregistryholdsastrong
	referencetotheswarm, and``Swarm.__del__``previouslydidnotcall
	``self.dm.destroy()``. Thismattersfortransientswarmsusedinside
	time-steppingloops (globalexpressionevaluation, checkpointreads,
	meshadaptationtransfers).
	canremainallocatedlongerthanintended. Historicallythiswas
	exacerbatedwhenthemodelregistrykeptstrongreferencestoswarms,
	and``Swarm.__del__``didnotcall``self.dm.destroy()``. Thismatters
	fortransientswarmsusedinsidetime-steppingloops (globalexpression
	evaluation, checkpointreads, meshadaptationtransfers).

[Copilot]

test_swarm_del_calls_dm_destroy can pass even if the actual self.dm.destroy() call is removed, because the substring check also matches the docstring (which contains self.dm.destroy()). Consider parsing the function body (e.g., via ast) and asserting there is a Call to .destroy on self.dm, or otherwise ensuring the match can’t come from the docstring/comments.

Suggested change

	importinspect
	src=inspect.getsource(uw.swarm.Swarm.__del__)
	assert"dm.destroy"insrc, (
	"Swarm.__del__ no longer contains dm.destroy() — §1a regression"
	importast
	importinspect
	importtextwrap
	src=textwrap.dedent(inspect.getsource(uw.swarm.Swarm.__del__))
	tree=ast.parse(src)
	defis_self_dm_destroy_call(node):
	return (
	isinstance(node, ast.Call)
	andisinstance(node.func, ast.Attribute)
	andnode.func.attr=="destroy"
	andisinstance(node.func.value, ast.Attribute)
	andnode.func.value.attr=="dm"
	andisinstance(node.func.value.value, ast.Name)
	andnode.func.value.value.id=="self"
	)
	assertany(is_self_dm_destroy_call(node) fornodeinast.walk(tree)), (
	"Swarm.__del__ no longer contains a call to self.dm.destroy() — §1a regression"

[Copilot]

This leak test is skipped unless psutil is installed (pytest.importorskip('psutil')), but the repo doesn’t appear to declare psutil in its test dependencies. If this is meant to guard the regression in CI, consider adding psutil to the test extras/CI environment or rewriting the test to avoid the external dependency so it actually runs by default.

[Copilot]

Swarm.__del__ unregisters via self._model_ref(), but here __init__ registers via uw.get_default_model() again. If the default model changes between storing _model_ref and this registration (e.g. reset_default_model() in tests/interactive sessions), the swarm can end up registered in a different model than _model_ref points to, and __del__ would unregister the wrong registry. Registering with the same model instance used to set _model_ref (or using _model_ref() consistently) would avoid that mismatch and potential registry/variable leaks.