JIT C-source generation is non-deterministic across MPI ranks (flaky hard abort at np>1)

[lmoresi]

_jitextension.py guards against generate_c_source producing different C on different ranks, and raises a hard error when it does. That guard fires intermittently — about half of runs — on a Stokes solve with a power-law transversely isotropic viscosity under the consistent tangent.

RuntimeError: JIT C-source hash differs across MPI ranks:
{'471c0c17f9aba995', '5e640ad7875d8e9f'}. This indicates non-determinism in
generate_c_source — likely a set or dict whose iteration order leaks into the C
output. Treating this as a hard error since cache reuse would be unsound.

src/underworld3/utilities/_jitextension.py:680

Repro

On feature/adjoint-rotated-bc (PR #751):

pixi run -e amr-dev mpirun -n 2 python -m pytest \
  --config-file=tests/pytest.ini tests/test_0022_rotated_adjoint.py -q

Roughly 1 in 2 runs aborts; exit codes over six attempts: 1, 0, 1 and, with PYTHONHASHSEED=0 pinned on both ranks, 1, 0, 0. When it aborts it does so in about 6.7 s, in the fixture's first forward solve.

The setup is an annulus, TransverseIsotropicFlowModel, shear_viscosity_0 = (0.01 + eII)**(-1/3), shear_viscosity_1 = η₁·shear_viscosity_0 with η₁ a uw.expression, radial director, consistent_jacobian = True.

What we ruled out

• Not hash randomisation. Pinning PYTHONHASHSEED=0 on both ranks still aborts (1 of 3).
• Not a general regression. tests/test_0021_adjoint_misfit_through_gradient.py, which also uses TransverseIsotropicFlowModel, is stable at np=2 (3/3). Its viscosity is constant rather than power-law.
• Not the expression on its own. Building the same solver outside pytest and solving never diverges (3/3), across four variants that isolate the features: isotropic power-law; TI constant; TI power-law with a plain sympy.Float coefficient; TI power-law with a uw.expression coefficient. All ok.
• Not the adjoint. The abort lands in the first forward solve, before any adjoint code runs. The same file passes 7/7 serially.

So the trigger needs the pytest process context — several solvers and meshes built in one interpreter — on top of an expression of this shape. The guard's own hypothesis (a set or dict whose iteration order leaks into the emitted C) is still the likely mechanism, but it is not driven by PYTHONHASHSEED, which points at ordering derived from something like object identity or insertion order that differs between ranks rather than from string hashing.

Why it is not currently visible

CI's parallel pass collects only tests/parallel/test_*.py (scripts/test.sh:199), and every file that exercises this expression shape sits in the serial band. Any solve of this kind at np>1 in user code is exposed.

Impact

The guard is doing its job — failing loudly rather than letting ranks compile different kernels — so this is a correctness-preserving abort, not a wrong answer. But it makes a legitimate parallel run fail at random, and the underlying non-determinism would be unsound if the guard were ever relaxed.

Underworld development team with AI support from Claude Code

[lmoresi]

Root cause found: sympy prints Mul factors in stored order, not canonical order

We dumped both ranks' canonical_source on a failing run and diffed them.

12 lines differ out of 225. Every one has an identical token multiset and identical length. The only difference is the order of factors in a commutative product:

rank0:  petsc_x[1]*pow(petsc_x[0], 3)        rank0:  petsc_x[1]*petsc_x[0]
rank1:  pow(petsc_x[0], 3)*petsc_x[1]        rank1:  petsc_x[0]*petsc_x[1]

Same mathematics, same operations, same constants (constants[0]/constants[1] are identically indexed on both ranks). Only the print order changes.

So the guard's own hypothesis — "likely a set or dict whose iteration order leaks into the C output" — points at the wrong layer. Nothing in UW3's dictionaries is reordering anything. sympy is printing Mul args in stored order.

Mechanism

StrPrinter._print_Mul, which C99CodePrinter inherits, opens with an "unevaluated Mul" branch that prints expr.args as stored, deliberately bypassing as_ordered_factors():

args = expr.args
if args[0] is S.One or any(
        isinstance(a, Number) or
        a.is_Pow and all(ai.is_Integer for ai in a.args)
        for a in args[1:]):
    # ... prints in stored arg order, never consulting self.order

Outside that branch it does canonicalise. Confirmed with plain symbols: both x0**3*x1 and x1*x0**3 print as pow(x0, 3)*x1 — which is exactly rank1's form. rank1 is the canonically sorted output; rank0 is the unsorted one.

Note the printer is constructed as c_code_printers["c99"]({"user_functions": custom_functions}) with order unset (printer.order is None), but setting order would not help: this branch never consults it.

Stored Mul.args order is not a sympy invariant across processes whose construction history differs — and two ranks genuinely differ (different cell counts, different boundary nodes, different uw.function.evaluate point sets feeding sympy's cache).

That explains every negative result in the original report:

• not PYTHONHASHSEED — no string hashing is involved;
• not reproducible in a fresh script — no divergent construction history;
• test_0021 escapes it — its constant viscosity never builds products of this shape.

The guard is right to fire

real_modname = f"fn_ptr_ext_{source_hash}" and the C symbol prefix both derive from the hash. Ranks that disagree compile and load differently named modules, which breaks the rank-0-compiles/others-load protocol. This is not a cosmetic difference the guard over-reacts to. What is wrong is that the hash is taken over text sympy never promised to make canonical.

Candidate fix (tested)

Every rank's source is mathematically correct, so the cheapest correct repair is to stop requiring agreement and instead establish it — adopt rank 0's source on mismatch, rehash, warn:

canonical_codeguys = underworld3.mpi.comm.bcast(canonical_codeguys, root=0)
canonical_source = "\n".join(e[1] for e in canonical_codeguys)
source_hash = hashlib.sha256(
    (canonical_source + "\n---\n" + _abi_salt()).encode("utf-8")
).hexdigest()[:16]

Measured: 6 runs, 6 clean, against a ~50% baseline failure rate (tests/test_0022_rotated_adjoint.py at np=2). Every rank then compiles and loads the same artefact.

Caveat: the bcast is a collective, so it must sit where all ranks reach it. In the patch above they do — the mismatch branch is entered collectively, since the allgather result is the same on every rank.

The deeper fix

Canonicalising factor order before printing (rebuilding each Mul with sorted args, or pre-canonicalising the expressions) is better hygiene and would make the emitted C reproducible rather than merely agreed-upon. But it changes every generated kernel and so invalidates every cached hash in the wild, on disk and in memory. Worth doing deliberately rather than as a hotfix.

Suggested order: broadcast now to stop the flaky aborts, canonicalise later as a separate, deliberate cache-invalidating change.

Underworld development team with AI support from Claude Code

[lmoresi]

Correction: the printer is not at fault, and an ordering fix does not work

My earlier comment attributed this to StrPrinter._print_Mul's unevaluated-Mul fast path printing expr.args in stored order. That is wrong, and I have disproved it. Correcting the record, because it would send the next person down a dead end.

The printer sorts, and forcing it changes nothing

The C printers do not use StrPrinter._print_Mul at all — C89CodePrinter._print_Mul delegates to CodePrinter._print_Mul (sympy/printing/codeprinter.py), which has no such fast path:

if self.order not in ('old', 'none'):
    args = expr.as_ordered_factors()

printer.order is None, so it does sort. And sorting is deterministic for our objects — measured on the real UWCoordinate base scalars:

Mul._from_args((N.y, N.x**3)).as_ordered_factors()  ->  [N.x**3, N.y]
Mul._from_args((N.x**3, N.y)).as_ordered_factors()  ->  [N.x**3, N.y]

Two independent checks confirm the printer is innocent:

1. A _print_Mul override that forces as_ordered_factors() produced byte-identical C to the stock printer across all five kernel bundles of the failing case, including the power-law TI ones.
2. With that override installed, a run still aborted (1 of 7).

It is not commutative ordering at all

Recursively sorting every Mul/Add argument by default_sort_key before hashing still leaves the two ranks' expressions different (2 of 3 runs diverged, 2 equations differing). So the divergence is structural, not an ordering artefact — even though the emitted C looks like pure factor reordering.

What is actually established

Dumping sympy.srepr of each kernel expression per rank on a failing run: of 15 equations, 14 hash identically and eqn_7 differs. So the divergence is in the expression handed to the printer, not the printing.

And the solver's own symbolic blocks all agree across ranks. Hashing them directly at np=2:

c_tensor AGREE   flux AGREE   _u_F1 AGREE   _uu_G0 AGREE   _uu_G1 AGREE
_uu_G2 AGREE     _uu_G3 AGREE   _up_G0 AGREE   _pu_G1 AGREE

...on runs where the JIT then aborted with divergent C.

Inputs to the JIT agree; the JIT's output diverges. The non-determinism is inside the lowering pipeline in getext/generate_c_source — the phases between the solver's pointwise functions and printer.doprint: _reveal_constants, unwrap, the constants_subs_map xreplace, to_matrix(mesh.N), the free_syms walk, and the cse path. Some of those consume .atoms() (a set) directly; _stable_sorted is already used in places, which suggests this class of bug has been hit before there.

So the guard's original wording — "likely a set or dict whose iteration order leaks into the C output" — was closer to right than my correction was. It just needs to be read as pointing at the lowering, not the printing.

Consequences for the fix

• The ordering fix is not available. There is no ordering to canonicalise at the printer level; both the printer override and a full recursive Mul/Add sort were measured not to work.
• The broadcast fix still works (6 clean runs against a ~50% baseline) but is a workaround: it makes ranks agree rather than making generation deterministic.
• The real fix is to find the set/dict iteration in the lowering pipeline. eqn_7 in tests/test_0022_rotated_adjoint.py's fixture is a reproducer, at roughly one abort in two runs at np=2.

A diagnostic worth adding regardless: on mismatch, have the guard dump each rank's source (or their first differing line) under an env var. Comparing token multisets per line is what showed this was a reordering rather than different mathematics, and it took minutes instead of reading 90 KB expressions.

Underworld development team with AI support from Claude Code

[lmoresi]

Mitigation up as PR #753; and the trigger is branch-specific

PR #753 makes the ranks agree instead of refusing: adopt rank 0's source, rehash, warn. 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. It does not close this issue; the lowering is still non-deterministic.

One finding that changes where to look, and corrects something I said earlier.

The trigger does not reproduce on development

Running the same script at np=2:

branch	divergences
development	0 of 12
feature/discrete-adjoint	2 of 4

I earlier wrote that this was "latent JIT non-determinism rather than anything this branch introduced", reasoning from test_0021 being stable at np=2. That was the wrong comparison — test_0021 is on the adjoint branch too. The branch-vs-development comparison is the sharper test, and it says the adjoint branch is what exposes it.

I checked the obvious suspect and ruled it out: the Newton-alpha machinery is already on development (12 references in petsc_generic_snes_solvers.pyx), so it is not the differentiator. What on that branch changes the lowering's input enough to expose this is not yet identified — and worth knowing, because it may point straight at the non-deterministic step.

Caveat on the comparison: the two checkouts differ in worktree state as well as in branch (JIT cache, .meshes), so "branch" is the likely explanation but not a controlled one. Anyone chasing this should re-run the pair in fresh worktrees before treating it as settled.

Where the bug still is

Unchanged by the above, and the most useful part of the diagnosis: the solver's symbolic blocks hash identically across ranks on runs that then abort, while sympy.srepr of the kernel expressions shows exactly one of 15 differing. Inputs to the JIT agree; its output diverges. The non-determinism is in the lowering — _reveal_constants, unwrap, the constants_subs_map xreplace, to_matrix, the free_syms walk, cse — several of which consume .atoms(), a set.

Underworld development team with AI support from Claude Code

[lmoresi]

Correction: this is NOT branch-specific. It is pre-existing on development.

My previous comment reported development 0/12 versus feature/discrete-adjoint 2/4 and concluded the adjoint branch exposes this. That was an uncontrolled comparison and the conclusion is wrong. The two checkouts differed in JIT cache and .meshes state as well as in branch — I flagged that caveat and should have resolved it before drawing the conclusion rather than after.

Controlled re-run: one worktree, one pixi environment, one set of caches, branch as the only variable. Detection is the #753 repair firing, which counts divergences without aborting.

arm	repair fired
development	2 / 10
feature/discrete-adjoint	2 / 10

Identical. There is no branch regression. The earlier 0/12 was a low-probability sample, not a different behaviour: at a ~20% rate, 0 in 12 has probability ~0.07.

The rate itself is environment-sensitive — 2/4 and 4/6 in a warm-cache worktree, 2/10 in both controlled arms — so any future comparison here needs the controlled form: same worktree, same caches, one variable. Rates from different checkouts are not comparable, which is exactly the trap I fell into.

What stands

Everything about the localisation is unaffected, since it was measured within single runs:

• the two ranks' C differs only by the order of factors in commutative products — identical token multisets, identical length;
• the solver's symbolic blocks (c_tensor, flux, _u_F1, _uu_G0–_uu_G3, _up_G0, _pu_G1) hash identically across ranks on runs that abort;
• sympy.srepr per kernel shows exactly one of 15 equations differing.

Inputs to the JIT agree; its output diverges. The non-determinism is in the lowering — _reveal_constants, unwrap, the constants_subs_map xreplace, to_matrix, the free_syms walk, cse — several of which consume .atoms(), a set.

What changes is the urgency and the blast radius: this is a latent bug that any sufficiently complex expression can hit at np>1 on development today, not something arriving with a feature branch. It is not a merge blocker for the adjoint work, and it is more worth fixing properly than I implied.

Underworld development team with AI support from Claude Code