Make a stalled JAX compile report itself

[Jammy2211]

Closes #1516. Phase 1 of 3 in the jax-compile-stall epic (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_test delaunay (#245), autogalaxy_workspace_test multi_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 is

autofit.non_linear.jax_compile - INFO - JAX jit compiling vectorized (vmap)
    likelihood function, could take seconds or minutes...

and then nothing at all until a cap kills it.

log_on_first_compile also does two very different things under that one log line: func(*args, **kwargs) (tracing, lowering, XLA compilation) and then jax.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._grad and the batched latent computation in analysis/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.

	Env var	Default
Heartbeat — still compiling <desc>, Ns elapsed on an interval	PYAUTOFIT_JAX_COMPILE_HEARTBEAT_SECS	30, 0 disables
faulthandler watchdog — the process dumps its own traceback if the compile overruns	PYAUTOFIT_JAX_COMPILE_DUMP_SECS	300 when CI is set, 0 otherwise
Compile wait and execution wait timed and logged separately	—	always on

The existing complete in {n} seconds summary 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 the CI branch, 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):

19:56:50 ... - JAX jit compiling vectorized (vmap) likelihood function, could take seconds or minutes...
19:56:52 ... - JAX jit still compiling vectorized (vmap) likelihood function, 2s elapsed...
19:56:54 ... - JAX jit still compiling vectorized (vmap) likelihood function, 4s elapsed...
Timeout (0:00:05)!
Thread 0x00007f2460870080 (most recent call first):
  File ".../stall_proof.py", line 7 in <lambda>
  File ".../autofit/non_linear/jax_compile.py", line 221 in wrapper
...
[exit 137 — SIGKILL]

That is precisely the evidence three quarantines did not have. And on a real successful jax.vmap(jax.jit(...)) compile:

JAX jit compilation of vectorized (vmap) likelihood function: traced, lowered
  and compiled in 0.1 seconds, result materialized in 0.0 seconds.
JAX jit compilation of vectorized (vmap) likelihood function complete in 0.1 seconds.

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

1. Fitness._vmap builds jax.vmap(jax.jit(self.call)) — vmap of jit, the inverted ordering — while analysis/latent.py builds jax.jit(jax.vmap(...)), the conventional one. The path that stalls is exactly the vmap path; 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.
2. Both NEEDS_FIX stalls post-date the persistent-compilation-cache default (PyAutoConf#128, merged 2026-07-17); the eight SLOW-marked entries predate it. Cache-lock contention is a live hypothesis alongside the version-interaction one.

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-heart is 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