diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 3a985a1..1af30c3 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -2,8 +2,8 @@ # # Cheap, fast (<5 min), CPU-only. Runs on every PR and on push to main. # Does NOT produce real result artifacts — the smoke step short-circuits -# every profile script via AUTOLENS_PROFILING_SMOKE=1 (lands as part of -# Phase 5; see scripts under likelihood/, simulators/, searches/nautilus/). +# every profile script via AUTOLENS_PROFILING_SMOKE=1 (see scripts under +# likelihood_runtime/, likelihood_breakdown/, simulators/, searches/nautilus/). # # Profile re-runs that actually produce results live in `profile.yml`. @@ -60,12 +60,15 @@ jobs: - name: Install runtime dependencies run: | - python -m pip install --upgrade pip - pip install numpy scipy matplotlib jax jaxlib - pip install nautilus-sampler==1.0.5 - pip install numba - # autoconf / autofit / autoarray / autogalaxy / autolens are imported - # from the sibling checkouts via PYTHONPATH below — no install needed. + python -m pip install --upgrade pip setuptools wheel + # Install each sibling checkout in dependency order (the PyAutoHeart + # lib-tests.yml pattern): satisfies the full third-party dep chain + # (pyyaml, astropy, jax, numba, ...) without hand-maintaining a + # package list. PYTHONPATH below still resolves the PyAuto* imports + # from source, so the checkouts stay authoritative. + for r in PyAutoConf PyAutoFit PyAutoArray PyAutoGalaxy PyAutoLens; do + pip install "./$r[optional]" + done - name: Install lint tools run: | @@ -92,9 +95,11 @@ jobs: working-directory: autolens_profiling run: | # Conservative scope: only README.md files in this repo, exclude - # GitHub UI links that lychee tends to false-positive on. + # GitHub UI links that lychee tends to false-positive on, and the + # Slack invite (302 -> 403 to any non-browser client). lychee \ --exclude '^https://github\.com/.*/(issues|pull|commit|releases)/' \ + --exclude '^https://join\.slack\.com/' \ --no-progress \ --accept '200..=299,429' \ --max-redirects 5 \ @@ -111,6 +116,7 @@ jobs: # Each script reads AUTOLENS_PROFILING_SMOKE at module top and # exits 0 immediately after the import + setup section. Catches # import-graph breakage without running the full profile. - python likelihood/imaging/mge.py + python likelihood_runtime/imaging/mge.py + python likelihood_breakdown/imaging/mge.py python simulators/imaging.py - python searches/nautilus/simple.py + python searches/nautilus/imaging/mge.py diff --git a/.github/workflows/profile.yml b/.github/workflows/profile.yml index 45b7006..99e1dfa 100644 --- a/.github/workflows/profile.yml +++ b/.github/workflows/profile.yml @@ -92,7 +92,7 @@ jobs: pip install nautilus-sampler==1.0.5 pip install numba - - name: Run likelihood/ scripts + - name: Run likelihood_runtime/ scripts if: ${{ inputs.sections == '' || contains(inputs.sections, 'likelihood') }} working-directory: autolens_profiling continue-on-error: true @@ -102,8 +102,8 @@ jobs: PYTHONPATH: ${{ github.workspace }}/PyAutoConf:${{ github.workspace }}/PyAutoFit:${{ github.workspace }}/PyAutoArray:${{ github.workspace }}/PyAutoGalaxy:${{ github.workspace }}/PyAutoLens run: | set +e # don't abort the whole step on one script's failure - for script in likelihood/imaging/*.py likelihood/interferometer/*.py \ - likelihood/point_source/*.py likelihood/datacube/delaunay.py; do + for script in likelihood_runtime/imaging/*.py likelihood_runtime/interferometer/*.py \ + likelihood_runtime/point_source/*.py likelihood_runtime/datacube/delaunay.py; do echo "::group::$script" python "$script" || echo "::warning::$script failed — see log; dashboard cell will show ERR" echo "::endgroup::" @@ -163,7 +163,7 @@ jobs: # gitignored; future hardware-tier work may flip a subset to # tracked, in which case this `add -f` becomes redundant but # harmless. - git add -f results/ likelihood/README.md likelihood/*/README.md \ + git add -f results/ likelihood_runtime/README.md likelihood_breakdown/README.md \ simulators/README.md searches/README.md \ searches/nautilus/README.md README.md git commit -m "chore: refresh profile artifacts + dashboard tables [skip ci] diff --git a/.gitignore b/.gitignore index 5a9d230..6e3c4f2 100644 --- a/.gitignore +++ b/.gitignore @@ -46,12 +46,10 @@ dataset/datacube/sim_*/ # that's purely an auto-cache. dataset/**/lensed_source.fits -# Profiling result artifacts — written by scripts under likelihood/, simulators/, searches/. -# Each run produces fresh versioned JSON+PNG; Phase 4's dashboard will explicitly select -# which artifacts to commit. Until then, results stay local. +# Retired legacy results section (pre likelihood_runtime/breakdown split); +# nothing writes here since the phase-2 output retarget. Kept ignored so old +# local checkouts don't surface stale artifacts as untracked. results/likelihood/ -results/simulators/ -results/searches/ -# Latent profiling sweep outputs — regenerated by latent/sweep.py, not committed. +# Retired legacy latent output home (latent now writes results/latent/, in-repo). latent/results/ diff --git a/_profile_cli.py b/_profile_cli.py index e23e0bb..934c6bb 100644 --- a/_profile_cli.py +++ b/_profile_cli.py @@ -257,3 +257,50 @@ def auto_simulate_if_missing( ], check=True, ) + + +def check_pinned(got, expected, *, label: str, rtol: float = 1e-4): + """Compare a computed likelihood/evidence against its pinned baseline. + + Profiling runs **record and flag** drift; they never adjudicate library + correctness (that is autolens_workspace_test's remit — see + ``results/notes/design_lock_in.md``). Returns ``None`` when ``got`` is + within ``rtol`` of ``expected``; otherwise prints a loud warning and + returns a drift record for the result JSON, which PyAutoHeart's vitals + scan picks up. Never raises — a changed computation must not kill a + profiling job (the timing numbers are still data; the flag marks them + non-comparable to the pinned baseline). + """ + got_f = float(got) + rel = abs(got_f - expected) / max(abs(expected), 1e-300) + if rel <= rtol: + return None + print( + f" WARNING: PINNED-VALUE DRIFT [{label}] — got {got_f!r}, " + f"pinned {expected!r} (rel diff {rel:.3e} > rtol {rtol:g}). " + f"Timings from this run are NOT comparable to the pinned baseline; " + f"file a bug / check autolens_workspace_test before trusting trends." + ) + return { + "label": label, + "expected": expected, + "got": got_f, + "rel_diff": rel, + "rtol": rtol, + } + + +def record_pinned_check(json_path, expected, drift_records) -> None: + """Merge the pinned-value check outcome into an already-written result JSON. + + Adds ``pinned_expected`` (the baseline value, or ``None`` when the + instrument has no pin) and ``pinned_drift`` (list of drift records — + empty means every compared value matched). The guard blocks run after + the summary JSON is written, so this rewrites the file in place. + """ + import json + + data = json.loads(json_path.read_text()) + data["pinned_expected"] = expected + data["pinned_drift"] = drift_records + json_path.write_text(json.dumps(data, indent=2)) diff --git a/dataset/imaging/euclid/data.fits b/dataset/imaging/euclid/data.fits new file mode 100644 index 0000000..474839a Binary files /dev/null and b/dataset/imaging/euclid/data.fits differ diff --git a/dataset/imaging/euclid/noise_map.fits b/dataset/imaging/euclid/noise_map.fits new file mode 100644 index 0000000..c8fa54c Binary files /dev/null and b/dataset/imaging/euclid/noise_map.fits differ diff --git a/dataset/imaging/euclid/positions.json b/dataset/imaging/euclid/positions.json new file mode 100644 index 0000000..81705cc --- /dev/null +++ b/dataset/imaging/euclid/positions.json @@ -0,0 +1,28 @@ +{ + "type": "instance", + "class_path": "autoarray.structures.grids.irregular_2d.Grid2DIrregular", + "arguments": { + "values": { + "type": "ndarray", + "array": [ + [ + 0.27109374999999997, + 1.6729626159565016 + ], + [ + -1.49921875, + 0.1953067707493031 + ], + [ + 1.49921875, + -0.1953067707493031 + ], + [ + -0.27109374999999997, + -1.6729626159565019 + ] + ], + "dtype": "float64" + } + } +} \ No newline at end of file diff --git a/dataset/imaging/euclid/psf.fits b/dataset/imaging/euclid/psf.fits new file mode 100644 index 0000000..1f7babf Binary files /dev/null and b/dataset/imaging/euclid/psf.fits differ diff --git a/dataset/imaging/euclid/tracer.json b/dataset/imaging/euclid/tracer.json new file mode 100644 index 0000000..d218053 --- /dev/null +++ b/dataset/imaging/euclid/tracer.json @@ -0,0 +1,110 @@ +{ + "type": "instance", + "class_path": "autolens.lens.tracer.Tracer", + "arguments": { + "galaxies": { + "type": "list", + "values": [ + { + "type": "instance", + "class_path": "autogalaxy.galaxy.galaxy.Galaxy", + "arguments": { + "redshift": 0.5, + "label": null, + "bulge": { + "type": "instance", + "class_path": "autogalaxy.profiles.light.standard.sersic.Sersic", + "arguments": { + "centre": { + "type": "tuple", + "values": [ + 0.0, + 0.0 + ] + }, + "sersic_index": 3.0, + "intensity": 2.0, + "ell_comps": { + "type": "tuple", + "values": [ + 0.05263157894736841, + 3.2227547345982974e-18 + ] + }, + "effective_radius": 0.6 + } + }, + "mass": { + "type": "instance", + "class_path": "autogalaxy.profiles.mass.total.isothermal.Isothermal", + "arguments": { + "einstein_radius": 1.6, + "ell_comps": { + "type": "tuple", + "values": [ + 0.05263157894736841, + 3.2227547345982974e-18 + ] + }, + "centre": { + "type": "tuple", + "values": [ + 0.0, + 0.0 + ] + } + } + }, + "shear": { + "type": "instance", + "class_path": "autogalaxy.profiles.mass.sheets.external_shear.ExternalShear", + "arguments": { + "gamma_1": 0.05, + "gamma_2": 0.05 + } + } + } + }, + { + "type": "instance", + "class_path": "autogalaxy.galaxy.galaxy.Galaxy", + "arguments": { + "redshift": 1.0, + "label": null, + "bulge": { + "type": "instance", + "class_path": "autogalaxy.profiles.light.standard.sersic_core.SersicCore", + "arguments": { + "gamma": 0.25, + "centre": { + "type": "tuple", + "values": [ + 0.0, + 0.0 + ] + }, + "radius_break": 0.025, + "sersic_index": 1.0, + "intensity": 4.0, + "alpha": 3.0, + "ell_comps": { + "type": "tuple", + "values": [ + 0.0962250448649376, + -0.05555555555555551 + ] + }, + "effective_radius": 0.1 + } + } + } + } + ] + }, + "cosmology": { + "type": "instance", + "class_path": "autogalaxy.cosmology.model.Planck15", + "arguments": {} + } + } +} \ No newline at end of file diff --git a/hpc/batch_gpu/submit_probe_delaunay_a100 b/hpc/batch_gpu/submit_probe_delaunay_a100 new file mode 100644 index 0000000..ed5e64c --- /dev/null +++ b/hpc/batch_gpu/submit_probe_delaunay_a100 @@ -0,0 +1,55 @@ +#!/bin/bash -l +# +# Probe-only sweep — DELAUNAY cells (10-30 min per XLA compile; scipy +# Delaunay via pure_callback bloats the graph, so probes are single-point +# per the vram/README.md methodology). +# +# Re-validates the vram/config.py batch-size table for the +# PreOptimizationTimes campaign (autolens_profiling#54). Each probe writes +# results/runtime//delaunay/vmap_probe_delaunay[_sparse].json +# and exits before any timing loop. See submit_probe_fast_a100 for the +# fast cells and vram/README.md "Probe-only submits" for the ingest step. + +#SBATCH -J probe_delaunay_a100 +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64gb +#SBATCH --time=8:00:00 +#SBATCH -o output/output.%A.out +#SBATCH -e error/error.%A.err +#SBATCH --mail-type=END,FAIL +#SBATCH --mail-user=james.w.nightingale@durham.ac.uk + +export AP_ROOT=/mnt/ral/jnightin/autolens_profiling +source $AP_ROOT/activate.sh + +export JAX_PLATFORM_NAME=cuda +export JAX_PLATFORMS=cuda,cpu +export XLA_PYTHON_CLIENT_PREALLOCATE=false +export JAX_ENABLE_X64=True +export NUMBA_CACHE_DIR=/tmp/numba_cache +export MPLCONFIGDIR=/tmp/matplotlib + +nvidia-smi +cd $AP_ROOT + +run_probe () { + echo "=== probe: $1 --instrument $2 $3 ===" + date + python3 "$1" --instrument "$2" --vmap-probe $3 || echo "!!! probe FAILED: $1 $2 $3" +} + +# Imaging delaunay — campaign instruments, dense + sparse. +for inst in ao jwst hst; do + run_probe likelihood_runtime/imaging/delaunay.py "$inst" "" + run_probe likelihood_runtime/imaging/delaunay.py "$inst" "--sparse" +done + +# Interferometer delaunay — campaign anchor (alma_high) + the cheap sanity cell. +run_probe likelihood_runtime/interferometer/delaunay.py sma "" +run_probe likelihood_runtime/interferometer/delaunay.py alma_high "" + +echo "Finished." +date diff --git a/hpc/batch_gpu/submit_probe_fast_a100 b/hpc/batch_gpu/submit_probe_fast_a100 new file mode 100644 index 0000000..4a5251e --- /dev/null +++ b/hpc/batch_gpu/submit_probe_fast_a100 @@ -0,0 +1,61 @@ +#!/bin/bash -l +# +# Probe-only sweep — FAST cells (mge / pixelization; ~10 s per XLA compile). +# +# Re-validates the vram/config.py batch-size table for the +# PreOptimizationTimes campaign (autolens_profiling#54): every campaign +# vmap cell is re-probed against current source, dense + sparse imaging +# variants. Each probe writes +# results/runtime///vmap_probe_[_sparse].json +# and exits before any timing loop (--vmap-probe). Runtime cells then +# prefer these fresh JSONs over the static table automatically +# (vram.resolve_vmap_batch), and the table rows get updated from them at +# ingest — see vram/README.md "Probe-only submits". +# +# Pair with submit_probe_delaunay_a100 (slow cells, separate job). + +#SBATCH -J probe_fast_a100 +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64gb +#SBATCH --time=1:00:00 +#SBATCH -o output/output.%A.out +#SBATCH -e error/error.%A.err +#SBATCH --mail-type=END,FAIL +#SBATCH --mail-user=james.w.nightingale@durham.ac.uk + +export AP_ROOT=/mnt/ral/jnightin/autolens_profiling +source $AP_ROOT/activate.sh + +export JAX_PLATFORM_NAME=cuda +export JAX_PLATFORMS=cuda,cpu +export XLA_PYTHON_CLIENT_PREALLOCATE=false +export JAX_ENABLE_X64=True +export NUMBA_CACHE_DIR=/tmp/numba_cache +export MPLCONFIGDIR=/tmp/matplotlib + +nvidia-smi +cd $AP_ROOT + +run_probe () { + echo "=== probe: $1 --instrument $2 $3 ===" + date + python3 "$1" --instrument "$2" --vmap-probe $3 || echo "!!! probe FAILED: $1 $2 $3" +} + +# Imaging mge / pixelization — campaign instruments, dense + sparse. +for inst in ao jwst hst; do + run_probe likelihood_runtime/imaging/mge.py "$inst" "" + run_probe likelihood_runtime/imaging/mge.py "$inst" "--sparse" + run_probe likelihood_runtime/imaging/pixelization.py "$inst" "" + run_probe likelihood_runtime/imaging/pixelization.py "$inst" "--sparse" +done + +# Interferometer mge / pixelization — sma only (alma+ intentionally blocked). +run_probe likelihood_runtime/interferometer/mge.py sma "" +run_probe likelihood_runtime/interferometer/pixelization.py sma "" + +echo "Finished." +date diff --git a/likelihood_runtime/datacube/delaunay.py b/likelihood_runtime/datacube/delaunay.py index b3061c1..fb1907e 100644 --- a/likelihood_runtime/datacube/delaunay.py +++ b/likelihood_runtime/datacube/delaunay.py @@ -120,14 +120,16 @@ # Sweep-driver CLI args (--config-name / --output-dir / --use-mixed-precision). # Tolerates extra/unknown args via parse_known_args inside the helper. -from _profile_cli import ( # noqa: E402 +from _profile_cli import ( auto_simulate_if_missing, + check_pinned, device_info_dict, parse_profile_cli, + record_pinned_check, # noqa: E402 resolve_output_paths, ) from simulators.interferometer import INSTRUMENTS # noqa: E402 -from vram import ProbeResult, vmap_batch_for, write_probe_json # noqa: E402 +from vram import ProbeResult, write_probe_json # noqa: E402 _cli = parse_profile_cli() @@ -405,7 +407,7 @@ def _build_transformer(uv_wavelengths, real_space_mask): if _cli.vmap_probe: probe_path = ( _cli.output_dir or (_workspace_root / "results" / "runtime" / "datacube" / "delaunay") - ) / "vmap_probe.json" + ) / "vmap_probe_delaunay.json" probe_path.parent.mkdir(parents=True, exist_ok=True) import json @@ -595,6 +597,9 @@ def full_cube_pipeline_from_params(params_tree): # Identical channels = exact N × single-channel log-evidence (for "sma"). # For "hannah" the per-channel literal isn't pinned yet, so the assertion is # skipped until the value below is filled in from a clean run. +_pinned_drift: list = [] +_pinned_expected = None + EXPECTED_LOG_EVIDENCE_PER_CHANNEL = { "sma": None, "alma": None, @@ -603,6 +608,7 @@ def full_cube_pipeline_from_params(params_tree): _per_channel = EXPECTED_LOG_EVIDENCE_PER_CHANNEL.get(instrument) expected_cube_log_evidence = n_channels * _per_channel if _per_channel is not None else None +_pinned_expected = expected_cube_log_evidence if expected_cube_log_evidence is None: print( @@ -613,24 +619,19 @@ def full_cube_pipeline_from_params(params_tree): f"EXPECTED_LOG_EVIDENCE_PER_CHANNEL[{instrument!r}]." ) else: - np.testing.assert_allclose( - cube_log_evidence_ref, - expected_cube_log_evidence, - rtol=1e-4, - err_msg=( - f"datacube/delaunay[{instrument}]: regression — eager cube log_evidence " - f"drifted (got {cube_log_evidence_ref}, expected {expected_cube_log_evidence})" - ), - ) - print( - f"\n Eager cube regression assertion PASSED: log_evidence matches " - f"{expected_cube_log_evidence:.6f}" - ) + _rec = check_pinned(cube_log_evidence_ref, _pinned_expected, label="eager", rtol=1e-4) + if _rec is not None: + _pinned_drift.append(_rec) if full_cube_result is not None: - np.testing.assert_allclose( - float(full_cube_result), - expected_cube_log_evidence, - rtol=1e-3, - err_msg=f"datacube/delaunay[{instrument}]: regression — full cube log_evidence drifted", - ) - print(" Full-pipeline cube regression assertion PASSED") + _rec = check_pinned(float(full_cube_result), _pinned_expected, label="full", rtol=1e-3) + if _rec is not None: + _pinned_drift.append(_rec) + + +# Pinned-value outcome -> result JSON: profiling records and flags drift, +# never adjudicates library correctness (autolens_workspace_test's remit; +# boundary rule in results/notes/design_lock_in.md). PyAutoHeart's vitals +# scan reads the pinned_drift field. +record_pinned_check(dict_path, _pinned_expected, _pinned_drift) +if _pinned_expected is not None and not _pinned_drift: + print(" Pinned-value check PASSED (recorded in result JSON).") diff --git a/likelihood_runtime/imaging/delaunay.py b/likelihood_runtime/imaging/delaunay.py index 9d78abd..ef762ee 100644 --- a/likelihood_runtime/imaging/delaunay.py +++ b/likelihood_runtime/imaging/delaunay.py @@ -73,17 +73,19 @@ # Sweep-driver CLI args (--config-name / --output-dir / --use-mixed-precision). # Tolerates extra/unknown args via parse_known_args inside the helper. -from _profile_cli import ( # noqa: E402 +from _profile_cli import ( auto_simulate_if_missing, + check_pinned, device_info_dict, parse_profile_cli, + record_pinned_check, # noqa: E402 resolve_output_paths, ) from simulators.imaging import INSTRUMENTS # noqa: E402 from vram import ( # noqa: E402 probe_vmap_memory, recommend_batch_size, - vmap_batch_for, + resolve_vmap_batch, write_probe_json, ) @@ -470,13 +472,25 @@ def full_pipeline_from_params(params_tree): print("\n--- vmap batched evaluation ---") -batch_size = vmap_batch_for("imaging", "delaunay", instrument) or 3 +_batch_resolved, _batch_source = resolve_vmap_batch( + "imaging", + "delaunay", + instrument, + output_dir=_cli.output_dir + or (_workspace_root / "results" / "runtime" / "imaging" / "delaunay"), + path="sparse" if _cli.use_sparse_operator else "dense", + backend=jax.default_backend(), +) +print(f" vmap batch_size: {_batch_resolved} (source: {_batch_source})") +batch_size = _batch_resolved or 3 -# Skip vmap if vmap_batch_for explicitly returns None for this cell. -_vmap_skipped = vmap_batch_for("imaging", "delaunay", instrument) is None +# Skip vmap if batch resolution (probe JSON / table) returned None for this cell. +_vmap_skipped = _batch_resolved is None if _vmap_skipped: - print(" SKIPPED: vmap_batch_for() returned None for this (cell, instrument).") + print( + f" SKIPPED: batch resolution returned None for this (cell, instrument) — source: {_batch_source}." + ) vmap_per_call = None vmap_speedup = None vmap_batch_time = None @@ -585,7 +599,7 @@ def full_pipeline_from_params(params_tree): "inversion_path": "sparse" if _cli.use_sparse_operator else "dense", }, "full_pipeline_single_jit": full_pipeline_per_call, - "vmap": "SKIPPED — vmap_batch_for() returned None for this (cell, instrument)" + "vmap": "SKIPPED — batch resolution returned None for this (cell, instrument)" if _vmap_skipped else { "batch_size": batch_size, @@ -615,31 +629,45 @@ def full_pipeline_from_params(params_tree): # matches imaging/pixelization (adaptive meshes amplify fp drift through # Cholesky / log_det). vmap result asserted only when DELAUNAY_VMAP=1 # (vmap compile takes 20+ min). -EXPECTED_LOG_EVIDENCE_HST = ( - 29110.92085793 # 1500-pixel Hilbert/Delaunay, MGE-60 lens, adapt_image=lensed_source -) +# Pinned empirically per instrument (mirrors interferometer cells); missing +# means "skip the assertion and print the value so it can be pasted in here +# on a clean run". The old module-level EXPECTED_LOG_EVIDENCE_HST constant +# was asserted for every --instrument, which failed any non-hst run by +# construction (phase-2 local validation sweep, autolens_profiling#54). +# The eager ≡ JIT ≡ vmap cross-consistency assertions still guard every +# instrument. +_pinned_drift: list = [] +_pinned_expected = None + +EXPECTED_LOG_EVIDENCE = { + "hst": 29110.92085793 # 1500-pixel Hilbert/Delaunay, MGE-60 lens, adapt_image=lensed_source +} -np.testing.assert_allclose( - log_evidence_ref, - EXPECTED_LOG_EVIDENCE_HST, - rtol=1e-4, - err_msg=( - f"imaging/delaunay[{instrument}]: regression — eager log_evidence drifted " - f"(got {log_evidence_ref}, expected {EXPECTED_LOG_EVIDENCE_HST})" - ), -) -print(f" Eager regression assertion PASSED: log_evidence matches {EXPECTED_LOG_EVIDENCE_HST:.6f}") -np.testing.assert_allclose( - float(full_result), - EXPECTED_LOG_EVIDENCE_HST, - rtol=1e-3, - err_msg=f"imaging/delaunay[{instrument}]: regression — full log_evidence drifted", -) -if not _vmap_skipped: - np.testing.assert_allclose( - np.array(result_vmap), - EXPECTED_LOG_EVIDENCE_HST, - rtol=1e-3, - err_msg=f"imaging/delaunay[{instrument}]: regression — vmap log_evidence drifted", +expected_log_evidence = EXPECTED_LOG_EVIDENCE.get(instrument) +_pinned_expected = expected_log_evidence + +if expected_log_evidence is None: + print( + f" Regression assertion SKIPPED for {instrument} " + f"(no pinned value). Eager log_evidence = {log_evidence_ref}" ) -print(f" Regression assertion PASSED: log_evidence matches {EXPECTED_LOG_EVIDENCE_HST:.6f}") +else: + _rec = check_pinned(log_evidence_ref, _pinned_expected, label="eager", rtol=1e-4) + if _rec is not None: + _pinned_drift.append(_rec) + _rec = check_pinned(float(full_result), _pinned_expected, label="full", rtol=1e-3) + if _rec is not None: + _pinned_drift.append(_rec) + if not _vmap_skipped: + _rec = check_pinned(np.array(result_vmap), _pinned_expected, label="vmap", rtol=1e-3) + if _rec is not None: + _pinned_drift.append(_rec) + + +# Pinned-value outcome -> result JSON: profiling records and flags drift, +# never adjudicates library correctness (autolens_workspace_test's remit; +# boundary rule in results/notes/design_lock_in.md). PyAutoHeart's vitals +# scan reads the pinned_drift field. +record_pinned_check(dict_path, _pinned_expected, _pinned_drift) +if _pinned_expected is not None and not _pinned_drift: + print(" Pinned-value check PASSED (recorded in result JSON).") diff --git a/likelihood_runtime/imaging/mge.py b/likelihood_runtime/imaging/mge.py index 9d6c519..9464691 100644 --- a/likelihood_runtime/imaging/mge.py +++ b/likelihood_runtime/imaging/mge.py @@ -80,17 +80,19 @@ # Sweep-driver CLI args (--config-name / --output-dir / --use-mixed-precision). # Tolerates extra/unknown args via parse_known_args inside the helper. sys.path.insert(0, str(Path(__file__).resolve().parents[2])) -from _profile_cli import ( # noqa: E402 +from _profile_cli import ( auto_simulate_if_missing, + check_pinned, device_info_dict, parse_profile_cli, + record_pinned_check, # noqa: E402 resolve_output_paths, ) from simulators.imaging import INSTRUMENTS # noqa: E402 from vram import ( # noqa: E402 probe_vmap_memory, recommend_batch_size, - vmap_batch_for, + resolve_vmap_batch, write_probe_json, ) @@ -437,7 +439,16 @@ def full_pipeline_from_params(params_tree): print("\n--- vmap batched evaluation ---") -batch_size = vmap_batch_for("imaging", "mge", instrument) or 3 +_batch_resolved, _batch_source = resolve_vmap_batch( + "imaging", + "mge", + instrument, + output_dir=_cli.output_dir or (_workspace_root / "results" / "runtime" / "imaging" / "mge"), + path="sparse" if _cli.use_sparse_operator else "dense", + backend=jax.default_backend(), +) +print(f" vmap batch_size: {_batch_resolved} (source: {_batch_source})") +batch_size = _batch_resolved or 3 # Build the batched pytree: every leaf gets a fresh leading batch axis. No # flat-vector reshaping required — JAX walks the pytree via the registration @@ -568,30 +579,44 @@ def full_pipeline_from_params(params_tree): # default centre/ell_comps priors) put the evaluation point at the # physically-meaningful truth operating point. Eager, JIT, and vmap all # agree to ~1e-11 precision. -EXPECTED_LOG_LIKELIHOOD_HST = 27379.38890685539 +# Pinned empirically per instrument (mirrors interferometer/mge.py); +# ``None``/missing means "skip the assertion and print the value so it can +# be pasted in here on a clean run". Only hst has a pinned value — the old +# module-level EXPECTED_LOG_LIKELIHOOD_HST constant was asserted for every +# --instrument, which failed any non-hst run by construction (found by the +# phase-2 local validation sweep, autolens_profiling#54). The eager ≡ JIT ≡ +# vmap cross-consistency assertions above still guard every instrument. +_pinned_drift: list = [] +_pinned_expected = None + +EXPECTED_LOG_LIKELIHOOD = { + "hst": 27379.38890685539, +} -np.testing.assert_allclose( - log_likelihood_ref, - EXPECTED_LOG_LIKELIHOOD_HST, - rtol=1e-4, - err_msg=( - f"imaging/mge[{instrument}]: regression — eager log_likelihood drifted " - f"(got {log_likelihood_ref}, expected {EXPECTED_LOG_LIKELIHOOD_HST})" - ), -) -print( - f" Eager regression assertion PASSED: log_likelihood matches {EXPECTED_LOG_LIKELIHOOD_HST:.6f}" -) -np.testing.assert_allclose( - float(full_result), - EXPECTED_LOG_LIKELIHOOD_HST, - rtol=1e-4, - err_msg=f"imaging/mge[{instrument}]: regression — full log_likelihood drifted", -) -np.testing.assert_allclose( - np.array(result_vmap), - EXPECTED_LOG_LIKELIHOOD_HST, - rtol=1e-4, - err_msg=f"imaging/mge[{instrument}]: regression — vmap log_likelihood drifted", -) -print(f" Regression assertion PASSED: log_likelihood matches {EXPECTED_LOG_LIKELIHOOD_HST:.6f}") +expected_log_likelihood = EXPECTED_LOG_LIKELIHOOD.get(instrument) +_pinned_expected = expected_log_likelihood + +if expected_log_likelihood is None: + print( + f" Regression assertion SKIPPED for {instrument} " + f"(no pinned value). Eager log_likelihood = {log_likelihood_ref}" + ) +else: + _rec = check_pinned(log_likelihood_ref, _pinned_expected, label="eager", rtol=1e-4) + if _rec is not None: + _pinned_drift.append(_rec) + _rec = check_pinned(float(full_result), _pinned_expected, label="full", rtol=1e-4) + if _rec is not None: + _pinned_drift.append(_rec) + _rec = check_pinned(np.array(result_vmap), _pinned_expected, label="vmap", rtol=1e-4) + if _rec is not None: + _pinned_drift.append(_rec) + + +# Pinned-value outcome -> result JSON: profiling records and flags drift, +# never adjudicates library correctness (autolens_workspace_test's remit; +# boundary rule in results/notes/design_lock_in.md). PyAutoHeart's vitals +# scan reads the pinned_drift field. +record_pinned_check(dict_path, _pinned_expected, _pinned_drift) +if _pinned_expected is not None and not _pinned_drift: + print(" Pinned-value check PASSED (recorded in result JSON).") diff --git a/likelihood_runtime/imaging/pixelization.py b/likelihood_runtime/imaging/pixelization.py index 57715b4..a87d3e7 100644 --- a/likelihood_runtime/imaging/pixelization.py +++ b/likelihood_runtime/imaging/pixelization.py @@ -66,17 +66,19 @@ # Sweep-driver CLI args (--config-name / --output-dir / --use-mixed-precision). # Tolerates extra/unknown args via parse_known_args inside the helper. -from _profile_cli import ( # noqa: E402 +from _profile_cli import ( auto_simulate_if_missing, + check_pinned, device_info_dict, parse_profile_cli, + record_pinned_check, # noqa: E402 resolve_output_paths, ) from simulators.imaging import INSTRUMENTS # noqa: E402 from vram import ( # noqa: E402 probe_vmap_memory, recommend_batch_size, - vmap_batch_for, + resolve_vmap_batch, write_probe_json, ) @@ -466,7 +468,17 @@ def full_pipeline_from_params(params_tree): print("\n--- vmap batched evaluation ---") -batch_size = vmap_batch_for("imaging", "pixelization", instrument) or 3 +_batch_resolved, _batch_source = resolve_vmap_batch( + "imaging", + "pixelization", + instrument, + output_dir=_cli.output_dir + or (_workspace_root / "results" / "runtime" / "imaging" / "pixelization"), + path="sparse" if _cli.use_sparse_operator else "dense", + backend=jax.default_backend(), +) +print(f" vmap batch_size: {_batch_resolved} (source: {_batch_source})") +batch_size = _batch_resolved or 3 vmap_batch_time = None vmap_per_call = None vmap_speedup = None @@ -632,31 +644,45 @@ def full_pipeline_from_params(params_tree): # 1581x1581 mapping matrix relative to the non-adaptive baseline (which # previously matched at 1e-4). The 1e-3 envelope is still tight enough to # catch real numerical regressions while accommodating the adaptive path. -EXPECTED_LOG_EVIDENCE_HST = ( - 28370.27770182 # 39x39 = 1521 source pixels, MGE-60 lens light, adapt_image=lensed_source -) +# Pinned empirically per instrument (mirrors interferometer cells); missing +# means "skip the assertion and print the value so it can be pasted in here +# on a clean run". The old module-level EXPECTED_LOG_EVIDENCE_HST constant +# was asserted for every --instrument, which failed any non-hst run by +# construction (phase-2 local validation sweep, autolens_profiling#54). +# The eager ≡ JIT ≡ vmap cross-consistency assertions still guard every +# instrument. +_pinned_drift: list = [] +_pinned_expected = None + +EXPECTED_LOG_EVIDENCE = { + "hst": 28370.27770182 # 39x39 = 1521 source pixels, MGE-60 lens light, adapt_image=lensed_source +} -np.testing.assert_allclose( - log_evidence_ref, - EXPECTED_LOG_EVIDENCE_HST, - rtol=1e-4, - err_msg=( - f"imaging/pixelization[{instrument}]: regression — eager log_evidence drifted " - f"(got {log_evidence_ref}, expected {EXPECTED_LOG_EVIDENCE_HST})" - ), -) -print(f" Eager regression assertion PASSED: log_evidence matches {EXPECTED_LOG_EVIDENCE_HST:.6f}") -np.testing.assert_allclose( - float(full_result), - EXPECTED_LOG_EVIDENCE_HST, - rtol=1e-3, - err_msg=f"imaging/pixelization[{instrument}]: regression — full log_evidence drifted", -) -if result_vmap is not None: - np.testing.assert_allclose( - np.array(result_vmap), - EXPECTED_LOG_EVIDENCE_HST, - rtol=1e-3, - err_msg=f"imaging/pixelization[{instrument}]: regression — vmap log_evidence drifted", +expected_log_evidence = EXPECTED_LOG_EVIDENCE.get(instrument) +_pinned_expected = expected_log_evidence + +if expected_log_evidence is None: + print( + f" Regression assertion SKIPPED for {instrument} " + f"(no pinned value). Eager log_evidence = {log_evidence_ref}" ) -print(f" Regression assertion PASSED: log_evidence matches {EXPECTED_LOG_EVIDENCE_HST:.6f}") +else: + _rec = check_pinned(log_evidence_ref, _pinned_expected, label="eager", rtol=1e-4) + if _rec is not None: + _pinned_drift.append(_rec) + _rec = check_pinned(float(full_result), _pinned_expected, label="full", rtol=1e-3) + if _rec is not None: + _pinned_drift.append(_rec) + if result_vmap is not None: + _rec = check_pinned(np.array(result_vmap), _pinned_expected, label="vmap", rtol=1e-3) + if _rec is not None: + _pinned_drift.append(_rec) + + +# Pinned-value outcome -> result JSON: profiling records and flags drift, +# never adjudicates library correctness (autolens_workspace_test's remit; +# boundary rule in results/notes/design_lock_in.md). PyAutoHeart's vitals +# scan reads the pinned_drift field. +record_pinned_check(dict_path, _pinned_expected, _pinned_drift) +if _pinned_expected is not None and not _pinned_drift: + print(" Pinned-value check PASSED (recorded in result JSON).") diff --git a/likelihood_runtime/interferometer/delaunay.py b/likelihood_runtime/interferometer/delaunay.py index fc108e4..bd2a8a0 100644 --- a/likelihood_runtime/interferometer/delaunay.py +++ b/likelihood_runtime/interferometer/delaunay.py @@ -114,17 +114,19 @@ # Sweep-driver CLI args (--config-name / --output-dir / --use-mixed-precision). # Tolerates extra/unknown args via parse_known_args inside the helper. -from _profile_cli import ( # noqa: E402 +from _profile_cli import ( auto_simulate_if_missing, + check_pinned, device_info_dict, parse_profile_cli, + record_pinned_check, # noqa: E402 resolve_output_paths, ) from simulators.interferometer import INSTRUMENTS # noqa: E402 from vram import ( # noqa: E402 probe_vmap_memory, recommend_batch_size, - vmap_batch_for, + resolve_vmap_batch, write_probe_json, ) @@ -455,7 +457,11 @@ def full_pipeline_from_params(params_tree): recommended = recommend_batch_size(probe) probe_path = ( _cli.output_dir or (_workspace_root / "results" / "runtime" / "interferometer" / "delaunay") - ) / "vmap_probe.json" + ) / ( + "vmap_probe_delaunay_sparse.json" + if _cli.use_sparse_operator + else "vmap_probe_delaunay.json" + ) write_probe_json(probe, recommended, probe_path) print(f"\n vmap_probe samples: {probe.samples}") print(f" per_replica: {probe.per_replica_mb:.1f} MB / replica") @@ -469,10 +475,20 @@ def full_pipeline_from_params(params_tree): print("\n--- vmap batched evaluation ---") -batch_size = vmap_batch_for("interferometer", "delaunay", instrument) or 3 +_batch_resolved, _batch_source = resolve_vmap_batch( + "interferometer", + "delaunay", + instrument, + output_dir=_cli.output_dir + or (_workspace_root / "results" / "runtime" / "interferometer" / "delaunay"), + path="sparse" if _cli.use_sparse_operator else "dense", + backend=jax.default_backend(), +) +print(f" vmap batch_size: {_batch_resolved} (source: {_batch_source})") +batch_size = _batch_resolved or 3 -# Skip vmap if vmap_batch_for explicitly returns None for this cell. -_vmap_skipped = vmap_batch_for("interferometer", "delaunay", instrument) is None +# Skip vmap if batch resolution (probe JSON / table) returned None for this cell. +_vmap_skipped = _batch_resolved is None vmap_batch_time = None vmap_per_call = None @@ -483,7 +499,9 @@ def full_pipeline_from_params(params_tree): _n_leaves = len(jax.tree_util.tree_leaves(params_tree)) if _vmap_skipped: - print(" SKIPPED: vmap_batch_for() returned None for this (cell, instrument).") + print( + f" SKIPPED: batch resolution returned None for this (cell, instrument) — source: {_batch_source}." + ) elif _n_leaves == 0: print( " SKIPPED: model has 0 free parameters (all fixed to truth); " @@ -588,7 +606,7 @@ def full_pipeline_from_params(params_tree): if vmap_per_call is None: if _vmap_skipped: - vmap_payload = "SKIPPED — vmap_batch_for() returned None for this (cell, instrument)" + vmap_payload = "SKIPPED — batch resolution returned None for this (cell, instrument)" else: vmap_payload = "SKIPPED — model has 0 free parameters (all fixed to truth)" else: @@ -647,6 +665,9 @@ def full_pipeline_from_params(params_tree): # print the value so it can be pasted in here on a clean run". sma was # bumped to mask_radius=3.5 in 2026-05-21's INSTRUMENTS refactor — the # old mask_radius=3.0 value no longer applies and needs re-measuring. +_pinned_drift: list = [] +_pinned_expected = None + EXPECTED_LOG_EVIDENCE = { "sma": None, "alma": None, @@ -654,6 +675,7 @@ def full_pipeline_from_params(params_tree): } expected_log_evidence = EXPECTED_LOG_EVIDENCE.get(instrument) +_pinned_expected = expected_log_evidence if expected_log_evidence is None: print( @@ -662,28 +684,22 @@ def full_pipeline_from_params(params_tree): f"and paste it into EXPECTED_LOG_EVIDENCE[{instrument!r}]." ) else: - np.testing.assert_allclose( - figure_of_merit_ref, - expected_log_evidence, - rtol=1e-4, - err_msg=( - f"interferometer/delaunay[{instrument}]: regression — eager log_evidence " - f"drifted (got {figure_of_merit_ref}, expected {expected_log_evidence})" - ), - ) - print(f" Eager regression assertion PASSED: log_evidence matches {expected_log_evidence:.6f}") - np.testing.assert_allclose( - float(full_result), - expected_log_evidence, - rtol=1e-3, - err_msg=f"interferometer/delaunay[{instrument}]: regression — full log_evidence drifted", - ) - print(" Full-pipeline regression assertion PASSED") + _rec = check_pinned(figure_of_merit_ref, _pinned_expected, label="eager", rtol=1e-4) + if _rec is not None: + _pinned_drift.append(_rec) + _rec = check_pinned(float(full_result), _pinned_expected, label="full", rtol=1e-3) + if _rec is not None: + _pinned_drift.append(_rec) if result_vmap is not None: - np.testing.assert_allclose( - np.array(result_vmap), - expected_log_evidence, - rtol=1e-3, - err_msg=f"interferometer/delaunay[{instrument}]: regression — vmap log_evidence drifted", - ) - print(" vmap regression assertion PASSED") + _rec = check_pinned(np.array(result_vmap), _pinned_expected, label="vmap", rtol=1e-3) + if _rec is not None: + _pinned_drift.append(_rec) + + +# Pinned-value outcome -> result JSON: profiling records and flags drift, +# never adjudicates library correctness (autolens_workspace_test's remit; +# boundary rule in results/notes/design_lock_in.md). PyAutoHeart's vitals +# scan reads the pinned_drift field. +record_pinned_check(dict_path, _pinned_expected, _pinned_drift) +if _pinned_expected is not None and not _pinned_drift: + print(" Pinned-value check PASSED (recorded in result JSON).") diff --git a/likelihood_runtime/interferometer/mge.py b/likelihood_runtime/interferometer/mge.py index d55e5f7..a63c3f4 100644 --- a/likelihood_runtime/interferometer/mge.py +++ b/likelihood_runtime/interferometer/mge.py @@ -77,17 +77,19 @@ # Plus this script's own --use-dft override to compare NUFFT against the # historical DFT baseline on SMA. sys.path.insert(0, str(Path(__file__).resolve().parents[2])) -from _profile_cli import ( # noqa: E402 +from _profile_cli import ( auto_simulate_if_missing, + check_pinned, device_info_dict, parse_profile_cli, + record_pinned_check, # noqa: E402 resolve_output_paths, ) from simulators.interferometer import INSTRUMENTS # noqa: E402 from vram import ( # noqa: E402 probe_vmap_memory, recommend_batch_size, - vmap_batch_for, + resolve_vmap_batch, write_probe_json, ) @@ -352,7 +354,7 @@ def full_pipeline_from_params(params_tree): recommended = recommend_batch_size(probe) probe_path = ( _cli.output_dir or (_workspace_root / "results" / "runtime" / "interferometer" / "mge") - ) / "vmap_probe.json" + ) / ("vmap_probe_mge_sparse.json" if _cli.use_sparse_operator else "vmap_probe_mge.json") write_probe_json(probe, recommended, probe_path) print(f"\n vmap_probe samples: {probe.samples}") print(f" per_replica: {probe.per_replica_mb:.1f} MB / replica") @@ -366,7 +368,17 @@ def full_pipeline_from_params(params_tree): print("\n--- vmap batched evaluation ---") -batch_size = vmap_batch_for("interferometer", "mge", instrument) or 3 +_batch_resolved, _batch_source = resolve_vmap_batch( + "interferometer", + "mge", + instrument, + output_dir=_cli.output_dir + or (_workspace_root / "results" / "runtime" / "interferometer" / "mge"), + path="sparse" if _cli.use_sparse_operator else "dense", + backend=jax.default_backend(), +) +print(f" vmap batch_size: {_batch_resolved} (source: {_batch_source})") +batch_size = _batch_resolved or 3 parameters = jax.tree_util.tree_map( lambda leaf: jnp.broadcast_to(leaf, (batch_size, *leaf.shape)), @@ -557,6 +569,9 @@ def full_pipeline_from_params(params_tree): # print the value so it can be pasted in here on a clean run". sma was # bumped to mask_radius=3.5 in 2026-05-21's INSTRUMENTS refactor — the # old mask_radius=3.0 value no longer applies and needs re-measuring. +_pinned_drift: list = [] +_pinned_expected = None + EXPECTED_LOG_LIKELIHOOD = { "sma": None, "alma": None, @@ -564,6 +579,7 @@ def full_pipeline_from_params(params_tree): } expected_log_likelihood = EXPECTED_LOG_LIKELIHOOD.get(instrument) +_pinned_expected = expected_log_likelihood if expected_log_likelihood is None: print( @@ -576,26 +592,23 @@ def full_pipeline_from_params(params_tree): # the rtol when the user asked for mixed precision. _regression_rtol = 1e-3 if _cli.use_mixed_precision else 1e-4 transformer_label = "DFT" if USE_DFT else "NUFFT" - np.testing.assert_allclose( - log_likelihood_ref, - expected_log_likelihood, - rtol=_regression_rtol, - err_msg=( - f"interferometer/mge[{instrument}, {transformer_label}]: " - f"regression — eager log_likelihood drifted " - f"(got {log_likelihood_ref}, expected {expected_log_likelihood})" - ), - ) - np.testing.assert_allclose( - float(full_result), - expected_log_likelihood, - rtol=_regression_rtol, - err_msg=f"interferometer/mge[{instrument}]: regression — full log_likelihood drifted", + _rec = check_pinned(log_likelihood_ref, _pinned_expected, label="eager", rtol=_regression_rtol) + if _rec is not None: + _pinned_drift.append(_rec) + _rec = check_pinned(float(full_result), _pinned_expected, label="full", rtol=_regression_rtol) + if _rec is not None: + _pinned_drift.append(_rec) + _rec = check_pinned( + np.array(result_vmap), _pinned_expected, label="vmap", rtol=_regression_rtol ) - np.testing.assert_allclose( - np.array(result_vmap), - expected_log_likelihood, - rtol=_regression_rtol, - err_msg=f"interferometer/mge[{instrument}]: regression — vmap log_likelihood drifted", - ) - print(f" Regression assertion PASSED: log_likelihood matches {expected_log_likelihood:.6f}") + if _rec is not None: + _pinned_drift.append(_rec) + + +# Pinned-value outcome -> result JSON: profiling records and flags drift, +# never adjudicates library correctness (autolens_workspace_test's remit; +# boundary rule in results/notes/design_lock_in.md). PyAutoHeart's vitals +# scan reads the pinned_drift field. +record_pinned_check(dict_path, _pinned_expected, _pinned_drift) +if _pinned_expected is not None and not _pinned_drift: + print(" Pinned-value check PASSED (recorded in result JSON).") diff --git a/likelihood_runtime/interferometer/pixelization.py b/likelihood_runtime/interferometer/pixelization.py index 4827bd9..22175bf 100644 --- a/likelihood_runtime/interferometer/pixelization.py +++ b/likelihood_runtime/interferometer/pixelization.py @@ -75,17 +75,19 @@ # Sweep-driver CLI args (--config-name / --output-dir / --use-mixed-precision). # Tolerates extra/unknown args via parse_known_args inside the helper. -from _profile_cli import ( # noqa: E402 +from _profile_cli import ( auto_simulate_if_missing, + check_pinned, device_info_dict, parse_profile_cli, + record_pinned_check, # noqa: E402 resolve_output_paths, ) from simulators.interferometer import INSTRUMENTS # noqa: E402 from vram import ( # noqa: E402 probe_vmap_memory, recommend_batch_size, - vmap_batch_for, + resolve_vmap_batch, write_probe_json, ) @@ -378,7 +380,11 @@ def full_pipeline_from_params(params_tree): probe_path = ( _cli.output_dir or (_workspace_root / "results" / "runtime" / "interferometer" / "pixelization") - ) / "vmap_probe.json" + ) / ( + "vmap_probe_pixelization_sparse.json" + if _cli.use_sparse_operator + else "vmap_probe_pixelization.json" + ) write_probe_json(probe, recommended, probe_path) print(f"\n vmap_probe samples: {probe.samples}") print(f" per_replica: {probe.per_replica_mb:.1f} MB / replica") @@ -396,7 +402,17 @@ def full_pipeline_from_params(params_tree): print("\n--- vmap batched evaluation ---") -batch_size = vmap_batch_for("interferometer", "pixelization", instrument) or 3 +_batch_resolved, _batch_source = resolve_vmap_batch( + "interferometer", + "pixelization", + instrument, + output_dir=_cli.output_dir + or (_workspace_root / "results" / "runtime" / "interferometer" / "pixelization"), + path="sparse" if _cli.use_sparse_operator else "dense", + backend=jax.default_backend(), +) +print(f" vmap batch_size: {_batch_resolved} (source: {_batch_source})") +batch_size = _batch_resolved or 3 vmap_batch_time = None vmap_per_call = None vmap_speedup = None @@ -617,6 +633,9 @@ def full_pipeline_from_params(params_tree): # print the value so it can be pasted in here on a clean run". sma was # bumped to mask_radius=3.5 in 2026-05-21's INSTRUMENTS refactor — the # old mask_radius=3.0 value no longer applies and needs re-measuring. +_pinned_drift: list = [] +_pinned_expected = None + EXPECTED_LOG_EVIDENCE = { "sma": None, "alma": None, @@ -624,6 +643,7 @@ def full_pipeline_from_params(params_tree): } expected_log_evidence = EXPECTED_LOG_EVIDENCE.get(instrument) +_pinned_expected = expected_log_evidence if expected_log_evidence is None: print( @@ -631,27 +651,22 @@ def full_pipeline_from_params(params_tree): f"(no pinned value). Eager log_evidence = {figure_of_merit_ref}" ) else: - np.testing.assert_allclose( - figure_of_merit_ref, - expected_log_evidence, - rtol=1e-4, - err_msg=( - f"interferometer/pixelization[{instrument}]: regression — " - f"eager log_evidence drifted " - f"(got {figure_of_merit_ref}, expected {expected_log_evidence})" - ), - ) - np.testing.assert_allclose( - float(full_result), - expected_log_evidence, - rtol=1e-3, - err_msg=f"interferometer/pixelization[{instrument}]: regression — full log_evidence drifted", - ) + _rec = check_pinned(figure_of_merit_ref, _pinned_expected, label="eager", rtol=1e-4) + if _rec is not None: + _pinned_drift.append(_rec) + _rec = check_pinned(float(full_result), _pinned_expected, label="full", rtol=1e-3) + if _rec is not None: + _pinned_drift.append(_rec) if result_vmap is not None: - np.testing.assert_allclose( - np.array(result_vmap), - expected_log_evidence, - rtol=1e-3, - err_msg=f"interferometer/pixelization[{instrument}]: regression — vmap log_evidence drifted", - ) - print(f" Regression assertion PASSED: log_evidence matches {expected_log_evidence:.6f}") + _rec = check_pinned(np.array(result_vmap), _pinned_expected, label="vmap", rtol=1e-3) + if _rec is not None: + _pinned_drift.append(_rec) + + +# Pinned-value outcome -> result JSON: profiling records and flags drift, +# never adjudicates library correctness (autolens_workspace_test's remit; +# boundary rule in results/notes/design_lock_in.md). PyAutoHeart's vitals +# scan reads the pinned_drift field. +record_pinned_check(dict_path, _pinned_expected, _pinned_drift) +if _pinned_expected is not None and not _pinned_drift: + print(" Pinned-value check PASSED (recorded in result JSON).") diff --git a/likelihood_runtime/point_source/image_plane.py b/likelihood_runtime/point_source/image_plane.py index ef63bd4..2318eac 100644 --- a/likelihood_runtime/point_source/image_plane.py +++ b/likelihood_runtime/point_source/image_plane.py @@ -62,17 +62,19 @@ # Sweep-driver CLI args (--config-name / --output-dir / --use-mixed-precision). # Tolerates extra/unknown args via parse_known_args inside the helper. sys.path.insert(0, str(Path(__file__).resolve().parents[2])) -from _profile_cli import ( # noqa: E402 +from _profile_cli import ( auto_simulate_if_missing, + check_pinned, device_info_dict, parse_profile_cli, + record_pinned_check, # noqa: E402 resolve_output_paths, ) from simulators.point_source import INSTRUMENTS # noqa: E402 from vram import ( # noqa: E402 probe_vmap_memory, recommend_batch_size, - vmap_batch_for, + resolve_vmap_batch, write_probe_json, ) @@ -302,7 +304,11 @@ def full_pipeline_from_params(params_tree): probe_path = ( _cli.output_dir or (_workspace_root / "results" / "runtime" / "point_source" / "image_plane") - ) / "vmap_probe.json" + ) / ( + "vmap_probe_image_plane_sparse.json" + if _cli.use_sparse_operator + else "vmap_probe_image_plane.json" + ) write_probe_json(probe, recommended, probe_path) print(f"\n vmap_probe samples: {probe.samples}") print(f" per_replica: {probe.per_replica_mb:.1f} MB / replica") @@ -316,7 +322,17 @@ def full_pipeline_from_params(params_tree): print("\n--- vmap batched evaluation ---") -batch_size = vmap_batch_for("point_source", "image_plane", instrument) or 3 +_batch_resolved, _batch_source = resolve_vmap_batch( + "point_source", + "image_plane", + instrument, + output_dir=_cli.output_dir + or (_workspace_root / "results" / "runtime" / "point_source" / "image_plane"), + path="sparse" if _cli.use_sparse_operator else "dense", + backend=jax.default_backend(), +) +print(f" vmap batch_size: {_batch_resolved} (source: {_batch_source})") +batch_size = _batch_resolved or 3 batched_params = jax.tree_util.tree_map( lambda leaf: jnp.broadcast_to(leaf, (batch_size, *leaf.shape)), @@ -363,7 +379,6 @@ def full_pipeline_from_params(params_tree): f"eager={log_likelihood_ref} vs JIT={float(full_result)}" ), ) -print(" Tier 1: eager ≡ JIT assertion PASSED") np.testing.assert_allclose( np.array(result_vmap), @@ -371,7 +386,6 @@ def full_pipeline_from_params(params_tree): rtol=1e-4, err_msg="point_source/image_plane: JIT vs vmap mismatch", ) -print(" Tier 2: JIT ≡ vmap assertion PASSED") # =================================================================== @@ -497,42 +511,29 @@ def full_pipeline_from_params(params_tree): # regenerated; the new value reflects the current truth-aligned # evaluation against the dataset committed in # autolens_workspace_developer@f8a5cef. -EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE = 7.196577317761017 +_pinned_drift: list = [] +_pinned_expected = None -np.testing.assert_allclose( - log_likelihood_ref, - EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE, - rtol=1e-4, - err_msg=( - f"point_source/image_plane: regression — eager log_likelihood drifted " - f"(got {log_likelihood_ref}, expected {EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE})" - ), -) -print( - f" Eager regression assertion PASSED: log_likelihood matches " - f"{EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE:.6f}" -) -np.testing.assert_allclose( - float(full_result), - EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE, - rtol=1e-4, - err_msg=( - f"point_source/image_plane: regression — JIT log_likelihood drifted " - f"(got {float(full_result)}, expected {EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE})" - ), -) -np.testing.assert_allclose( - np.array(result_vmap), - EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE, - rtol=1e-4, - err_msg=( - f"point_source/image_plane: regression — vmap log_likelihood drifted " - f"(expected {EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE})" - ), -) -print( - f" Tier 3: regression assertion PASSED: log_likelihood matches " - f"{EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE:.6f}" -) +EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE = 7.196577317761017 +_pinned_expected = EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE + +_rec = check_pinned(log_likelihood_ref, _pinned_expected, label="eager", rtol=1e-4) +if _rec is not None: + _pinned_drift.append(_rec) +_rec = check_pinned(float(full_result), _pinned_expected, label="JIT", rtol=1e-4) +if _rec is not None: + _pinned_drift.append(_rec) +_rec = check_pinned(np.array(result_vmap), _pinned_expected, label="vmap", rtol=1e-4) +if _rec is not None: + _pinned_drift.append(_rec) timer.summary() + + +# Pinned-value outcome -> result JSON: profiling records and flags drift, +# never adjudicates library correctness (autolens_workspace_test's remit; +# boundary rule in results/notes/design_lock_in.md). PyAutoHeart's vitals +# scan reads the pinned_drift field. +record_pinned_check(dict_path, _pinned_expected, _pinned_drift) +if _pinned_expected is not None and not _pinned_drift: + print(" Pinned-value check PASSED (recorded in result JSON).") diff --git a/likelihood_runtime/point_source/source_plane.py b/likelihood_runtime/point_source/source_plane.py index f833627..f323889 100644 --- a/likelihood_runtime/point_source/source_plane.py +++ b/likelihood_runtime/point_source/source_plane.py @@ -44,17 +44,19 @@ # Sweep-driver CLI args (--config-name / --output-dir / --use-mixed-precision). # Tolerates extra/unknown args via parse_known_args inside the helper. sys.path.insert(0, str(Path(__file__).resolve().parents[2])) -from _profile_cli import ( # noqa: E402 +from _profile_cli import ( auto_simulate_if_missing, + check_pinned, device_info_dict, parse_profile_cli, + record_pinned_check, # noqa: E402 resolve_output_paths, ) from simulators.point_source import INSTRUMENTS # noqa: E402 from vram import ( # noqa: E402 probe_vmap_memory, recommend_batch_size, - vmap_batch_for, + resolve_vmap_batch, write_probe_json, ) @@ -304,7 +306,11 @@ def full_pipeline_from_params(params_tree): probe_path = ( _cli.output_dir or (_workspace_root / "results" / "runtime" / "point_source" / "source_plane") - ) / "vmap_probe.json" + ) / ( + "vmap_probe_source_plane_sparse.json" + if _cli.use_sparse_operator + else "vmap_probe_source_plane.json" + ) write_probe_json(probe, recommended, probe_path) print(f"\n vmap_probe samples: {probe.samples}") print(f" per_replica: {probe.per_replica_mb:.1f} MB / replica") @@ -356,7 +362,17 @@ def ray_trace_to_source_plane(params_tree, positions_raw): print("\n--- vmap over ray-trace prefix ---") -batch_size = vmap_batch_for("point_source", "source_plane", instrument) or 3 +_batch_resolved, _batch_source = resolve_vmap_batch( + "point_source", + "source_plane", + instrument, + output_dir=_cli.output_dir + or (_workspace_root / "results" / "runtime" / "point_source" / "source_plane"), + path="sparse" if _cli.use_sparse_operator else "dense", + backend=jax.default_backend(), +) +print(f" vmap batch_size: {_batch_resolved} (source: {_batch_source})") +batch_size = _batch_resolved or 3 batched_params = jax.tree_util.tree_map( lambda leaf: jnp.broadcast_to(leaf, (batch_size, *leaf.shape)), @@ -550,35 +566,28 @@ def ray_trace_to_source_plane(params_tree, positions_raw): # log-likelihood. The refreshed value reflects the current truth-aligned # evaluation against the dataset committed in # autolens_workspace_developer@f8a5cef. +_pinned_drift: list = [] +_pinned_expected = None + EXPECTED_LOG_LIKELIHOOD_SOURCE_PLANE = -33788.35731127962 +_pinned_expected = EXPECTED_LOG_LIKELIHOOD_SOURCE_PLANE -np.testing.assert_allclose( - log_likelihood_ref, - EXPECTED_LOG_LIKELIHOOD_SOURCE_PLANE, - rtol=1e-4, - err_msg=( - f"point_source/source_plane: regression — eager log_likelihood drifted " - f"(got {log_likelihood_ref}, expected {EXPECTED_LOG_LIKELIHOOD_SOURCE_PLANE})" - ), -) -print( - f" Eager regression assertion PASSED: log_likelihood matches " - f"{EXPECTED_LOG_LIKELIHOOD_SOURCE_PLANE:.6f}" -) +_rec = check_pinned(log_likelihood_ref, _pinned_expected, label="eager", rtol=1e-4) +if _rec is not None: + _pinned_drift.append(_rec) if full_pipeline_jits: - np.testing.assert_allclose( - float(full_result), - EXPECTED_LOG_LIKELIHOOD_SOURCE_PLANE, - rtol=1e-4, - err_msg=( - f"point_source/source_plane: regression — JIT log_likelihood drifted " - f"(got {float(full_result)}, expected {EXPECTED_LOG_LIKELIHOOD_SOURCE_PLANE})" - ), - ) - print( - f" JIT regression assertion PASSED: log_likelihood matches " - f"{EXPECTED_LOG_LIKELIHOOD_SOURCE_PLANE:.6f}" - ) + _rec = check_pinned(float(full_result), _pinned_expected, label="JIT", rtol=1e-4) + if _rec is not None: + _pinned_drift.append(_rec) timer.summary() + + +# Pinned-value outcome -> result JSON: profiling records and flags drift, +# never adjudicates library correctness (autolens_workspace_test's remit; +# boundary rule in results/notes/design_lock_in.md). PyAutoHeart's vitals +# scan reads the pinned_drift field. +record_pinned_check(dict_path, _pinned_expected, _pinned_drift) +if _pinned_expected is not None and not _pinned_drift: + print(" Pinned-value check PASSED (recorded in result JSON).") diff --git a/results/notes/design_lock_in.md b/results/notes/design_lock_in.md index cbbf580..80c8651 100644 --- a/results/notes/design_lock_in.md +++ b/results/notes/design_lock_in.md @@ -93,6 +93,26 @@ Rule of thumb: **cross-release trend** questions read versioned summaries; - Baselines are **append-only**: never edited after the campaign closes; a post-optimization campaign snapshots a new name next to it. +## Pinned values: record-and-flag, never adjudicate (added phase 2) + +The runtime cells carry pinned log-likelihood/evidence values. Their job here +is **benchmark comparability** — a timing is only comparable to a baseline if +it measures the same computation — *not* library regression testing (that is +`autolens_workspace_test`'s remit). Accordingly: + +- Drift against a pin **never crashes a profiling run**: `check_pinned` + (in `_profile_cli.py`) prints a loud warning and records + `pinned_expected` + `pinned_drift` into the result JSON + (`record_pinned_check`). The numbers still land; the flag marks them + non-comparable to the pinned baseline. +- The **eager ≡ JIT ≡ vmap cross-consistency assertions stay hard** — a + mismatch there invalidates the measurement itself. +- On a drift flag the follow-up is "file a bug / check workspace_test", + never "debug the library inside this repo". +- **PyAutoHeart pairing**: the `pinned_drift` JSON field is the machine + surface Heart's vitals scan reads (tracked as its own Mind task, + `feature/pyautoheart/profiling_drift_check.md`). + ## Deliberately out of scope here Profiling runs (phases 2–4); searches profiling; `point_source` cells (in the diff --git a/results/notes/vram_validation_2026_07_08.md b/results/notes/vram_validation_2026_07_08.md new file mode 100644 index 0000000..e86ace5 --- /dev/null +++ b/results/notes/vram_validation_2026_07_08.md @@ -0,0 +1,67 @@ +# Phase-2 vram-first validation sweep — 2026-07-08 + +Local CPU validation of every PreOptimizationTimes campaign cell **before** +A100 time is spent, per the campaign's vram-first rule. Tracked by +[autolens_profiling#54](https://github.com/PyAutoLabs/autolens_profiling/issues/54); +design context: [`design_lock_in.md`](./design_lock_in.md). PyAutoLens +2026.7.6.649, JAX CPU backend, cheapest instrument per cell; sweep artifacts +were scratch-tier (validation, not campaign data) — this note is the durable +record. + +## Verdict + +**All 9 campaign cells run clean end-to-end on current source** (after the +fixes below). The A100 probe leg (`hpc/batch_gpu/submit_probe_{fast,delaunay}_a100`) +can be dispatched without expecting functional surprises. + +| Cell | Instrument | Path | Wall time (CPU) | +|------|-----------|------|-----------------| +| imaging/mge | euclid | dense | 47 s | +| imaging/pixelization | euclid | dense | 401 s | +| imaging/pixelization | euclid | sparse | 194 s | +| imaging/delaunay | euclid | dense | 381 s | +| imaging/delaunay | euclid | sparse | 213 s | +| interferometer/mge | sma | dense | 23 s | +| interferometer/pixelization | sma | dense | 70 s | +| interferometer/delaunay | sma | dense | 124 s | +| datacube/delaunay | sma | — | 265 s | + +Also validated live: the `--vmap-probe` pipeline end-to-end on CPU, and the +probe-JSON-over-table resolution (a fresh probe in the cell output dir won, +`source: probe (vmap_probe_mge.json)` in the log). + +## Defects flushed (all fixed on the phase-2 branch) + +1. **A100 batch table applied on every backend.** Runtime cells vmapped at + table batch (64) on laptop CPU → OOM-kill (exit 137, pixelization dense at + 97 s); the probe's 65 GB budget default is A100-specific too, so a CPU + probe *also* recommended 64. Fix: `resolve_vmap_batch` takes the running + backend and clamps to 3 on non-GPU; probes record their backend and + resolution ignores probes from a different backend. +2. **Instrument-naive pins.** `imaging/{mge,pixelization,delaunay}` asserted + an HST-only pinned likelihood/evidence for every `--instrument` — any + non-HST run failed by construction (mge at 62 s; pixelization sparse died + at 2499 s *after* 41 min of honest compute, at the final assert). + Converted to the per-instrument dict pattern `interferometer/mge` had. +3. **Pins were hard asserts** (design defect, user-directed redesign): they + guard benchmark comparability, not library correctness, and must never + kill a profiling job — see the record-and-flag section of + [`design_lock_in.md`](./design_lock_in.md). + +## Empirical bonus + +Pixelization **sparse at batch 64 on CPU did not OOM where dense did** (it +survived to the final assert), directly confirming the sparse path's smaller +per-replica footprint — the basis for `VMAP_BATCH_SPARSE`'s +fall-back-to-dense-is-conservative rule. + +## What phase 3 consumes + +- Batch sizes: `vram.resolve_vmap_batch` (probe JSON → table → clamp), sparse + rows via `VMAP_BATCH_SPARSE` once the A100 probes are ingested. +- The dense table's provenance is 2026-05-24 (`vram.config.PROVENANCE`) — + ingest the probe-only A100 submits' JSONs before trusting it for the + campaign (workflow: `vram/README.md` "Probe-only submits + ingest"). +- No pins exist yet for euclid/ao/jwst imaging or any interferometer + instrument — first clean campaign runs should pin them (values are printed + by every unpinned run). diff --git a/results/simulators/imaging_euclid_summary_v2026.7.6.649.json b/results/simulators/imaging_euclid_summary_v2026.7.6.649.json new file mode 100644 index 0000000..c4a2db1 --- /dev/null +++ b/results/simulators/imaging_euclid_summary_v2026.7.6.649.json @@ -0,0 +1,45 @@ +{ + "autolens_version": "2026.7.6.649", + "type": "imaging", + "instrument": "euclid", + "configuration": { + "grid_shape": [ + 71, + 71 + ], + "pixel_scales": 0.1, + "mask_radius": 3.5, + "psf_shape": [ + 21, + 21 + ], + "psf_sigma": 0.1, + "over_sampling_sub_sizes": [ + 32, + 8, + 2 + ] + }, + "phases": { + "setup_grids": 0.061925699999846984, + "setup_psf_simulator": 0.0005141000001458451, + "setup_galaxies": 8.140000136336312e-05, + "setup_tracer": 1.1600001016631722e-05, + "image_2d_eager": 0.3962143999997352, + "image_2d_jit_lower": 0.47080600000117556, + "image_2d_jit_compile": 0.3903323000013188, + "image_2d_jit_first_call": 0.018089100001816405, + "image_2d_jit_steady_x10": 0.16007420000096317, + "via_tracer_from": 0.6722835999971721, + "solver_build": 8.159999924828298e-05, + "solver_solve_eager": 0.05072459999792045, + "output_fits": 0.01322970000182977, + "output_lensed_source": 0.03972860000067158, + "output_json": 0.002425700000458164 + }, + "key_timings": { + "image_2d_eager_s": 0.3962143999997352, + "via_tracer_from_s": 0.6722835999971721, + "solver_solve_eager_s": 0.05072459999792045 + } +} \ No newline at end of file diff --git a/results/simulators/imaging_euclid_summary_v2026.7.6.649.png b/results/simulators/imaging_euclid_summary_v2026.7.6.649.png new file mode 100644 index 0000000..438aeaf Binary files /dev/null and b/results/simulators/imaging_euclid_summary_v2026.7.6.649.png differ diff --git a/simulators/README.md b/simulators/README.md index bd80fed..108c4ae 100644 --- a/simulators/README.md +++ b/simulators/README.md @@ -1,6 +1,6 @@ # simulators -Run-time profiling for the PyAutoLens simulators — the scripts that produce mock lensed datasets (imaging, interferometer, point-source, cluster-scale, group-scale, multi-wavelength) used as inputs to fitting pipelines. These scripts profile **simulation cost**, not likelihood cost (which lives under [`likelihood/`](../likelihood/README.md)). They time grid construction, ray-tracing, PSF convolution, point-solver work, FITS / JSON output, and whatever else dominates each science case. +Run-time profiling for the PyAutoLens simulators — the scripts that produce mock lensed datasets (imaging, interferometer, point-source, cluster-scale, group-scale, multi-wavelength) used as inputs to fitting pipelines. These scripts profile **simulation cost**, not likelihood cost (which lives under [`likelihood_runtime/`](../likelihood_runtime/README.md) / [`likelihood_breakdown/`](../likelihood_breakdown/README.md)). They time grid construction, ray-tracing, PSF convolution, point-solver work, FITS / JSON output, and whatever else dominates each science case. ## Why profile the simulators? @@ -43,7 +43,9 @@ Old versions are retained alongside new ones so cross-release trends stay visibl Auto-generated by `scripts/build_readme.py` from the latest `*_summary_v.json` artifacts under `results/simulators/`. Cell value is the **total wall time** of the simulator run-to-completion (sum of all `timer.section` phases), not per-likelihood. Hardware tier is CPU only today. -_No data yet — run a simulator under `simulators/` to populate. See section README._ +| Script | Total wall time | PyAutoLens version | +|--------|-----------------|--------------------| +| `imaging_euclid.py` | 2.28 s | v2026.7.6.649 | ## Running a script diff --git a/test/test_vram_config.py b/test/test_vram_config.py new file mode 100644 index 0000000..b7d1a1f --- /dev/null +++ b/test/test_vram_config.py @@ -0,0 +1,191 @@ +"""Unit tests for ``vram.config`` batch-size lookup + resolution. + +Covers the sparse-path fallback rule and the probe-JSON-over-table +resolution order added for the PreOptimizationTimes campaign (phase 2, +autolens_profiling#54). No JAX dependency. + +Run:: + + cd autolens_profiling + python -m pytest test/test_vram_config.py +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +# Make `vram` importable without installing the repo. +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from vram.config import ( # noqa: E402 + VMAP_BATCH, + VMAP_BATCH_SPARSE, + resolve_vmap_batch, + vmap_batch_for, +) + + +def _write_probe( + tmp_path: Path, + *, + model: str = "pixelization", + sparse: bool = False, + dataset: str = "imaging", + instrument: str = "hst", + recommended=24, +) -> Path: + suffix = "_sparse" if sparse else "" + p = tmp_path / f"vmap_probe_{model}{suffix}.json" + p.write_text( + json.dumps( + { + "dataset": dataset, + "model": model, + "instrument": instrument, + "recommended_batch_size": recommended, + } + ) + ) + return p + + +# --------------------------------------------------------------------------- +# vmap_batch_for — path axis +# --------------------------------------------------------------------------- + + +def test_dense_lookup_unchanged(): + assert ( + vmap_batch_for("imaging", "pixelization", "hst") + == VMAP_BATCH[("imaging", "pixelization", "hst")] + ) + + +def test_sparse_falls_back_to_dense_when_unprobed(): + key = ("imaging", "pixelization", "hst") + assert key not in VMAP_BATCH_SPARSE # precondition: sparse row unprobed + assert vmap_batch_for(*key, path="sparse") == VMAP_BATCH[key] + + +def test_sparse_row_wins_when_present(monkeypatch): + key = ("imaging", "pixelization", "hst") + monkeypatch.setitem(VMAP_BATCH_SPARSE, key, 32) + assert vmap_batch_for(*key, path="sparse") == 32 + assert vmap_batch_for(*key) == VMAP_BATCH[key] # dense unaffected + + +def test_intentional_none_preserved(): + assert vmap_batch_for("datacube", "delaunay", "sma") is None + assert vmap_batch_for("interferometer", "mge", "alma") is None + + +# --------------------------------------------------------------------------- +# resolve_vmap_batch — probe JSON over table +# --------------------------------------------------------------------------- + + +def test_resolve_without_output_dir_uses_table(): + batch, source = resolve_vmap_batch("imaging", "pixelization", "hst") + assert batch == VMAP_BATCH[("imaging", "pixelization", "hst")] + assert source == "table" + + +def test_resolve_prefers_matching_probe(tmp_path): + _write_probe(tmp_path, recommended=24) + batch, source = resolve_vmap_batch("imaging", "pixelization", "hst", output_dir=tmp_path) + assert batch == 24 + assert source.startswith("probe") + + +def test_resolve_sparse_probe_selected_by_path(tmp_path): + _write_probe(tmp_path, recommended=24) # dense probe + _write_probe(tmp_path, sparse=True, recommended=40) + batch, source = resolve_vmap_batch( + "imaging", "pixelization", "hst", output_dir=tmp_path, path="sparse" + ) + assert batch == 40 + assert "sparse" in source + + +def test_resolve_ignores_cell_mismatch(tmp_path): + _write_probe(tmp_path, instrument="euclid") # probe for a different instrument + batch, source = resolve_vmap_batch("imaging", "pixelization", "hst", output_dir=tmp_path) + assert batch == VMAP_BATCH[("imaging", "pixelization", "hst")] + assert "mismatch" in source + + +def test_resolve_ignores_invalid_recommendation(tmp_path): + _write_probe(tmp_path, recommended=None) + batch, source = resolve_vmap_batch("imaging", "pixelization", "hst", output_dir=tmp_path) + assert batch == VMAP_BATCH[("imaging", "pixelization", "hst")] + assert "no valid recommendation" in source + + +def test_resolve_survives_corrupt_json(tmp_path): + (tmp_path / "vmap_probe_pixelization.json").write_text("{not json") + batch, source = resolve_vmap_batch("imaging", "pixelization", "hst", output_dir=tmp_path) + assert batch == VMAP_BATCH[("imaging", "pixelization", "hst")] + assert "unreadable" in source + + +def test_resolve_missing_probe_uses_table(tmp_path): + batch, source = resolve_vmap_batch("imaging", "pixelization", "hst", output_dir=tmp_path) + assert batch == VMAP_BATCH[("imaging", "pixelization", "hst")] + assert source == "table" + + +# --------------------------------------------------------------------------- +# backend awareness (non-GPU clamp + probe backend matching) +# --------------------------------------------------------------------------- + + +def test_non_gpu_backend_clamps_table_value(): + batch, source = resolve_vmap_batch("imaging", "pixelization", "euclid", backend="cpu") + assert batch == 3 + assert "clamped" in source + + +def test_gpu_backend_not_clamped(): + batch, source = resolve_vmap_batch("imaging", "pixelization", "euclid", backend="gpu") + assert batch == VMAP_BATCH[("imaging", "pixelization", "euclid")] + assert "clamped" not in source + + +def test_intentional_none_not_clamped_to_int(): + batch, _ = resolve_vmap_batch("datacube", "delaunay", "sma", backend="cpu") + assert batch is None + + +def test_probe_from_other_backend_ignored(tmp_path): + p = _write_probe(tmp_path, recommended=64) + data = json.loads(p.read_text()) + data["backend"] = "cpu" + p.write_text(json.dumps(data)) + batch, source = resolve_vmap_batch( + "imaging", "pixelization", "hst", output_dir=tmp_path, backend="gpu" + ) + assert batch == VMAP_BATCH[("imaging", "pixelization", "hst")] + assert "mismatch" in source + + +def test_probe_same_backend_accepted_and_clamped_on_cpu(tmp_path): + p = _write_probe(tmp_path, recommended=64) + data = json.loads(p.read_text()) + data["backend"] = "cpu" + p.write_text(json.dumps(data)) + batch, source = resolve_vmap_batch( + "imaging", "pixelization", "hst", output_dir=tmp_path, backend="cpu" + ) + assert batch == 3 # probe wins, then the non-GPU cap applies + assert source.startswith("probe") and "clamped" in source + + +def test_legacy_probe_without_backend_accepted(tmp_path): + _write_probe(tmp_path, recommended=24) # no backend field + batch, source = resolve_vmap_batch( + "imaging", "pixelization", "hst", output_dir=tmp_path, backend="gpu" + ) + assert batch == 24 + assert source.startswith("probe") diff --git a/vram/README.md b/vram/README.md index e89397d..4090df3 100644 --- a/vram/README.md +++ b/vram/README.md @@ -30,7 +30,8 @@ From `vram`: | Name | Purpose | |------|---------| -| `vmap_batch_for(dataset, model, instrument)` | Return per-cell batch_size (or `None` if vmap is intentionally skipped or the cell hasn't been probed). | +| `vmap_batch_for(dataset, model, instrument, path="dense")` | Return per-cell batch_size (or `None` if vmap is intentionally skipped or the cell hasn't been probed). `path="sparse"` consults `VMAP_BATCH_SPARSE`, falling back to the dense row while unprobed. | +| `resolve_vmap_batch(dataset, model, instrument, output_dir=None, path="dense")` | The runtime cells' entry point: prefer a fresh, cell-matching `vmap_probe_[_sparse].json` in `output_dir` over the table. Returns `(batch_size, source)`. | | `probe_vmap_memory(func, args, batch_sizes=(2, 4))` | JIT-vmap `func` at each batch, read `compiled.memory_analysis()`, return a `ProbeResult`. | | `recommend_batch_size(probe, vram_budget_gb=70.0, max_batch=64)` | Linear extrapolation → max batch fitting in budget. | | `write_probe_json(probe, recommended, path)` | Serialise probe + recommendation to JSON. | @@ -80,16 +81,45 @@ XLA recompiles for each new batch_size (no cache reuse). Compile cost varies: `batch_sizes=(1,)` and accept linear extrapolation. If the chosen batch OOMs at run time, manually re-probe at a smaller batch. +## Batch-size resolution order (what a runtime cell actually uses) + +1. **Fresh probe JSON** — `results/runtime///vmap_probe_[_sparse].json`, + written by a `--vmap-probe` run (Phase A of the A100 submits). Used only if + its `dataset`/`model`/`instrument` fields match the cell, so a stale table + can never OOM a job whose submit re-probed. +2. **The curated table** — `VMAP_BATCH` (dense) / `VMAP_BATCH_SPARSE` + (w-tilde sparse-operator rows; unprobed sparse cells fall back to the dense + value, which is conservative — the sparse path's per-replica footprint has + been smaller on every cell measured). +3. `None` → the cell script skips vmap with a logged reason. + +Each script logs which source won. Table provenance lives in +`vram.config.PROVENANCE` — check it before a campaign trusts the numbers. + +## Probe-only submits + ingest + +`hpc/batch_gpu/submit_probe_fast_a100` (mge/pixelization, ~minutes) and +`hpc/batch_gpu/submit_probe_delaunay_a100` (delaunay, hours) re-probe every +PreOptimizationTimes campaign vmap cell — dense + sparse imaging variants — +against current source. Ingest after the jobs finish: + +1. Commit the probe JSONs from the HPC checkout (they land under + `results/runtime///`). +2. Copy each `recommended_batch_size` into `VMAP_BATCH` / + `VMAP_BATCH_SPARSE` in `vram/config.py` (keep the per-row MB/replica + comments) and update `PROVENANCE`. +3. Re-run the regular profile to confirm the chosen batch holds at steady + state — the cuFFT-scratch halvings noted in the table comments came from + exactly this step, static analysis alone missed them. + ## Adding a new instrument 1. Add an entry to the appropriate `INSTRUMENTS` dict in `simulators/{imaging,interferometer}.py`. -2. Create a probe SLURM script at - `z_projects/profiling/hpc/batch_gpu/probe_vmap__` - (clone any existing one). -3. Run the probe; pull the resulting JSON to - `output/runtime///vmap_probe.json`. -4. Read `recommended_batch_size` and add the row to `VMAP_BATCH` in - `vram/config.py`. +2. Add the cell to a probe-only submit in `hpc/batch_gpu/` (clone a + `run_probe` line). +3. Run the probe; the JSON lands at + `results/runtime///vmap_probe_[_sparse].json`. +4. Ingest per "Probe-only submits + ingest" above. 5. Re-run the regular profile to confirm the chosen batch holds at steady state (vmap completes, doesn't OOM). diff --git a/vram/__init__.py b/vram/__init__.py index 2724d07..3dbe861 100644 --- a/vram/__init__.py +++ b/vram/__init__.py @@ -14,7 +14,13 @@ See ``vram/README.md`` for methodology and how to extend. """ -from vram.config import VMAP_BATCH, vmap_batch_for +from vram.config import ( + PROVENANCE, + VMAP_BATCH, + VMAP_BATCH_SPARSE, + resolve_vmap_batch, + vmap_batch_for, +) from vram.probe import ( ProbeResult, probe_vmap_memory, @@ -23,7 +29,10 @@ ) __all__ = [ + "PROVENANCE", "VMAP_BATCH", + "VMAP_BATCH_SPARSE", + "resolve_vmap_batch", "vmap_batch_for", "ProbeResult", "probe_vmap_memory", diff --git a/vram/config.py b/vram/config.py index 67fbe19..bc10c8e 100644 --- a/vram/config.py +++ b/vram/config.py @@ -123,9 +123,34 @@ } -def vmap_batch_for(dataset: str, model: str, instrument: str) -> int | None: +# Sparse-operator (w-tilde) rows. The sparse inversion path has a different +# per-replica memory profile than dense mapping, so it gets its own table +# rather than a fourth key element (keeps VMAP_BATCH's shape stable for the +# many existing readers). Unprobed sparse cells fall back to the dense row — +# conservative, since the sparse path's per-replica footprint is smaller for +# every cell measured so far. Populated by the probe-only SLURM submits +# (``hpc/batch_gpu/submit_probe_*``) as campaign probes come in. +VMAP_BATCH_SPARSE: dict[tuple[str, str, str], int | None] = {} + + +# Where each table's numbers came from — checked before a campaign trusts +# them. Update whenever probe results are ingested. +PROVENANCE: dict[str, str] = { + "VMAP_BATCH": ( + "probed 2026-05-24 on RAL A100 80GB PCIe (PyAutoLens ~2026.5.x); " + "manual halvings for cuFFT scratch failures noted per row" + ), + "VMAP_BATCH_SPARSE": "unprobed — falls back to dense rows", +} + + +def vmap_batch_for(dataset: str, model: str, instrument: str, path: str = "dense") -> int | None: """Return the per-(dataset, model, instrument) vmap batch_size for A100. + ``path`` selects the inversion path: ``"dense"`` (mapping-matrix, the + default) or ``"sparse"`` (w-tilde sparse operator). Sparse cells without + a probed row fall back to the dense value. + Returns ``None`` when vmap is intentionally skipped (cube cells), blocked (interferometer mge/pix at ALMA+), or when the cell hasn't been probed yet. Callers should default to a small fallback (typically @@ -133,4 +158,88 @@ def vmap_batch_for(dataset: str, model: str, instrument: str) -> int | None: entirely when ``None`` is returned for a known-blocked or known-skipped cell. """ - return VMAP_BATCH.get((dataset, model, instrument)) + key = (dataset, model, instrument) + if path == "sparse" and key in VMAP_BATCH_SPARSE: + return VMAP_BATCH_SPARSE[key] + return VMAP_BATCH.get(key) + + +# On non-GPU backends the vmap row is a correctness/timing sample, not a +# production configuration — the table and the probe budget are both +# A100-oriented, and batch 64 on a laptop CPU exhausts host RAM (found by +# the phase-2 local validation sweep, autolens_profiling#54). +_NON_GPU_BATCH_CAP = 3 + + +def resolve_vmap_batch( + dataset: str, + model: str, + instrument: str, + output_dir=None, + path: str = "dense", + backend: str | None = None, +) -> tuple[int | None, str]: + """Resolve the vmap batch_size, preferring a fresh probe over the table. + + Resolution order (returns ``(batch_size, source)`` where ``source`` + says which step won, for the run log): + + 1. ``/vmap_probe_[_sparse].json`` — written by a + ``--vmap-probe`` run (Phase A of the A100 submit scripts). Used only + if its ``dataset``/``model``/``instrument`` fields match this cell, + so a stale table can never OOM a job whose submit re-probed. + 2. The curated table (``vmap_batch_for``, including the sparse + fallback-to-dense rule). + + ``backend`` is the running JAX backend (``jax.default_backend()``). + The table and the probe budget are A100-oriented, so on any non-``gpu`` + backend the resolved batch is clamped to ``_NON_GPU_BATCH_CAP`` (an + intentional ``None`` stays ``None``). + + A probe JSON that cannot be read or does not match is ignored (with the + mismatch reported in ``source``), never fatal — the table is the + fallback, not the probe. + """ + import json + from pathlib import Path + + def _clamped(batch: int | None, source: str) -> tuple[int | None, str]: + if ( + backend is not None + and backend != "gpu" + and isinstance(batch, int) + and batch > _NON_GPU_BATCH_CAP + ): + return _NON_GPU_BATCH_CAP, f"{source}, clamped to {_NON_GPU_BATCH_CAP} on {backend}" + return batch, source + + if output_dir is not None: + suffix = "_sparse" if path == "sparse" else "" + probe_path = Path(output_dir) / f"vmap_probe_{model}{suffix}.json" + if probe_path.exists(): + try: + data = json.loads(probe_path.read_text()) + except (OSError, ValueError) as exc: + return _clamped( + vmap_batch_for(dataset, model, instrument, path=path), + f"table (probe JSON unreadable: {exc})", + ) + probe_backend = data.get("backend") + matches = ( + data.get("dataset") == dataset + and data.get("model") == model + and data.get("instrument") == instrument + # A probe measured on another backend is not evidence for + # this one (CPU memory analysis says nothing about A100 and + # vice versa). Probes predating the backend field pass. + and (probe_backend is None or backend is None or probe_backend == backend) + ) + recommended = data.get("recommended_batch_size") + if matches and isinstance(recommended, int) and recommended >= 1: + return _clamped(recommended, f"probe ({probe_path.name})") + return _clamped( + vmap_batch_for(dataset, model, instrument, path=path), + f"table (probe JSON ignored: " + f"{'cell mismatch' if not matches else 'no valid recommendation'})", + ) + return _clamped(vmap_batch_for(dataset, model, instrument, path=path), "table") diff --git a/vram/probe.py b/vram/probe.py index 2bf77bb..99c6725 100644 --- a/vram/probe.py +++ b/vram/probe.py @@ -192,8 +192,11 @@ def write_probe_json( extra: dict | None = None, ) -> None: """Serialise probe samples + recommendation to JSON.""" + import jax + output_path.parent.mkdir(parents=True, exist_ok=True) payload = { + "backend": jax.default_backend(), "dataset": probe.dataset, "model": probe.model, "instrument": probe.instrument,