fix: per-script timeout budget so full-dataset jax_grad scripts are measurable

[Jammy2211]

Overview

Three autolens_workspace_test JAX-gradient scripts TIMEOUT at the 300s per-script
smoke cap (PyAutoHeart workspace-smoke run 30858578587, 2026-08-03T22:25Z):
imaging/jax_grad/pixelization.py, imaging/jax_grad/regularization.py,
point_source/jax_grad/gradient.py. They are killed, not errored.

The question asked was whether this is a real slowdown (JIT compile regression, lost
caching) or a cap that was always marginal. The answer is mixed, and the current
tooling cannot fully settle it — which is itself the primary defect to fix.

What is established

point_source/jax_grad/gradient.py is definitively a budget problem. Run locally
under the resolved smoke profile it PASSES in 665.9s — 2.2x the cap, with no hang
and no correctness failure. Nothing about it fits in 300s.

Its stored baseline describes a different script. Heart's 39.24s entry is slugged
autolens_test__scripts__jax_grad__point_source.json, i.e. the pre-#216 path,
measured 2026-07-30. Commit 50f1c33 (2026-07-31) then added solver blocks C-F: a
pixel_scale_precision=1e-5 solver (100x finer than production) and 6-step FD sweeps
across four further models (+281 lines). It grew by design.

regularization.py has never completed in any recorded run. Its earlier
FAIL (103.5s) was the tfp-nightly/bessel_kve import gap
(workspace-validation.yml:247-254), i.e. time-to-failure, not time-to-success. It
loops over 4 mesh/regularization variants; a local partial run reached an early
variant at 131.5s, extrapolating to ~310-560s in CI. There is no passing baseline to
regress from.

Root cause of the cap mismatch. These scripts declare ENV: jax full_datasets.
Precisely: full_datasets unsets PYAUTO_SMALL_DATASETS and jax unsets
PYAUTO_DISABLE_JAX; PYAUTO_TEST_MODE remains 2 (verified by resolving the
profile, not by reading docs). But PYAUTO_TEST_MODE=2 only bypasses
AbstractSearch.fit, and these scripts run no search — they build Fitness and call
jax.value_and_grad/finite differences directly. So test mode reduces nothing here:
they do full-resolution FD work in smoke exactly as in release, where the cap is
already 1800s with a written rationale (workspace-validation.yml:319-326). Smoke was
left at 300s on the premise that smoke work is reduced; for this class that premise
does not hold.

What is NOT established (deliberately not overclaimed)

Timings across the two runs that day:

script	06:31Z	22:25Z	delta
lp.py	40.0s	39.6s	-1%
mge.py	42.1s	41.8s	-0.7%
delaunay.py	84.6s	92.4s	+9%
knn.py	175.8s	200.0s	+14%
pixelization.py	244.8s	>300s	>= +23%

The slowdown scales monotonically with runtime. That is compatible with runner
contention — but equally with a mesh/inversion/XLA compile regression whose effect
grows with graph size, under which flat lp.py/mge.py timings are exactly what
you would expect. So the flat cheap scripts do not discriminate between the two
hypotheses, and the earlier claim that they did was too strong.

An earlier draft cited the profiling prompt's ">400s on 2026-07-13" as corroboration
that pixelization.py was always marginal. Withdrawn: commit 74673c8
(2026-07-23) retired variants and rewrote both gradient scripts, so that figure
describes an ancestor, not the current script. pixelization.py therefore has exactly
one successful observation (244.8s), which cannot establish a distribution.

No library change explains a targeted slowdown: the only inversion-adjacent commit in
the window, PyAutoArray #431 (5006f347), is gated behind PYAUTO_SMALL_DATASETS=1,
which these scripts unset — a no-op for them.

The blocker is that the evidence to settle this does not exist. build_util.py
discards subprocess.TimeoutExpired.stdout/stderr, recording only
"Timed out after Ns". A timed-out script therefore reveals nothing about which
variant or block was executing. Raising the budget is not symptom-suppression here: a
process killed at 300s produces no measurement at all, so a realistic budget plus
preserved output is the instrument that makes the regression question answerable.

Plan

• Preserve timeout stdout/stderr so a TIMEOUT says where it died.
• Add parent-side per-script timeout resolution (the value is already reachable).
• Give the jax_grad class a realistic budget, from measured durations.
• Re-run and compare pixelization.py against its 244.8s baseline to settle the
  regression question with data.
• Leave smoke's 300s default untouched for every other script.

