Python Environments on amplitUDE

Python is essential for AI/ML workloads on amplitUDE. This guide covers setting up isolated, reproducible Python environments optimized for GPU-accelerated scientific computing.

Overview

amplitUDE provides multiple Python options:

System Python 3.9 - Native installation, basic packages only
Anaconda module (anaconda3/2023) - Python 3.11.7 with scientific stack
Custom Conda/Mamba environments - Recommended for AI/ML projects
Apptainer containers - Maximum reproducibility (see Apptainer guide)

Recommendation: Use the anaconda3/2023 module as a base, then create project-specific environments.

Available Python Installations

Check System Python

# System Python (native installation)
python3 --version
# Output: Python 3.9.21

# Check available modules
module avail anaconda
python --version
# Output Python 3.11.7

Load Anaconda Module

# Load Anaconda (Python 3.11.7)
module load anaconda3/2023

python3 --version
# Output: Python 3.11.7

# Check pre-installed packages
conda list | head -20

Choosing the Right Python

Use Case	Python to Use	Why
AI/ML, Deep Learning	`anaconda3/2023`	Pre-configured scientific stack, Python 3.11
Quick scripts	System Python 3.9	No module loading needed
Reproducible research	Custom conda env	Full dependency control
Legacy code (Python <3.9)	Custom conda env	Install specific version

Using the Anaconda Module

Basic Usage

# Load module
module load anaconda3/2023

# Verify Python version
python3 --version
# Python 3.11.7

# List some pre-installed packages
conda list | grep -E "numpy|scipy|pandas|scikit|matplotlib"

What’s Included?

The anaconda3/2023 module includes:

Core Scientific Computing:

NumPy, SciPy, Pandas
Matplotlib, Seaborn
Scikit-learn
IPython, Jupyter

Check specific package:

module load anaconda3/2023
python -c "import numpy; print(f'NumPy {numpy.__version__}')"

Creating Custom Environments

Why Custom Environments?

Even with the anaconda module, create project-specific environments for:

Different Python versions (3.9, 3.10, 3.11, 3.12)
Specific package versions (PyTorch 2.0 vs 2.3)
Avoiding dependency conflicts
Reproducible research
Clean project separation

Basic Workflow

# Load anaconda module first
module load anaconda3/2023

# Create environment
conda create -n my-project python=3.11

# Activate environment
conda activate my-project

# Install packages
conda install numpy pandas matplotlib

# Deactivate when done
conda deactivate

Example: Data Science Environment

module load anaconda3/2023

conda create -n datascience python=3.11 \
    numpy scipy pandas matplotlib seaborn \
    scikit-learn jupyter \
    -c conda-forge

conda activate datascience

Installing Mamba (Recommended)

Mamba is a much faster drop-in replacement for Conda.

Why Mamba?

10-30x faster environment solving
Compatible with all conda commands
Better dependency resolution
Same package repositories

Install Mamba

module load anaconda3/2023

# Install mamba (one-time setup)
conda install mamba -c conda-forge

# Verify
mamba --version

Use Mamba

# All conda commands work with mamba
mamba create -n test-env python=3.11 numpy scipy pandas

# Activate works the same
conda activate test-env

# Install packages with mamba (much faster!)
mamba install scikit-learn matplotlib

Speed Comparison:

conda create (20 packages):  2-5 minutes ⏳
mamba create (20 packages):  10-30 seconds ⚡

AI/ML Environments

PyTorch (GPU-enabled)

module load anaconda3/2023

# Create environment with PyTorch + CUDA
mamba create -n pytorch-gpu python=3.11 \
    pytorch torchvision torchaudio pytorch-cuda=12.1 \
    -c pytorch -c nvidia

conda activate pytorch-gpu

Verify GPU support (on GPU node):

# Request GPU node first
salloc --partition=GPU-H200 --gres=gpu:1 --time=00:30:00
srun --pty bash

# Load environment
module load anaconda3/2023
conda activate pytorch-gpu

