synchronised_array_update: collective dirty-variable flush replaces rank-local event replay (#379 item 2)

[lmoresi]

Third of the #379 series, stacked on #381 (base is bugfix/canonical-callback-guard; retarget to development once #381 merges). Independent of #382.

The problem

uw.synchronised_array_update queued every write event per rank and replayed the queue at context exit, between blanket barriers. The queue contents are rank-dependent (each rank queues what it locally wrote), while mesh-variable callbacks contain collective operations (the ghost synchronisation). So ranks that wrote unevenly replayed different collectives — a masked write on one rank hung every other rank — and the barriers around the batch could not fix collectives inside it. Per-item exceptions were swallowed into a logger warning, the same swallow that hid #376.

The fix

Writes inside the context now mark their canonical variable dirty instead of queueing events. At exit:

• Mesh variables (collective packs) are flushed by cross-rank agreement. Every variable gets a creation-order registration id — construction is collective, so the ids line up on every rank by construction. (Names cannot serve as the key: temporary variables embed rank-local id() values in their names.) The union of dirty ids is allgathered — unconditionally, empty sets included, so ranks stay matched — and every rank flushes the union in id order. Ranks that wrote nothing pack their unchanged values, which is what keeps the per-variable collective matched.
• Swarm variables flush rank-locally (their packs touch only local particle storage); particle migration stays rank-agreed through the existing _needs_migration reduction at solve entry.
• Exceptions propagate. If the context body raises, the flush is skipped entirely — ranks unwind exceptions asymmetrically, and replaying collectives during unwinding was itself a hang vector. The written values remain in the local arrays. (This is a deliberate behaviour change from replay-during-unwind.)

Two adjacent traps, fixed because the rank-uneven promise is broken without them

1. The .array views called pack(sync=True) directly on every setitem — a collective per write that bypassed the delay context entirely (and model.py's persistent-variable transfer does exactly this inside the context). They now route through the canonical array, so .array writes follow the same path as .data writes: immediate pack outside a context, one agreed flush inside.
2. The first .data access performed collective vec setup lazily, so a first touch from rank-conditional code deadlocked. The canonical array is now created eagerly at variable construction (which is collective) — the same reasoning as the existing VE Stokes hangs on first solve at specific (np, mesh) combinations #130 coordinate-cache pre-population.

Also: the two duplicate delay-context classes collapse into one; swarm.access() / mesh.access() compatibility shims inherit the new exit semantics automatically.

Verification

• tests/test_0140_synchronised_updates.py rewritten from a print-script that could not fail into real tests (level_1 / tier_a, written first): one flush per variable at exit, silence during the context, masked-write canonical marking, per-level nested flushing, exception-skips-flush, legacy untagged replay preserved.
• New tests/parallel/test_0758_synchronised_update_collective_flush.py (level_2 / tier_a): rank-uneven writes, opposite per-rank write order, helper-created variables, mixed swarm+mesh, zero-size slice on one rank. 5 passed at np=2 and np=4. The rank-uneven case hangs on the old machinery.
• Serial quick gate level_1 and tier_a: 411 passed, 1 pre-existing xfail. Swarm-write and mesh-deform MPI smokes at np=2: passed.

Flagged, not fixed (out of scope): _vector_stats/_tensor_stats create variables named _temp_mag_{id(self)} — rank-asymmetric names (TODO in a follow-up issue); a latent np NameError on the unit-aware tensor write path is marked with a TODO(BUG) at the site.

Underworld development team with AI support from Claude Code

[lmoresi]

Adversarial review — round 1

One reviewer attacked the flush design with live np=2 repros (faulthandler-verified hang tracebacks). The creation-order id mechanism itself survived: no rank-conditional variable-construction path exists in src/ (every _BaseMeshVariable.__init__ contains world collectives, so a violation fails fast rather than desynchronising the counter), nested symmetric contexts scope correctly, and in-context reads are not stale. But four real findings, two of them blocking:

1. BLOCKING — a caught rank-local exception now deadlocks. The raising rank skips the exit allgather entirely while the other ranks run it; the allgather then pairs with an unrelated collective. Verified at np=2: symmetric writes, rank 0 raises and catches a ValueError outside the context → rank 1 blocks in the allgather, rank 0 in the next barrier. The old machinery replayed unconditionally, so this exact shape used to exit cleanly. Fix: the exit must run one rank-agreement collective on BOTH paths — allgather an aborted flag first, and skip the flushes everywhere if any rank aborted.

2. BLOCKING — v.array = values (attribute assignment) bypasses the deferral. The views were rerouted, but the property setter still packs collectively per write. Verified hang at np=2 for a rank-conditional v.array = vals inside the context. Not a regression (it hung before too), but it falsifies the docstring's rank-uneven-writes promise for the most natural spelling of a full-field write. Fix: route the setter through the view path.

3. Should-fix — the two most important legacy callbacks contain collectives, violating the stated legacy-replay invariant: mesh.X.coords writes replay a full collective deform rank-locally at exit (verified hang at np=2), and the deprecated swarm.points callback runs collective migrate() (removed by Retire the swarm.points writable stack: read-only snapshot, coords is the write path (#379 item 1) #382). Fix: mesh coordinates become a first-class collective-flush citizen — register the coords callback through the canonical guard and give the mesh a registration id, so coordinate writes inside the context join the same rank-agreed flush.

4. Should-fix — dirty_local's strong refs can flush a stale canonical. swarm.migrate() inside the context invalidates and rebuilds the canonical; the flush then fires the pinned pre-migration array (verified serial repro: user reads and PETSc disagree afterwards; in parallel, uneven size changes raise rank-locally before the allgather — finding 1's shape). Fix: store the owner weakly and re-resolve its live canonical at flush time, as the collective branch already does.

Also recorded: the missing-id RuntimeError is nearly unreachable in practice (the orchestration model's registry pins every registered variable strongly — pre-existing), and counter misalignment would be undetectable in principle, resting on "construction always contains a world collective" — true today, worth a cheap debug assertion later. Test gaps to close with the fixes: exception path, .array view path in parallel, the setter, and coordinate writes inside the context.

Fixes follow on this branch.

Underworld development team with AI support from Claude Code

[lmoresi]

Round 1 — response (fixed; branch rebased onto the updated #381)

All four findings addressed:

1. Caught-exception deadlock — the exit allgather now runs on both paths, carrying an aborted flag; any rank aborting makes every rank skip the flushes together, and all rank-local flushing moved after the agreement. Pinned by test_caught_rank_local_exception_exits_everywhere (np=2/np=4 — this shape hung before the fix).
2. v.array = setter bypass — attribute assignment now routes through the view write path, so it defers like every other write. One deliberate consequence: under active units the setter now enforces the same unit-aware-input requirement as view writes, instead of silently packing plain arrays. Pinned by test_rank_uneven_array_setter.
3. Legacy coords callback ran a rank-local collective deform — mesh coordinates are now first-class flush citizens: the coords callback registers through the canonical guard and the mesh carries a registration id, so mesh.X.coords writes inside the context defer to one rank-agreed deform at exit. Pinned by test_coords_write_inside_context_defers_deform. (The other offending legacy callback, swarm.points, is deleted by Retire the swarm.points writable stack: read-only snapshot, coords is the write path (#379 item 1) #382.)
4. Stale-canonical flush — swarm-variable dirt is recorded as a weak owner reference and the LIVE canonical is re-resolved at flush time. A migrate() issued inside the context forfeits unflushed prior writes but leaves user reads and PETSc consistent — documented on the flush target.

Also from the review: the near-unreachable missing-id error and the in-principle-undetectable counter misalignment are accepted as documented residual risk ("construction always contains a world collective" — true today, cheap debug assertion left for a follow-up), and the rank-uneven-.array-view shape got its own test.

9/9 parallel tests at np=2 and np=4; quick gate 415 passed.

Underworld development team with AI support from Claude Code

[lmoresi]

Round 2 — verification and one blocking find (fixed in 58091e5)

The round-2 verifier confirmed all round-1 fixes closed (collective counts traced through every nesting/abort combination; setter and both view classes route through the canonical; flush targets re-resolve live). It found one blocking defect introduced by the coords canonical registration: _deform_mesh re-wraps the coordinate array after a rebuild and copied the callback list verbatim — the canonical dispatch is bound (by weakref) to the OLD array's identity, so on the replacement it classified every write as a foreign copy and silently returned. A second mesh.X.coords write did nothing: no deform, no version bump, no error. This is the same bug class the #381 round-1 review caught in sync_data, now closed the same way: the re-wrap re-homes canonical callbacks via their recorded original function. Regression test writes coordinates twice and checks the DM.

Also hardened from a round-2 note: rank-local flush phases can raise rank-locally; when a collective flush follows, ranks now agree on local-phase success before entering the per-variable collectives, so no rank unwinds while others block.

Residual (accepted, documented): a rank-local exception in the context body followed by other user collectives after the raising point is inherent to rank-local exceptions and outside this machinery's reach.

Serial battery 21 passed; 9/9 parallel at np=2; gate re-running on the stack tip.

Underworld development team with AI support from Claude Code