|
| 1 | +# JAX vmap materialisation hang — root-caused to XLA CPU's Eigen thread pool, all seven quarantines cleared (jax-compile-stall phase 3) |
| 2 | + |
| 3 | +- **Issue:** PyAutoFit#1528 (closed) · autolens_workspace_test#245 (closed) · **PRs:** PyAutoFit#1529 (`4130b22`), PyAutoHands#269 (`ad64e12`), autolens_workspace_test#281 (`69ee97c`), autogalaxy_workspace_test#114 (`adf5ffe`) — all merged 2026-08-27 |
| 4 | +- **Repos:** PyAutoFit (`non_linear/jax_compile.py`), PyAutoHands (`autohands/build_util.py`, `autohands/env_config.py`), both `*_workspace_test` (`config/build/profile_{smoke,release}.yaml`, `no_run.yaml`, `smoke_tests.txt`) |
| 5 | +- **Epic:** `jax-compile-stall`, phase 3 of 3 — **CLOSES the epic.** Phase 1 shipped the instrumentation, phase 2 measured and stopped short, phase 3 found the cause and restored the coverage. |
| 6 | +- **Task slug:** `jax-stall-block-until-ready` (the `active.md` entry and the branch name; this record is filed under the corrected name, since the task slug still carries the "stall" misnomer) |
| 7 | +- **Status: SHIPPED, with the root cause deliberately unfinished** — see "What is still not known". |
| 8 | + |
| 9 | +## The headline: the epic's name was wrong for a month |
| 10 | + |
| 11 | +Every marker, prompt and issue called this an *"intermittent XLA compile stall"*. **It is not a compile stall.** Compilation completes in ~12-18s; what never returns is **materialising the vmap result**. Captured at both doors, in the same run: |
| 12 | + |
| 13 | +``` |
| 14 | +jax/_src/api.py 2764 try_to_block jax/_src/array.py 642 _value |
| 15 | +jax/_src/api.py 2781 block_until_ready jax/_src/profiler.py 420 wrapper |
| 16 | +jax_compile.py 249 _block_until_ready jax/_src/array.py 391 __str__ |
| 17 | +jax_compile.py 339 wrapper mge_group.py 319 <module> |
| 18 | +mge_group.py 313 <module> |
| 19 | +``` |
| 20 | + |
| 21 | +One failure, two doors. The line-313 door is the first `_vmap` call; the line-319 door is `print()` of the second, which short-circuits `log_on_first_compile` and so logs nothing at all. |
| 22 | + |
| 23 | +## The trigger: XLA CPU's multithreaded Eigen thread pool |
| 24 | + |
| 25 | +`XLA_FLAGS=--xla_cpu_multi_thread_eigen=false`, in `profile_smoke.yaml` **and** `profile_release.yaml` of both test workspaces so gate, mega-run and release-integrate cannot disagree. ABAB on `imaging/jax_likelihood/mge_group.py`, 3 repeats x 2 legs, 300s cap: |
| 26 | + |
| 27 | +| Arm | Result | Run | |
| 28 | +|---|---|---| |
| 29 | +| default | 4/4 HANG | 33076408637 | |
| 30 | +| compilation cache disabled | 6/6 HANG | 33078033016 | |
| 31 | +| flag + `host_platform_device_count=1` + `OMP_NUM_THREADS=1` | 6/6 PASS, 63.2s | 33080638510 | |
| 32 | +| **flag alone** | **6/6 PASS, 62.3s** | 33082896277 | |
| 33 | +| default again | 4/6 HANG, verdict `STALL` | 33085024710 | |
| 34 | + |
| 35 | +**12 passes / 0 hangs with the flag; 2 passes / 14 hangs without. Fisher exact p ~ 3e-6.** The other two variables did nothing — 62.3s alone vs 63.2s together. Cost ~15%; reaches only JAX/XLA scripts since nothing else reads `XLA_FLAGS`. |
| 36 | + |
| 37 | +## Coverage restored — the point of the campaign |
| 38 | + |
| 39 | +Family re-times 33087267785 (al) + 33087271848 (ag): **42/42 completions, `NEITHER` on every entry and every leg**, slowest 20% of cap. |
| 40 | + |
| 41 | +| Repo | Entry | 3.12 | 3.13 | was | |
| 42 | +|---|---|---|---|---| |
| 43 | +| al | `imaging/jax_likelihood/mge_group.py` | 61.2s | 61.3s | NEEDS_FIX 08-24, **16/16 lifetime cap hits, zero completions** | |
| 44 | +| al | `multi_dataset/jax_likelihood/delaunay.py` | 24.4s | 23.2s | NEEDS_FIX 08-01 (#245) | |
| 45 | +| al | `multi_dataset/jax_likelihood/mge.py` | 16.0s | 14.5s | smoke_tests.txt 08-22 | |
| 46 | +| al | `multi_dataset/jax_likelihood/shared_preloads.py` | 43.9s | 43.8s | smoke_tests.txt 07-22 | |
| 47 | +| ag | `multi_dataset/jax_likelihood/rectangular.py` | 20.0s | 21.9s | NEEDS_FIX 08-01 | |
| 48 | +| ag | `imaging/jax_likelihood/mge_group.py` | 32.8s | 32.6s | NEEDS_FIX 08-23 | |
| 49 | +| ag | `imaging/jax_likelihood/rectangular_mge.py` | 22.9s | 24.1s | NEEDS_FIX 08-23 | |
| 50 | + |
| 51 | +Both exclusion-list disagreements resolved, and they ran opposite ways: al excluded `shared_preloads.py` from `smoke_tests.txt` as "300s+, the heaviest entry" while `no_run.yaml` admitted it to the weekly sweep (where it burned a cap on 08-25) — it measures **43.9s**, so that 300s+ was the stall, not slowness. ag had `rectangular.py` NEEDS_FIX in `no_run.yaml` but still live in `smoke_tests.txt`. |
| 52 | + |
| 53 | +## The finding that unlocked everything: the evidence was being thrown away |
| 54 | + |
| 55 | +**Nothing set `PYTHONUNBUFFERED`**, and the runners capture through `subprocess.PIPE`. A pipe is not a tty, so a child's `print()` is block-buffered and flushed only at exit; `logging` goes to stderr and arrives immediately. A script SIGKILLed at its cap loses its **entire stdout buffer**. |
| 56 | + |
| 57 | +So every *"and then silence"* tail in this campaign — the 08-01, 08-23, 08-24 and 08-25 markers — was silence of **stderr alone**. Five scripts were quarantined on evidence truncated before anyone read it. Proven from the *passing* leg, not inferred: its logging lines interleave by timestamp and then all eight prints arrive in one block at exit. |
| 58 | + |
| 59 | +## Traps recorded |
| 60 | + |
| 61 | +- **A phase-blind heartbeat is worse than none.** It said `"still compiling"` in both halves, so a stalled run produced *positive evidence for the wrong cause*. Five markers were written against it. Instrumentation that can be confidently wrong is a liability, not a diagnostic. |
| 62 | +- **A summary line emitted after both halves finish describes only the runs that did not fail.** The compile/execute split never appeared on a stalled run. |
| 63 | +- **`faulthandler` under SIGABRT works while a C extension holds the GIL** — that is why it is the right tool, and why SIGKILL alone had produced nothing for a month. It shows Python frames only, so XLA's own threads stay opaque. |
| 64 | +- **A watchdog scoped to one function cannot see a hang outside it.** PyAutoFit's watchdog is disarmed the moment the first compile returns; the general fix belongs in the runner (`kill_group`), not the library. |
| 65 | +- **`env_config.apply_profile` does `env[key] = str(value)`** — an unquoted empty YAML value becomes the literal string `"None"`, i.e. a truthy cache directory named `None`. Quote empty values. |
| 66 | +- **Every pre-existing test in `test_script_timeout.py` passed `flush=True`**, so the suite proved output survives in the one case real scripts never hit. A test suite can be uniformly blind in exactly the dimension that matters. |
| 67 | +- **The tenant firewall is real and will catch you.** Citing `PyAutoLabs/PyAutoFit#NNNN` in a `.py` under an organ fails CI. Fix by removing the instance fact, never by growing `FIREWALL_ALLOWLIST` — each entry is another file an adopting fork must rewrite. Organ names are fine (framework identity). |
| 68 | +- **Sufficiency is not necessity.** The hang rate wanders: this script passed 2/2 at 02:18 and hung 2/2 at 13:33 on the identical commit. Twelve passes with a flag prove nothing without re-measuring the control — hence ABAB, not before/after. |
| 69 | + |
| 70 | +## Hypotheses refuted (so nobody re-runs them) |
| 71 | + |
| 72 | +- **Persistent compilation cache** — 6/6 hangs with it disabled. Verified applied: no "cache has been enabled" line, compiles dropped 15.7s → ~12s. |
| 73 | +- **jax/jaxlib version** — the 2026-08-24 run, hanging 3/3 on *both* legs, had byte-identical jax `0.11.1` / jaxlib `0.11.1` / ml_dtypes `0.6.0` / numpy `2.5.2` / scipy `1.17.1`. No bisect needed. |
| 74 | +- **3.12 vs 3.13** — 3.13 passed 2/2 then hung 2/2 on the same commit hours apart. Not causal. |
| 75 | +- **`vmap(jit)` ordering** — untested at power here; phase 2's p=0.070 stands, and the flag result makes it moot. |
| 76 | + |
| 77 | +## What is still not known |
| 78 | + |
| 79 | +**This is a workaround, not a root-cause fix.** We know *where* it hangs and *what* avoids it. We do not know **why** XLA's CPU thread pool wedges on these graphs — `faulthandler` reports Python frames only, so the wedged worker threads are still opaque. Getting further needs a native stack (gdb / py-spy) and a minimal jaxpr reproducer outside the workspace, then an upstream JAX/XLA report. Filed as a follow-up research prompt. |
| 80 | + |
| 81 | +Both `no_run.yaml` block comments record this, and say that removing the flag brings the entries back. |
| 82 | + |
| 83 | +## Heart |
| 84 | + |
| 85 | +**Not consulted** — `pyauto-heart` unreachable from the `web-github` environment, as on phases 1 and 2. Every merge was on the human's explicit `/prm`. |
| 86 | + |
| 87 | +## Provenance |
| 88 | + |
| 89 | +Planned, implemented, measured and shipped by one `web-github` session on 2026-08-27 (`claude/jax-compile-stall-root-cause-d33hi0`), against direct clones rather than worktrees. Eleven CI dispatches of the phase-2 `retime.yml` harness did the measuring. |
| 90 | + |
| 91 | +## Original prompt |
| 92 | + |
| 93 | +# Phase 3: root-cause the XLA vmap compile stall and clear every NEEDS_FIX it caused |
| 94 | + |
| 95 | +Type: bug |
| 96 | +Target: ci |
| 97 | +Repos: |
| 98 | +- @PyAutoFit |
| 99 | +- @autogalaxy_workspace_test |
| 100 | +- @autolens_workspace_test |
| 101 | +Difficulty: large |
| 102 | +Autonomy: supervised |
| 103 | +Priority: high |
| 104 | +Status: formalised |
| 105 | +Epic: jax-compile-stall |
| 106 | +Phase: 3 |
| 107 | +Campaign: bug/ci/jax_vmap_jit_compile_stall.md (Phase 3 — the fix; blocked on phases 1 and 2) |
| 108 | +Filed: 2026-08-23 |
| 109 | +Issued: 2026-08-27 |
| 110 | + |
| 111 | +## Blocked on phase 1 |
| 112 | + |
| 113 | +Do not start this before phase 1's watchdog has shipped and a CI stall has |
| 114 | +actually dumped a traceback. Diagnosing a hang that leaves no evidence is what |
| 115 | +produced three quarantines and no root cause; repeating it without the |
| 116 | +instrumentation would be a fourth. |
| 117 | + |
| 118 | +## Reproduce deliberately |
| 119 | + |
| 120 | +Loop `imaging/jax_likelihood/mge_group.py` under its declared CI env profile |
| 121 | +until it hangs, rather than waiting for CI to hit it. Attach `py-spy dump` to |
| 122 | +the hung process as well as reading phase 1's own `faulthandler` output — the |
| 123 | +two see different things, and `py-spy` can read native frames the in-process |
| 124 | +dump cannot. |
| 125 | + |
| 126 | +## Hypotheses, in the order they are cheapest to test |
| 127 | + |
| 128 | +1. **`vmap` of `jit`, the inverted ordering.** `Fitness._vmap` |
| 129 | + (`autofit/non_linear/fitness.py`) builds `jax.vmap(jax.jit(self.call))`; |
| 130 | + `autofit/non_linear/analysis/latent.py` builds |
| 131 | + `jax.jit(jax.vmap(compute_latent_for_model))`, the conventional order. The |
| 132 | + stalling path is exactly the `vmap` path, and the `_jit`-only scripts in the |
| 133 | + same directories do not stall. One-line A/B — try it first. |
| 134 | +2. **Persistent compilation cache contention.** `JAX_COMPILATION_CACHE_DIR` has |
| 135 | + defaulted on since PyAutoConf#128 (merged 2026-07-17). Both NEEDS_FIX stalls |
| 136 | + post-date it; the eight SLOW entries predate it. A/B with the cache dir set |
| 137 | + to empty (which disables it) and see whether the stall probability moves. |
| 138 | +3. **JAX/XLA version interaction.** This repo has form: `delaunay_mge.py` is |
| 139 | + disabled outright because `jax 0.7` removed |
| 140 | + `jax.interpreters.xla.pytype_aval_mappings`, and the smoke installer once |
| 141 | + clobbered a working `tfp-nightly`. Pin-bisect jax/jaxlib across a run set. |
| 142 | +4. **Runner CPU contention.** `complete/2026/07/jax-compile-time-research.md` |
| 143 | + records that XLA compiles on **host** CPUs and that compile timing is |
| 144 | + load-sensitive by up to 7×, which is why a hosted runner is the place this |
| 145 | + reproduces and a workstation is not. |
| 146 | +5. **Graph size in the vmap trace.** The affected set has a shape: plain `mge.py` |
| 147 | + passes in 9.4s, while the *composite* variants — group, rectangular-MGE, |
| 148 | + delaunay-MGE — stall or are already out. Complexity-driven compile blowup was |
| 149 | + argued against by the autolens_profiling#71 research ("compile cost is |
| 150 | + op-pattern-driven, not complexity-driven"), so treat this as the hypothesis |
| 151 | + of last resort, not the first. |
| 152 | + |
| 153 | +## Then restore the coverage |
| 154 | + |
| 155 | +The point of the campaign. Quarantining removes exactly the heaviest JAX paths, |
| 156 | +which are the ones most worth testing. |
| 157 | + |
| 158 | +1. Clear the NEEDS_FIX markers this campaign inherits — including the |
| 159 | + 2026-08-01 `multi_dataset/jax_likelihood/rectangular.py` one and the |
| 160 | + `autolens_workspace_test` `delaunay.py` entry citing #245. |
| 161 | +2. Re-enable `multi_dataset/jax_likelihood/mge.py` and `shared_preloads.py` in |
| 162 | + `autolens_workspace_test`'s `smoke_tests.txt` (folded in from the superseded |
| 163 | + 2026-08-22 filing). |
| 164 | +3. Anything that stays out after the fix stays out with a **recorded deliberate |
| 165 | + reason**, not as an accumulated one-off. |
| 166 | + |
| 167 | +## Acceptance |
| 168 | + |
| 169 | +- A stated root cause, or an explicit recorded decision that it is an |
| 170 | + infrastructure limit to be worked around rather than fixed. Not another |
| 171 | + quarantine. |
| 172 | +- Every entry marked NEEDS_FIX for this signature restored to its suite, or |
| 173 | + re-marked with the real reason phase 2 established. |
| 174 | +- The `multi_dataset/jax_likelihood/` family back under CI coverage in both test |
| 175 | + workspaces, or its absence recorded as a deliberate choice. |
0 commit comments