diff --git a/complete/2026/08/jax-vmap-materialisation-hang.md b/complete/2026/08/jax-vmap-materialisation-hang.md new file mode 100644 index 00000000..46e47d20 --- /dev/null +++ b/complete/2026/08/jax-vmap-materialisation-hang.md @@ -0,0 +1,175 @@ +# JAX vmap materialisation hang — root-caused to XLA CPU's Eigen thread pool, all seven quarantines cleared (jax-compile-stall phase 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 +- **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`) +- **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. +- **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) +- **Status: SHIPPED, with the root cause deliberately unfinished** — see "What is still not known". + +## The headline: the epic's name was wrong for a month + +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: + +``` +jax/_src/api.py 2764 try_to_block jax/_src/array.py 642 _value +jax/_src/api.py 2781 block_until_ready jax/_src/profiler.py 420 wrapper +jax_compile.py 249 _block_until_ready jax/_src/array.py 391 __str__ +jax_compile.py 339 wrapper mge_group.py 319 +mge_group.py 313 +``` + +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. + +## The trigger: XLA CPU's multithreaded Eigen thread pool + +`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: + +| Arm | Result | Run | +|---|---|---| +| default | 4/4 HANG | 33076408637 | +| compilation cache disabled | 6/6 HANG | 33078033016 | +| flag + `host_platform_device_count=1` + `OMP_NUM_THREADS=1` | 6/6 PASS, 63.2s | 33080638510 | +| **flag alone** | **6/6 PASS, 62.3s** | 33082896277 | +| default again | 4/6 HANG, verdict `STALL` | 33085024710 | + +**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`. + +## Coverage restored — the point of the campaign + +Family re-times 33087267785 (al) + 33087271848 (ag): **42/42 completions, `NEITHER` on every entry and every leg**, slowest 20% of cap. + +| Repo | Entry | 3.12 | 3.13 | was | +|---|---|---|---|---| +| al | `imaging/jax_likelihood/mge_group.py` | 61.2s | 61.3s | NEEDS_FIX 08-24, **16/16 lifetime cap hits, zero completions** | +| al | `multi_dataset/jax_likelihood/delaunay.py` | 24.4s | 23.2s | NEEDS_FIX 08-01 (#245) | +| al | `multi_dataset/jax_likelihood/mge.py` | 16.0s | 14.5s | smoke_tests.txt 08-22 | +| al | `multi_dataset/jax_likelihood/shared_preloads.py` | 43.9s | 43.8s | smoke_tests.txt 07-22 | +| ag | `multi_dataset/jax_likelihood/rectangular.py` | 20.0s | 21.9s | NEEDS_FIX 08-01 | +| ag | `imaging/jax_likelihood/mge_group.py` | 32.8s | 32.6s | NEEDS_FIX 08-23 | +| ag | `imaging/jax_likelihood/rectangular_mge.py` | 22.9s | 24.1s | NEEDS_FIX 08-23 | + +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`. + +## The finding that unlocked everything: the evidence was being thrown away + +**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**. + +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. + +## Traps recorded + +- **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. +- **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. +- **`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. +- **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. +- **`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. +- **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. +- **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). +- **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. + +## Hypotheses refuted (so nobody re-runs them) + +- **Persistent compilation cache** — 6/6 hangs with it disabled. Verified applied: no "cache has been enabled" line, compiles dropped 15.7s → ~12s. +- **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. +- **3.12 vs 3.13** — 3.13 passed 2/2 then hung 2/2 on the same commit hours apart. Not causal. +- **`vmap(jit)` ordering** — untested at power here; phase 2's p=0.070 stands, and the flag result makes it moot. + +## What is still not known + +**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. + +Both `no_run.yaml` block comments record this, and say that removing the flag brings the entries back. + +## Heart + +**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`. + +## Provenance + +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. + +## Original prompt + +# Phase 3: root-cause the XLA vmap compile stall and clear every NEEDS_FIX it caused + +Type: bug +Target: ci +Repos: +- @PyAutoFit +- @autogalaxy_workspace_test +- @autolens_workspace_test +Difficulty: large +Autonomy: supervised +Priority: high +Status: formalised +Epic: jax-compile-stall +Phase: 3 +Campaign: bug/ci/jax_vmap_jit_compile_stall.md (Phase 3 — the fix; blocked on phases 1 and 2) +Filed: 2026-08-23 +Issued: 2026-08-27 + +## Blocked on phase 1 + +Do not start this before phase 1's watchdog has shipped and a CI stall has +actually dumped a traceback. Diagnosing a hang that leaves no evidence is what +produced three quarantines and no root cause; repeating it without the +instrumentation would be a fourth. + +## Reproduce deliberately + +Loop `imaging/jax_likelihood/mge_group.py` under its declared CI env profile +until it hangs, rather than waiting for CI to hit it. Attach `py-spy dump` to +the hung process as well as reading phase 1's own `faulthandler` output — the +two see different things, and `py-spy` can read native frames the in-process +dump cannot. + +## Hypotheses, in the order they are cheapest to test + +1. **`vmap` of `jit`, the inverted ordering.** `Fitness._vmap` + (`autofit/non_linear/fitness.py`) builds `jax.vmap(jax.jit(self.call))`; + `autofit/non_linear/analysis/latent.py` builds + `jax.jit(jax.vmap(compute_latent_for_model))`, the conventional order. The + stalling path is exactly the `vmap` path, and the `_jit`-only scripts in the + same directories do not stall. One-line A/B — try it first. +2. **Persistent compilation cache contention.** `JAX_COMPILATION_CACHE_DIR` has + defaulted on since PyAutoConf#128 (merged 2026-07-17). Both NEEDS_FIX stalls + post-date it; the eight SLOW entries predate it. A/B with the cache dir set + to empty (which disables it) and see whether the stall probability moves. +3. **JAX/XLA version interaction.** This repo has form: `delaunay_mge.py` is + disabled outright because `jax 0.7` removed + `jax.interpreters.xla.pytype_aval_mappings`, and the smoke installer once + clobbered a working `tfp-nightly`. Pin-bisect jax/jaxlib across a run set. +4. **Runner CPU contention.** `complete/2026/07/jax-compile-time-research.md` + records that XLA compiles on **host** CPUs and that compile timing is + load-sensitive by up to 7×, which is why a hosted runner is the place this + reproduces and a workstation is not. +5. **Graph size in the vmap trace.** The affected set has a shape: plain `mge.py` + passes in 9.4s, while the *composite* variants — group, rectangular-MGE, + delaunay-MGE — stall or are already out. Complexity-driven compile blowup was + argued against by the autolens_profiling#71 research ("compile cost is + op-pattern-driven, not complexity-driven"), so treat this as the hypothesis + of last resort, not the first. + +## Then restore the coverage + +The point of the campaign. Quarantining removes exactly the heaviest JAX paths, +which are the ones most worth testing. + +1. Clear the NEEDS_FIX markers this campaign inherits — including the + 2026-08-01 `multi_dataset/jax_likelihood/rectangular.py` one and the + `autolens_workspace_test` `delaunay.py` entry citing #245. +2. Re-enable `multi_dataset/jax_likelihood/mge.py` and `shared_preloads.py` in + `autolens_workspace_test`'s `smoke_tests.txt` (folded in from the superseded + 2026-08-22 filing). +3. Anything that stays out after the fix stays out with a **recorded deliberate + reason**, not as an accumulated one-off. + +## Acceptance + +- A stated root cause, or an explicit recorded decision that it is an + infrastructure limit to be worked around rather than fixed. Not another + quarantine. +- Every entry marked NEEDS_FIX for this signature restored to its suite, or + re-marked with the real reason phase 2 established. +- The `multi_dataset/jax_likelihood/` family back under CI coverage in both test + workspaces, or its absence recorded as a deliberate choice. diff --git a/complete/index.md b/complete/index.md index 237a3b56..661d076a 100644 --- a/complete/index.md +++ b/complete/index.md @@ -6,7 +6,7 @@ Token-light navigation over the finished-work records (schema: only then grep a dated bucket. Curators: edit the band between the CURATED markers; everything below GENERATED is rebuilt. -1154 records across 7 buckets. +1155 records across 7 buckets. ## Highlights @@ -126,6 +126,7 @@ _(curate hard-won records here — survives regeneration.)_ - [jax-likelihood-smoke-pins-stale](2026/08/jax-likelihood-smoke-pins-stale.md) - [jax-stall-shared-preloads-retime-refutation](2026/08/jax-stall-shared-preloads-retime-refutation.md) - [jax-traceback-filtering-release-harness](2026/08/jax-traceback-filtering-release-harness.md) +- [jax-vmap-materialisation-hang](2026/08/jax-vmap-materialisation-hang.md) — jax-compile-stall phase 3 - [knowledge-board](2026/08/knowledge-board.md) — auto-closed on merge - [lacosmic-cr-option-and-star-pass-decoupling](2026/08/lacosmic-cr-option-and-star-pass-decoupling.md) — both deliberately OPEN — closure is the default-flip decision, not this ship - [latex-docstrings-invalid-escape-warnings](2026/08/latex-docstrings-invalid-escape-warnings.md) diff --git a/dashboard.html b/dashboard.html index ed2588ff..dd53f52f 100644 --- a/dashboard.html +++ b/dashboard.html @@ -269,7 +269,7 @@

