gh-1088: Add realistic benchmark tests

.gitignore

@@ -16,3 +16,4 @@ dist
 styles
 tests/regression/**/results
 uv.lock
+healpy-data

benchmarks/README.md

@@ -0,0 +1,28 @@
+# GLASS Benchmarks
+
+These benchmarks are intended to allow benchmarking GLASS on various machines,
+architectures and node configurations.
+
+Before running any benchmarks, you must first setup your venv with the required
+dependencies. For all backends, this will require installing the dependency
+group `benchmarks`, like so,
+
+```sh
+uv sync --group benchmarks
+```
+
+However, it may also be necessary to install other dependencies / dependency
+groups. For example, read the
+[prerequisites for Archer2](./archer2/README.md#prerequisites)
+
+##  healpy-data
+
+glass depends on data from the healpy-data repo. If not found locally, glass
+downloads this from the internet. Since most clusters will not allow internet
+access from a worker node. We must provide a local copy of this data before
+submitting a job. To download this data we can use git:
+
+```sh
+git clone --depth 1 https://github.com/healpy/healpy-data
+export HEALPY_DATAPATH="$(pwd)/healpy-data"
+```

benchmarks/__init__.py

@@ -0,0 +1 @@
+"""Benchmarks from GLASS."""

benchmarks/archer2/README.md

@@ -0,0 +1,90 @@
+# Archer2
+
+This document contains instructions specific to running the glass benchmarks on
+Archer2
+
+## Prerequisites
+
+As Archer2 is a managed cluster, there are several steps to setting up your
+environment. First you should [setup uv](#setting-up-uv-on-archer2). Then you
+will need to install you python virtual environment. Finally, you will need to
+load specific modules for the GPU benchmark.
+
+Note that for the GPU benchmarks Archer2 only support rocm up to v0.6.x.
+Therefore, we are restricted to using `jax==0.4.35`. This in turn restricts us
+to using python 3.12. Thus, to produce a useful CPU vs GPU comparison create
+your venv using the following command
+
+```sh
+uv venv --python 3.12
+```
+
+### CPU prerequisites
+
+To setup your python environment for running the cpu benchmark on Archer2, run
+the following commands.
+
+```sh
+uv sync --group benchmarks
+```
+
+### GPU prerequisites
+
+For the gpu benchmark, there is an additional dependency group `archer2-gpu`:
+
+```sh
+uv sync --group benchmarks --group archer2-gpu
+```
+
+Once your python environment is setup you must load the relevant modules via the
+provided script [setup-gpu-env.sh](./setup-gpu-env.sh).
+
+> Note that setup-gpu-env.sh is called automatically by the submission script
+> [submit-gpu.sh](./submit-gpu.sh)
+
+## Running the benchmarks
+
+Benchmarks should be submitted as a batch job to slurm via the provided script.
+
+For example to benchmark using jax with amd/rocm, run the following from the
+root of the glass repo on Archer2:
+
+```sh
+sbatch benchmarks/archer2/submit-gpu.sh -d "$(pwd)" -x jax --healpy-datapath "$HEALPY_DATAPATH"
+```
+
+> To understand what HEALPY_DATAPATH is, read an explanation in
+> [benchmarks/README.md#healpy-data](../README.md#healpy-data)
+
+## Setting up UV on Archer2
+
+Firstly, install uv via curl onto the `/work` partition
+
+```sh
+cd "${HOME/home/work}"
+curl -LsSf https://astral.sh/uv/install.sh | sh
+```
+
+Then we must make sure uv is available on the login node and the worker node. To
+do this we can update our start up scripts (`.profile`) on both partitions.
+Therefore, execute the following
+
+```sh
+cat <<'EOF' >> "$HOME/.profile"
+WORK_DIR="${HOME/home/work}"
+cd "$WORK_DIR"
+source "$WORK_DIR/.profile"
+EOF
+```
+
+and similarly
+
+```sh
+cat <<'EOF' >> "${HOME/home/work}/.profile"
+export HOME="${HOME/home/work}"
+source "$HOME/.local/bin/env"
+EOF
+```
+
+Now when you next login to archer2, uv will be in your path and you will be on
+the `/work` partition as your `HOME` dir.

benchmarks/archer2/setup-gpu-env.sh

@@ -0,0 +1,11 @@
+#!/bin/bash --login
+
+# Load GPU modules
+module load PrgEnv-amd/8.6.0
+module load rocm
+module load craype-accel-amd-gfx90a
+module load craype-x86-milan
+
+# Ensure the rocm library and build is know to jax
+export LD_LIBRARY_PATH="/opt/rocm/lib:$LD_LIBRARY_PATH"
+export ROCM_PATH="/opt/rocm"

benchmarks/archer2/submit-cpu.sh

@@ -0,0 +1,109 @@
+#!/bin/bash --login
+# shellcheck disable=SC1091
+
+#SBATCH --job-name=glass_benchmark_cpu
+#SBATCH --output=%x-%j.out
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --cpus-per-task=1
+#SBATCH --time=0:30:0
+#SBATCH --partition=standard
+#SBATCH --qos=standard
+
+GLASS_DIR=""
+ARRAY_BACKEND="numpy"
+HEALPY_DATAPATH=""
+RUN_PROFILE="false"
+
+help() {
+  echo "Usage:"
+  echo "    $0 -d <glass/dir> [-x <array_backend>] [--healpy-datapath <healpy-datapath>] [-h|--help]"
+  echo ""
+  echo "ARGS:"
+  echo "    -h | --help                          Display this help message."
+  echo "    -d | --glass-dir <glass/dir>         Path to the cloned glass directory."
+  echo "    -x | --array-backend <array_backend> The array backend to use for the benchmarks."
+  echo "                                         Defaults to NumPy."
+  echo "    --healpy-datapath <healpy-datapath>  The path to the healpy-data repo to allow"
+  echo "                                         running offline. Defaults to <glass/dir>/healpy-data"
+}
+
+# Ensure uv is available
+source "${HOME/home/work}/.profile" # HOME starts as /home/... but uv needs to be on /work/...
+
+# Stop unintentional multi-threading within software libraries
+export OMP_NUM_THREADS=1
+# Ensure the cpus-per-task option is propagated to srun commands
+export SRUN_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK
+
+# check for no input arguments and show help
+if [ $# -eq 0 ];
+then
+    help
+    exit 1
+fi
+
+while [ $# -gt 0 ] ; do
+    case $1 in
+        -h | --help)
+            help
+            exit 0
+            ;;
+        -d | --glass-dir)
+            GLASS_DIR="$2"
+            shift 2
+            continue
+            ;;
+        -x | --array-backend)
+            ARRAY_BACKEND="$2"
+            shift 2
+            continue
+            ;;
+        -p | --profile)
+            RUN_PROFILE="true"
+            shift 1
+            continue
+            ;;
+        --healpy-datapath)
+            HEALPY_DATAPATH="$2"
+            shift 2
+            continue
+            ;;
+        *)
+            echo "Invalid option: $1" >&2;
+            help
+            exit 1
+            ;;
+    esac
+    shift 1
+done
+
+# Ensure GLASS_DIR is provided
+if [[ "$GLASS_DIR" == "" ]]
+then
+  echo "GLASS_DIR must be provided"
+  help
+  exit 1
+fi
+
+# Set HEALPY_DATAPATH default
+if [[ "$HEALPY_DATAPATH" == "" ]]; then
+    HEALPY_DATAPATH="$GLASS_DIR/healpy-data"
+fi
+
+# Get execution method
+PYTHON_COMMAND="$GLASS_DIR/.venv/bin/python"
+if [[ -n "$SLURM_JOB_ID" || -n "$SLURM_BATCH_SCRIPT" ]]; then
+  echo "Running under SLURM (batch). SLURM_JOB_ID=${SLURM_JOB_ID:-unknown}"
+  PYTHON_COMMAND="srun $PYTHON_COMMAND"
+else
+  echo "Running directly from CLI"
+fi
+
+PROFILE_ENV_VAR=""
+if [[ "$RUN_PROFILE" == "true" ]]
+then
+    PROFILE_ENV_VAR="RUN_PROFILE='$RUN_PROFILE'"
+fi
+
+ARRAY_BACKEND="$ARRAY_BACKEND" HEALPY_DATAPATH="$HEALPY_DATAPATH"  "$PROFILE_ENV_VAR" $PYTHON_COMMAND benchmarks/lensing.py

