fix: make a stalled JAX vmap compile report itself

[Jammy2211]

Overview

Intermittent XLA compile stalls in the JAX vmap likelihood path have been quarantined three separate times across autolens_workspace_test and autogalaxy_workspace_test without ever being diagnosed. The reason is that a stalled run leaves no evidence: the last line it emits is JAX jit compiling vectorized (vmap) likelihood function... from autofit.non_linear.jax_compile, and then silence until the cap kills it.

This is phase 1 of the jax-compile-stall epic and it is PyAutoFit-only: give the compile wrapper a heartbeat, a faulthandler watchdog, and a compile-vs-execute timing split, so the next CI stall diagnoses itself. Phases 2 (SLOW-vs-stall audit) and 3 (root cause + un-quarantine) are both blocked on the evidence this phase creates.

Surfaced by autogalaxy_workspace_test#109, where a newly-added 300s per-script cap turned what had been four 6-hour silent Actions cancellations into one TIMEOUT (300s) imaging/jax_likelihood/mge_group.py with the compiling-step tail attached.

Plan

• Give log_on_first_compile a heartbeat so a long compile keeps reporting liveness with elapsed seconds, instead of going silent.
• Arm a faulthandler watchdog around the first call so a compile that overruns a threshold dumps its own traceback to stderr before anything kills it.
• Time the trace/lower/compile half separately from the jax.block_until_ready execution half, so the log names which of the two is stuck.
• Default the watchdog on under CI (via the CI environment variable) and off locally, both overridable — so the next CI stall self-diagnoses with no workspace edit and no runner edit.
• Keep it entirely library-side. 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.

Detailed implementation plan

Work Classification

Library — PyAutoFit source only. No workspace follow-on for this phase.

Affected Repositories

• PyAutoFit (primary)

Branch Survey

Repository	Current Branch	Dirty?
PyAutoFit	main	clean (fresh clone)

worktree_check_conflict jax-compile-stall-evidence PyAutoFit → exit 0, no conflict. No existing PyAutoFit branch touches this area.

Suggested branch: feature/jax-compile-stall-evidence

Worktree root: ~/Code/PyAutoLabs-wt/jax-compile-stall-evidence/ (created by /start_library in a local-dev session; this task was started in a web-github session against a direct clone).

The defect in the current instrumentation

log_on_first_compile(func, description) does two very different things under one log line on the first call:

1. result = func(*args, **kwargs) — tracing, lowering, XLA compilation;
2. jax.block_until_ready(result) — execution, since JAX dispatches asynchronously.

then logs one complete in {n} seconds summary. A hang in either half is externally identical, and a merely-slow compile is indistinguishable from a stopped one. Nothing reports liveness in between.

Implementation Steps

1. autofit/non_linear/jax_compile.py — heartbeat: while the first call is in flight, log still compiling {description}, {n}s elapsed on an interval. Daemon thread, stopped in the existing finally so it can never hold the process open. Interval from PYAUTOFIT_JAX_COMPILE_HEARTBEAT_SECS, default 30, 0 disables.
2. autofit/non_linear/jax_compile.py — watchdog: faulthandler.dump_traceback_later(secs, repeat=True, exit=False) before the first call, cancel_dump_traceback_later() in the finally. Threshold from PYAUTOFIT_JAX_COMPILE_DUMP_SECS, defaulting to 300 when the CI env var is set and 0 otherwise.
3. autofit/non_linear/jax_compile.py — timing split: time func(...) and jax.block_until_ready(result) separately and log both. Keep the existing single complete in {n} seconds summary line unchanged so nothing reading it breaks.
4. test_autofit/non_linear/test_jax_compile.py — extend with cases that need no JAX import: heartbeat fires for a slow fake callable; watchdog armed and cancelled (monkeypatched faulthandler); env-var defaults including the CI branch; threads are daemon and joined.

All four call sites pick this up automatically: Fitness._vmap, Fitness._jit, Fitness._grad in autofit/non_linear/fitness.py, and the batched latent computation in autofit/non_linear/analysis/latent.py.

Key Files

• autofit/non_linear/jax_compile.py — log_on_first_compile, the whole change.
• test_autofit/non_linear/test_jax_compile.py — existing test module, extended.
• autofit/non_linear/fitness.py — call sites (_vmap, _jit, _grad); read-only for this phase.
• autofit/non_linear/analysis/latent.py — fourth call site; read-only for this phase.

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 exactly the fork phase 3 needs, so the limitation does not undermine the phase.

Two findings recorded for phase 3 (not acted on here)

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

