gh-1088: Add realistic benchmark tests

[connoraird]

Description

• Adds a more realistic benchmark to be used for understanding likely performance of GLASS on various machines and configurations

Closes: #1088 #1084

Changelog entry

Added: Realistic benchmark test for benchmarking glass on various machines

Checks

• Is your code passing linting?
• Is your code passing tests?
• Have you added additional tests (if required)?
• Have you modified/extended the documentation (if required)?
• Have you added a one-liner changelog entry above (if required)?

[connoraird]

@paddyroddy I'm not sure why readthedocs is failing. Any ideas?

[paddyroddy]

@paddyroddy I'm not sure why readthedocs is failing. Any ideas?

I've been trying to fix it, no luck yet

[paddyroddy]

Have gone through. Fixed typos. And tried to fix readthedocs.

[paddyroddy]

This whole file seems very similar to tests/fixtures/array_backends.py. Is there a way we can reduce the duplication?

[connoraird]

I think the only similarity is the if blocks, i.e.:

if not ARRAY_BACKEND or ARRAY_BACKEND == "numpy":
    xp_available_backends["numpy"] = np
	...

Everything else is new, no?

[paddyroddy]

I guess all of this?

glass/tests/fixtures/array_backends.py

Lines 22 to 96 in faec065

	# environment variable to specify array backends for testing
	# can be:
	# - a particular array library (numpy, jax, array_api_strict, ...)
	# - all (try finding every supported array library available in the environment)
	ARRAY_BACKEND: str = os.environ.get("ARRAY_BACKEND", "")
	def _check_version(lib: str, array_api_compliant_version: str) -> None:
	"""
	Check if installed library's version is compliant with the array API standard.
	Parameters
	----------
	lib
	name of the library.
	array_api_compliant_version
	version of the library compliant with the array API standard.
	Raises
	------
	ImportError
	If the installed version is not compliant with the array API standard.
	"""
	lib_version = packaging.version.Version(importlib.metadata.version(lib))
	if lib_version < packaging.version.Version(array_api_compliant_version):
	msg = f"{lib} must be >= {array_api_compliant_version}; found {lib_version}"
	raise ImportError(msg)
	def _import_and_add_numpy(xp_available_backends: dict[str, ModuleType]) -> None:
	"""Add numpy to the backends dictionary."""
	_check_version("numpy", "2.2.6")
	xp_available_backends["numpy"] = np
	def _import_and_add_array_api_strict(
	xp_available_backends: dict[str, ModuleType],
	) -> None:
	"""Add array_api_strict to the backends dictionary."""
	import array_api_strict # noqa: PLC0415
	_check_version("array_api_strict", "2.3.1")
	xp_available_backends["array_api_strict"] = array_api_strict
	array_api_strict.set_array_api_strict_flags(api_version="2025.12")
	def _import_and_add_jax(xp_available_backends: dict[str, ModuleType]) -> None:
	"""Add jax to the backends dictionary."""
	import jax # noqa: PLC0415
	_check_version("jax", "0.6.2")
	xp_available_backends["jax.numpy"] = jax.numpy
	# enable 64 bit numbers
	jax.config.update("jax_enable_x64", val=True)
	# a dictionary with all array backends to test
	xp_available_backends: dict[str, ModuleType] = {}
	# if no backend passed, use numpy by default
	if not ARRAY_BACKEND or ARRAY_BACKEND == "numpy":
	_import_and_add_numpy(xp_available_backends)
	elif ARRAY_BACKEND == "array_api_strict":
	_import_and_add_array_api_strict(xp_available_backends)
	elif ARRAY_BACKEND == "jax":
	_import_and_add_jax(xp_available_backends)
	# if all, try importing every backend
	elif ARRAY_BACKEND == "all":
	_import_and_add_numpy(xp_available_backends)
	_import_and_add_array_api_strict(xp_available_backends)
	_import_and_add_jax(xp_available_backends)
	else:
	msg = f"unsupported array backend: {ARRAY_BACKEND}"
	raise ValueError(msg)

[connoraird]

I decided not to do that here and just always import numpy, jax, array_api_strict, etc. In fact I can probably remove array_api_strict as I doubt we care about benchmarking that.

[paddyroddy]

This seems very unusual. Surely we go function_to_benchmark, xp, args, kwargs?

[connoraird]

I know what you mean. The reason for this is I want to ensure xp is always provided as a named variable. If I have xp, *args, **kwargs then the first positional arg will be assumed to be xp but it doesn't have to be.

[paddyroddy]

Does xp have to be named?

How about:

def run_benchmark(
    function_to_benchmark: FunctionType,
    xp: ModuleType,
    *,
    args: tuple[Any, ...] = (),
    kwargs: dict[str, Any] | None = None,
):

[connoraird]

Doesn't have to be. However, if I then call run_benchmark like this:

x: FloatArray

# some logic to set `x` to and define `some_benchmark`

run_benchmark(
        some_benchmark,
		x,
)

Then xp in run_benchmark will be x.

[paddyroddy]

Clearly. Is that not a fault with run_benchmark? Not entirely sure what you're trying to do here.

