Make every MPI rank compile the same generated C, rather than refuse

src/underworld3/utilities/_jitextension.py

@@ -688,20 +688,36 @@ def getext(
         (canonical_source + "\n---\n" + _abi_salt()).encode("utf-8")
     ).hexdigest()[:16]
 
-    # Determinism check: all ranks must agree on the hash. A mismatch means
-    # generate_c_source isn't deterministic across ranks (typically caused
-    # by set/dict-iteration order leaking into the emitted C). Caching
-    # cannot work correctly if ranks disagree, so fail loudly rather than
-    # let stale entries propagate.
-    if underworld3.mpi.size > 1:
-        all_hashes = underworld3.mpi.comm.allgather(source_hash)
-        if any(h != source_hash for h in all_hashes):
-            raise RuntimeError(
-                f"JIT C-source hash differs across MPI ranks: {set(all_hashes)}. "
-                f"This indicates non-determinism in generate_c_source — likely "
-                f"a set or dict whose iteration order leaks into the C output. "
-                f"Treating this as a hard error since cache reuse would be unsound."
-            )
+    # All ranks must end up compiling and loading the SAME module: the module
+    # name and the C symbol prefix are both derived from `source_hash` below, so
+    # ranks that disagree would build disjoint artefacts and the
+    # rank-0-compiles/others-load protocol would break.
+    #
+    # Agreement used to be REQUIRED here, and a mismatch was a hard error. It
+    # fires in practice: the lowering above is not yet deterministic across
+    # ranks (#752), and a Stokes solve with a power-law transversely isotropic
+    # viscosity trips it in roughly half of np=2 runs. What we measured there
+    # matters for why this is safe to repair rather than refuse:
+    #
+    #   * the sources differ only in the ORDER of factors in commutative
+    #     products — identical token multisets, identical length, identical
+    #     mathematics. Every rank's source is a correct kernel for the same
+    #     equation;
+    #   * the solver's own symbolic blocks (constitutive tensor, flux, every
+    #     Jacobian block) hash IDENTICALLY across ranks on the runs that abort.
+    #     What differs is produced inside this function, not handed to it.
+    #
+    # So the disagreement is about which of several correct spellings to
+    # compile, and adopting one of them is enough. Rank 0's is taken, and every
+    # rank rehashes from it, which restores the one invariant that matters: one
+    # source, one hash, one module.
+    #
+    # This is a REPAIR, not a fix. The non-determinism upstream is still a bug
+    # and still worth finding, which is why it is said out loud rather than
+    # papered over silently.
+    canonical_codeguys, canonical_source, source_hash = _agree_source_across_ranks(
+        canonical_codeguys, canonical_source, source_hash
+    )
 
     # Derive the real modname/randstr from the hash — same source ⇒ same
     # compiled artefact, different sources ⇒ disjoint symbol namespaces.
@@ -862,6 +878,41 @@ def _aux_component_offsets(mesh):
     return offsets
 
 
