gh-1084: Add config to run on archer2 with GPU

.shellcheckrc

@@ -0,0 +1,2 @@
+disable=SC1091 # openBinaryFile: does not exist (No such file or directory)
+disable=SC1090 # Can't follow non-constant source. Use a directive to specify location

noxfile.py

@@ -237,6 +237,8 @@ def benchmarks(session: nox.Session) -> None:
     _check_revision_count(session.posargs, expected_count=1)
     revision = session.posargs[0]
 
+    _setup_array_backend(session)
+
     # overwrite current package with specified revision
     session.install(f"git+{GLASS_REPO_URL}@{revision}")
     session.run(

pyproject.toml

@@ -7,6 +7,9 @@ requires = [
 ]
 
 [dependency-groups]
+archer2-gpu = [
+    "jax[rocm7-local]",
+]
 build = [
     "build",
 ]

tests/benchmarks/archer2/README.md

@@ -0,0 +1,47 @@
+# GLASS benchmarks on Archer2
+
+For a more consistent benchmarking and regression testing environment we have
+been trialing using the UCL machine [Archer2](https://www.archer2.ac.uk/).
+
+## Setting up
+
+1. **Install uv:** 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, save the following into `$HOME/.profile`
+
+   ```sh
+   WORK_DIR="${HOME/home/work}"
+   cd "$WORK_DIR"
+   source "$WORK_DIR/.profile"
+   ```
+
+   and then save the following into `${HOME/home/work}/.profile`.
+
+   ```sh
+   export HOME="${HOME/home/work}"
+   source "$HOME/.local/env"
+   ```
+
+   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.
+
+2. **Clone GLASS:** Clone the glass repo into the `/work` partition of Archer2 -
+   `/work/<budget-code>/<budget-code>/<user-id>`
+
+   ```sh
+   cd "${HOME/home/work}"
+   git clone https://github.com/glass-dev/glass.git
+   ```
+
+3. **Run the benchmarks script:** Now we have cloned glass, we can run the
+   script [run_regression_test.sh](./run_regression_test.sh) which will setup
+   the required environments and submit regression test job to slurm. A help
+   message is provided. Just run `run_regression_test.sh -h` from the root of
+   the GLASS repo.

tests/benchmarks/archer2/run_regression_test.sh

@@ -0,0 +1,137 @@
+#!/bin/bash -l
+
+GLASS_DIR=""
+GLASS_REPO_URL="https://github.com/glass-dev/glass"
+CPU_OR_GPU="cpu"
+START_REF=""
+END_REF=""
+START_VENV=".venv-start"
+END_VENV=".venv-end"
+ACCOUNT=""
+SETUP_ENVS="true"
+
+help() {
+  echo "Usage:"
+  echo "    $0 "
+  echo ""
+  echo "ARGS:"
+  echo "    -h | --help                      Display this help message."
+  echo "    -d | --glass-dir <glass/dir>     Path to the cloned glass directory."
+  echo "    -s | --start-ref <start_ref>     The git ref to be used as the initial state."
+  echo "    -e | --end-ref <end_ref>         The git ref to be used as the final state."
+  echo "    -a | --account <archer2_account> The archer2 account code to run jobs against."
+  echo "    -g | --gpu                       Flag to state the gpu benchmark should be ran."
+  echo "    --skip-setup                     Flag to state if the setup (installation of "
+  echo "                                     dependencies) can be skipped. Good for simply"
+  echo "                                     re-submitting"
+}
+
+# check for no input arguments and show help
+if [ $# -eq 0 ];
+then
+    help
+    exit 1
+fi
+
+# parse input arguments
+while [ $# -gt 0 ] ; do
+    case $1 in
+        -h | --help)
+            help
+            exit 0
+            ;;
+        -d | --glass-dir)
+            GLASS_DIR=$2
+            shift 2
+            continue
+            ;;
+        -s | --start-ref)
+            START_REF=$2
+            shift 2
+            continue
+            ;;
+        -e | --end-ref)
+            END_REF=$2
+            shift 2
+            continue
+            ;;
+        -a | --account)
+            ACCOUNT=$2
+            shift 2
+            continue
+            ;;
+        -g | --gpu)
+            CPU_OR_GPU="gpu"
+            shift 1
+            continue
+            ;;
+        --skip-setup)
+            SETUP_ENVS="false"
+            shift 1
+            continue
+            ;;
+        *)
+            echo "Invalid option: $1" >&2;
+            help
+            exit 1
+            ;;
+    esac
+    shift 1
+done
+
+if [[ "$START_REF" == "" || "$END_REF" == "" ]]
+then
+  echo "START_REF and END_REF must be provided"
+  help
+  exit 1
+fi
+
+if [[ "$GLASS_DIR" == "" ]]
+then
+  echo "GLASS_DIR must be provided"
+  help
+  exit 1
+fi
+
+if [[ "$ACCOUNT" == "" ]]
+then
+  echo "ACCOUNT must be provided"
+  help
+  exit 1
+fi
+
+BENCHMARKS_DIR="$GLASS_DIR/tests/benchmarks"
+
+
+if [[ "$SETUP_ENVS" == "true" ]]
+then
+  # Setup base environment
+  rm -rf "${GLASS_DIR:?}/$START_VENV" # Cleanup old venv
+  uv venv  "$GLASS_DIR/$START_VENV"
+  source "$GLASS_DIR/$START_VENV/bin/activate"
+  if [ $CPU_OR_GPU = "gpu" ]; then
+    uv sync --active --group test --group archer2-gpu
+  else
+    uv sync --active --group test
+  fi
+  uv pip uninstall glass -y # Make sure no installation of glass already exists
+  uv pip install "git+$GLASS_REPO_URL@$START_REF"
+
+  # Setup head environment
+  rm -rf "${GLASS_DIR:?}/$END_VENV" # Cleanup old venv
+  uv venv "$GLASS_DIR/$END_VENV"
+  source "$GLASS_DIR/$END_VENV/bin/activate"
+  if [ $CPU_OR_GPU = "gpu" ]; then
+    uv sync --active --group test --group archer2-gpu
+  else
+    uv sync --active --group test
+  fi
+  uv pip uninstall glass -y # Make sure no installation of glass already exists
+  uv pip install "git+$GLASS_REPO_URL@$END_REF"
+
+  # Remove old benchmark results
+  rm -rf "$BENCHMARKS_DIR/outputs"
+fi
+
+# Submit job
+sbatch --account="$ACCOUNT" "$BENCHMARKS_DIR/archer2/submission_script_$CPU_OR_GPU.sh" "$GLASS_DIR"