Testing

pytest test_autofit/non_linear/test_jax_compile.py, then the PyAutoFit suite.

Epic

Phase 1 of 3 in the jax-compile-stall epic. Ledger: PyAutoMind/draft/bug/ci/jax_vmap_jit_compile_stall.md.

• Phase 1 (this issue) — evidence: heartbeat + faulthandler + compile/execute split. PyAutoFit.
• Phase 2 — are the eight SLOW-marked jax_likelihood/jax_grad entries genuinely slow, or is this stall mislabelled? A slow script has a tight timing distribution; a stalling one is bimodal.
• Phase 3 — root cause and fix, then clear every NEEDS_FIX marker for this signature and restore the quarantined coverage.

Original Prompt

Click to expand starting prompt

Phase 1: make a stalled JAX compile report itself (heartbeat + faulthandler + compile/execute split)

Type: bug
Target: ci
Repos:

• @PyAutoFit
  Difficulty: small
  Autonomy: supervised
  Priority: high
  Status: formalised
  Epic: jax-compile-stall
  Phase: 1
  Campaign: bug/ci/jax_vmap_jit_compile_stall.md (Phase 1 — the enabler; phases 2 and 3 are blocked on this)
  Filed: 2026-08-23

Why this is phase 1

The stall's whole cost is that it produces no evidence. The last line any
killed run emits is

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

and then silence until the cap kills it. Three separate quarantines
(autolens_workspace_test delaunay #245, autogalaxy_workspace_test
multi_dataset/.../rectangular.py 2026-08-01, imaging/.../mge_group.py
2026-08-23) produced no diagnosis between them, because there was nothing to
diagnose from. Phases 2 and 3 of this campaign both consume evidence this
phase creates.

What is wrong with the current instrumentation

log_on_first_compile(func, description) in
autofit/non_linear/jax_compile.py wraps the jax.jit / jax.vmap /
jax.grad callables so the "this is compiling" line lands where the user
actually waits — on the first call. Inside that first call it does two very
different things under one log line:

1. result = func(*args, **kwargs) — tracing, lowering and XLA compilation;
2. jax.block_until_ready(result) — execution, because JAX dispatches
   asynchronously.

Then it logs one complete in {n} seconds summary. So a hang anywhere in
either half looks identical from the outside, and a compile that is merely
slow looks identical to one that has stopped. Nothing reports liveness in
between.

Task

All of this is library-side in @PyAutoFit. Do not touch the workspace
scripts — they are user-facing documentation, and a per-script workaround is
the quarantine pattern this campaign exists to stop.

1. Heartbeat. While the first call is in flight, log
   still compiling {description}, {n}s elapsed on an interval. Daemon thread,
   stopped in the existing finally so it can never hold the process open.
   Interval from PYAUTOFIT_JAX_COMPILE_HEARTBEAT_SECS, default 30, 0
   disables.
2. Watchdog. Arm faulthandler.dump_traceback_later(secs, repeat=True, exit=False) before the first call and cancel_dump_traceback_later() in the
   finally, so a compile that overruns dumps its own traceback to stderr
   before anything kills it. Threshold from PYAUTOFIT_JAX_COMPILE_DUMP_SECS.
3. Default it on under CI. Default the threshold to 300 when the CI
   environment variable is set and 0 (off) otherwise, both overridable. This
   is what makes the next CI stall self-diagnosing with no workspace edit and
   no runner edit — the alternative, wiring an env var into each workspace's
   config/build/env_vars_*.yaml, is a second repo touch for the same effect.
4. Split the timing. Time func(...) and jax.block_until_ready(result)
   separately and log both, so the record says which half is stuck. Keep the
   existing single complete in {n} seconds summary line unchanged.

Applies automatically to all four call sites: Fitness._vmap, Fitness._jit,
Fitness._grad (autofit/non_linear/fitness.py) and the batched latent
computation in autofit/non_linear/analysis/latent.py.

Known limitation, to be stated in the PR

A Python traceback taken during XLA compilation parks at the pybind boundary —
it 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 exactly the fork phase 3 needs, so the
limitation does not undermine the phase.

Acceptance

• A stalled first compile emits periodic liveness lines with elapsed time.
• A stalled first compile leaves a traceback behind in CI without any workspace
  or runner change.
• The log distinguishes the compile wait from the execution wait.
• Covered by tests in test_autofit/non_linear/test_jax_compile.py that need no
  JAX import: heartbeat fires, watchdog is armed and cancelled, env defaults
  including the CI branch.
• No workspace script and no CI runner is modified by this phase.