Skip to content

Python

Upcoming changes to Cori GPU modulefiles

Cori GPU modulefiles are moving to a new location. Workflows on this system must be adjusted to account for this new change. Please see here for more information.

There are many options for using Python on GPUs, each with their own set of pros/cons. We have developed some advice for porting your Python code to GPU here.

On the Cori GPU nodes, we recommend that users build a custom conda environment for the Python GPU framework they would like to use. You can find instructions for building a custom conda environment here.

In all cases you'll need to:

  1. Make sure you have sourced your conda environment via source activate mypythonenv
  2. Run your code with the general format srun -n 1 python yourscript.py args...

CuPy

  • module load python cuda
  • Build a custom conda environment
  • pip install cupy into your environment following the directions here
  • Note that your CuPy version and CUDA version must match

Numba CUDA

  • module load python cuda
  • Build a custom conda environment
  • conda install numba and cudatoolkit into your environment following the directions here Note that your CUDA and cudatoolkit versions must match.

PyOpenCL

  • module load python
  • Build a custom conda environment with conda install -n pyoencl-env -c conda-forge pyopencl cudatoolkit
  • And then create a link (only once) to the CUDA OpenCL vendor driver via
    ln -s /etc/OpenCL/vendors/nvidia.icd ~/.conda/envs/pyopencl-env/etc/OpenCL/vendors/nvidia.icd
    

PyCUDA

  • module load python cuda
  • build a custom conda environment
  • pip install pycuda

JAX

  • module load python cuda
  • Build a custom conda environment
  • JAX installation is somewhat complex. Check here for the most up-to-date information. Note that your jaxlib version must match your CUDA version.

RAPIDS

RAPIDS moves too quickly for us to provide an up-to-date distribution. We suggest you build a conda environment and install the latest version yourself. You can use the installation helper to generate the conda install command you need.

If you intend to use RAPIDS via scripts/command line, you're ready to go. If you would like to create your own RAPIDS kernel to use in Jupyter, you'll need to conda install ipykernel and python -m ipykernel install --user --name rapids_env --display-name rapids

You'll need to restart your Jupyter server. When you log in, you should now see your rapids kernel as an option.

If you need other libraries like matplotlib, you may want to install them during your original conda install command OR you may want to install later via pip with --user. These will help you avoid dependency problems.

For more information about how to use NVIDIA RAPIDS, please see our Examples page.

MPI4py

You can build mpi4py and install it into a conda environment on Cori to be used with one of the MPI implementations available for use with the GPU nodes.

In a conda environment, you can install mpi4py via:

module load cgpu gcc cuda openmpi python
conda create -n test python=3.8 -y
source activate test
MPICC="$(which mpicc)" pip install --no-binary mpi4py mpi4py

Or use our legacy directions:

user@cgpu12:~> conda create -n mpi4pygpu python=3.8
user@cgpu12:~> source activate mpi4pygpu
(mpi4pygpu) user@cgpu12:~> module load gcc cuda openmpi/4.0.3  # (or pgi/intel instead of gcc)
(mpi4pygpu) user@cgpu12:~> wget https://bitbucket.org/mpi4py/mpi4py/downloads/mpi4py-3.0.3.tar.gz
(mpi4pygpu) user@cgpu12:~> tar zxvf mpi4py-3.0.3.tar.gz
(mpi4pygpu) user@cgpu12:~> cd mpi4py-3.0.3
(mpi4pygpu) user@cgpu12:~> python setup.py build --mpicc="mpicc -B /usr/bin"
(mpi4pygpu) user@cgpu12:~> python setup.py install

Deep Learning Software

Tensorflow:

module load tensorflow/gpu-1.13.1-py36

PyTorch:

module load pytorch/v1.1.0-gpu