Detailed implementation plan

Affected Repositories

• PyAutoHands (primary) — runner changes
• autolens_workspace_test — profile override
• PyAutoHeart — comment correction

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoHands	main	clean
./PyAutoHeart	main	clean
./autolens_workspace_test	main	clean

Suggested branch: feature/jax-grad-smoke-timeout-budget

Phase 1 — PyAutoHands

1. Preserve timeout output. build_util.py:391 (execute_script) and the
   notebook path (execute_notebook, ~line 246) catch subprocess.TimeoutExpired
   and drop e.stdout/e.stderr. Capture them into ScriptResult.error_message
   (truncated tail). Do this FIRST — it is what makes the next measurement legible.
2. Per-script timeout, parent-side. build_util.py:12 reads
   TIMEOUT_SECS = int(os.environ.get("BUILD_SCRIPT_TIMEOUT", "300")) once at import
   in the parent and applies it as subprocess.run(..., timeout=TIMEOUT_SECS), while
   the per-script env from env_config.build_env_for_script is passed only to the
   child — so a per-script value is silently ignored today.
   Resolve the effective timeout from the env dict execute_script /
   execute_notebook already receive.
   Use the existing set: syntax, not a new schema key — validate_env_profiles.py:60
   defines ALLOWED_OVERRIDE_KEYS = {"pattern", "set", "unset"}, so a timeout: key
   would be rejected by profile validation. set: {BUILD_SCRIPT_TIMEOUT: "1800"}
   validates today and needs no validator change.
3. Precedence. Because apply_profile writes over the copied ambient env, the
   per-script value wins over the ambient global. Document this explicitly:
   per-script profile value > ambient BUILD_SCRIPT_TIMEOUT > 300 default.
   Note the trap that motivates it: run_all.py:256 exports BUILD_SCRIPT_TIMEOUT
   unconditionally, even when 300 was only the implicit CLI default — so an
   "explicit global always wins" rule would make per-script values work under Heart's
   workflow but be silently ignored under local run_all. Reject that rule for this
   reason and record it.
4. Keep the SLOW banner honest. slow_skip_check.py:198 quotes
   build_util.TIMEOUT_SECS as "the" cap in the mega-run report. With per-script
   budgets that figure is wrong for overridden scripts — the same
   understated-cap failure that previously biased park/un-park decisions
   (the "60s cap myth"). Surface the in-force per-script value or qualify the text.
5. Record provenance. Include the effective timeout in the result JSON so a future
   TIMEOUT is self-describing.
6. Tests. Matching vs non-matching path; script AND notebook paths; malformed /
   zero / negative values; precedence vs ambient global; that release's 1800s export
   still applies where no override matches.

Phase 2 — measure, then set the budget (after phase 1 lands)

7. Run the three scripts with a generous provisional budget and record ACTUAL CI
   durations. Local reference: gradient.py 665.9s PASS; calibration knn.py local
   141.6s vs CI 200.0s (CI/local ~1.41), so gradient.py projects to ~940s in CI.
8. autolens_workspace_test/config/build/profile_smoke.yaml — add the override with
   the measured basis in a comment.
9. Compare pixelization.py's completed duration against its 244.8s baseline. A
   result near 245-300s indicates contention plus an inadequate cap; substantially
   above indicates a genuine mesh/inversion regression, which then routes to
   draft/feature/profiling/profiling_agent_jax_compile_time_scope.md rather than
   being absorbed here.
10. PyAutoHeart/.github/workflows/workspace-validation.yml:319-326 — correct the
    "smoke mode is unchanged (still 300s)" comment.

Scope decision — pattern breadth

_pattern_matches (env_config.py:413) substring-matches patterns containing /, so
jax_grad/ matches 8 scripts: the 6 imaging ones, point_source/jax_grad/gradient.py,
and interferometer/jax_grad/gradient.py (scripts/weak/jax_grad.py does NOT match —
no trailing slash). Family-wide is chosen deliberately over three explicit patterns:
knn.py is already at 200.0s (67% of cap and rising) and would be the next failure.
Noted: interferometer/jax_grad/gradient.py is SLOW-skipped in no_run.yaml:47 for
flaking at the 1800s cap, so it will not run regardless — but it is direct evidence
that this family genuinely needs release-scale budgets.

Key Files

