From fd9b0397e1bf20a7a394911f1cd2a433106a680d Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 23 Apr 2026 18:37:01 +1000 Subject: [PATCH 1/2] Fix VE_Stokes first-solve parallel deadlock (#130) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit mesh._get_coords_for_basis contains PETSc collectives (DMClone, createInterpolation, globalToLocal). It was populated lazily from MeshVariable.coords_nd, which is accessed from rank-local code paths such as rbf_interpolate inside global_evaluate_nd's per-particle loop. On partition geometries where some ranks had only interior migrated particles, those ranks skipped the RBF path while others entered it — leaving the collective waiting forever. Pre-populate the coord cache at the end of _BaseMeshVariable.__init__ (a collective context) so every variable's basis is filled by all ranks together. Also repopulate 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, 16x8 failure case. Verified against the full (np, mesh) matrix from the issue: previously hanging configs (np=3 with 16x8/32x16, np=4 with 16x8, np=16 with 16x8) all complete; previously OK configs still OK; level_1+tier_a suite passes. Underworld development team with AI support from Claude Code --- .../discretisation/discretisation_mesh.py | 9 +++ .../discretisation_mesh_variables.py | 12 ++++ .../test_0780_ve_stokes_first_solve_mpi.py | 72 +++++++++++++++++++ 3 files changed, 93 insertions(+) create mode 100644 tests/parallel/test_0780_ve_stokes_first_solve_mpi.py diff --git a/src/underworld3/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index ea7a7124c..ad59e40fe 100644 --- a/src/underworld3/discretisation/discretisation_mesh.py +++ b/src/underworld3/discretisation/discretisation_mesh.py @@ -1653,6 +1653,15 @@ def nuke_coords_and_rebuild( # Invalidate projected boundary normals (rebuilt lazily on access) self._projected_normals = None + # TODO(BUG): issue #130 — refill the coord cache for every already- + # registered variable. Variables created before this rebuild would + # otherwise have their cache entry (from __init__) wiped above and + # refill lazily from rank-local code paths (rbf_interpolate), which + # deadlocks when the collectives inside _get_coords_for_basis are + # reached by only a subset of ranks. + for _var in list(self.vars.values()): + self._get_coords_for_var(_var) + if verbose and uw.mpi.rank == 0: print( f"Mesh Spatial Discretisation Complete", diff --git a/src/underworld3/discretisation/discretisation_mesh_variables.py b/src/underworld3/discretisation/discretisation_mesh_variables.py index d5b69aabe..b159cecb8 100644 --- a/src/underworld3/discretisation/discretisation_mesh_variables.py +++ b/src/underworld3/discretisation/discretisation_mesh_variables.py @@ -388,6 +388,18 @@ def __init__( self.mesh.vars[self.clean_name] = self self._setup_ds() + # TODO(BUG): issue #130 — pre-populate the mesh's coordinate cache + # for this variable's basis. mesh._get_coords_for_basis contains MPI + # collectives (DMClone, createInterpolation, globalToLocal) that + # deadlock when triggered lazily from rank-local code paths (e.g. + # rbf_interpolate inside global_evaluate_nd's per-particle loop): + # ranks with no exterior points skip the call, while ranks with + # exterior points enter the collective and wait forever. Variable + # construction is collective, so filling the cache here ensures all + # ranks populate it together and subsequent rank-local lookups are + # cache hits. + self.mesh._get_coords_for_var(self) + # Setup public view of data - using NDArray_With_Callback self._array_cache = None # Will be created lazily when first accessed self._data_cache = None # Will be created lazily when first accessed diff --git a/tests/parallel/test_0780_ve_stokes_first_solve_mpi.py b/tests/parallel/test_0780_ve_stokes_first_solve_mpi.py new file mode 100644 index 000000000..913e0ebc5 --- /dev/null +++ b/tests/parallel/test_0780_ve_stokes_first_solve_mpi.py @@ -0,0 +1,72 @@ +""" +MPI regression test for VE_Stokes first-solve deadlock (issue #130). + +The bug: the first ``VE_Stokes.solve()`` call on a fresh in-memory mesh +deadlocked at specific ``(np, mesh)`` partition geometries (e.g. np=4 +with a 16x8 StructuredQuadBox → 4x2 rank partition). Root cause was +lazy invocation of ``mesh._get_coords_for_basis`` (DMClone + +createInterpolation + globalToLocal collectives) from rank-local code +paths in ``global_evaluate_nd``: ranks whose migrated particles were +all interior skipped the RBF path and never entered the collective, +while ranks with exterior particles did — deadlocking forever. + +The fix pre-populates each variable's coordinate cache at the end of +``_BaseMeshVariable.__init__`` (a collective context), so subsequent +rank-local lookups always hit the cache. + +This test runs the canonical failure case at np=4, 16x8 and fails fast +via the pytest timeout rather than blocking indefinitely. +""" + +import sympy +import pytest +import underworld3 as uw +from underworld3.function import expression + + +pytestmark = [ + pytest.mark.level_2, + pytest.mark.tier_a, + pytest.mark.mpi(min_size=2), + pytest.mark.timeout(60), +] + + +@pytest.mark.mpi(min_size=2) +def test_ve_stokes_first_solve_does_not_deadlock(): + """ + First VE_Stokes.solve() must complete under MPI on a partition-sensitive + mesh (16x8 -> 4x2 rank partition at np=4). Before the fix for issue #130, + this hung at the first solve indefinitely. + """ + mesh = uw.meshing.StructuredQuadBox(elementRes=(16, 8)) + v = uw.discretisation.MeshVariable("V", mesh, mesh.dim, degree=2) + p = uw.discretisation.MeshVariable("P", mesh, 1, degree=1) + + stokes = uw.systems.VE_Stokes( + mesh, velocityField=v, pressureField=p, order=2 + ) + stokes.constitutive_model = ( + uw.constitutive_models.ViscoElasticPlasticFlowModel + ) + stokes.constitutive_model.Parameters.shear_viscosity_0 = 1.0 + stokes.constitutive_model.Parameters.shear_modulus = 1.0 + stokes.constitutive_model.Parameters.dt_elastic = 0.02 + + V_top = expression("V_top", 0.5, "Top BC") + stokes.add_dirichlet_bc((V_top, 0.0), "Top") + stokes.add_dirichlet_bc((0.0, 0.0), "Bottom") + stokes.add_dirichlet_bc((sympy.oo, 0.0), "Left") + stokes.add_dirichlet_bc((sympy.oo, 0.0), "Right") + + stokes.solve(timestep=0.02) + + # Velocity must carry the top BC. If the solve returned without diverging + # (pre-fix deadlock would hit the pytest timeout) but with zero velocity, + # something else went wrong. + import numpy as np + v_max = float(np.abs(v.data).max()) if v.data.size else 0.0 + gathered = uw.mpi.comm.allgather(v_max) + assert max(gathered) > 1.0e-6, ( + f"Velocity field is effectively zero after solve; max|v|={gathered}" + ) From 7b9fdc75d028bb8a2967d2feeed51b4960cc59d9 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 23 Apr 2026 20:52:09 +1000 Subject: [PATCH 2/2] Address Copilot review on #136 - _get_coords_for_basis: destroy matInterp/vecScale after use so the per-variable eager cache fill and the nuke_coords_and_rebuild refill loop don't leak PETSc objects in long adapt runs. - test_0780_ve_stokes_first_solve_mpi: require mpi(min_size=4). The bug only reproduces at the 4x2 partition geometry; min_size=2 would let CI pass with a partition that never hit the deadlock. - Relabel the two bugfix comment blocks from TODO(BUG) to BUGFIX(#130) so they don't read as unfinished work. Underworld development team with AI support from Claude Code --- .../discretisation/discretisation_mesh.py | 14 ++++++++++---- .../discretisation_mesh_variables.py | 4 ++-- .../test_0780_ve_stokes_first_solve_mpi.py | 4 ++-- 3 files changed, 14 insertions(+), 8 deletions(-) diff --git a/src/underworld3/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index ad59e40fe..14f85a105 100644 --- a/src/underworld3/discretisation/discretisation_mesh.py +++ b/src/underworld3/discretisation/discretisation_mesh.py @@ -1653,10 +1653,10 @@ def nuke_coords_and_rebuild( # Invalidate projected boundary normals (rebuilt lazily on access) self._projected_normals = None - # TODO(BUG): issue #130 — refill the coord cache for every already- - # registered variable. Variables created before this rebuild would - # otherwise have their cache entry (from __init__) wiped above and - # refill lazily from rank-local code paths (rbf_interpolate), which + # BUGFIX(#130): refill the coord cache for every already-registered + # variable. Variables created before this rebuild would otherwise + # have their cache entry (from __init__) wiped above and refill + # lazily from rank-local code paths (rbf_interpolate), which # deadlocks when the collectives inside _get_coords_for_basis are # reached by only a subset of ranks. for _var in list(self.vars.values()): @@ -2733,6 +2733,12 @@ def _get_coords_for_basis(self, degree, continuous): dmnew.restoreGlobalVec(coordsNewG) dmnew.restoreLocalVec(coordsNewL) + # Clean up the PETSc interpolation objects built above. Without this + # they accumulate until Python GC runs — noticeable in long adapt + # loops that re-fill the coord cache per variable. + matInterp.destroy() + if vecScale is not None: + vecScale.destroy() dmnew.destroy() dmfe.destroy() diff --git a/src/underworld3/discretisation/discretisation_mesh_variables.py b/src/underworld3/discretisation/discretisation_mesh_variables.py index b159cecb8..133590732 100644 --- a/src/underworld3/discretisation/discretisation_mesh_variables.py +++ b/src/underworld3/discretisation/discretisation_mesh_variables.py @@ -388,8 +388,8 @@ def __init__( self.mesh.vars[self.clean_name] = self self._setup_ds() - # TODO(BUG): issue #130 — pre-populate the mesh's coordinate cache - # for this variable's basis. mesh._get_coords_for_basis contains MPI + # BUGFIX(#130): pre-populate the mesh's coordinate cache for this + # variable's basis. mesh._get_coords_for_basis contains MPI # collectives (DMClone, createInterpolation, globalToLocal) that # deadlock when triggered lazily from rank-local code paths (e.g. # rbf_interpolate inside global_evaluate_nd's per-particle loop): diff --git a/tests/parallel/test_0780_ve_stokes_first_solve_mpi.py b/tests/parallel/test_0780_ve_stokes_first_solve_mpi.py index 913e0ebc5..bc0c39831 100644 --- a/tests/parallel/test_0780_ve_stokes_first_solve_mpi.py +++ b/tests/parallel/test_0780_ve_stokes_first_solve_mpi.py @@ -27,12 +27,12 @@ pytestmark = [ pytest.mark.level_2, pytest.mark.tier_a, - pytest.mark.mpi(min_size=2), + pytest.mark.mpi(min_size=4), pytest.mark.timeout(60), ] -@pytest.mark.mpi(min_size=2) +@pytest.mark.mpi(min_size=4) def test_ve_stokes_first_solve_does_not_deadlock(): """ First VE_Stokes.solve() must complete under MPI on a partition-sensitive