Planned

latent-nan-guard-honest-run — planned 2026-07-22

Backlog markdown version

-

135 filed prompts, not started — sorted most-pickable first (priority, then size). 25 of them belong to an epic and are listed only under Epics below.

+

135 filed prompts, not started — sorted most-pickable first (priority, then size). 24 of them belong to an epic and are listed only under Epics below.

feature — 28

Numba CPU likelihood phase 1: batched MGE convolution + operated-matrix caching✨ featureautoarraymediumsupervisedhigh

@@ -302,10 +302,11 @@

Backlog

Teach repos_sync --write to stamp organ config surfaces✨ featurepyautomindhardsupervisedlow

-research — 14 +research — 15

Deep research: Can we speed up Delaunay in PyAutoArray?🔬 researchautoarraytoo-largesupervisedhigh

Delaunay-family JAX modules never hit the persistent compilation cache🔬 researchautoarraymediumsupervisedmedium

Quick-update plotting cost — minutes per update, and it is not JAX…🔬 researchautolensmediumsupervisedmedium

+

Why does XLA CPU's Eigen thread pool wedge on the multi_dataset vmap…🔬 researchcimediumsupervisedmedium

Use readthedocs or migrate to GitHub docs🔬 researchpyautohandssmallsupervisednormal

Re-baseline the slacs0008 acceptance parity after the HAP-dedupe fix🔬 researchpyautoreducesmallsupervisednormal