benchmarks/archer2/submit-gpu.sh

@@ -0,0 +1,98 @@
+#!/bin/bash --login
+# shellcheck disable=SC1091
+
+#SBATCH --job-name=glass_benchmark_gpu
+#SBATCH --output=%x-%j.out
+#SBATCH --cpus-per-task=1
+#SBATCH --nodes=1
+#SBATCH --gpus=1
+#SBATCH --time=0:30:0
+#SBATCH --partition=gpu
+#SBATCH --qos=gpu-shd
+
+# Recommended environment settings
+# Stop unintentional multi-threading within software libraries
+export OMP_NUM_THREADS=1
+# Ensure the cpus-per-task option is propagated to srun commands
+export SRUN_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK
+
+GLASS_DIR=""
+ARRAY_BACKEND="numpy"
+HEALPY_DATAPATH=""
+
+help() {
+  echo "Usage:"
+  echo "    $0 -d <glass/dir> [-x <array_backend>] [--healpy-datapath <healpy-datapath>] [-h|--help]"
+  echo ""
+  echo "ARGS:"
+  echo "    -h | --help                          Display this help message."
+  echo "    -d | --glass-dir <glass/dir>         Path to the cloned glass directory."
+  echo "    -x | --array-backend <array_backend> The array backend to use for the benchmarks."
+  echo "                                         Defaults to NumPy."
+  echo "    --healpy-datapath <healpy-datapath>  The path to the healpy-data repo to allow"
+  echo "                                         running offline. Defaults to <glass/dir>/healpy-data"
+}
+
+# Ensure uv is available
+source "${HOME/home/work}/.profile" # HOME starts as /home/... but uv needs to be on /work/...
+
+# Stop unintentional multi-threading within software libraries
+export OMP_NUM_THREADS=1
+# Ensure the cpus-per-task option is propagated to srun commands
+export SRUN_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK
+
+# check for no input arguments and show help
+if [ $# -eq 0 ];
+then
+    help
+    exit 1
+fi
+
+while [ $# -gt 0 ] ; do
+    case $1 in
+        -h | --help)
+            help
+            exit 0
+            ;;
+        -d | --glass-dir)
+            GLASS_DIR="$2"
+            shift 2
+            continue
+            ;;
+        -x | --array-backend)
+            ARRAY_BACKEND="$2"
+            shift 2
+            continue
+            ;;
+        --healpy-datapath)
+            HEALPY_DATAPATH="$2"
+            shift 2
+            continue
+            ;;
+        *)
+            echo "Invalid option: $1" >&2;
+            help
+            exit 1
+            ;;
+    esac
+    shift 1
+done
+
+# Ensure GLASS_DIR is provided
+if [[ "$GLASS_DIR" == "" ]]
+then
+  echo "GLASS_DIR must be provided"
+  help
+  exit 1
+fi
+
+# Set HEALPY_DATAPATH default
+if [[ "$HEALPY_DATAPATH" == "" ]]; then
+    HEALPY_DATAPATH="$GLASS_DIR/healpy-data"
+fi
+
+# Setup environment
+source "$GLASS_DIR/benchmarks/archer2/setup-gpu-env.sh"
+
+# Run benchmark via slurm
+HEALPY_DATAPATH="$HEALPY_DATAPATH" ARRAY_BACKEND="$ARRAY_BACKEND" srun "$GLASS_DIR/.venv/bin/python" benchmarks/lensing.py