+def _agree_source_across_ranks(canonical_codeguys, canonical_source, source_hash):
+    """Make every rank compile the SAME generated C, and say so if they did not.
+
+    Returns the (possibly replaced) ``(codeguys, source, hash)``. Serial runs and
+    runs where the ranks already agree are returned untouched, so the common path
+    costs one ``allgather`` of a 16-character string.
+
+    See the call site for why adopting one rank's source is a sound repair rather
+    than papering over a wrong answer. Separated out so the repair can be tested
+    directly — forcing a real disagreement through the JIT means reproducing a
+    non-deterministic bug, which is not a test.
+    """
+    import hashlib          # module-local in generate_c_source too
+
+    if underworld3.mpi.size <= 1:
+        return canonical_codeguys, canonical_source, source_hash
+
+    all_hashes = underworld3.mpi.comm.allgather(source_hash)
+    if all(h == source_hash for h in all_hashes):
+        return canonical_codeguys, canonical_source, source_hash
+
+    canonical_codeguys = underworld3.mpi.comm.bcast(canonical_codeguys, root=0)
+    canonical_source = "\n".join(entry[1] for entry in canonical_codeguys)
+    source_hash = hashlib.sha256(
+        (canonical_source + "\n---\n" + _abi_salt()).encode("utf-8")
+    ).hexdigest()[:16]
+    underworld3.mpi.pprint(
+        f"[jit] WARNING: generated C differed across ranks "
+        f"({sorted(set(all_hashes))}); adopted rank 0's source so every rank "
+        f"compiles the same module. The kernels are mathematically identical — "
+        f"see issue #752 for the upstream non-determinism."
+    )
+    return canonical_codeguys, canonical_source, source_hash
+
+
 def generate_c_source(
     name,
     mesh: underworld3.discretisation.Mesh,

tests/parallel/test_0022_jit_rank_source_agreement.py

@@ -0,0 +1,121 @@
+"""Every rank must compile the SAME generated C.
+
+The JIT derives both the module name and the C symbol prefix from a hash of the
+generated source, so ranks that generate different text build disjoint artefacts
+and the rank-0-compiles/others-load protocol breaks. Agreement used to be
+REQUIRED — a mismatch raised — and it fires in practice: the lowering in
+``generate_c_source`` is not deterministic across ranks (#752), and a Stokes solve
+with a power-law transversely isotropic viscosity trips it in roughly half of np=2
+runs.
+
+Two measurements say why adopting one rank's source is a sound repair rather than
+a way of ignoring a wrong answer:
+
+  * the sources differ only in the ORDER of factors in commutative products —
+    identical token multisets, identical length, identical mathematics — so every
+    rank's source is a correct kernel for the same equation;
+  * the solver's own symbolic blocks (constitutive tensor, flux, every Jacobian
+    block) hash IDENTICALLY across ranks on the runs that abort. What differs is
+    produced inside ``generate_c_source``, not handed to it.
+
+This file tests the repair directly, by handing
+``_agree_source_across_ranks`` a disagreement built on purpose. Reproducing a
+real one means reproducing a non-deterministic bug at about one run in two, which
+is not a test — it would pass half the time with the repair removed.
+
+The repair is not a fix. #752 is still open, and when the lowering is made
+deterministic the disagreement branch should stop being reached — these tests
+keep working either way, because they construct the disagreement themselves.
+"""
+
+import hashlib
+
+import pytest
+
+import underworld3 as uw
+from underworld3.utilities._jitextension import (
+    _abi_salt,
+    _agree_source_across_ranks,
+)
+
+pytestmark = [pytest.mark.timeout(300)]
+
+
+def _hash_of(codeguys):
+    """The same hash ``generate_c_source`` derives the module name from."""
+    source = "\n".join(entry[1] for entry in codeguys)
+    return hashlib.sha256(
+        (source + "\n---\n" + _abi_salt()).encode("utf-8")
+    ).hexdigest()[:16]
+
+
+@pytest.mark.skipif(uw.mpi.size < 2, reason="needs at least two ranks")
+def test_ranks_that_disagree_all_adopt_rank_zero_source():
+    """The contract. Each rank arrives with different text; all leave with rank
+    0's, and with the hash that text actually produces — not merely with hashes
+    that happen to match each other."""
+    mine = [["eqn_0", f"/* generated on rank {uw.mpi.rank} */"]]
+    agreed, source, digest = _agree_source_across_ranks(
+        mine, "\n".join(e[1] for e in mine), _hash_of(mine)
+    )
+
+    expected_source = "/* generated on rank 0 */"
+    assert source == expected_source
+    assert agreed == [["eqn_0", expected_source]]
+    assert digest == _hash_of([["eqn_0", expected_source]])
+
+    # and every rank really did land on the same one
+    assert len(set(uw.mpi.comm.allgather(digest))) == 1
+
+
+@pytest.mark.skipif(uw.mpi.size < 2, reason="needs at least two ranks")
+def test_agreeing_ranks_are_left_exactly_as_they_were():
+    """The common path must be a no-op, not a broadcast. Ranks that already agree
+    keep their own objects: a repair that rewrote the source on every call would
+    hide the upstream bug rather than report it, and would pay a collective on
+    every kernel."""
+    same = [["eqn_0", "/* identical on every rank */"]]
+    before = _hash_of(same)
+    agreed, source, digest = _agree_source_across_ranks(
+        same, "\n".join(e[1] for e in same), before
+    )
+
+    assert digest == before
+    assert agreed is same
+    assert source == "/* identical on every rank */"
+
+
+@pytest.mark.skipif(uw.mpi.size < 2, reason="needs at least two ranks")
+def test_a_real_kernel_generates_the_same_hash_on_every_rank():
+    """End to end: the symptom #752 actually presents as. A Stokes solve with a
+    power-law transversely isotropic viscosity is the case that trips the
+    lowering; it must complete rather than abort, whichever way the lowering
+    happens to fall on each rank."""
+    import sympy
+
+    mesh = uw.meshing.Annulus(radiusInner=0.5, radiusOuter=1.0,
+                              cellSize=0.25, qdegree=3)
+    x, y = mesh.X
+    r = sympy.sqrt(x**2 + y**2)
+    unit_r = sympy.Matrix([[x / r, y / r]])
+    th = sympy.atan2(y, x)
+    v = uw.discretisation.MeshVariable("v_jit", mesh, 2, degree=2)
+    p = uw.discretisation.MeshVariable("p_jit", mesh, 1, degree=1)
+
+    stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p)
+    edot = mesh.vector.strain_tensor(v.sym)
+    eII = sympy.sqrt(sympy.Rational(1, 2) * (edot[0, 0] ** 2 + edot[1, 1] ** 2)
+                     + edot[0, 1] ** 2)
+    eta_0 = (sympy.Float(0.01) + eII) ** sympy.Rational(-1, 3)
+    stokes.constitutive_model = uw.constitutive_models.TransverseIsotropicFlowModel
+    stokes.constitutive_model.Parameters.shear_viscosity_0 = eta_0
+    stokes.constitutive_model.Parameters.shear_viscosity_1 = 0.2 * eta_0
+    stokes.constitutive_model.Parameters.director = unit_r
+    stokes.bodyforce = 1.0e2 * sympy.cos(3 * th) * (r - 0.5) / 0.5 * unit_r
+    stokes.add_dirichlet_bc((0.0, 0.0), "Lower")
+    stokes.add_rotated_freeslip_bc(0.0, "Upper")
+    stokes.consistent_jacobian = True
+    stokes.tolerance = 1.0e-9
+
+    stokes.solve(zero_init_guess=True)   # raised here before the repair
+    assert stokes.snes is not None