Kernel-CDF bandwidth defaults — config-dependent quality, investigate adaptivity🔬 researchautoarraymediumsupervisednormal

@@ -413,6 +414,12 @@

Backlog 2026-08-27 filed +Why does XLA CPU's Eigen thread pool wedge on the multi_dataset vmap… + + + +2026-08-27 +filed One construction path for plane-bound lensing quantities @@ -464,7 +471,7 @@

Backlog autocti_workspace has no Navigator Check, so its CI can never roll… - + 2026-08-24 filed Un-park multi_galaxy/features/scaling_relation/slam once a capped run… @@ -704,12 +711,6 @@

Backlog HowToLens ch4 tutorial 3: mask overlay is never actually drawn - -2026-08-03 -filed -Un-park imaging/features/scaling_relation/slam — the PyAutoArray#431… - -

Epics markdown version

@@ -734,10 +735,9 @@

Epics

Source & Cluster arc — magnification science, PointSolver trust, cluster extended sources✨ featureautolenstoo-largesupervisedhigh

-Intermittent XLA compile stall in the JAX vmap likelihood path — 2 queued prompt(s), in order -

Intermittent XLA compile stall in the JAX vmap likelihood path — ledger: draft/bug/ci/jax_vmap_jit_compile_stall.md — CLOSED AS PARTIAL 2026-08-23 — record complete/2026/08/jax-compile-stall-slow-vs-stall-audit.md

-

Phase 3: root-cause the XLA vmap compile stall and clear every NEEDS_FIX…🐛 bugcilargesupervisedhigh

-

Intermittent XLA compile stall in JAX vmap likelihood scripts — third repo…🐛 bugcitoo-largesupervisedhigh

+JAX vmap result never materialises (was: "intermittent XLA compile stall" — the name was wrong) — 1 queued prompt(s), in order +

JAX vmap result never materialises (was: "intermittent XLA compile stall" — the name was wrong) — ledger: draft/bug/ci/jax_vmap_jit_compile_stall.md — SHIPPED 2026-08-27 — all 3 phases done; record

+

JAX vmap result never materialises — campaign map…🐛 bugcitoo-largesupervisedhigh

