From b534a29b89983e9febc24ec2febd4abbeb8d9822 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 17 Sep 2026 09:14:54 +1000 Subject: [PATCH] fix: make every MPI rank compile the same generated C, rather than refuse 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 was REQUIRED, and a mismatch raised. 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 under the consistent tangent trips it in roughly half of np=2 runs - measured 2 of 4 on feature/discrete-adjoint. The abort lands in the forward solve. Two measurements say 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. Diffing both ranks' output on a failing run: 12 of 225 lines differ, every one with an identical token multiset and identical length. Every rank's source is a correct kernel for the same equation. * the solver's own symbolic blocks - constitutive tensor, flux, and every Jacobian block - hash IDENTICALLY across ranks on the runs that abort. What differs is produced inside generate_c_source, not handed to it. So the disagreement is over which of several correct spellings to compile, and adopting one is enough. Rank 0's is taken and every rank rehashes from it, which restores the invariant that actually matters: one source, one hash, one module. Measured on the reproducing branch: the repair fired in 4 of 6 runs and all 6 succeeded, against a baseline of 2 aborts in 4. This is a REPAIR, not a fix. It is said out loud - one warning naming the issue - rather than papering over the upstream non-determinism, which is still a bug and still worth finding. The common path is untouched: ranks that already agree keep their own objects and pay one allgather of a 16-character string. Two dead ends recorded so they are not retried. The C printers do NOT use StrPrinter._print_Mul - C89CodePrinter delegates to CodePrinter._print_Mul, which has no stored-order fast path and does sort, and forcing that sort produces byte-identical C while a run still aborts. And it is not commutative ordering at all: recursively sorting every Mul/Add argument by default_sort_key still leaves the ranks differing. The repair is extracted as _agree_source_across_ranks so it can be tested directly. Reproducing a real disagreement 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 tests build the disagreement themselves, so they keep working once the lowering is made deterministic. Underworld development team with AI support from Claude Code --- src/underworld3/utilities/_jitextension.py | 79 ++++++++++-- .../test_0022_jit_rank_source_agreement.py | 121 ++++++++++++++++++ 2 files changed, 186 insertions(+), 14 deletions(-) create mode 100644 tests/parallel/test_0022_jit_rank_source_agreement.py diff --git a/src/underworld3/utilities/_jitextension.py b/src/underworld3/utilities/_jitextension.py index baa3ccf6..96b82fef 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/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, diff --git a/tests/parallel/test_0022_jit_rank_source_agreement.py b/tests/parallel/test_0022_jit_rank_source_agreement.py new file mode 100644 index 00000000..b0f079e1 --- /dev/null +++ b/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