tests/benchmarks/archer2/submission_script_cpu.sh

@@ -0,0 +1,41 @@
+#!/bin/bash --login
+
+#SBATCH --job-name=glass_reg_test_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
+
+# Ensure uv is available
+source "${HOME/home/work}/.profile" # HOME starts as /home/... but uv needs to be on /work/...
+
+# 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
+
+# Set path to class and select base or head
+GLASS_DIR="$1"
+BENCHMARKS_DIR="$GLASS_DIR/tests/benchmarks"
+BENCHMARKS_SHARED_FLAGS=(
+  "--benchmark-storage=file://$BENCHMARKS_DIR/outputs"
+  "--benchmark-calibration-precision=1000"
+  "--benchmark-columns=mean,stddev,rounds"
+  "--benchmark-max-time=5.0"
+  "--benchmark-sort=name"
+  "--benchmark-timer=time.process_time"
+)
+
+# Generate the base report for comparison later
+source "$GLASS_DIR/.venv-start/bin/activate"
+srun uv run --active pytest "$BENCHMARKS_DIR" --benchmark-autosave "${BENCHMARKS_SHARED_FLAGS[@]}"
+deactivate
+
+# Run the stable and unstable benchmarks and compare to the base ref
+source "$GLASS_DIR/.venv-end/bin/activate"
+srun uv run --active pytest "$BENCHMARKS_DIR" -m stable --benchmark-compare=0001 --benchmark-compare-fail=mean:5% "${BENCHMARKS_SHARED_FLAGS[@]}"
+srun uv run --active pytest "$BENCHMARKS_DIR" -m unstable --benchmark-compare=0001 --benchmark-compare-fail=mean:0.0005 "${BENCHMARKS_SHARED_FLAGS[@]}"

tests/benchmarks/archer2/submission_script_gpu.sh