• PyAutoHands/autohands/build_util.py — TIMEOUT_SECS:12, execute_script:367, execute_notebook:219
• PyAutoHands/autohands/env_config.py — build_env_for_script:293, apply_profile:327, _pattern_matches:413
• PyAutoHands/autohands/validate_env_profiles.py — ALLOWED_OVERRIDE_KEYS:60
• PyAutoHands/autohands/run_all.py:256 — unconditional export
• PyAutoHands/autohands/slow_skip_check.py:198 — banner figure
• autolens_workspace_test/config/build/profile_smoke.yaml
• PyAutoHeart/.github/workflows/workspace-validation.yml:319-326

Out of scope

Other failures in the same run: autolens_test, interferometer; autofit_test, jax_assertions; howtogalaxy, chapter_4_pixelizations. The durable speed-up remains
the Profiling Agent prompt; this task does not duplicate it.

Follow-ups found while investigating (file separately)

• Heart script_timing slugs are path-derived, so the fix: validate the project before generate.py deletes notebooks/ #216 restructure orphaned every
  jax_grad baseline; no new-layout slug exists, so Heart accumulates no history.
• Every stored timing history is one value repeated 7x (a rolling-window fill, not 7
  observations), so the "median of 7" regression ratio is a single-observation compare.
• Local runs of lp.py fail assertions that pass in CI (identical JAX 0.10.2; numpy
  2.2.6 local vs 2.4.6 CI). Local jax_grad assertion failures are therefore NOT
  trustworthy evidence of source defects — worth its own investigation.

Original Prompt

Click to expand starting prompt

Original request (verbatim)

4. Three JAX-gradient script timeouts

In the PyAutoLabs workspace, three scripts TIMEOUT at the 300s smoke cap:

autolens_workspace_test scripts/imaging/jax_grad/pixelization.py
autolens_workspace_test scripts/imaging/jax_grad/regularization.py
autolens_workspace_test scripts/point_source/jax_grad/gradient.py

Evidence: PyAutoHeart workspace-smoke run 30858578587
(2026-08-03T22:25Z), jobs "smoke / run_scripts (3.12, autolens_test,
imaging)" and "(3.12, autolens_test, point_source)". All three are
TIMEOUT, not error — they never finished.

Work out whether this is real slowdown (JIT compile-time regression,
lost caching) or just a cap that was always marginal for these three:
check whether they ever passed in smoke and how close they were. Compare
against Heart's script_timing baselines if there's history. Fix the
cause, or — with an explicit reason — reclassify them; don't just raise
the cap silently. Route through start_dev.

[Jammy2211]

Phase 1 shipped — PR #227

Runner-side only. No script's cap has moved yet — phase 2 sets the actual jax_grad/ budget once it can be measured.

What landed: build_util.timeout_for(env) resolves the cap parent-side (the per-script value was already reachable — it just went only to the child), and the TimeoutExpired handlers now preserve a truncated output tail plus the cap in force.

No profile schema change: profiles use the existing set: key, since ALLOWED_OVERRIDE_KEYS is {pattern, set, unset} and a new timeout: key would have been rejected by validate_env_profiles.

Tests: 25, control-tested against the unfixed behaviour first — 11 fail without the change, so they are not vacuous. Full suite 281 pass / 2 fail, both verified pre-existing on unmodified main 9601f18.

Correction to this issue's own diagnosis

Two claims in the issue body are weaker than first written, and adversarial review caught both:

1. The ">400s on 2026-07-13" corroboration for pixelization.py is withdrawn — commit 74673c8 (2026-07-23) rewrote both gradient scripts, so that figure describes an ancestor.
2. "Not a regression" is too strong. The slowdown scales monotonically with runtime (−1% at 40s → +23% at 245s), which fits runner contention and a graph-size-dependent XLA regression equally well. Flat lp.py/mge.py do not discriminate between them.

What remains solid: gradient.py passes locally in 665.9s (2.2× the cap) — no hang, no correctness failure — and regularization.py has never completed in any recorded run.

Phase 2 gate

The budget must not be set from local numbers. The lp.py control fails locally while passing in CI, and pixelization.py/regularization.py abort early locally, so their local durations are truncated rather than total.

Phase 2 needs a workspace-smoke dispatch with #227 merged, reading actual CI durations from the now-preserved timeout output. Comparing pixelization.py's completed duration against its 244.8s baseline is what finally settles the open regression question.

[Jammy2211]

Measured — the question is settled: a marginal cap, not a slowdown

PyAutoHeart workspace-smoke run 30938311069, the first run in which these scripts could finish. All three now PASS, and both autolens_test jobs are green.

