pypomp on greatlakes

The greatlakes GPUs are a convenient resource for high-performance computing at University of Michigan. The greatlakes gpu partition has 52 Nvidia Tesla V100 GPUs.

Testing in an interactive session

To obtain access to a GPU via an interactive seesion (i.e., a terminal prompt) once can do

salloc --account=ionides0 --partition=gpu --gpus=v100:1 --cpus-per-gpu=1 --mem=8000

I set up the Python environment on the login node

module load python
python -m venv ~/opt/py3.12
source ~/opt/py3.12/bin/activate
pip install jax[cuda12]
pip install pypomp pytest tqdm
git clone git@github.com:pypomp/pypomp

and built jax with cuda:

pip install -U "jax[cuda12]"

Then install everything else needed,

pip install pypomp pytest tqdm

At this point, running salloc, as above, starts a machine with a GPU.
Then check we have access to the GPUs:

salloc --account=ionides0 --partition=gpu --gpus=v100:1 --cpus-per-gpu=1 --mem=8000
nvidia-smi

It is necessary to restart the Python environment for the new machine. A simple JAX test is working. So are the pypomp tests.

module load python
source ~/opt/py3.12/bin/activate
python -c "import jax.numpy as np; print(np.ones((3,3)))"
pytest pypomp/test

The tests took almost 6 minutes to run.

An earlier test with pympomp 0.0.2 failed, perhaps because of some difficulty with tensorflow (which is not needed for the current pypomp, but was installed anyhow).

An earlier test with pypomp 0.0.3 failed, perhaps because of using anaconda. There is a rumor that the anaconda distribution can have problems with jax[cuda12].

Testing a batch job

Once the Python environment is set up as above, one can replace the interactive session above by a batch job. The following is put in greatlakes.sbat

#!/bin/bash

#SBATCH --job-name=gpu-pypomp-test
#SBATCH --mail-user=ionides@umich.edu
#SBATCH --mail-type=BEGIN,END,FAIL
#SBATCH --account=ionides0
#SBATCH --partition=gpu
#SBATCH --gpus=v100:1
#SBATCH --cpus-per-gpu=1
#SBATCH --mem=4GB
### wall time hours:minutes:seconds
#SBATCH --time=00:10:00

###   Load software modules
module load python/3.12.1
module list

####  Commands your job should run follow this line
echo "Running on $SLURM_JOB_NODELIST"
echo "Running in $(pwd)"
source ~/opt/py3.12/bin/activate
pytest ~/pypomp/test

Then, sbatch greatlakes.sbat runs the tests.