@@ -0,0 +1,46 @@
+#!/bin/bash --login
+
+#SBATCH --job-name=glass_reg_test_gpu
+#SBATCH --output=%x-%j.out
+#SBATCH --cpus-per-task=1
+#SBATCH --gpus=1
+#SBATCH --time=0:30:0
+#SBATCH --partition=gpu
+#SBATCH --qos=gpu-shd
+
+# Load GPU modules
+module load PrgEnv-amd/8.4.0
+module load rocm
+module load craype-accel-amd-gfx90a
+module load craype-x86-milan
+
+# Ensure uv is available
+source "${HOME/home/work}/.profile" # HOME starts as /home/... but uv need to be on /work/...
+
+# 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
+
+# Set path to class and select base or head
+GLASS_DIR="$1"
+BENCHMARKS_DIR="$GLASS_DIR/tests/benchmarks"
+BENCHMARKS_SHARED_FLAGS=(
+  "--benchmark-storage=file://$BENCHMARKS_DIR/outputs"
+  "--benchmark-calibration-precision=1000"
+  "--benchmark-columns=mean,stddev,rounds"
+  "--benchmark-max-time=5.0"
+  "--benchmark-sort=name"
+  "--benchmark-timer=time.process_time"
+)
+
+# Generate the base report for comparison later
+source "$GLASS_DIR/.venv-start/bin/activate"
+srun uv run --active pytest "$BENCHMARKS_DIR" --benchmark-autosave "${BENCHMARKS_SHARED_FLAGS[@]}"
+deactivate
+
+# Run the stable and unstable benchmarks and compare to the base ref
+source "$GLASS_DIR/.venv-end/bin/activate"
+srun uv run --active pytest "$BENCHMARKS_DIR" -m stable --benchmark-compare=0001 --benchmark-compare-fail=mean:5% "${BENCHMARKS_SHARED_FLAGS[@]}"
+srun uv run --active pytest "$BENCHMARKS_DIR" -m unstable --benchmark-compare=0001 --benchmark-compare-fail=mean:0.0005 "${BENCHMARKS_SHARED_FLAGS[@]}"

tests/benchmarks/myriad/run_regression_test.sh

@@ -1,9 +1,9 @@
 #!/bin/bash -l
 
-# Request ten minutes of wallclock time (format hours:minutes:seconds).
+# Request three hours of wallclock time (format hours:minutes:seconds).
 #$ -l h_rt=3:0:0
 
-# Request 1 gigabyte of RAM (must be an integer followed by M, G, or T)
+# Request 4 gigabyte of RAM (must be an integer followed by M, G, or T)
 #$ -l mem=4G
 
 # Request exclusive access to a node

tests/benchmarks/test_fields.py

@@ -125,7 +125,6 @@ def function_to_benchmark() -> list[Any]:
 @pytest.mark.stable
 def test_cls2cov(
     benchmark: BenchmarkFixture,
-    compare: Compare,
     generator_consumer: GeneratorConsumer,
     xpb: ModuleType,
 ) -> None:
@@ -148,10 +147,6 @@ def function_to_benchmark() -> list[Any]:
     assert cov.shape == (nl, nc + 1)
     assert cov.dtype == xpb.float64
 
-    compare.assert_allclose(cov[:, 0], xpb.asarray([1.0, 1.5, 2.0]))
-    compare.assert_allclose(cov[:, 1], xpb.asarray([1.5, 2.0, 2.5]))
-    compare.assert_allclose(cov[:, 2], 0)
-
 
 @pytest.mark.stable
 @pytest.mark.parametrize("use_rng", [False, True])
@@ -164,7 +159,7 @@ def test_generate_grf(
     use_rng: bool,  # noqa: FBT001
 ) -> None:
     """Benchmarks for glass.fields._generate_grf with positional arguments only."""
-    gls: AngularPowerSpectra = [urngb.random(1_000)]
+    gls: AngularPowerSpectra = [urngb.random(1_00)]
     nside = 4
 
     def function_to_benchmark() -> list[Any]:

tests/fixtures/array_backends.py

@@ -109,9 +109,7 @@ def xp(request: pytest.FixtureRequest) -> ModuleType:
 
 @pytest.fixture(
     params=[
-        xp
-        for name, xp in xp_available_backends.items()
-        if name not in {"array_api_strict", "jax.numpy"}
+        xp for name, xp in xp_available_backends.items() if name != "array_api_strict"
     ],
     scope="session",
 )