script	06:31Z	22:25Z	now
lp.py	40.0s	39.6s	41.6s
mge.py	42.1s	41.8s	41.2s
delaunay.py	84.6s	92.4s	89.6s
knn.py	175.8s	200.0s	186.5s
pixelization.py	244.8s	TIMEOUT	259.8s PASS
regularization.py	FAIL (import)	TIMEOUT	301.7s PASS
gradient.py	TIMEOUT	TIMEOUT	568.2s PASS

The regression hypothesis is dead

This issue deliberately left open whether a genuine mesh/inversion XLA regression sat underneath the cap problem — the 2026-08-03 slowdown scaled monotonically with runtime (−1% at 40s → +23% at 245s), which fits a graph-size-dependent regression as well as contention, and the flat cheap scripts did not discriminate.

It resolves as contention:

• pixelization.py returned 259.8s against its 244.8s baseline — +6%, ordinary variance;
• knn.py 200.0 → 186.5 and delaunay.py 92.4 → 89.6 both came back down.

A graph-size-dependent regression does not undo itself. The 22:25Z run was simply a slow run. Nothing routes to the Profiling Agent.

Per script

• pixelization.py — always marginal, exactly as suspected. At 82% of cap it had no headroom to absorb a slow runner.
• regularization.py — 301.7s: it exceeded the 300s cap by 1.7 seconds. It was never far too slow; it missed by a hair, and the tfp import gap masked that until the install was fixed.
• gradient.py — 568.2s, 1.9× the cap. Grew by design in 50f1c33.

A projection of mine was wrong

I estimated gradient.py at ~940s in CI from a local measurement (665.9s) scaled by the knn.py CI/local ratio (~1.41x). The real figure is 568.2s — CI is faster than local for that script, so the ratio did not transfer between scripts. The conclusion held; the extrapolated number did not. autolens_workspace_test#250 replaces the projected basis with the measured one.

Follow-up

The 1800s budget was explicitly provisional. With a measured max of 568.2s it is 3.2× the slowest script, so autolens_workspace_test#250 tightens it to 900s (~1.6× headroom; the largest swing ever seen, +23%, lands at ~699s).

[Jammy2211]

Closing — shipped, measured, and cleaned up

Verdict: a marginal cap, not a slowdown. All seven live jax_grad scripts pass on PyAutoHeart workspace-smoke run 30938311069: gradient.py 568.2s, regularization.py 301.7s, pixelization.py 259.8s, knn.py 186.5s, delaunay.py 89.6s, lp.py 41.6s, mge.py 41.2s.

Merged

PR
PyAutoHands#227 → 52408a84	timeout_for(env) resolves the cap parent-side; TimeoutExpired handlers preserve the output tail + cap in force; 25 control-tested tests
PyAutoHeart#138 → 1bdf8188	corrects the release cap's "smoke is unchanged (still 300s)" premise
autolens_workspace_test#249 → a430ea61	the jax_grad/ budget, plus a second runner (run_smoke.py) carrying the identical parent-side bug
autolens_workspace_test#250 → 527c1669	tightens 1800s → 900s on measured durations (~1.6x headroom over 568.2s)

Heart effect

The workspace-validation reason went from 3 failed, 3 timeout naming these scripts to 2 failed naming only autolens (not _test) interferometer likelihood_function — pre-existing and outside this task.

Corrections made during the work

• The ">400s on 2026-07-13" corroboration was withdrawn — 74673c8 rewrote both gradient scripts, so it described an ancestor.
• "Not a regression" was too strong and was held as not established until the measurement settled it.
• The ~940s CI projection for gradient.py was wrong (real: 568.2s); the local/CI ratio did not transfer between scripts. fix: notebook + env parsers silently mis-handle raw-string (r""") docstrings #250 records "measure, do not extrapolate".

Follow-ups filed separately (not absorbed)

• Heart script_timing slugs are path-derived, so the fix: validate the project before generate.py deletes notebooks/ #216 restructure orphaned every jax_grad baseline — no history has accumulated since 2026-07-24 — and every stored history is one value repeated 7x, making the "median of 7" ratio a single-observation compare.
• Local jax_grad runs fail assertions that pass in CI (lp.py control demonstrates it); suspect numpy 2.2.6 vs 2.4.6.

Record: PyAutoMind/complete/2026/08/jax-grad-smoke-timeout-budget.md. Worktree and all four branches removed.