Expectation propagation (EP) campaign — 9 queued prompt(s), in order diff --git a/dashboard.md b/dashboard.md index 26688199..1c351a5e 100644 --- a/dashboard.md +++ b/dashboard.md @@ -247,7 +247,7 @@ Scoped but not started; some are not yet prompt files. Full detail in [`planned. ## Backlog -**135** filed prompts, not started. Each section is sorted most-pickable first (priority, then size). **25** of them belong to an epic and are listed only under [Epics](#epics) below. +**135** filed prompts, not started. Each section is sorted most-pickable first (priority, then size). **24** of them belong to an epic and are listed only under [Epics](#epics) below.
feature — 28 @@ -479,7 +479,7 @@ Scoped but not started; some are not yet prompt files. Full detail in [`planned.
-research — 14 +research — 15
📋 Deep research: Can we speed up Delaunay in PyAutoArray? — autoarray · too-large · supervised · high @@ -505,6 +505,14 @@ Scoped but not started; some are not yet prompt files. Full detail in [`planned.
+
📋 Why does XLA CPU's Eigen thread pool wedge on the multi_dataset vmap… — ci · medium · supervised · medium + +``` +/start_dev draft/research/ci/xla_cpu_eigen_pool_deadlock.md +``` + +
+
📋 Use readthedocs or migrate to GitHub docs — pyautohands · small · supervised · normal ``` @@ -1180,6 +1188,7 @@ The 50 newest things to happen to the work in hand, newest first — issued, par | Date | Event | Task | |------|-------|------| +| 2026-08-27 | filed | Why does XLA CPU's Eigen thread pool wedge on the multi_dataset vmap… | | 2026-08-27 | filed | One construction path for plane-bound lensing quantities | | 2026-08-27 | filed | Multi-plane time delays | | 2026-08-27 | filed | LensCalc NumPy Hessian step is too coarse for multi-plane tracers | @@ -1189,12 +1198,12 @@ The 50 newest things to happen to the work in hand, newest first — issued, par | 2026-08-24 | filed | wiki-currency's --check-version gate rots on every library main merge | | 2026-08-24 | filed | interferometer/jax_grad/gradient.py: eager and jitted likelihoods… | | 2026-08-24 | filed | autocti_workspace has no Navigator Check, so its CI can never roll… | -| 2026-08-24 | filed | Un-park multi_galaxy/features/scaling_relation/slam once a capped run… |
… 10 more (40 left) | Date | Event | Task | |------|-------|------| +| 2026-08-24 | filed | Un-park multi_galaxy/features/scaling_relation/slam once a capped run… | | 2026-08-24 | filed | Induct PyAutoReduce into the PyAutoHands release machinery (date… | | 2026-08-24 | filed | Heart's local smoke runner cannot run any CTI workspace — no autocti… | | 2026-08-23 | filed | pynufft removal: unswept downstream residue (1 hard break + stale… | @@ -1204,12 +1213,12 @@ The 50 newest things to happen to the work in hand, newest first — issued, par | 2026-08-22 | filed | Untrack the generated FITS test artifacts in autoarray | | 2026-08-22 | filed | The reconstruction noise map describes a different estimator than the… | | 2026-08-22 | filed | Point-source JSON datasets record no resolution regime | -| 2026-08-22 | filed | Is Intel macOS a supported platform, and what is the numpy-only… |
… 10 more (30 left) | Date | Event | Task | |------|-------|------| +| 2026-08-22 | filed | Is Intel macOS a supported platform, and what is the numpy-only… | | 2026-08-21 | filed | Rectangular mesh split: Bilinear (fast CPU default) vs RTU… | | 2026-08-20 | filed | Numba CPU likelihood phase 2: kernel-CDF numba fast path (the 49-88%… | | 2026-08-20 | filed | Numba CPU likelihood phase 1: batched MGE convolution +… | @@ -1219,12 +1228,12 @@ The 50 newest things to happen to the work in hand, newest first — issued, par | 2026-08-19 | filed | RTD organism docs currency: Nerves page, organ-count drift, hands.md… | | 2026-08-19 | filed | Deduplicate repos_sync.py's check/write pairs | | 2026-08-19 | filed | Bug in autocti_workspace: the dataset_1d results/database example… | -| 2026-08-18 | parked | single-source-density-design |
… 10 more (20 left) | Date | Event | Task | |------|-------|------| +| 2026-08-18 | parked | single-source-density-design | | 2026-08-18 | parked | prior-message-collapse-design | | 2026-08-18 | filed | @PyAutoFit TransformedMessage.logpdf/pdf omit the transform… | | 2026-08-17 | filed | Which other searches need prior-support handling — coverage audit… | @@ -1234,12 +1243,12 @@ The 50 newest things to happen to the work in hand, newest first — issued, par | 2026-08-09 | found | isothermal-ell-sph-oversampling-at-the-cusp | | 2026-08-08 | parked | pyautoreduce-slacs1430-acs-comparison | | 2026-08-08 | filed | Regenerate autolens_workspace markdown/ so the MGE pages show… | -| 2026-08-07 | filed | Regenerate setup_notebook-drifted notebooks in… |
… 10 more (10 left) | Date | Event | Task | |------|-------|------| +| 2026-08-07 | filed | Regenerate setup_notebook-drifted notebooks in… | | 2026-08-06 | filed | Triage: Convolver "No blurring_image provided" warning in canonical… | | 2026-08-06 | filed | Rewrite PyAutoCTI docs/api — 55 of 89 autosummary entries are dead | | 2026-08-06 | filed | Dependency-cap refresh 2026-08: safe bumps, astropy 8 decision, two… | @@ -1249,7 +1258,6 @@ The 50 newest things to happen to the work in hand, newest first — issued, par | 2026-08-04 | filed | autolens_workspace_developer: broad stale-API rot (56 symbols, no CI) | | 2026-08-04 | filed | Nightly release has been blocked 8 nights running — triage the streak | | 2026-08-04 | filed | HowToLens ch4 tutorial 3: mask overlay is never actually drawn | -| 2026-08-03 | filed | Un-park imaging/features/scaling_relation/slam — the PyAutoArray#431… |
@@ -1399,25 +1407,17 @@ Continue the 'Cluster strong lensing — Source & Cluster arc' epic. Its canonic
-Intermittent XLA compile stall in the JAX vmap likelihood path — 2 queued prompt(s), in order - -
📋 Intermittent XLA compile stall in the JAX vmap likelihood path — ledger: `draft/bug/ci/jax_vmap_jit_compile_stall.md` — CLOSED AS PARTIAL 2026-08-23 — record complete/2026/08/jax-compile-stall-slow-vs-stall-audit.md - -``` -Continue the 'Intermittent XLA compile stall in the JAX vmap likelihood path' epic. Its canonical state lives in draft/bug/ci/jax_vmap_jit_compile_stall.md — read that ledger (and any DECISIONS/RESULTS files beside it) first. Cross-check this epic's entry in PyAutoMind/epics.md, any related rows in PyAutoMind/active.md, and the referenced repos' open issues and PRs, to work out the last completed phase and what is currently in flight. Then pick the next logical step and continue it through the normal workflow (/start_dev — filing the phase's prompt first if none exists), updating the ledger as the work advances. Note: phase 1 (watchdog) shipped in full; phases 2/3 stopped deliberately at a measured-but-not-root-caused state. The stall is instrumented and characterised (>100x bimodality inside one compile step; vmap-of-jit contributory at p=0.070 but NOT causal; the compile-cache hypothesis never tested) and NOTHING was un-quarantined. Resume via draft/research/ci/smoke_timing_and_profiling.md, which is where this gets dug up. Superseded complete/2026/08/multi-dataset-jax-likelihood-xla-stall.md (was draft/bug/autolens_workspace_test/multi_dataset_jax_likelihood_xla_stall.md). -``` - -
+JAX vmap result never materialises (was: "intermittent XLA compile stall" — the name was wrong) — 1 queued prompt(s), in order -
📋 Phase 3: root-cause the XLA vmap compile stall and clear every NEEDS_FIX… — ci · large · supervised · high +
📋 JAX vmap result never materialises (was: "intermittent XLA compile stall" — the name was wrong) — ledger: `draft/bug/ci/jax_vmap_jit_compile_stall.md` — SHIPPED 2026-08-27 — all 3 phases done; record ``` -/start_dev draft/bug/ci/jax_compile_stall_3_root_cause.md +Continue the 'JAX vmap result never materialises (was: "intermittent XLA compile stall" — the name was wrong)' epic. Its canonical state lives in draft/bug/ci/jax_vmap_jit_compile_stall.md — read that ledger (and any DECISIONS/RESULTS files beside it) first. Cross-check this epic's entry in PyAutoMind/epics.md, any related rows in PyAutoMind/active.md, and the referenced repos' open issues and PRs, to work out the last completed phase and what is currently in flight. Then pick the next logical step and continue it through the normal workflow (/start_dev — filing the phase's prompt first if none exists), updating the ledger as the work advances. Note: phase 1 (watchdog) shipped in full; phases 2/3 stopped deliberately at a measured-but-not-root-caused state. The stall is instrumented and characterised (>100x bimodality inside one compile step; vmap-of-jit contributory at p=0.070 but NOT causal; the compile-cache hypothesis never tested) and NOTHING was un-quarantined. Resumed 2026-08-27 as phase 3 (PyAutoFit#1528) — NOT via draft/research/ci/smoke_timing_and_profiling.md, ```
-
📋 Intermittent XLA compile stall in JAX vmap likelihood scripts — third repo… — ci · too-large · supervised · high +
📋 JAX vmap result never materialises — campaign map… — ci · too-large · supervised · high ``` /start_dev draft/bug/ci/jax_vmap_jit_compile_stall.md diff --git a/draft/bug/ci/jax_compile_stall_3_root_cause.md b/draft/bug/ci/jax_compile_stall_3_root_cause.md deleted file mode 100644 index fbcb1de4..00000000 --- a/draft/bug/ci/jax_compile_stall_3_root_cause.md +++ /dev/null @@ -1,82 +0,0 @@ -# Phase 3: root-cause the XLA vmap compile stall and clear every NEEDS_FIX it caused - -Type: bug -Target: ci -Repos: -- @PyAutoFit -- @autogalaxy_workspace_test -- @autolens_workspace_test -Difficulty: large -Autonomy: supervised -Priority: high -Status: formalised -Epic: jax-compile-stall -Phase: 3 -Campaign: bug/ci/jax_vmap_jit_compile_stall.md (Phase 3 — the fix; blocked on phases 1 and 2) -Filed: 2026-08-23 - -## Blocked on phase 1 - -Do not start this before phase 1's watchdog has shipped and a CI stall has -actually dumped a traceback. Diagnosing a hang that leaves no evidence is what -produced three quarantines and no root cause; repeating it without the -instrumentation would be a fourth. - -## Reproduce deliberately - -Loop `imaging/jax_likelihood/mge_group.py` under its declared CI env profile -until it hangs, rather than waiting for CI to hit it. Attach `py-spy dump` to -the hung process as well as reading phase 1's own `faulthandler` output — the -two see different things, and `py-spy` can read native frames the in-process -dump cannot. - -## Hypotheses, in the order they are cheapest to test - -1. **`vmap` of `jit`, the inverted ordering.** `Fitness._vmap` - (`autofit/non_linear/fitness.py`) builds `jax.vmap(jax.jit(self.call))`; - `autofit/non_linear/analysis/latent.py` builds - `jax.jit(jax.vmap(compute_latent_for_model))`, the conventional order. The - stalling path is exactly the `vmap` path, and the `_jit`-only scripts in the - same directories do not stall. One-line A/B — try it first. -2. **Persistent compilation cache contention.** `JAX_COMPILATION_CACHE_DIR` has - defaulted on since PyAutoConf#128 (merged 2026-07-17). Both NEEDS_FIX stalls - post-date it; the eight SLOW entries predate it. A/B with the cache dir set - to empty (which disables it) and see whether the stall probability moves. -3. **JAX/XLA version interaction.** This repo has form: `delaunay_mge.py` is - disabled outright because `jax 0.7` removed - `jax.interpreters.xla.pytype_aval_mappings`, and the smoke installer once - clobbered a working `tfp-nightly`. Pin-bisect jax/jaxlib across a run set. -4. **Runner CPU contention.** `complete/2026/07/jax-compile-time-research.md` - records that XLA compiles on **host** CPUs and that compile timing is - load-sensitive by up to 7×, which is why a hosted runner is the place this - reproduces and a workstation is not. -5. **Graph size in the vmap trace.** The affected set has a shape: plain `mge.py` - passes in 9.4s, while the *composite* variants — group, rectangular-MGE, - delaunay-MGE — stall or are already out. Complexity-driven compile blowup was - argued against by the autolens_profiling#71 research ("compile cost is - op-pattern-driven, not complexity-driven"), so treat this as the hypothesis - of last resort, not the first. - -## Then restore the coverage - -The point of the campaign. Quarantining removes exactly the heaviest JAX paths, -which are the ones most worth testing. - -1. Clear the NEEDS_FIX markers this campaign inherits — including the - 2026-08-01 `multi_dataset/jax_likelihood/rectangular.py` one and the - `autolens_workspace_test` `delaunay.py` entry citing #245. -2. Re-enable `multi_dataset/jax_likelihood/mge.py` and `shared_preloads.py` in - `autolens_workspace_test`'s `smoke_tests.txt` (folded in from the superseded - 2026-08-22 filing). -3. Anything that stays out after the fix stays out with a **recorded deliberate - reason**, not as an accumulated one-off. - -## Acceptance - -- A stated root cause, or an explicit recorded decision that it is an - infrastructure limit to be worked around rather than fixed. Not another - quarantine. -- Every entry marked NEEDS_FIX for this signature restored to its suite, or - re-marked with the real reason phase 2 established. -- The `multi_dataset/jax_likelihood/` family back under CI coverage in both test - workspaces, or its absence recorded as a deliberate choice. diff --git a/draft/bug/ci/jax_vmap_jit_compile_stall.md b/draft/bug/ci/jax_vmap_jit_compile_stall.md index 629cd9e5..db9bfc85 100644 --- a/draft/bug/ci/jax_vmap_jit_compile_stall.md +++ b/draft/bug/ci/jax_vmap_jit_compile_stall.md @@ -1,4 +1,4 @@ -# Intermittent XLA compile stall in JAX vmap likelihood scripts — third repo, still unfixed +# JAX vmap result never materialises — campaign map (SHIPPED 2026-08-27; the "XLA compile stall" name was wrong) Type: bug Target: ci @@ -146,6 +146,26 @@ of this file, and every marker calling this an "intermittent XLA compile stall", inherit a guess made before there was any evidence. See the record's final section. +## SHIPPED — 2026-08-27 + +All three phases are done. Record: +[`complete/2026/08/jax-vmap-materialisation-hang.md`](../../../complete/2026/08/jax-vmap-materialisation-hang.md). + +Phase 3 established that the hang is **materialising the vmap result**, not +compiling it — captured at `jax.block_until_ready`/`try_to_block` and at +`jax.Array._value`, with compilation finished ~12-18s earlier. The trigger is +XLA CPU's multithreaded Eigen thread pool; the workaround is +`XLA_FLAGS=--xla_cpu_multi_thread_eigen=false` in both test workspaces' smoke +and release profiles (ABAB, 12 passes/0 hangs with vs 2 passes/14 hangs +without, Fisher exact p ~ 3e-6). All seven quarantined entries are restored, +42/42 completions. + +**Why the pool wedges is still unknown** — a workaround, not a root-cause fix. +The follow-up lives under `draft/research/ci/`. + +The section below is kept as the 2026-08-23 state, for the history of how this +was worked. + ## CLOSED AS PARTIAL — 2026-08-23 Phase 1 shipped. Phases 2 and 3 were taken to a deliberate stopping point and the diff --git a/draft/research/ci/xla_cpu_eigen_pool_deadlock.md b/draft/research/ci/xla_cpu_eigen_pool_deadlock.md new file mode 100644 index 00000000..45f5267a --- /dev/null +++ b/draft/research/ci/xla_cpu_eigen_pool_deadlock.md @@ -0,0 +1,77 @@ +# Why does XLA CPU's Eigen thread pool wedge on the multi_dataset vmap graphs? + +Type: research +Target: ci +Repos: +- @PyAutoFit +- @autolens_workspace_test +- @autogalaxy_workspace_test +Difficulty: medium +Autonomy: supervised +Priority: medium +Status: formalised +Filed: 2026-08-27 + +## Why this exists + +`complete/2026/08/jax-vmap-materialisation-hang.md` shipped a **workaround, not +a root cause**. We know *where* it hangs and *what* avoids it: + +- the hang is materialising a vmap result — stacks at + `jax.block_until_ready`/`try_to_block` and at `jax.Array._value` via + `__str__`, with compilation finished ~12-18s earlier; +- `XLA_FLAGS=--xla_cpu_multi_thread_eigen=false` avoids it, ABAB, 12 passes / + 0 hangs against 2 passes / 14 hangs (Fisher exact p ~ 3e-6). + +We do **not** know why the pool wedges. `faulthandler` reports Python frames +only, so the wedged XLA worker threads are exactly the part still invisible. + +Two costs of leaving it here. Every JAX script in both test workspaces now runs +single-threaded Eigen, ~15% slower on the heaviest; and the flag is a load-bearing +line whose removal silently brings back seven quarantines. + +## The questions, in order + +1. **A native stack from the wedged threads.** `faulthandler` cannot see them. + Attach `gdb -p` (`thread apply all bt`) or `py-spy dump --native` to a hung + process. Where in XLA's runtime are the workers parked — a condition variable + in the Eigen pool, a work-queue steal loop, a barrier? +2. **A minimal reproducer outside the workspace.** Strip to the smallest jaxpr + that still wedges: how much of the `multi_dataset` composition is needed, and + does batch size or vmap axis count move the probability? A standalone script + is what an upstream report needs. +3. **CPU count versus cgroup quota.** The likeliest mechanism, and untested. A + hosted runner advertises more CPUs than the container can schedule, so a pool + sized from `nproc` can oversubscribe and deadlock. Compare + `os.cpu_count()` / `os.sched_getaffinity` / `/sys/fs/cgroup/cpu.max` on a + runner, and A/B `--xla_cpu_multi_thread_eigen=true` with the pool sized to + the real quota. If this is it, the fix is sizing the pool correctly rather + than disabling threading, and the ~15% comes back. +4. **Upstream.** File with JAX/XLA once 1-3 give a reproducer and a native + stack. Link the report from both workspaces' `no_run.yaml` block comments. + +## Do not repeat + +Refuted in phase 3, with evidence in the record: the persistent compilation +cache; the jax/jaxlib version (identical `0.11.1` stack on the day it hung 3/3 +on both legs); the 3.12-vs-3.13 split. Do not re-run these. + +## Method notes that cost a day to learn + +- **The hang rate wanders.** The same script passed 2/2 and hung 2/2 on the + identical commit hours apart. Any claim needs ABAB against a re-measured + control, never before/after. +- The dispatchable harness is `retime.yml` in either workspace; its `STALL` / + `NEITHER` verdicts are the classifier to quote. +- CI picks up library branches by **matching branch name** across the dependency + chain, so an experiment needs the same branch name in every repo it touches. + +## Acceptance + +- A native stack naming where the XLA workers are parked, or a recorded finding + that it could not be obtained and why. +- Question 3 answered either way, since it decides whether the ~15% is + recoverable. +- An upstream issue filed, or a recorded decision not to with the reason. +- If a better fix lands, the flag is removed from all four profiles and the + family re-timed before anything is called done. diff --git a/epics.md b/epics.md index 0898eada..0ceac917 100644 --- a/epics.md +++ b/epics.md @@ -29,9 +29,15 @@ epic, never picked standalone. - notes: 12 phased prompts under draft/; issue phases ONE at a time as predecessors near shipping — no bulk issue queues. ## jax-compile-stall -- title: Intermittent XLA compile stall in the JAX vmap likelihood path +- title: JAX vmap result never materialises (was: "intermittent XLA compile stall" — the name was wrong) - ledger: draft/bug/ci/jax_vmap_jit_compile_stall.md -- status: CLOSED AS PARTIAL 2026-08-23 — record complete/2026/08/jax-compile-stall-slow-vs-stall-audit.md +- status: SHIPPED 2026-08-27 — all 3 phases done; record + complete/2026/08/jax-vmap-materialisation-hang.md. Root cause is XLA CPU's multithreaded Eigen + thread pool; workaround XLA_FLAGS=--xla_cpu_multi_thread_eigen=false in both test workspaces' + smoke AND release profiles (ABAB: 12 pass/0 hang with vs 2 pass/14 hang without, Fisher p~3e-6). + All 7 quarantined entries restored, 42/42 completions. PyAutoFit#1528, PRs PyAutoFit#1529, + PyAutoHands#269, autolens_workspace_test#281, autogalaxy_workspace_test#114. + NOT a root-cause fix: why the pool wedges is still unknown — follow-up filed under draft/research/ci/. - NEW EVIDENCE 2026-08-25: multi_dataset/jax_likelihood/shared_preloads.py stalled at TIMEOUT (300s) in PyAutoHeart Workspace Smoke run 32902243623 — one day after the 2026-08-24 retime refuted its SLOW marker and returned it to mega-run coverage. N=5 per leg measures the fast mode of a bimodal failure @@ -43,7 +49,8 @@ epic, never picked standalone. - CORRECTION (post-close-out): the captured stack shows the hang is in jax.block_until_ready, NOT in compilation. The epic's name and every marker calling this an "XLA compile stall" are wrong. Resume from "why does block_until_ready never return", not from compiler behaviour. -- notes: phase 1 (watchdog) shipped in full; phases 2/3 stopped deliberately at a measured-but-not-root-caused state. The stall is instrumented and characterised (>100x bimodality inside one compile step; vmap-of-jit contributory at p=0.070 but NOT causal; the compile-cache hypothesis never tested) and NOTHING was un-quarantined. Resume via draft/research/ci/smoke_timing_and_profiling.md, which is where this gets dug up. Superseded complete/2026/08/multi-dataset-jax-likelihood-xla-stall.md (was draft/bug/autolens_workspace_test/multi_dataset_jax_likelihood_xla_stall.md). +- notes: phase 1 (watchdog) shipped in full; phases 2/3 stopped deliberately at a measured-but-not-root-caused state. The stall is instrumented and characterised (>100x bimodality inside one compile step; vmap-of-jit contributory at p=0.070 but NOT causal; the compile-cache hypothesis never tested) and NOTHING was un-quarantined. Resumed 2026-08-27 as phase 3 (PyAutoFit#1528) — NOT via draft/research/ci/smoke_timing_and_profiling.md, + which the 2026-08-23 close-out named as the resume door but which was never written. Superseded complete/2026/08/multi-dataset-jax-likelihood-xla-stall.md (was draft/bug/autolens_workspace_test/multi_dataset_jax_likelihood_xla_stall.md). ## graphical-ep - title: Expectation propagation (EP) campaign