[paddyroddy]

What about

def run_benchmark(
      function_to_benchmark: Callable[..., Any],
      *,
      xp: ModuleType,
      args: tuple[Any, ...] = (),
      kwargs: dict[str, Any] | None = None,
  ) -> None:

[paddyroddy]

I've been trying to fix it, no luck yet

I've pushed another attempted fix. Will request a Copilot review also.

[Copilot]

Pull request overview

This PR adds a new “realistic” benchmarking workflow (scripts + docs) to measure GLASS performance across local machines and HPC clusters (e.g., Archer2/Cosma8/AAC6), and makes a few array-backend compatibility tweaks in core lensing/shell code to support those benchmarks.

Changes:

• Introduces a new benchmarks/ suite with a lensing benchmark, shared utilities, and cluster submission scripts/docs.
• Updates dependency groups for benchmark + GPU environments and adjusts example execution to install glass-ext-camb.
• Improves array-backend compatibility in lensing/shear code and shells tests.

Reviewed changes

Copilot reviewed 20 out of 21 changed files in this pull request and generated 18 comments.

Show a summary per file

File	Description
tests/regression/archer2/README.md	Repoints setup instructions to shared benchmark docs (link needs correction).
tests/core/test_shells.py	Adjusts assertion construction for backend compatibility.
pyproject.toml	Adds benchmark/GPU dependency groups and lint ignores; modifies examples extra.
noxfile.py	Ensures glass-ext-camb is installed for running example notebooks.
glass/shells.py	Forces RadialWindow.zeff to be a Python float.
glass/lensing.py	Moves shear computation utilities toward backend-agnostic array ops.
benchmarks/README.md	Adds benchmark runner documentation (contains a small typo).
benchmarks/lensing.py	Adds a realistic lensing benchmark script.
benchmarks/cosma8/submit-gpu.sh	Adds Cosma8 GPU Slurm submission script (path handling needs tightening).
benchmarks/cosma8/submit-cpu.sh	Adds Cosma8 CPU Slurm submission script (path handling needs tightening).
benchmarks/cosma8/README.md	Adds Cosma8 benchmark instructions (contains a broken relative link).
benchmarks/benchmark_utils.py	Adds backend selection + timing utilities and a Cosmology wrapper (contains a bug).
benchmarks/archer2/submit-gpu.sh	Adds Archer2 GPU Slurm submission script (path handling needs tightening).
benchmarks/archer2/submit-cpu.sh	Adds Archer2 CPU submission script (path handling needs tightening).
benchmarks/archer2/setup-gpu-env.sh	Adds Archer2 GPU module/environment setup script (minor typo).
benchmarks/archer2/README.md	Adds Archer2 benchmark instructions (contains a broken relative link).
benchmarks/aac6/submit-gpu.sh	Adds AAC6 GPU Slurm submission script (path handling needs tightening).
benchmarks/aac6/README.md	Adds AAC6 benchmark instructions (contains a broken relative link).
benchmarks/init.py	Adds benchmarks package marker.
.typos.toml	Adds a shell-word exception for “HSA”.
.gitignore	Ignores healpy-data directory used for offline runs.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[paddyroddy]

Had another pass

[paddyroddy]

Suggested change

	For the gpu benchmark, there is an additional dependency group `aac6-gpu` which
	For the GPU benchmark, there is an additional dependency group `aac6-gpu` which

[paddyroddy]

Suggested change

	For the gpu benchmark, there is an additional dependency group `archer2-gpu`
	For the GPU benchmark, there is an additional dependency group `archer2-gpu`

[paddyroddy]

Suggested change

	job to the archer2 amd gpu testbed queue.
	job to the ARCHER2 AMD GPU testbed queue.

[paddyroddy]

Suggested change

	# Flags to maximise jax gpu performance
	# Flags to maximise JAX GPU performance

[paddyroddy]

Suggested change

	For the gpu benchmark, there is an additional dependency group `cosma8-gpu`
	For the GPU benchmark, there is an additional dependency group `cosma8-gpu`

[paddyroddy]

Suggested change

	# Ensure the rocm library and build is known to jax
	# Ensure the ROCm library and build is known to JAX

[paddyroddy]

Suggested change

	example to benchmark using jax with amd/rocm, run the following from the root of
	example to benchmark using JAX with AMD/ROCm, run the following from the root of

[paddyroddy]

Suggested change

	Note that for the GPU benchmarks ARCHER2 only support rocm up to v0.6.x.
	Note that for the GPU benchmarks ARCHER2 only support ROCm up to v0.6.x.

[paddyroddy]

Any idea why we need the ty: ignore[unresolved-import] statements in this file?

[paddyroddy]

I've realised that benchmarks aren't currently being checked by ty. I'm going to push adding this config and fixing the few warnings.

[paddyroddy]

Are we installing these dependencies?

[paddyroddy]

I wonder if the issues we were having with glass-ext-camb was because previously they were only included as optional-dependencies, whereas now we have them as dependency-group.

[connoraird]

Making a note that there seems to be an issue with OOM error on Archer2 exclusive nodes with JAX. This page contains potential solutions