feature: env-profile migration steps 4-8 — derive the release JAX set

[Jammy2211]

Overview

Finish the env-profile migration (build-chain campaign Phase 3, umbrella #155, research + design on #161; design doc docs/env_profile_redesign.md, merged #162; steps 1–3 shipped: validator #163, single resolver, scrubbed baseline env #165). In the ag/al release profiles JAX is re-enabled by a hand-enumerated list of ~10 path patterns — a list that silently rots: 15 scripts literally named for JAX (jax_assertions/, hessian_jax.py, tracer_jax.py, …) run NumPy in every nightly while asserting JAX behaviour. Step 4 replaces enumeration with the design's derivation rule (JAX-on in release mode iff a path segment/stem matches jax_* / *_jax / *_jit); steps 5–8 make the validator binding, rename the profiles, fold the duplicate library env read, and (human-gated) decide the use_jax sentinel.

This issue covers steps 4–8, executed in phases. This session: step 4 only.

Plan

Step 4 — the derivation rule (this session, three PRs, PyAutoHands first):

• Phase A (PyAutoHands): move is_jax_marked from the validator into autohands/env_config.py (one owner of the rule); resolver applies the derivation when a profile sets a new derive_jax_markers: true top-level key (explicit key, not a filename heuristic); validator accepts the key, resolves through the shared logic, and skips non-runnable __init__.py files; unit tests. Old profiles + new resolver = unchanged behaviour, so this merges first.
• Phase B (autolens_workspace_test): rename the two mid-stem files imaging/modeling_visualization_jit_delaunay.py → modeling_visualization_delaunay_jit.py (and _rectangular) so the _jit suffix carries the marker (human decision 2026-07-23); triage database/scrape/ on the NumPy path first — only fall back to a scrape_jax/ folder rename if JAX is genuinely required (human decision 2026-07-23; no exception list, ever); delete all 10 PYAUTO_DISABLE_JAX overrides from env_vars_release.yaml and add the key; triage the 10 runnable turns-ON scripts individually under their exact new JAX-on env (none has ever run JAX in the nightly).
• Phase C (autogalaxy_workspace_test): same minus renames/scrape — delete 8 overrides, add the key, triage the 3 runnable jax_assertions/ scripts.
• Verification (binding): resolved-env diff via resolve_clean() over every script, old profile vs new profile+rule, read line by line — must be exactly the triaged set and nothing else. Never smoke (smoke cannot see config-only changes). Prove against the real CI env (workspace-validation.yml), not a synthetic clean shell. The release path is only exercised by the nightly mega-run — mega-run verification accepted and stated on each PR.

Later phases (same issue):

• Step 5 — flip --strict-derivation / --strict-markers to errors; wire the validator into the three *_workspace_test PR gates.
• Step 6 — rename env_vars.yaml → profile_smoke.yaml, env_vars_release.yaml → profile_release.yaml (last; re-sweep for old names after every merge from main).
• Step 7 — fold al.AnalysisDataset's duplicate PYAUTO_DISABLE_JAX read into base af.Analysis (own plan; PyAutoFit + PyAutoLens, library-first; currently queued behind the testmode-env-drift PyAutoFit claim).
• Step 8 — use_jax: Optional[bool] = None sentinel: human decision required before any code; may be declined.

Detailed implementation plan

Affected Repositories

• PyAutoHands (primary — the rule + validator)
• autolens_workspace_test
• autogalaxy_workspace_test
• (steps 7–8 later: PyAutoFit, PyAutoLens)
• autofit_workspace_test needs no work — validator clean (0 warnings, the Skip EmceePlotter / NautilusPlotter under PYAUTO_TEST_MODE=2 #47 shape)

Branch Survey (2026-07-23)

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

No worktree conflicts (worktree_check_conflict clean for all three).

Suggested branch: feature/env-profile-derivation
Worktree root: ~/Code/PyAutoLabs-wt/env-profile-derivation/
Work classification: Both (PyAutoHands source first, workspace_test follow-ups)

Measured baseline (validator, 2026-07-23 — live re-measurement, not the stale 07-17 counts)

repo	override-enumeration	JAX-marked but resolves NumPy	unmarked but resolves JAX
autolens_workspace_test	10	11	9
autogalaxy_workspace_test	8	4	0
autofit_workspace_test	0	0	0

Nuances found in planning: imaging/modeling_visualization_jit_delaunay.py / _rectangular.py are already in no_run.yaml (SLOW, 2026-05-07) — no immediate nightly impact, but post-migration the strict validator would see them as consistent NumPy scripts while their content asserts jax.Array, hence the rename. 4 of 7 database/scrape/ files are also no_run'd; only general.py + scaling_relation.py actually run. The 15 turns-ON scripts are not no_run'd — they genuinely run tonight on NumPy, so flipping them is a real nightly change and each is triaged individually.

Implementation Steps (Phase A)

1. autohands/env_config.py — move in is_jax_marked (from validate_env_profiles.py:46); in build_env_for_script, when env_config.get("derive_jax_markers") is truthy, after defaults+overrides: if is_jax_marked(file): env["PYAUTO_DISABLE_JAX"] = "0".
2. autohands/validate_env_profiles.py — import is_jax_marked from env_config (delete local copy); add derive_jax_markers to ALLOWED_TOP_KEYS; resolve_clean applies the same derivation via shared logic; skip __init__.py when collecting scripts.
3. Unit tests (rule, derivation on/off, validator key, __init__.py skip).

Implementation Steps (Phase B — al)

1. git mv the two renames; update config/build/no_run.yaml:26-27; sweep for other references.
2. Scrape triage: run database/scrape/general.py + scaling_relation.py under the resolved NumPy release env. Pass → drop override, state on PR. Fail (JAX-rooted) → git mv scripts/database/scrape → scrape_jax + sweep ~8 refs (smoke profile ×2, no_run ×4, release profile, AGENTS.md, folder CLAUDE.md).
3. config/build/env_vars_release.yaml — delete all PYAUTO_DISABLE_JAX overrides; add derive_jax_markers: true; rewrite header comment.
4. Turns-ON triage (each under its exact new resolved env): jax_assertions/ ×7, hessian_jax.py, profiles_jit.py, tracer_jax.py. Fail → trivial fix-forward, else no_run.yaml with dated reason, surfaced on the PR.

Implementation Steps (Phase C — ag)

As B steps 3–4: delete 8 overrides, add key, triage jax_assertions/ ×3.

Key Files

• PyAutoHands/autohands/env_config.py — the single resolver; gains the rule + derivation
• PyAutoHands/autohands/validate_env_profiles.py — PR-time gate; resolve_clean() is the verification harness
• autolens_workspace_test/config/build/env_vars_release.yaml — 10 overrides to delete
• autogalaxy_workspace_test/config/build/env_vars_release.yaml — 8 overrides to delete
• PyAutoHeart/.github/workflows/workspace-validation.yml:295,322 — how the nightly invokes the release profile (read-only; parity check)
• PyAutoHands/docs/env_profile_redesign.md — the decided design (§2 policy, §3 rule)

Config traps (from prior sweeps)

overrides: order is load-bearing (later entries win); an empty overrides: block loads as None, not []; verify config-only changes by resolved-env diff, never smoke.

Planning notes

Brain Feature Agent scored this too-large (score 25, repo-count proxy) and proposed a generic 4-phase split — overridden with the prompt's own step phasing (step 4 alone; 5+6 pair; 7 independent; 8 decision-first), recorded here per standing guidance.

Original Prompt

Click to expand starting prompt

Finish the env-profile migration: steps 4-8 (derivation, strict gate, rename, one-reader, sentinel)

Type: feature
Target: workspaces
Repos:

• autolens_workspace_test
• autogalaxy_workspace_test
• autofit_workspace_test
• PyAutoHands
• PyAutoFit
• PyAutoLens
  Difficulty: large
  Autonomy: supervised
  Priority: normal
  Status: formalised

Build-chain campaign Phase 3, the last remaining work from the umbrella closed
2026-07-22 (epic PyAutoHands#155). Tracked on PyAutoHands#161. Design
merged: PyAutoHands/docs/env_profile_redesign.md (#162). Steps 1-3 done:
validator (#163), single resolver, scrubbed baseline env (#165 + profiles).

What this is about

Each *_workspace_test repo has two env profiles under config/build/:
env_vars.yaml (smoke — the per-PR gate) and env_vars_release.yaml
(release — the nightly mega-run). PyAutoHands/autohands/env_config.py
resolves them: a defaults: block applied to every script, then an ordered
overrides: list of path patterns that flip vars for specific scripts.

In the ag/al release profiles PYAUTO_DISABLE_JAX: "1" is the default and
JAX is re-enabled by a hand-enumerated list of ~10 folder patterns. That
enumeration is the failure mode the whole campaign exists to remove: a
hand-maintained list of paths silently rots as scripts are added or renamed,
and the failure is quiet — a jax_grad/ script that falls off the list runs
NumPy and asserts a gradient that is all zeros, or tests nothing at all.

The decided policy (design doc §2, human 2026-07-16) is NOT "release =
production defaults" — that was rejected on measured JIT compile cost
(autolens_profiling#71). mode=release validates every script on the NumPy
path plus a derived, loud, bounded JAX set. Step 4 is what makes that set
derived instead of enumerated.

Measure first — the 2026-07-17 numbers are stale

Re-run the validator before planning; do not carry old counts into the work.
Measured 2026-07-22 (after the autobuild→autohands rename and the
dead-entry purge, both of which moved these numbers):

python3 PyAutoHands/autohands/validate_env_profiles.py <workspace_test_root>

repo	override-enumeration	JAX-marked but resolves NumPy	unmarked but resolves JAX
autolens_workspace_test	10	11	9
autogalaxy_workspace_test	8	4	0
autofit_workspace_test	0	0	0 (clean — DISABLE_JAX "0", no overrides, the #47 shape)

The 3 dead patterns the old brief mentioned (ag quantity/) are gone —
purged by the no_run/env_vars dead-entry campaign. autofit_workspace_test
needs no work at all; this task is ag + al.

Step 4 — the derivation rule (do this first; RELEASE-SURFACE RISK)

Replace the enumerated PYAUTO_DISABLE_JAX overrides with one rule in
autohands/env_config.py: JAX-on in mode=release iff a path segment matches
jax_* / *_jax / *_jit.

This is not a config tidy — it changes which backend scripts run on. Every
script below either starts or stops using JAX, and none of the "starts" has
ever run under JAX. Triage each individually: does it pass under JAX, and is it
fast enough for the nightly? Three groups, from the measured lists:

1. Turns ON (marked, currently NumPy — 15 scripts). al: jax_assertions/
   (8 files, incl. __init__.py), hessian_jax.py, profiles_jit.py,
   tracer_jax.py. ag: jax_assertions/ (4 files). Highest value — files
   literally named for JAX that have never been JAX-validated — and the
   highest chance of a red nightly.
2. database/scrape/ (7 files, al). The one place the rule does not match:
   these need JAX on but carry no marker. Decide rename vs one
   documented parity entry. The design says no exception list, so a rename is
   the honest option, but it touches paths that no_run.yaml,
   copy_files.yaml and docs may reference — sweep for references first.
3. Turns OFF — a regression risk the old brief missed.
   imaging/modeling_visualization_jit_delaunay.py and
   ..._jit_rectangular.py resolve JAX-on today only because the override
   pattern imaging/modeling_visualization_jit substring-matches them. Under
   the derivation rule their stems end _delaunay / _rectangular, so they
   would stop getting JAX and silently no-op (the very failure the rule
   exists to prevent). Resolve this before writing the rule — it is
   evidence the *_jit suffix match is the wrong shape, and a
   segment-contains form may be what the design actually needs.

Also decide what the rule does with __init__.py (3 of the counted files are
package inits, not runnable scripts) — probably skip them in the validator
rather than special-case the rule.

Verification (binding). Verify by resolved-env diff, never by smoke:
run resolve_clean() (autohands/validate_env_profiles.py:56) over every
script, old profile vs new rule, and read the diff line by line — it must be
exactly the triaged set above and nothing else. Traps from prior config
sweeps: overrides: order is load-bearing (later entries win; defaults:
order is not), and an empty overrides: block loads as None, not []. And
per the step-3 near-miss: prove it against the real CI env, not a synthetic
clean one — a synthetic run passed 65/65 while the real env would have gone
red, because the smoke profiles relied on a PYAUTO_ var injected by the
workflow rather than by the profile.

The release path is exercised only by the mega-run, not the per-PR gate. Land
behind a rehearsal, or accept mega-run verification and say so on the PR.

Step 5 — make the validator binding

Flip --strict-derivation / --strict-markers from warnings to errors and
wire validate_env_profiles into the three *_workspace_test PR gates.
Depends on step 4 (today it would fail on 42 warnings across ag+al).
Acceptance: a PR adding a JAX-marked script outside the derived set fails its
own gate.

Step 6 — rename the profiles

env_vars.yaml → profile_smoke.yaml, env_vars_release.yaml →
profile_release.yaml. Update autohands/env_config.py discovery,
run_python.py, run_smoke.py, workspace-validation.yml, docs. Wide but
mechanical; do it last, and re-sweep for the old names after every merge
from main — git rename-detection is blind to new files using the old name,
which is how the autobuild→autohands rename leaked.

Step 7 — library one-reader fold

al.AnalysisDataset reads PYAUTO_DISABLE_JAX a second time, before
super().__init__() (design doc failure mode 8). Fold it into base
af.Analysis so there is exactly one reader. Own plan; PyAutoFit + PyAutoLens,
so library-first then workspace. Behaviour-preserving — prove it with a unit
test asserting both construction orders resolve identically.

Step 8 — use_jax: Optional[bool] = None sentinel (HUMAN-GATED)

Only if the human wants "explicitly requested" distinguishable from
"defaulted". Cross-repo API change; may be declined — ask before building.

Suggested phasing

Step 4 is a session on its own (the triage is the work). Steps 5+6 can pair.
Step 7 is independent of 4-6 and could go first if a session is short. Step 8
needs a decision before any code.

[Jammy2211]

Library PR created (Phase A): #182 — the derive_jax_markers derivation rule + single apply_profile resolution core + validator updates. 157 tests pass; behaviour-preserving until a profile adopts the key (live-repo validator identical minus the 3 __init__.py warnings). Shipped to PR-open under human RED acknowledgement (3 pre-existing reasons, recorded on the PR). Next: /start_workspace for Phases B (autolens_workspace_test: renames + scrape triage + profile) and C (autogalaxy_workspace_test: profile) — merge #182 first.

[Jammy2211]

Step 4 shipped to PR-open — 2026-07-23

PRs (merge in this order)

1. feat: JAX-marker derivation in the env resolver (#161 step 4) #182 — the derivation rule (derive_jax_markers, single apply_profile resolver, validator updates). 157 tests pass.
2. feat: derive the release JAX set from name markers (#181 step 4) autolens_workspace_test#203 — Phase B: 10 overrides deleted, 6 renames (2 mid-stem _jit_ files; 4 explicit-use_jax=True scrape scripts → *_jax). Smoke 17/17.
3. feat: derive the release JAX set from name markers (#181 step 4) autogalaxy_workspace_test#85 — Phase C: 8 overrides deleted. Smoke 26/26.

Verification (binding, per plan)

Resolved-env diff over every script (154 al + 71 ag), old profile vs new: exactly the triaged set — al 14 changes (11 turn ON, 3 scrape turn OFF), ag 4 (jax_assertions ON), only PYAUTO_DISABLE_JAX, renamed scripts resolve identically. All 13 runnable turns-ON scripts pass locally under their exact new JAX-on release env (2–12 s each — no no_run additions needed); both runnable turns-OFF scrape scripts pass under NumPy. Validator: 0 errors 0 warnings in both repos including under --strict-derivation --strict-markers — step 5's gate flip is unblocked. CI-env parity checked (workflow injects only PYAUTO_SKIP_WORKSPACE_VERSION_CHECK, which the profiles pin themselves).

Triage findings worth recording

• The scrape folder's JAX override existed for the four no_run'd scripts (explicit use_jax=True in code), not the two that actually run — general.py/scaling_relation.py have no JAX usage and are now validated NumPy, per policy.
• Both mid-stem _jit_ files were already no_run'd (SLOW, 2026-05-07), so the turns-OFF risk flagged in the prompt had no live nightly impact — but the renames close the trap for any future un-skip.
• Full effect lands only via the nightly mega-run (release path is not in the per-PR gate) — mega-run verification accepted and stated on each PR.

Remaining on this issue

Step 5 (flip validator strict flags to errors + wire into the 3 PR gates — now unblocked), step 6 (profile rename, last), step 7 (library one-reader fold; queued behind testmode-env-drift's PyAutoFit claim), step 8 (use_jax sentinel — human decision first, may be declined).

Heart RED (3 pre-existing reasons) acknowledged by human for all three PR-opens; recorded on each PR body.

[Jammy2211]

Step 4 merged + step 5 at PR-open — 2026-07-23

Step 4 merged (order respected): PyAutoHands#182 → autolens_workspace_test#203 + autogalaxy_workspace_test#85. Post-merge validator on live main: all three _test repos 0 errors / 0 warnings under --strict-derivation --strict-markers.

Step 5 shipped: PyAutoLabs/PyAutoHeart#100 — the strict validation step lands centrally in Heart's reusable smoke workflow (all nine calling repos gate on merge; guarded on env_vars_release.yaml presence so user workspaces/HowTo* skip). Acceptance verified: an injected PYAUTO_DISABLE_JAX override fails the gate with exit 1.

Step 6 re-scoped — bigger than the brief. The live-reference sweep for env_vars.yaml/env_vars_release.yaml spans ~14 repos / 40+ files: user workspaces also carry env_vars_release.yaml; all ten workspace repos vendor a run_smoke.py with the hardcoded name; plus autohands modules+tests, Heart workflow+docs+test, AGENTS/README files, Brain agent docs. This is ~10 mechanical-rename PRs with the wide-rename re-sweep traps — deferring to its own session (needs a discovery-with-fallback shape: prefer profile_smoke.yaml, fall back to env_vars.yaml during the migration window).

Remaining: step 6 (own session, scoped above), step 7 (one-reader fold — unblocks when testmode-env-drift releases its PyAutoFit claim), step 8 (use_jax sentinel — human decision first).

[Jammy2211]

Step 6 stages 1+2 COMPLETE — 2026-07-23

All 14 rename PRs merged today: stage 1 (dual-name tooling) PyAutoHands#183 + PyAutoHeart#101; stage 2 (the renames) autofit_workspace#113, autogalaxy_workspace#149, autolens_workspace#323, autofit_workspace_test#69, autogalaxy_workspace_test#86, autolens_workspace_test#204, HowToFit#30, HowToGalaxy#40, HowToLens#50, autocti_workspace_test#7, euclid_strong_lens_modeling_pipeline#34, PyAutoBrain#154.

Post-merge verification on live mains: branch + main smoke CI green in all 10 CI-carrying repos; strict validator 0/0 in all three _test repos; renamed profiles resolve byte-identically to the pre-rename files over all 288 _test scripts; global sweep shows only the intentional run_smoke.py legacy-fallback lines and dated migration notes.

Remaining — stage 3 (follow-up): drop the legacy env_vars*.yaml fallbacks from autohands (PROFILE_NAMES), the 10 vendored run_smoke.py copies, and Heart's workflows; make the validator error on legacy filenames so the old names cannot creep back. One small PR per repo, any time — the dual-name window is stable meanwhile.

Then: step 7 (one-reader fold — unblocks when testmode-env-drift releases PyAutoFit), step 8 (use_jax sentinel — human decision, may be declined).

[Jammy2211]

Steps 6 (complete) + 7 MERGED — 2026-07-23

Step 6 stage 3 (12 PRs): legacy env_vars*.yaml names are dead — canonical-only discovery, validator errors on a legacy-named file, workflow fallbacks and all 10 vendored run_smoke.py fallbacks removed. Step 7 (PyAutoLens#643): PYAUTO_DISABLE_JAX has exactly one reader (af.Analysis.__init__); AnalysisLens receives the base-resolved value at all three call sites. This fixed a live bug — AnalysisPoint and the weak analysis silently ignored the env downgrade (red-green proven). Nightly impact: exactly one script flips to NumPy (point_source/visualization.py — verified passing, 8 s); user workspaces pin DISABLE_JAX: "0" and are untouched.

Post-merge: branch CI green everywhere pre-merge and all main runs concluded success (11 repos, incl. PyAutoLens Tests+Docs). Live-main sweeps: zero legacy-name references; zero PYAUTO_DISABLE_JAX readers in autolens.

Remaining on this issue: step 8 only — the use_jax: Optional[bool] = None sentinel. Human decision; may be declined. Tonight's mega-run is the standing final verification of the release surface for everything merged today.

[Jammy2211]

Step 8 declined (human decision, 2026-07-23). The use_jax: Optional[bool] = None sentinel will not be built: after steps 4–7 the release surface is structurally protected (marker derivation + strict PR-gate validator + single env reader), CI scrubs ambient env, and the residual exposure (explicit use_jax=True overridden by a user's own shell var) is narrow and arguably intended. Recorded here so the option isn't re-derived later; reopen only if author-intent-vs-env loudness becomes a real need.

All other steps complete — closing. Steps 1–3 (validator, single resolver, scrubbed baseline), 4 (derivation), 5 (strict gate), 6 (rename + legacy names now errors), 7 (one-reader fold + AnalysisPoint/weak bug fix) are merged and verified; completion record lives in PyAutoMind.