jax-compile-stall epic: shipped phase 1, closed the rest as partial - #278
Merged
Conversation
Split the intermittent XLA vmap compile stall into the jax-compile-stall epic: phase 1 (PyAutoFit compile watchdog) is issued as PyAutoFit#1516, phases 2 (SLOW-vs-stall audit) and 3 (root cause + un-quarantine) stay in draft and are issued one at a time. The 2026-08-23 ci filing becomes the campaign ledger. Supersedes draft/bug/autolens_workspace_test/multi_dataset_jax_likelihood_xla_stall.md (same defect, filed 2026-08-22, never issued); its two unique asks are folded into phases 1 and 3. Records two source-level findings from the start_dev read of PyAutoFit that phase 3 starts from: Fitness._vmap builds jax.vmap(jax.jit(...)), the inverted ordering, on exactly the path that stalls; and log_on_first_compile covers both the compile wait and the block_until_ready execution wait with one log line and no heartbeat, which is why three quarantines produced no diagnosis. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015qk7hoavMnFyPtW4toYn8K
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015qk7hoavMnFyPtW4toYn8K
Phase 2 of the jax-compile-stall epic, issued as autolens_workspace_test#271. Corrects the census from the actual marker files: 21 SLOW-marked JAX entries across the two test workspaces, not 8, and 5 stall quarantines, not 3 -- imaging/jax_likelihood/rectangular_mge.py was quarantined in ag_test on 2026-08-23, and multi_dataset/jax_likelihood/mge.py is disabled in al_test's smoke_tests.txt. 26 entries in scope, not 11. Records three things the marker text settles before any re-timing: a 27.8s script SLOW-marked for flaking at a 1800s cap; the same script path SLOW in one repo and a stall in the other; and every 2026-07-14 SLOW marker reading 'flakes' -- bimodal language -- with no timing recorded behind it. Parked blocked: the measurement leg needs a re-timing harness, which does not exist, and where it lives is an organ-boundary call between a Heart-owned runner input and a per-workspace duplicate. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015qk7hoavMnFyPtW4toYn8K
Heart-side harness chosen: PyAutoHeart#161 gives the reusable smoke workflow a runner input, and autolens_workspace_test#272 / autogalaxy_workspace_test#110 add retime.py + a workflow_dispatch caller to each test workspace. #161 merges first — a caller cannot pass an input that does not exist. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015qk7hoavMnFyPtW4toYn8K
Phase 1 of the jax-compile-stall epic is shipped: PyAutoFit's compile wrapper now emits a heartbeat, arms a faulthandler watchdog (defaulting on under CI) and times the compile wait separately from the block_until_ready execution wait. The next CI stall diagnoses itself with no workspace or runner change. Record folds in the two phase-3 leads: Fitness._vmap is jax.vmap(jax.jit(...)), the inverted ordering, on exactly the path that stalls; and both NEEDS_FIX stalls post-date the persistent compile-cache default. Heart NOT consulted (web session, pyauto-heart unreachable); merge was a human instruction. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015qk7hoavMnFyPtW4toYn8K
) PyAutoHeart#161, autolens_workspace_test#272 and autogalaxy_workspace_test#110 all merged. The measurement leg is now unblocked; the full 26-entry sweep is left undispatched on purpose — many runner-hours at the 1800s cap. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015qk7hoavMnFyPtW4toYn8K
…#271) 40 script executions across both test workspaces settle three things: the datacube shared_preloads SLOW marker is refuted (10/10 completions, 34s worst against the 1800s cap it supposedly flakes at); rectangular_mge is a measured STALL, bimodal on both Python legs at ~80% per-compile probability; and mge_group + multi_dataset mge need the 1800s pass. The run also exposed a defect in phase 1's own watchdog: its 300s CI default collided with the 300s smoke cap, so the runner's SIGKILL beat the dump in all 20 stalled runs — heartbeats but no stacks. PyAutoFit#1518 derives the threshold from BUILD_SCRIPT_TIMEOUT instead. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015qk7hoavMnFyPtW4toYn8K
…ies (#271) Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015qk7hoavMnFyPtW4toYn8K
James's call: stop widening the timing measurement. Four entries was enough to prove the SLOW markers unreliable; the remaining 17 entries and the marker rewrites move out of the epic into draft/research/ci/smoke_timing_and_profiling.md, a follow-up task for proper smoke timing and profiling that carries the phase-2 results forward so nothing is re-measured. Phase 3 starts with the vmap-of-jit ordering A/B: experiment/jax-vmap-jit-ordering in PyAutoFit swaps jax.vmap(jax.jit(...)) for jax.jit(jax.vmap(...)), and a same-named branch in autogalaxy_workspace_test makes the workspace CI's matching-branch clone pick it up. Control is the measured 4/5 stall rate on both legs for imaging/jax_likelihood/rectangular_mge.py. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015qk7hoavMnFyPtW4toYn8K
Control 8/10 stalls, experiment 3/10, Fisher p=0.070. The stall survives the swap, so jax.vmap(jax.jit(...)) raises the probability rather than causing it. Recorded so nobody reads 80->30 as a fix. Healthy compile of rectangular_mge measured at 3.1s against a >300s stall — 100x bimodality inside one step. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015qk7hoavMnFyPtW4toYn8K
Phase 1 shipped in full. Phases 2 and 3 stop at a deliberate point (James's call): the stall is instrumented and characterised but NOT root-caused, and nothing was un-quarantined. The remaining breadth lives in draft/research/ci/smoke_timing_and_profiling.md, where this record is meant to be dug up. Record complete/2026/08/jax-compile-stall-slow-vs-stall-audit.md carries the 60 measured executions, the two live leads (vmap-of-jit contributory at p=0.070 but not causal; the compile-cache A/B never run), the traps, and what was deliberately left undone. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015qk7hoavMnFyPtW4toYn8K
Three conflicts, all in generated files — complete/index.md, dashboard.md and dashboard.html — from main's self-heal pushes landing while this branch was open. Resolved by taking main's copies and regenerating with the repo's own tooling (lifecycle.py index --apply, pyauto-brain intake --apply dashboard) rather than hand-merging generated output. main's curated Highlights band in complete/index.md survives regeneration. lifecycle check OK, index --check OK. Both jax-compile-stall records present; active.md back to its single pre-existing HOLD entry. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015qk7hoavMnFyPtW4toYn8K
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The Mind side of the
jax-compile-stallepic: the epic registration, four prompt files, the phase-1 completion record, the epic close-out record, and the follow-up task that inherits the unfinished work.Closed as partial by design (James, 2026-08-23): the stall is instrumented and characterised but not root-caused, and nothing was un-quarantined. The remaining breadth moves to a dedicated smoke timing/profiling task rather than being ground out one dispatch at a time.
What merged elsewhere
faulthandlerwatchdog + compile/execute timing splitrunner/runner-args/script-timeoutinputs on the reusable smoke workflowretime.pyharness +workflow_dispatchcallerIssues PyAutoFit#1516 and autolens_workspace_test#271 both closed.
What this PR adds
epics.md— thejax-compile-stallentry, statusCLOSED AS PARTIALdraft/bug/ci/jax_vmap_jit_compile_stall.md— the campaign ledger, carrying the original filing plus the phase map and outcome banner. It arrived on an unmerged branch (claude/backport-per-script-timeout-r3w1sv) and is brought ontomainheredraft/bug/ci/— phase 1 folded into its record; phases 2 and 3 retainedcomplete/2026/08/jax-compile-stall-evidence.md— phase 1's recordcomplete/2026/08/jax-compile-stall-slow-vs-stall-audit.md— the epic close-out: 60 measured executions, both live leads, the traps, and what was deliberately left undonedraft/research/ci/smoke_timing_and_profiling.md— the follow-up task, carrying the four measured verdicts forward so nothing is re-measureddraft/bug/autolens_workspace_test/multi_dataset_jax_likelihood_xla_stall.md— the same defect, filed 2026-08-22, never issuedThe findings worth keeping
interferometer/datacube/shared_preloads.py— completed 10/10 runs at a worst case of 34.0s, 1.9% of the cap it supposedly flakes at.rectangular_mge.pyis 3.1s; a stalled one exceeds 300s, same commit and runner image.vmap(jit)is contributory, not causal. A/B: 8/10 stalls on main vs 3/10 swapped, Fisher exact p = 0.070. The stall survives the swap, so the ordering is not necessary for it.Verification
lifecycle.py check: OK.lifecycle.py index --check: OK.dashboard.md/dashboard.htmlregenerated.active.mdis back to its single pre-existing HOLD entry.One caveat
Two 1800s re-timing runs were still in flight at close-out (ag_test 32668061785, al_test 32668067325) and should carry the first
faulthandlerstack. The record names them as the first thing to read when this is resumed.Generated by Claude Code