Skip to content

research: point-source CPU phase 4b — step-0 containment route A/B (PyAutoArray#579) #330

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jammy2211 merged 10 commits into from
Sep 27, 2026
Merged

research: point-source CPU phase 4b — step-0 containment route A/B (PyAutoArray#579) #330

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
311c7e6
feat(point_source_image): --step0-route library A/B in solver_config_…
Jammy2211 Sep 27, 2026
767c391
results(point_source_image): laptop step-0 route A/B, not quotable (P…
Jammy2211 Sep 27, 2026
d9b3382
fix(point_source_image): route-aware step-0 materialisation prefix; i…
Jammy2211 Sep 27, 2026
64e486a
hpc(point_source_image): RAL CPU + A100 submits for the step-0 route …
Jammy2211 Sep 27, 2026
d745d67
hpc(point_source_image): gpu-partition CPU-only fallback submit for t…
Jammy2211 Sep 27, 2026
54c5294
results(point_source_image): step-0 route A/B on RAL A100 (357322) an…
Jammy2211 Sep 27, 2026
e78d11c
docs(notes): point-source CPU campaign phase 4b — step-0 containment …
Jammy2211 Sep 27, 2026
dbc43d5
solver_config_sweep: backend-aware fiducial gate (CPU …812, A100 …806)
Jammy2211 Sep 27, 2026
6deb83c
results: quotable RAL 8490H step-0 route A/B (job 357321) + ledger Ph…
Jammy2211 Sep 27, 2026
cc46e65
hpc(point_source_image): step-0 submits' WALL-BASIS rows use the run …
Jammy2211 Sep 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
154 changes: 154 additions & 0 deletions ...batch_cpu/submit_breakdown_point_source_image_solver_config_sweep_step0_gpu_node_cpu_fp64
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,154 @@
#!/bin/bash -l
#
# 2026-09 point-source CPU speed-up campaign, phase 4b: the step-0 containment
# route A/B (PyAutoArray#579) on RAL CPU.
#
# WHAT THIS MEASURES. scripts/point_source_image/likelihood_breakdown/
# solver_config_sweep.py --step0-route control,gather,components,structured:
# the production FitPositionsImagePairAllSolved likelihood at the production
# geometry (+-9.9" / 0.2 / 1e-3 / MCS 15) with autoarray's _STEP0_CONTAINMENT
# set per route for tracing only (fresh closure + jax.clear_caches() per route),
# interleaved 20 rounds x 20 calls, bootstrap 90 % CI, vmap 1/4/16, compile_s,
# XLA memory, the route-aware step-0 prefix split, the 200 prior + 200 stress
# completeness draws, and gates.step0_route_bit_identity (stream / fiducial /
# prior log L, prior + stress positions, image counts, grad, vmap bit-identical
# to the control). "control" is the branch's library default ("structured");
# "gather" is the pre-#579 general path, i.e. the "before".
#
# LIBRARY. The branch is unpushed, so its PyAutoArray is a git clone of the
# shared mirror with the branch fetched from a bundle, checked out detached at
# the measured commit (AA_BRANCH_SRC below), prepended to PYTHONPATH. Every
# other library is the shared mirror (/mnt/ral/jnightin/PyAuto, via
# activate.sh); PyAutoLens is mirror main 4487eb47, unchanged by #579. The job
# refuses to run unless autoarray imports from the clone and the rest from the
# mirror, and asserts the JSON source_revisions afterwards.
#
# THREADS. OMP/MKL/OPENBLAS pinned to 1 (BLAS family only); the JAX CPU backend
# sizes its Eigen pool from the 8 CPUs of --cpus-per-task=8, as phases 1-4a.
# XLA_FLAGS are the phase-1 row's, byte for byte.
#
# NODE. FALLBACK for the 8490H submit (..._ral_cpu_fp64): on 2026-09-27 the `ral`
# partition was saturated (euclid_dr1 arrays; 10-4 had its memory fully
# allocated at load ~155), so this leg runs CPU-only on an idle node of the `gpu`
# partition (AMD EPYC 7702), with NO --gres: JAX_PLATFORMS=cpu and no GPU is
# visible. Pinned to euclid-ral-gpu-2 so it does not share a node with the A100
# leg on gpu-1. Only in-job ratios compare with the 8490H rows of phase 4a.

# WALL-BASIS:
# RAL job 357335 (this cell, same 4 routes, EPYC 7702 euclid-ral-gpu-2 CPU-only,
# 8 CPUs, 200 + 200 draws, 20 x 20) ran 417 s (7:14 job wall).
# cell: point_source_image/solver_config_sweep/simple device: cpu precision: fp64
# source: measured-wall wall: 417 ref: ral-job-357335-point_source_image-solver_config_sweep headroom: 3.0

#SBATCH -J point_source_step0_p4b_epyc
#SBATCH --partition=gpu
#SBATCH --nodelist=euclid-ral-gpu-2
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=8
#SBATCH --mem=32gb
#SBATCH --time=1:30: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="$(cd "${SLURM_SUBMIT_DIR:-$(dirname "$0")}/../.." && pwd)"
echo "AP_ROOT=$AP_ROOT"
source $AP_ROOT/activate.sh

export AA_BRANCH_SRC="${AA_BRANCH_SRC:-/mnt/ral/jnightin/autolens_profiling_wt/PyAutoArray_pointsolver-step0-gather}"
if [ ! -f "$AA_BRANCH_SRC/autoarray/__init__.py" ]; then
echo "FATAL: $AA_BRANCH_SRC is not a PyAutoArray checkout (see LIBRARY)" >&2
exit 2
fi
export PYTHONPATH="$AA_BRANCH_SRC:$PYTHONPATH"

export OMP_NUM_THREADS=1
export MKL_NUM_THREADS=1
export OPENBLAS_NUM_THREADS=1
export NPROC=${SLURM_CPUS_PER_TASK:-8}
export JAX_PLATFORMS=cpu
export JAX_ENABLE_X64=1
export XLA_FLAGS="--xla_disable_hlo_passes=constant_folding --xla_gpu_autotune_level=0 --xla_gpu_enable_triton_gemm=false"
export NUMBA_CACHE_DIR=$AP_ROOT/output/numba_cache_${SLURM_JOB_ID:-local}
mkdir -p "$NUMBA_CACHE_DIR" "$MPLCONFIGDIR"
# No persistent compile cache: compile_s is a gate and must be cold per route.
unset JAX_COMPILATION_CACHE_DIR

CONFIG_NAME=hpc_ral_cpu_epyc7702_fp64
RESULTS_DIR=$AP_ROOT/results/breakdown/point_source_image
mkdir -p "$RESULTS_DIR"

echo "=========================================="
date
hostname
lscpu | grep 'Model name' || true
echo "partition: ${SLURM_JOB_PARTITION:-unset}"
echo "cpus-per-task: ${SLURM_CPUS_PER_TASK:-unset}"
echo "sched_affinity: $(python3 -c 'import os; print(len(os.sched_getaffinity(0)))')"
echo "loadavg: $(cat /proc/loadavg)"
echo "Cell: point_source_image/likelihood_breakdown/solver_config_sweep --step0-route"
echo "Thread env: OMP=$OMP_NUM_THREADS MKL=$MKL_NUM_THREADS OPENBLAS=$OPENBLAS_NUM_THREADS NPROC=$NPROC"
echo "XLA_FLAGS: $XLA_FLAGS"
echo "PYTHONPATH: $PYTHONPATH"

echo "=== source revisions ==="
for r in PyAutoNerves PyAutoFit PyAutoArray PyAutoGalaxy PyAutoLens; do
printf '%-14s %s\n' "$r (mirror)" "$(git -C /mnt/ral/jnightin/PyAuto/$r rev-parse HEAD 2>&1)"
done
echo "PyAutoArray (branch clone, IMPORTED): $(git -C $AA_BRANCH_SRC rev-parse HEAD 2>&1)"
echo "AP_ROOT rev: $(git -C $AP_ROOT rev-parse HEAD 2>&1)"

cd $AP_ROOT

echo "=== import provenance ==="
python3 - <<'PYEOF' || exit 3
import os
if os.environ.get("CUDA_VISIBLE_DEVICES") not in (None, "", "NoDevFiles"):
raise SystemExit(f"REFUSING: a GPU is visible ({os.environ['CUDA_VISIBLE_DEVICES']}) on the CPU leg")
import autoarray, autofit, autogalaxy, autolens, jax, jaxlib, numpy
aa = os.path.realpath(os.environ["AA_BRANCH_SRC"])
for m in (autoarray, autofit, autogalaxy, autolens):
print(f"{m.__name__:<11} {m.__version__:<12} {m.__file__}")
if not os.path.realpath(autoarray.__file__).startswith(aa + "/"):
raise SystemExit(f"REFUSING: autoarray from {autoarray.__file__}, not {aa}")
for m in (autofit, autogalaxy, autolens):
if not m.__file__.startswith("/mnt/ral/jnightin/PyAuto/"):
raise SystemExit(f"REFUSING: {m.__name__} from {m.__file__}, not the shared mirror")
from autoarray.structures.triangles import array as a
print("_STEP0_CONTAINMENT", a._STEP0_CONTAINMENT)
print("jax", jax.__version__, "jaxlib", jaxlib.__version__, "numpy", numpy.__version__, jax.devices())
PYEOF

echo "=========================================="
echo "RUN ${CONFIG_NAME} -> ${RESULTS_DIR}"
date
echo "loadavg: $(cat /proc/loadavg)"
PYTHONUNBUFFERED=1 python3 scripts/point_source_image/likelihood_breakdown/solver_config_sweep.py \
--config-name "$CONFIG_NAME" --output-dir "$RESULTS_DIR" \
--step0-route control,gather,components,structured \
--n-draws 200 --n-stress 200 --rounds 20 --calls 20

echo "=== source_revisions check (JSON vs imported HEADs) ==="
python3 - "$RESULTS_DIR/solver_config_sweep_step0_${CONFIG_NAME}.json" <<'PYEOF'
import json, os, subprocess, sys
path = sys.argv[1]
d = json.load(open(path))
revs = d["source_revisions"]
def head(p):
return subprocess.check_output(["git", "-C", p, "rev-parse", "HEAD"], text=True).strip()
want = {r: head(f"/mnt/ral/jnightin/PyAuto/{r}") for r in ("PyAutoNerves", "PyAutoFit", "PyAutoGalaxy", "PyAutoLens")}
want["PyAutoArray"] = head(os.environ["AA_BRANCH_SRC"])
for r, h in want.items():
if revs[r] != h:
raise SystemExit(f"MISMATCH {path}: {r} JSON {revs[r]} != {h}")
subprocess.check_call(["git", "-C", os.environ["AA_BRANCH_SRC"], "merge-base", "--is-ancestor",
"c13b2d733538bec86cff3a15eccd9a0940a68463", revs["PyAutoArray"]])
print("autoarray_imported_from", d.get("autoarray_imported_from"))
print("OK", path, "all_gates_pass", d.get("all_gates_pass"))
PYEOF

echo "=========================================="
echo "loadavg: $(cat /proc/loadavg)"
echo "Finished."
date
149 changes: 149 additions & 0 deletions hpc/batch_cpu/submit_breakdown_point_source_image_solver_config_sweep_step0_ral_cpu_fp64
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,149 @@
#!/bin/bash -l
#
# 2026-09 point-source CPU speed-up campaign, phase 4b: the step-0 containment
# route A/B (PyAutoArray#579) on RAL CPU.
#
# WHAT THIS MEASURES. scripts/point_source_image/likelihood_breakdown/
# solver_config_sweep.py --step0-route control,gather,components,structured:
# the production FitPositionsImagePairAllSolved likelihood at the production
# geometry (+-9.9" / 0.2 / 1e-3 / MCS 15) with autoarray's _STEP0_CONTAINMENT
# set per route for tracing only (fresh closure + jax.clear_caches() per route),
# interleaved 20 rounds x 20 calls, bootstrap 90 % CI, vmap 1/4/16, compile_s,
# XLA memory, the route-aware step-0 prefix split, the 200 prior + 200 stress
# completeness draws, and gates.step0_route_bit_identity (stream / fiducial /
# prior log L, prior + stress positions, image counts, grad, vmap bit-identical
# to the control). "control" is the branch's library default ("structured");
# "gather" is the pre-#579 general path, i.e. the "before".
#
# LIBRARY. The branch is unpushed, so its PyAutoArray is a git clone of the
# shared mirror with the branch fetched from a bundle, checked out detached at
# the measured commit (AA_BRANCH_SRC below), prepended to PYTHONPATH. Every
# other library is the shared mirror (/mnt/ral/jnightin/PyAuto, via
# activate.sh); PyAutoLens is mirror main 4487eb47, unchanged by #579. The job
# refuses to run unless autoarray imports from the clone and the rest from the
# mirror, and asserts the JSON source_revisions afterwards.
#
# THREADS. OMP/MKL/OPENBLAS pinned to 1 (BLAS family only); the JAX CPU backend
# sizes its Eigen pool from the 8 CPUs of --cpus-per-task=8, as phases 1-4a.
# XLA_FLAGS are the phase-1 row's, byte for byte.
#
# NODE. Pinned to the Xeon Platinum 8490H node of phase 4a (euclid-ral-compute-10-4,
# jobs 356365 / 356367). Check `sinfo -p ral -N -o "%N %T %C"` first: `idle*`
# nodes are unreachable.

# WALL-BASIS:
# RAL job 357321 (this cell, --step0-route control,gather,components,structured,
# 8490H euclid-ral-compute-10-4 under load ~200, 8 CPUs, 200 + 200 draws, 20 x 20)
# ran 480 s (8:08 job wall). Headroom 3 covers a quieter or busier node.
# cell: point_source_image/solver_config_sweep/simple device: cpu precision: fp64
# source: measured-wall wall: 480 ref: ral-job-357321-point_source_image-solver_config_sweep headroom: 3.0

#SBATCH -J point_source_step0_p4b
#SBATCH --partition=ral
#SBATCH --nodelist=euclid-ral-compute-10-4
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=8
#SBATCH --mem=32gb
#SBATCH --time=1:30: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="$(cd "${SLURM_SUBMIT_DIR:-$(dirname "$0")}/../.." && pwd)"
echo "AP_ROOT=$AP_ROOT"
source $AP_ROOT/activate.sh

export AA_BRANCH_SRC="${AA_BRANCH_SRC:-/mnt/ral/jnightin/autolens_profiling_wt/PyAutoArray_pointsolver-step0-gather}"
if [ ! -f "$AA_BRANCH_SRC/autoarray/__init__.py" ]; then
echo "FATAL: $AA_BRANCH_SRC is not a PyAutoArray checkout (see LIBRARY)" >&2
exit 2
fi
export PYTHONPATH="$AA_BRANCH_SRC:$PYTHONPATH"

export OMP_NUM_THREADS=1
export MKL_NUM_THREADS=1
export OPENBLAS_NUM_THREADS=1
export NPROC=${SLURM_CPUS_PER_TASK:-8}
export JAX_PLATFORMS=cpu
export JAX_ENABLE_X64=1
export XLA_FLAGS="--xla_disable_hlo_passes=constant_folding --xla_gpu_autotune_level=0 --xla_gpu_enable_triton_gemm=false"
export NUMBA_CACHE_DIR=$AP_ROOT/output/numba_cache_${SLURM_JOB_ID:-local}
mkdir -p "$NUMBA_CACHE_DIR" "$MPLCONFIGDIR"
# No persistent compile cache: compile_s is a gate and must be cold per route.
unset JAX_COMPILATION_CACHE_DIR

CONFIG_NAME=hpc_ral_cpu_fp64
RESULTS_DIR=$AP_ROOT/results/breakdown/point_source_image
mkdir -p "$RESULTS_DIR"

echo "=========================================="
date
hostname
lscpu | grep 'Model name' || true
echo "partition: ${SLURM_JOB_PARTITION:-unset}"
echo "cpus-per-task: ${SLURM_CPUS_PER_TASK:-unset}"
echo "sched_affinity: $(python3 -c 'import os; print(len(os.sched_getaffinity(0)))')"
echo "loadavg: $(cat /proc/loadavg)"
echo "Cell: point_source_image/likelihood_breakdown/solver_config_sweep --step0-route"
echo "Thread env: OMP=$OMP_NUM_THREADS MKL=$MKL_NUM_THREADS OPENBLAS=$OPENBLAS_NUM_THREADS NPROC=$NPROC"
echo "XLA_FLAGS: $XLA_FLAGS"
echo "PYTHONPATH: $PYTHONPATH"

echo "=== source revisions ==="
for r in PyAutoNerves PyAutoFit PyAutoArray PyAutoGalaxy PyAutoLens; do
printf '%-14s %s\n' "$r (mirror)" "$(git -C /mnt/ral/jnightin/PyAuto/$r rev-parse HEAD 2>&1)"
done
echo "PyAutoArray (branch clone, IMPORTED): $(git -C $AA_BRANCH_SRC rev-parse HEAD 2>&1)"
echo "AP_ROOT rev: $(git -C $AP_ROOT rev-parse HEAD 2>&1)"

cd $AP_ROOT

echo "=== import provenance ==="
python3 - <<'PYEOF' || exit 3
import os, autoarray, autofit, autogalaxy, autolens, jax, jaxlib, numpy
aa = os.path.realpath(os.environ["AA_BRANCH_SRC"])
for m in (autoarray, autofit, autogalaxy, autolens):
print(f"{m.__name__:<11} {m.__version__:<12} {m.__file__}")
if not os.path.realpath(autoarray.__file__).startswith(aa + "/"):
raise SystemExit(f"REFUSING: autoarray from {autoarray.__file__}, not {aa}")
for m in (autofit, autogalaxy, autolens):
if not m.__file__.startswith("/mnt/ral/jnightin/PyAuto/"):
raise SystemExit(f"REFUSING: {m.__name__} from {m.__file__}, not the shared mirror")
from autoarray.structures.triangles import array as a
print("_STEP0_CONTAINMENT", a._STEP0_CONTAINMENT)
print("jax", jax.__version__, "jaxlib", jaxlib.__version__, "numpy", numpy.__version__, jax.devices())
PYEOF

echo "=========================================="
echo "RUN ${CONFIG_NAME} -> ${RESULTS_DIR}"
date
echo "loadavg: $(cat /proc/loadavg)"
PYTHONUNBUFFERED=1 python3 scripts/point_source_image/likelihood_breakdown/solver_config_sweep.py \
--config-name "$CONFIG_NAME" --output-dir "$RESULTS_DIR" \
--step0-route control,gather,components,structured \
--n-draws 200 --n-stress 200 --rounds 20 --calls 20

echo "=== source_revisions check (JSON vs imported HEADs) ==="
python3 - "$RESULTS_DIR/solver_config_sweep_step0_${CONFIG_NAME}.json" <<'PYEOF'
import json, os, subprocess, sys
path = sys.argv[1]
d = json.load(open(path))
revs = d["source_revisions"]
def head(p):
return subprocess.check_output(["git", "-C", p, "rev-parse", "HEAD"], text=True).strip()
want = {r: head(f"/mnt/ral/jnightin/PyAuto/{r}") for r in ("PyAutoNerves", "PyAutoFit", "PyAutoGalaxy", "PyAutoLens")}
want["PyAutoArray"] = head(os.environ["AA_BRANCH_SRC"])
for r, h in want.items():
if revs[r] != h:
raise SystemExit(f"MISMATCH {path}: {r} JSON {revs[r]} != {h}")
subprocess.check_call(["git", "-C", os.environ["AA_BRANCH_SRC"], "merge-base", "--is-ancestor",
"c13b2d733538bec86cff3a15eccd9a0940a68463", revs["PyAutoArray"]])
print("autoarray_imported_from", d.get("autoarray_imported_from"))
print("OK", path, "all_gates_pass", d.get("all_gates_pass"))
PYEOF

echo "=========================================="
echo "loadavg: $(cat /proc/loadavg)"
echo "Finished."
date
Loading
Loading