# Test GPU
python << EOF
import torch
print(f"PyTorch: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"CUDA version: {torch.version.cuda}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
EOF

Expected output:

PyTorch: 2.1.2
CUDA available: True
CUDA version: 12.1
GPU: NVIDIA H200
Memory: 141.0 GB

TensorFlow (GPU-enabled)

module load anaconda3/2023

mamba create -n tensorflow-gpu python=3.11 \
    tensorflow cudatoolkit=11.8 \
    -c conda-forge

conda activate tensorflow-gpu

# Test
python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"

JAX (High-Performance Computing)

module load anaconda3/2023

mamba create -n jax-gpu python=3.11

conda activate jax-gpu

# Install JAX with CUDA
pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# Verify
python -c "import jax; print(jax.devices())"

Hugging Face Transformers (for LLMs)

module load anaconda3/2023

mamba create -n transformers python=3.11 \
    pytorch pytorch-cuda=12.1 \
    -c pytorch -c nvidia

conda activate transformers

# Install Hugging Face ecosystem
pip install transformers datasets accelerate evaluate peft bitsandbytes

# Optional: monitoring tools
pip install wandb tensorboard

Complete ML Environment

module load anaconda3/2023

mamba create -n ml-complete python=3.11 \
    pytorch pytorch-cuda=12.1 \
    numpy scipy pandas scikit-learn matplotlib seaborn \
    -c pytorch -c nvidia -c conda-forge

conda activate ml-complete

# Add ML tools
pip install transformers datasets accelerate \
    optuna wandb tensorboard \
    plotly

Managing Environments

List Environments

module load anaconda3/2023

# Show all environments
conda env list

# Output:
# base                  *  /path/to/anaconda3
# pytorch-gpu              /home/user/anaconda3/envs/pytorch-gpu
# tensorflow-gpu           /home/user/anaconda3/envs/tensorflow-gpu

Activate/Deactivate

# Activate
conda activate my-env

# Check active environment
echo $CONDA_DEFAULT_ENV

# Deactivate
conda deactivate

Clone Environment

# Create copy of existing environment
conda create --name new-env --clone existing-env

Remove Environment

# Delete environment (permanent!)
conda env remove -n old-environment

# Verify deletion
conda env list

Check Environment Size

# See disk usage
du -sh ~/anaconda3/envs/*

# Example output:
# 2.3G    /home/user/anaconda3/envs/pytorch-gpu
# 1.8G    /home/user/anaconda3/envs/tensorflow-gpu

Clean Up Disk Space

# Remove cached packages
conda clean --all

# Or selectively:
conda clean --packages  # Unused packages
conda clean --tarballs  # Package archives

Environment Reproducibility

Export Environment

# Activate environment
conda activate my-project

# Export complete environment
conda env export > environment.yml

# Export minimal (only explicitly installed packages)
conda env export --from-history > environment_minimal.yml

Example environment.yml:

name: my-project
channels:
  - pytorch
  - nvidia
  - conda-forge
  - defaults
dependencies:
  - python=3.11
  - pytorch=2.1.0
  - pytorch-cuda=12.1
  - numpy=1.24.3
  - pandas=2.0.3
  - pip:
    - transformers==4.35.0
    - datasets==2.14.0

Recreate Environment

# Create from YAML file
conda env create -f environment.yml

# Activate
conda activate my-project

Using Environments in Slurm Jobs

CPU Job Template

#!/bin/bash
#SBATCH --job-name=python-analysis
#SBATCH --partition=STD-l-12h
#SBATCH --time=02:00:00
#SBATCH --mem=32G
#SBATCH --cpus-per-task=8
#SBATCH --output=logs/%x-%j.out
#SBATCH --error=logs/%x-%j.err

mkdir -p logs

# Load anaconda module
module load anaconda3/2023

# Activate environment
conda activate my-project

# Print diagnostics
echo "Job started: $(date)"
echo "Hostname: $(hostname)"
echo "Python: $(which python3)"
echo "Python version: $(python3 --version)"
echo "Environment: $CONDA_DEFAULT_ENV"

# Run script
python3 analyze_data.py

echo "Job completed: $(date)"

GPU Job Template

#!/bin/bash
#SBATCH --job-name=gpu-training
#SBATCH --partition=GPU-H200
#SBATCH --gres=gpu:1
#SBATCH --time=08:00:00
#SBATCH --mem=64G
#SBATCH --cpus-per-task=8
#SBATCH --output=logs/%x-%j.out
#SBATCH --error=logs/%x-%j.err

mkdir -p logs

# Load module
module load anaconda3/2023

# Activate PyTorch environment
conda activate pytorch-gpu

# Verify GPU
echo "=== GPU Check ==="
nvidia-smi --query-gpu=name,memory.total --format=csv
echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES"

# Test PyTorch GPU
python3 << EOF
import torch
assert torch.cuda.is_available(), "ERROR: CUDA not available!"
print(f"✓ PyTorch sees GPU: {torch.cuda.get_device_name(0)}")
print(f"✓ GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
EOF

# Run training
python3 train.py --epochs 100 --batch-size 64

Multi-GPU Job

#!/bin/bash
#SBATCH --job-name=multi-gpu
#SBATCH --partition=GPU-H200
#SBATCH --gres=gpu:4              # 4 GPUs
#SBATCH --time=12:00:00
#SBATCH --mem=256G
#SBATCH --cpus-per-task=32

module load anaconda3/2023
conda activate pytorch-gpu

# PyTorch Distributed Data Parallel
torchrun --nproc_per_node=4 train_ddp.py

Interactive Session

# Request interactive GPU node
salloc --partition=GPU-H200 --gres=gpu:1 --time=02:00:00 --mem=64G

# Start shell on allocated node
srun --pty bash

# Load environment
module load anaconda3/2023
conda activate pytorch-gpu

# Check GPU
nvidia-smi
python3 -c "import torch; print(torch.cuda.get_device_name(0))"

# Run interactive Python
python3
>>> import torch
>>> torch.cuda.is_available()
True

Best Practices

Environment Management

1. Load module in all job scripts

# Always include these lines
module load anaconda3/2023
conda activate my-env

2. Name environments descriptively

✅ Good examples:

pytorch-thesis-2024
bert-sentiment-analysis
cv-resnet-training

❌ Bad examples:

env1
test
new
my_env

3. One environment per project

Don’t install everything in one giant environment. Create focused environments:

# Project 1: Computer Vision
conda create -n cv-project python=3.11 pytorch torchvision

# Project 2: NLP
conda create -n nlp-project python=3.11 transformers datasets

# Project 3: Data Analysis
conda create -n analysis python=3.11 pandas scikit-learn

4. Export for reproducibility

# After setting up environment
conda env export > environment.yml

# Commit to version control
git add environment.yml
git commit -m "Add Python environment config"

5. Clean up regularly

# Monthly maintenance
conda clean --all

# Check sizes
du -sh ~/anaconda3/envs/*

# Remove unused environments
conda env remove -n old-project

Package Installation

1. Prefer Mamba over Conda

# Faster and more reliable
mamba install package-name

# Instead of
conda install package-name

2. Install conda packages first, then pip

# Correct order:
mamba install pytorch numpy pandas
pip install transformers  # Only if not in conda

# This prevents conflicts

3. Use appropriate channels

# PyTorch from official channel
mamba install pytorch -c pytorch

# Scientific packages from conda-forge
mamba install scikit-learn -c conda-forge

# Specify channel order
mamba install package -c pytorch -c nvidia -c conda-forge

4. Pin important versions

# Pin specific versions for reproducibility
mamba install pytorch=2.1.0 pytorch-cuda=12.1

# Or in environment.yml:
dependencies:
  - pytorch=2.1.0
  - pytorch-cuda=12.1

Storage Best Practices

1. Environments in $HOME

# Default location (recommended)
~/anaconda3/envs/

# 0.5 TB quota
# Permanent storage
# Backed up

2. Datasets in $SCRATCH

# Large datasets go here
export SCRATCH=/lustre/scratch/$USER
mkdir -p $SCRATCH/datasets

3. Configure cache directories

Add to ~/.bashrc:

# Hugging Face cache
export HF_HOME=$SCRATCH/.cache/huggingface
export TRANSFORMERS_CACHE=$SCRATCH/.cache/transformers
export HF_DATASETS_CACHE=$SCRATCH/.cache/datasets

# PyTorch cache
export TORCH_HOME=$SCRATCH/.cache/torch

# Conda package cache
export CONDA_PKGS_DIRS=$SCRATCH/.conda/pkgs

4. Monitor disk usage

# Check home quota
quota -s

# Check environment sizes
du -sh ~/anaconda3/envs/*

# Check cache sizes
du -sh ~/.cache/*

Troubleshooting

`conda activate` doesn’t work in job

Symptom:

CommandNotFoundError: Your shell has not been properly configured to use 'conda activate'.

Solution 1: Use source activate

module load anaconda3/2023
source activate my-env

Solution 2: Initialize conda

module load anaconda3/2023

# Initialize for bash
eval "$(conda shell.bash hook)"

# Now activate works
conda activate my-env

CUDA not available in PyTorch

Important: GPUs only available on compute nodes, not login nodes!

Check on GPU node:

# Request GPU first
salloc --partition=GPU-H200 --gres=gpu:1 --time=00:30:00
srun --pty bash

# Load environment
module load anaconda3/2023
conda activate pytorch-gpu

# Check CUDA
python3 << EOF
import torch
print(f"PyTorch: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"CUDA version: {torch.version.cuda}")
print(f"cuDNN version: {torch.backends.cudnn.version()}")
EOF

# Check GPU visibility
nvidia-smi
echo $CUDA_VISIBLE_DEVICES

If still no CUDA:

# Reinstall PyTorch with correct CUDA version
mamba install pytorch pytorch-cuda=12.1 -c pytorch -c nvidia --force-reinstall

Common mistake:

# ❌ Wrong: testing on login node (no GPU!)
ssh login2
python -c "import torch; print(torch.cuda.is_available())"
# Output: False (expected - login nodes have no GPUs)

# ✅ Correct: test on GPU node
salloc --partition=GPU-H200 --gres=gpu:1
srun --pty bash
python -c "import torch; print(torch.cuda.is_available())"
# Output: True

Out of disk space / Quota exceeded

Check quota:

quota -s

Solutions:

1. Clean conda cache

# Remove all cached packages
conda clean --all

# Or selectively:
conda clean --packages
conda clean --tarballs

2. Move cache to scratch

# Add to ~/.bashrc
export CONDA_PKGS_DIRS=/lustre/scratch/$USER/.conda/pkgs

# Create directory
mkdir -p $CONDA_PKGS_DIRS

3. Remove large environments

# Find large environments
du -sh ~/anaconda3/envs/* | sort -h

# Remove unused
conda env remove -n old-environment

Environment creation is slow

Use Mamba instead:

# Install mamba once
conda install mamba -c conda-forge

# Use mamba (10-30x faster!)
mamba create -n my-env python=3.11 pytorch numpy scipy

Speed comparison:

conda create: 2-5 minutes ⏳
mamba create: 10-30 seconds ⚡

Package not found

# Search across channels
conda search package-name

# Install from specific channel
conda install -c conda-forge package-name

# Or use pip
pip install package-name

Import error in Python

# Verify environment is activated
conda activate my-env
echo $CONDA_DEFAULT_ENV  # Should print 'my-env'

# Check if package installed
conda list | grep package-name

# If not installed:
conda install package-name
# or
pip install package-name

# Verify Python location
which python3
# Should be: /home/user/anaconda3/envs/my-env/bin/python3

Conflicting dependencies

# Create fresh environment
conda create -n fresh-env python=3.11

# Install packages one at a time
conda activate fresh-env
conda install package1
conda install package2  # If this fails, package1 and package2 conflict

# Use mamba (better conflict resolution)
mamba install package1 package2

Quick Reference

Task	Command
Load Anaconda	`module load anaconda3/2023`
Check Python version	`python3 --version`
Create environment	`conda create -n name python=3.11`
Create with mamba	`mamba create -n name python=3.11`
Activate	`conda activate name`
Deactivate	`conda deactivate`
List environments	`conda env list`
Install package (conda)	`conda install package`
Install package (mamba)	`mamba install package`
Install package (pip)	`pip install package`
Search package	`conda search package`
List packages	`conda list`
Export environment	`conda env export > env.yml`
Create from file	`conda env create -f env.yml`
Remove environment	`conda env remove -n name`
Clone environment	`conda create -n new --clone old`
Clean cache	`conda clean --all`
Check disk usage	`du -sh ~/anaconda3/envs/*`

Complete Example: ML Project Setup

# 1. Load module
module load anaconda3/2023

# 2. Install mamba (one-time)
conda install mamba -c conda-forge

# 3. Create ML environment
mamba create -n image-classification python=3.11 \
    pytorch torchvision pytorch-cuda=12.1 \
    numpy pandas matplotlib scikit-learn \
    -c pytorch -c nvidia -c conda-forge

# 4. Activate
conda activate image-classification

# 5. Install additional tools
pip install tensorboard wandb torchmetrics

# 6. Create project structure
mkdir -p ~/projects/image-clf/{data,models,logs,scripts}
cd ~/projects/image-clf

# 7. Export for reproducibility
conda env export > environment.yml

# 8. Create training script
cat > scripts/train.py << 'EOF'
import torch
print(f"PyTorch {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
EOF

# 9. Create job script
cat > submit.sh << 'EOF'
#!/bin/bash
#SBATCH --job-name=train
#SBATCH --partition=GPU-H200
#SBATCH --gres=gpu:1
#SBATCH --time=04:00:00
#SBATCH --mem=64G

module load anaconda3/2023
conda activate image-classification
python scripts/train.py
EOF

# 10. Submit
sbatch submit.sh

Next Steps

AI Workloads on amplitUDE - ML training workflows
Fine-Tuning LLMs - LLM fine-tuning guide
Apptainer Containers - Container workflows

Further Resources

Last updated: May 2026