Creating a conda environment for GPU computing in Jupyter

December 29, 2024

After a few mis-steps, here is how I set up a conda environment to use in Jupyter with tensorflow, pytorch, and using the GPU.

As a note, I do this on the node with the GPU, so that things (hopefully) compile correctly!

1. Create an environment

First, create a new environment for tensorflow and friends, and activate it.

mamba create -n gpu_notebook cudatoolkit tensorflow nvidia::cuda pytorch torchvision torchaudio pytorch-cuda -c pytorch -c nvidia
mamba activate gpu_notebook

2. Install the python libraries

Install the usual suspect python packages that you will probably want to use. For convenience, I usually put these in a file in my Git repo called requirements.txt.

$ cat requirements.txt 
jupyter
matplotlib
natsort
numpy
pandas
scipy
scikit-learn
seaborn
statsmodels

pip install -r requirements.txt

3. Reame your jupyter kernel

When you open jupyter there is a list of kernels that you can connect to. (If you have a window open that list will be on the top right.) If you rename your jupyter kernel it makes it much easier to find the kernel associated with this conda environment. The default name is something like Python 3 which is not helpful if you have lots of them!

a. Find where your kernel is installed

This command shows your jupyter kernels

jupyter kernelspec list

You’ll see your kernel(s) and the locations of them. In the location listed there is a file called kernel.json.

b. Edit that file:

vi $HOME/miniconda3/envs/gpu_notebook/share/jupyter/kernels/python3/kernel.json

c. Change the `name` to be meaningful

Change the value associated with the display_name key. Set it to something meaningful so you can find it in your browser!

4. Set up the XLA_FLAGS environment variable.

This was essential for me to get tensorflow working. There is a directory somewhere in your conda environment with the libdevice library that is needed. For my installation that was in nvvm/libdevice/libdevice.10.bc. Of course you can find yours with:

find ~/miniconda3/ -name libdevice

You want to set the XLA_FLAGS variable to point to the base of the nvvm folder. This command sets it inside the conda environment so it is always set when the conda environment is activated, and unset when it is deactivated.

conda env config vars set XLA_FLAGS=--xla_gpu_cuda_data_dir=$HOME/miniconda3/envs/gpu_notebook

5. Activate the environment

Don’t forget to submit this to a node with GPU capabilities!

Rob Edwards