Fix VE_Stokes first-solve parallel deadlock (#130)

[lmoresi]

Summary

• Pre-populates the mesh coordinate cache at the end of _BaseMeshVariable.__init__ so _get_coords_for_basis PETSc collectives (DMClone, createInterpolation, globalToLocal) are always entered collectively.
• Defense-in-depth: also repopulates for all registered variables at the end of nuke_coords_and_rebuild, so mesh adaptation and DM rebuilds don't re-open the same deadlock window.
• Adds tests/parallel/test_0780_ve_stokes_first_solve_mpi.py exercising the canonical np=4, 16×8 failure case with a 60s timeout.

Fixes #130.

Root cause

global_evaluate_nd migrates particles across ranks and evaluates the expression rank-locally. In default mode, interior points use petsc_interpolate and exterior points use rbf_evaluate → MeshVariable.rbf_interpolate → self.coords_nd → mesh._get_coords_for_basis(degree, continuous).

_get_coords_for_basis contains PETSc collectives on the mesh's communicator. The coord cache was populated lazily on first access, so whichever rank hit the cache miss first entered the collective — but ranks whose migrated particles were all interior skipped rbf_evaluate entirely and never arrived. The collective never completed.

That explains the partition sensitivity reported in #130: on a 4×2 partition at 16×8, the migration distribution left some ranks without exterior particles; on a 2×2 partition at 16×16 every rank had some, so every rank reached the collective and it completed.

The fix moves cache population into _BaseMeshVariable.__init__, which is always collective (it modifies the mesh DM), ensuring all ranks fill the cache together and subsequent rank-local lookups are cache hits.

Stack traces from the investigation (rank 1 vs. ranks 0/2/3 on the np=4, 16×8 repro) are in the bugfix commit history and the issue thread.

Test plan

• Full (np, mesh) matrix from the issue
  • np=3 {16×8, 32×16}: was HANG 5/5 → 2.9s / 3.0s
  • np=4 16×8: was HANG 5/5 → 3.2s
  • np=4 {16×16, 32×16}, np=8 16×8: previously OK → still OK
  • np=12 16×8, np=16 16×8 (oversubscribed on 16 logical cores): complete
• New parallel regression test at np=4, 16×8 — passes in 5.5s
• pytest -m "level_1 and tier_a" — 56 passed, 3 skipped, no regressions

Underworld development team with AI support from Claude Code

[Copilot]

Pull request overview

Fixes an MPI deadlock in the VE_Stokes first solve() by ensuring mesh coordinate-basis caches are populated collectively (avoiding rank-asymmetric entry into PETSc collectives), and adds a parallel regression test for the previously hanging configuration.

Changes:

• Eagerly pre-populates per-variable mesh coordinate caches during MeshVariable construction to avoid lazy collective calls from rank-local evaluation paths.
• Re-populates coordinate caches for all registered variables after nuke_coords_and_rebuild() to prevent the same deadlock window after mesh/DM rebuilds.
• Adds an MPI regression test intended to reproduce the canonical np=4, 16×8 hang case with a timeout.

Reviewed changes

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

File	Description
src/underworld3/discretisation/discretisation_mesh_variables.py	Eagerly fills coordinate cache for each new variable to force collective entry.
src/underworld3/discretisation/discretisation_mesh.py	Refills coordinate cache for all existing variables after coordinate/DM rebuild.
tests/parallel/test_0780_ve_stokes_first_solve_mpi.py	Adds an MPI regression test for the VE_Stokes first-solve deadlock scenario.

---

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

[Copilot]

This test is described as the canonical np=4, 16×8 reproduction, but the MPI marker only enforces min_size=2. As written, CI runs with 2 ranks will execute the test and it won’t exercise the partition geometry that originally deadlocked. Consider requiring 4 ranks here (e.g., pytest.mark.mpi(min_size=4) or skip/assert when uw.mpi.size != 4) so the regression is actually covered.

Suggested change

	pytest.mark.mpi(min_size=2),
	pytest.mark.timeout(60),
	]
	@pytest.mark.mpi(min_size=2)
	pytest.mark.mpi(min_size=4),
	pytest.mark.timeout(60),
	]
	@pytest.mark.mpi(min_size=4)

[Copilot]

The new eager call to mesh._get_coords_for_var(self) will invoke PETSc createInterpolation paths. In discretisation_mesh._get_coords_for_basis, the returned matInterp/vecScale objects are not explicitly destroyed, so calling this for every variable construction can accumulate PETSc objects until GC runs (and may increase memory usage in long runs / adaptive loops). Consider updating _get_coords_for_basis to destroy matInterp (and vecScale when non-null) after use, similar to src/underworld3/function/field_projection.py which does explicit cleanup.

Suggested change

	# cache hits.
	self.mesh._get_coords_for_var(self)
	# cache hits. Force PETSc cleanup afterwards so temporary
	# interpolation objects created on this path do not accumulate until
	# Python garbage collection runs.
	try:
	self.mesh._get_coords_for_var(self)
	finally:
	PETSc.garbage_cleanup()

[Copilot]

These comments are labeled TODO(BUG), but this block is part of the actual fix path (not a TODO). Consider dropping the TODO tag (or replacing with a neutral rationale comment like “BUGFIX(#130): …”) so it doesn’t look like intentionally unfinished work.

[Copilot]

This loop can now trigger coordinate interpolation for every registered variable right after resetting _coord_array. Given _get_coords_for_basis currently doesn’t explicitly destroy the interpolation matrix/scale objects it creates, repeatedly calling nuke_coords_and_rebuild (e.g., during mesh adaptation) could cause avoidable PETSc memory growth. Consider ensuring _get_coords_for_basis cleans up matInterp/vecScale deterministically if this eager refill remains.