Add config to run GPU on ARCHER2 #1084

Activity

1. paddyroddy

   assigned

   connoraird

   on May 20, 2026
2. paddyroddy

   added

   enhancementNew feature or requestNew feature or request

   array-apiWork is related to the Array APIWork is related to the Array API

   benchmarksBenchmarking workBenchmarking work

   on May 20, 2026
3. connoraird commented on Jul 2, 2026

   

   connoraird

   on Jul 2, 2026

   Contributor

   More actions

   On an archer2 gpu partition

   srun --export=PATH --partition=gpu --qos=gpu-shd --gpus=1 --time=1:0:0 --pty /bin/bash

   using these modules

   module load PrgEnv-amd/8.6.0
   module load rocm
   module load craype-accel-amd-gfx90a
   module load craype-x86-milan

   I managed to get jax to run on the Archer2 gpu node (not with glass) using python 3.12 and these packages

   uv pip install jax==0.4.35 jax-rocm60-plugin==0.4.35 jax-rocm60-pjrt==0.4.35

   Test script

   import jax
   print(jax.devices())
   import jax.numpy as jnp
   x = jnp.arange(5)
   print(x)

   Output

   connoraird@nid200002:~/glass> .venv/bin/python test_jax.py
   [RocmDevice(id=0)]
   [0 1 2 3 4]
4. paddyroddy commented on Jul 2, 2026

   

   paddyroddy

   on Jul 2, 2026

   MemberAuthor

   More actions

   You could change this to:

   uv run --isolated --with "jax==0.4.35" --with "jax-rocm60-plugin==0.4.35" \
     --with "jax-rocm60-pjrt==0.4.35" python -c \
     "import jax; print(jax.devices()); import jax.numpy as jnp; x = jnp.arange(5); print(x)"
5. paddyroddy

   added a parent issue on Sep 21, 2026

   • Enable benchmarking a real GLASS workflow #1086
6. paddyroddy commented on Oct 5, 2026

   

   paddyroddy

   on Oct 5, 2026

   MemberAuthor

   More actions

   Finished in #1098
7. paddyroddy

   closed this as completedon Oct 5, 2026