Skip to content

jax-compile-stall epic: shipped phase 1, closed the rest as partial - #278

Merged
Jammy2211 merged 12 commits into
mainfrom
claude/jax-vmap-jit-stall-swz2tc
Aug 23, 2026
Merged

jax-compile-stall epic: shipped phase 1, closed the rest as partial#278
Jammy2211 merged 12 commits into
mainfrom
claude/jax-vmap-jit-stall-swz2tc

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

The Mind side of the jax-compile-stall epic: 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

PR
PyAutoFit#1517 heartbeat + faulthandler watchdog + compile/execute timing split
PyAutoFit#1518 fixes #1517 — the dump threshold now sits below the runner's cap
PyAutoHeart#161 runner / runner-args / script-timeout inputs on the reusable smoke workflow
autolens_workspace_test#272, autogalaxy_workspace_test#110 the retime.py harness + workflow_dispatch caller

Issues PyAutoFit#1516 and autolens_workspace_test#271 both closed.

What this PR adds

  • epics.md — the jax-compile-stall entry, status CLOSED AS PARTIAL
  • draft/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 onto main here
  • Three phase prompts under draft/bug/ci/ — phase 1 folded into its record; phases 2 and 3 retained
  • complete/2026/08/jax-compile-stall-evidence.md — phase 1's record
  • complete/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 undone
  • draft/research/ci/smoke_timing_and_profiling.md — the follow-up task, carrying the four measured verdicts forward so nothing is re-measured
  • Supersede banner on draft/bug/autolens_workspace_test/multi_dataset_jax_likelihood_xla_stall.md — the same defect, filed 2026-08-22, never issued

The findings worth keeping

  • A SLOW marker is not evidence of slowness. Every 2026-07-14 marker reads "flakes at the 1800s cap" and records no timing at all. The first one measured — 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.
  • The stall is a >100× bimodality inside one step. A healthy compile of rectangular_mge.py is 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.
  • The compile-cache hypothesis was never tested — the obvious next experiment.
  • A watchdog whose threshold equals the cap never fires. #1517's 300s default against a 300s cap produced 20 stalled runs with heartbeats and zero stacks.

Verification

lifecycle.py check: OK. lifecycle.py index --check: OK. dashboard.md / dashboard.html regenerated. active.md is 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 faulthandler stack. The record names them as the first thing to read when this is resumed.


Generated by Claude Code

James Nightingale and others added 12 commits August 23, 2026 19:49
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
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
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
@Jammy2211
Jammy2211 merged commit 2da60e8 into main Aug 23, 2026
1 of 2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant