Make a stalled JAX compile report itself - #1517
Merged
Merged
Conversation
The same intermittent XLA compile stall has been quarantined three times across autolens_workspace_test and autogalaxy_workspace_test without ever being diagnosed, because a stalled run leaves no evidence: the last line it emits is log_on_first_compile's "JAX jit compiling ..." and then silence until a cap kills it. That first call contains two different waits under one log line -- tracing, lowering and XLA compilation, then jax.block_until_ready execution -- so the captured tail could not say which half was stuck, or whether the process was alive at all. Three additions, all inside the compile wrapper so every call site picks them up (Fitness._vmap/_jit/_grad and the batched latent computation), and no workspace script or CI runner is touched: - a heartbeat logging "still compiling <desc>, Ns elapsed" on an interval (PYAUTOFIT_JAX_COMPILE_HEARTBEAT_SECS, default 30, 0 disables); - a faulthandler watchdog dumping the process's own traceback if the compile overruns (PYAUTOFIT_JAX_COMPILE_DUMP_SECS, defaulting to 300 under CI and off elsewhere, so the next CI stall self-diagnoses with nothing else edited); - separate timings for the compile wait and the execution wait, alongside the existing summary line, which is unchanged. Diagnostics never break a fit: an unstartable heartbeat thread, an unarmable dump and a malformed interval all fall back and continue. Closes #1516 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015qk7hoavMnFyPtW4toYn8K
This was referenced Aug 27, 2026
Jammy2211
pushed a commit
to trexfr-ops/PyAutoFit
that referenced
this pull request
Aug 27, 2026
The watchdog shipped in PyAutoLabs#1517 never actually dumped. Its CI default was a flat 300s and the workspace smoke cap is also 300s, so the dump timer and the runner's kill timer raced and the kill won -- a SIGKILLed process writes no traceback. Measured: 20 stalled runs across four CI legs produced heartbeats from every one of them and not a single stack (autolens_workspace_test#271). Derive the CI default from BUILD_SCRIPT_TIMEOUT -- the per-script cap the workspace runners and PyAutoHands both enforce -- at 80% of it, so the traceback has time to reach stderr before the kill lands. 300s cap dumps at 240s; 1800s cap dumps at 1440s. With no usable cap advertised, fall back to a flat 240s rather than computing a fraction of zero, which would silently disable the dump. An explicit PYAUTOFIT_JAX_COMPILE_DUMP_SECS still wins. Four new tests, one of them pinning the invariant directly: the threshold is strictly below the cap for every cap. Verified end-to-end: with a 10s cap the dump fires at 8s and the traceback is captured, where before nothing was written at all. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015qk7hoavMnFyPtW4toYn8K
Jammy2211
pushed a commit
to trexfr-ops/PyAutoFit
that referenced
this pull request
Aug 27, 2026
The first call to a jax.jit/vmap/grad wrapper contains two waits -- trace/lower/compile, then jax.block_until_ready execution -- and the instrumentation could not tell a stalled run which one it was in. Two defects, both of which made a stall report the wrong diagnosis: 1. The heartbeat hardcoded "still compiling". A stalled CI run therefore logged `JAX jit still compiling ... 1770s elapsed` while its captured faulthandler stack sat in jax.block_until_ready. That is positive evidence for the wrong cause, and five quarantine markers across the two test workspaces were written against it, all calling this an "intermittent XLA compile stall". It is not one: compilation completes, in ~16s, and execution never returns. 2. The compile/execute split line was emitted only after BOTH halves finished. A stalled run never finishes the second, so it reported neither number and the split characterised only the healthy case -- twenty stalled runs could not say that compilation had in fact completed. Now: the compile half is logged the instant func() returns, the heartbeat moves to naming the materialize half and repeats the compile time on every beat, and the materialize half is logged on its own when it lands. A run SIGKILLed at any point after tracing has already said on stderr which half it was in and how long the other took. Also extracts _block_until_ready() so the execution half is substitutable in tests without JAX installed, and so the broad `except Exception` guards only the optional jax import rather than the wait itself. No call site changes: all four (Fitness._vmap/_jit/_grad, analysis/latent) inherit this through log_on_first_compile. For PyAutoFit#1528 (jax-compile-stall phase 3); follows PyAutoLabs#1517/PyAutoLabs#1518. Co-Authored-By: Claude <noreply@anthropic.com>
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
1 participant
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.
Closes #1516. Phase 1 of 3 in the
jax-compile-stallepic (ledger:PyAutoMind/draft/bug/ci/jax_vmap_jit_compile_stall.md).Why
The same intermittent XLA compile stall has been quarantined three times —
autolens_workspace_testdelaunay (#245),autogalaxy_workspace_testmulti_dataset/.../rectangular.py(2026-08-01),imaging/.../mge_group.py(2026-08-23) — and diagnosed zero times. The reason is that a stalled run leaves no evidence. The last line it emits isand then nothing at all until a cap kills it.
log_on_first_compilealso does two very different things under that one log line:func(*args, **kwargs)(tracing, lowering, XLA compilation) and thenjax.block_until_ready(result)(execution, since JAX dispatches asynchronously). So even the captured tail could not say which half was stuck — or whether the process was alive.What changed
All inside the compile wrapper, so every call site picks it up automatically:
Fitness._vmap,Fitness._jit,Fitness._gradand the batched latent computation inanalysis/latent.py. No workspace script and no CI runner is touched — the workspace scripts are user-facing documentation, and per-script workarounds are the quarantine pattern this epic exists to stop.still compiling <desc>, Ns elapsedon an intervalPYAUTOFIT_JAX_COMPILE_HEARTBEAT_SECS30,0disablesfaulthandlerwatchdog — the process dumps its own traceback if the compile overrunsPYAUTOFIT_JAX_COMPILE_DUMP_SECS300whenCIis set,0otherwiseThe existing
complete in {n} secondssummary line is unchanged.The CI-conditional default is what makes the next stall self-diagnosing with nothing else edited. The alternative — threading an env var through each workspace's
config/build/env_vars_*.yaml— is more repos touched for the same effect.Diagnostics never break a fit: an unstartable heartbeat thread, an unarmable dump (a capturing harness can leave stderr without a real fd) and a malformed interval all fall back and continue.
Verification
Full suite: 2008 passed, 34 skipped (Python 3.12).
test_autofit/non_linear/test_jax_compile.py: 16 passed, 12 of them new — heartbeat fires / stops / disables, watchdog armed and cancelled including on a raising compile, nothing armed when disabled, env defaults including theCIbranch, malformed and negative intervals.End-to-end, a simulated stall in a fresh process killed at 12s exactly as a CI cap would kill it (
CI=true, heartbeat 2s, dump 5s):That is precisely the evidence three quarantines did not have. And on a real successful
jax.vmap(jax.jit(...))compile:Known limitation, stated up front
A Python traceback taken during XLA compilation parks at the pybind boundary and will not show XLA internals. It still separates in compile from in execution from blocked on a Python-level lock — for instance the persistent compilation cache (
JAX_COMPILATION_CACHE_DIR, on by default since PyAutoConf#128). That three-way split is the fork phase 3 needs, so it does not undermine the phase.Two findings recorded for phase 3, deliberately not acted on here
Fitness._vmapbuildsjax.vmap(jax.jit(self.call))—vmapofjit, the inverted ordering — whileanalysis/latent.pybuildsjax.jit(jax.vmap(...)), the conventional one. The path that stalls is exactly thevmappath; the_jit-only scripts in the same directories do not stall. Unproven as causal, but it is a one-line A/B and the first thing phase 3 should try.Neither belongs in a diagnostics PR — changing the transform ordering while trying to observe the stall would destroy the thing being observed.
Heart gate
pyauto-heartis not reachable from this session (web session, PyAutoHeart not checked out), so the readiness verdict was not consulted. Treat this PR as un-gated by Heart and run the vitals check before merge.Generated by Claude Code