Python is available on AImageLab-HPC both through the default Anaconda distribution and through dedicated module versions. This page explains how to set up and manage Python environments.

Available Python Versions

The anaconda3/2023.09-0-none-none module is loaded by default at login, providing a base Python installation. Additional standalone Python versions are also available as modules:

module load python/3.9.21-gcc-11.4.0
module load python/3.10.16-gcc-11.4.0
module load python/3.11.11-gcc-11.4.0

Virtual Environments

Although Anaconda is available by default, virtual environments created with venv are strongly preferred over Conda environments: Conda environments can easily grow to several gigabytes, while venv environments are lightweight and only contain the packages you explicitly install.

Creating a virtual environment

Load the Python version you need (or skip this step to use Anaconda’s Python):

module load python/3.11.11-gcc-11.4.0

Create the virtual environment:

python -m venv /path/to/my_env

Note: Choose an arbitrary name for your environment. Create your virtual environments inside /homes/<username> (e.g. /homes/<username>/envs/my_env). NFS handles the many small files that make up a venv better than BeeGFS (/work), which is optimised for large sequential reads.

Activating and installing packages

source /path/to/my_env/bin/activate
pip install <package>

Once activated, your shell prompt will show the environment name:

(my_env) user@ailb-login-02:~$

Deactivating

deactivate

Using Your Environment in SLURM Batch Jobs

Activate the virtual environment in your job script before running Python:

#!/bin/bash
#SBATCH --job-name my_job
#SBATCH --partition <partition>
#SBATCH --gres gpu:1

module load python/3.11.11-gcc-11.4.0   # same version used to create the env
source /path/to/my_env/bin/activate

python my_script.py

If your environment was created on top of Anaconda’s Python (no module load), initialise Conda first:

#!/bin/bash
#SBATCH --job-name my_job
#SBATCH --partition <partition>
#SBATCH --gres gpu:1

. /usr/local/anaconda3/etc/profile.d/conda.sh
source /path/to/my_env/bin/activate

python my_script.py

Jupyter Notebooks

To use your virtual environment as a Jupyter kernel, register it with ipykernel:

source /path/to/my_env/bin/activate
pip install ipykernel
python -m ipykernel install --user --name my_env --display-name "Python (my_env)"

For a full walkthrough of how to start a Jupyter session on a compute node and connect to it from your local browser, see Jupyter Notebook on a Compute Node.

Common Issues

• Running out of disk space: Prefer venv over Conda environments. If your environment is very large, consider removing unused packages rather than moving it to /work.

• Jupyter kernel not connecting: Make sure you registered the kernel from inside the correct activated environment, and that port forwarding is set up correctly.

• Model downloads filling up storage: When loading pretrained models (e.g. via Hugging Face from_pretrained()), set a custom cache directory to avoid filling your home:

  bash export HF_HOME=/path/to/large/storage/.cache/huggingface