VE Stokes hangs on first solve at specific (np, mesh) combinations

[lmoresi]

Summary

A VE Stokes solve with only Dirichlet BCs deadlocks on the first solve() under MPI for specific (np, mesh) combinations. The hang is deterministic (5/5 trials) for the cases marked below, distinct from #96 (which was a PETSc coordinate-DM bug, fixed), and reproducible against a clean checkout of origin/development.

Failure surface (5 trials per cell, 45s timeout)

np \ mesh	8×4	16×8	16×16	24×12	32×16
1	OK 5.6s	OK 4.4s	—	—	OK 4.9s
2	OK 4.8s	OK 5/5	—	—	OK 5.7s
3	OK 5.9s	HANG 5/5	—	—	HANG 5/5
4	OK 5/5	HANG 5/5	OK 3.0s	OK 3.2s	OK 5/5
5	—	OK 3.6s	—	—	—
6	OK 7.8s	OK 9.3s	—	—	OK 8.3s
7	—	OK 4.7s	—	—	—
8	OK 11.7s	OK 5/5	—	—	OK 12.0s
12	—	OK 8.9s	—	—	—
16	—	HANG 4/5 (1 trial finished at 43s, near the 45s timeout)	—	—	—

Cells without a 5/5 marker were tested once.

Notable observations:

• Doubling the mesh (16×8 → 32×16) at np=4 fixes the hang (it was the np=4 16×8 case specifically). So this is not a load-per-rank or memory-pressure issue.
• 16×16 (clean 2×2 partition at np=4) works; 16×8 (4×2 partition) hangs — looks geometry/aspect-ratio sensitive.
• np=8 is unaffected on every mesh tested.

Reproducer

# ve_stokes_parallel_hang_repro.py
import sys, time, sympy
import underworld3 as uw
from underworld3.function import expression

nx, ny = int(sys.argv[1]), int(sys.argv[2])
t0 = time.time()
if uw.mpi.rank == 0:
    print(f"[t=0.0s np={uw.mpi.size} mesh={nx}x{ny}] starting", flush=True)

mesh = uw.meshing.StructuredQuadBox(elementRes=(nx, ny))
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")

if uw.mpi.rank == 0:
    print(f"[t={time.time()-t0:.1f}s] about to solve", flush=True)
stokes.solve(timestep=0.02)
if uw.mpi.rank == 0:
    print(f"[t={time.time()-t0:.1f}s] DONE", flush=True)

mpirun -np 4 python -u ve_stokes_parallel_hang_repro.py 16 8     # HANGS
mpirun -np 4 python -u ve_stokes_parallel_hang_repro.py 32 16    # OK
mpirun -np 8 python -u ve_stokes_parallel_hang_repro.py 16 8     # OK

The hang is at the first stokes.solve(). [about to solve] prints, [DONE] does not.

Distinct from #96

#96 (BdIntegral parallel hang) was a PETSc coordinate-DM bug triggered by DMSetCoordinateDisc clearing the coordinate field cache; fixed via UW_DMForceCoordinateField and the createCoordinateSpace switch in nuke_coords_and_rebuild. The reproducer there involved boundary integrals and HDF5 mesh loading.

This bug:

• Triggers on a fresh in-memory mesh (no checkpointing).
• Uses only Dirichlet BCs — no add_natural_bc, no surface integrals.
• Hangs on the first solve, before any DM rebuild path is exercised.
• Completely independent of Parallel bug - Calculating a boundary integral causes the next Stokes solve to hang #96's coordinate-DM machinery.

Suggested next investigations

1. Stack-trace the hung ranks (lldb -p <pid> on macOS, or compile with py-spy/gdb-py). Are all ranks in the same MPI collective, or is one rank in a different code path?
2. Print the PETSc partitioning (-dm_view or mesh.dm.viewFromOptions) for the failing vs working configs — what does the rank topology look like at np=3 16×8 vs np=4 16×16?
3. Try a plain (non-VE) Stokes solver with the same BC layout to localise whether the trigger is in the VE / SemiLagrangian_DDt machinery or in the base Stokes assembly.
4. Bisect into _setup_solver — if the hang is during PETSc DS/DM setup, the rank-asymmetry will be near a MPI_Allreduce or MPI_Bcast that some ranks skip.

Environment

• macOS 15.2 (Apple Silicon)
• OpenMPI 5.0.10
• PETSc arch petsc-325-uw-openmpi
• underworld3 bugfix/jit-c-cache (PR JIT cache: key on generated C source (closes #121, #123) #129) — but the bug also reproduces on a clean origin/development checkout (verified by reverting _jitextension.py and petsc_generic_snes_solvers.pyx and rebuilding).

Underworld development team with AI support from Claude Code

[lmoresi]

Investigated and fixed in #136.

Root cause

mesh._get_coords_for_basis(degree, continuous) contains PETSc collectives (DMClone, createInterpolation, globalToLocal) but was populated lazily on first access via MeshVariable.coords_nd. That property is read from rank-local code paths — most importantly MeshVariable.rbf_interpolate inside global_evaluate_nd's per-particle RBF loop.

In global_evaluate_nd, after swarm migration, each rank evaluates the expression on its local migrated particles. In default mode (rbf=False), interior points go through petsc_interpolate and only exterior points trigger rbf_evaluate → rbf_interpolate → self.coords_nd → the collective. When a rank's migrated particles were all interior, it skipped the RBF path entirely and never entered the collective. Ranks with exterior particles did enter and waited forever.

That's why the bug was partition-sensitive: on the 4×2 partition at 16×8, migration distribution left some ranks without exterior particles; on the 2×2 partition at 16×16 every rank had some, so every rank reached the collective together and it completed.

Evidence

Python faulthandler.dump_traceback_later(20) was sufficient to expose the divergence without needing py-spy (which requires root on macOS). At the hang on np=4, 16×8:

• Ranks 0, 2, 3 — blocked at discretisation_mesh.py:2696 (dmnew = dmold.clone()), called via coords_nd → _get_coords_for_var → _get_coords_for_basis(degree=2, continuous=True) for variable V.
• Rank 1 — still in _global_evaluate_nd (Cython), never entered rbf_interpolate. Its migrated particles were all interior.

Post-migration particle counts at np=4, 16×8: 47 / 18 / 54 / 61. Rank 1's 18 particles happened to sit entirely inside the mesh domain.

The (degree=2, continuous=True) cache key for V was the trigger: the mesh's intrinsic key is (isSimplex, 1, True) and all the semi-Lagrangian psi_star / work vars are degree=1, so those cache entries were populated earlier in setup paths. The velocity's degree-2 key was the first cache miss to land in a rank-local context.

Fix

Pre-populate the coord cache at the end of _BaseMeshVariable.__init__ (a collective context, because it modifies the mesh DM). Plus defense-in-depth: re-populate for all registered variables at the end of nuke_coords_and_rebuild so mesh adaptation / DM rebuilds don't re-open the deadlock window.

Verification

Full (np, mesh) matrix from the report:

• np=3 {16×8, 32×16}, np=4 16×8, np=16 16×8 — previously HANG → now complete
• np=4 {16×16, 32×16}, np=8 16×8 — previously OK → still OK
• pytest -m "level_1 and tier_a" — 56 passed, no regressions

Regression test: tests/parallel/test_0780_ve_stokes_first_solve_mpi.py exercises the canonical np=4, 16×8 failure case with a 60s timeout.

Underworld development team with AI support from Claude Code

[lmoresi]

Fixed in #136 (merged to development at c514db3). Closing manually — GitHub auto-close via "Fixes #130" only fires on default-branch (main) merges.