Singularity.jax

Bootstrap: docker
From: nvidia/cuda:11.2.2-cudnn8-devel-ubuntu18.04

%environment 
  SHELL=/bin/bash
  export SHELL
  export HDF5_USE_FILE_LOCKING='FALSE'
  export PATH=/usr/local/cuda/bin:$PATH
  export CPATH=/usr/local/cuda/include:$CPATH
  export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
  
%setup
# runs on host
# the path to the image is $SINGULARITY_ROOTFS

%post
# post-setup script

# load environment variables
. /environment

# use bash as default shell
echo 'SHELL=/bin/bash' >> /environment

# make environment file executable
chmod +x /environment

# default mount paths
mkdir /scratch

# fix this issue: https://github.com/singularityware/singularity/issues/1182#issuecomment-381796545
touch /usr/bin/nvidia-smi
apt-get update && apt-get -y install locales
locale-gen en_US.UTF-8
apt-get -y --force-yes install vim wget bzip2 git
apt-get -y install build-essential
apt-get -y install software-properties-common
add-apt-repository ppa:deadsnakes/ppa
apt-get update
apt-get -y install python3.8
apt install python3-pip

pip install --upgrade pip
pip install Pillow 
pip install matplotlib 
pip install tqdm 
pip install h5py
pip install pandas
pip install comet_ml
pip install pyyaml
pip install simpleitk
pip install --upgrade jax jaxlib==0.1.64+cuda112 -f https://storage.googleapis.com/jax-releases/jax_releases.html

%runscript
# executes with the singularity run command
# delete this section to use existing docker ENTRYPOINT command
echo "Welcome to Singularity JAX pytorch ubuntu 18.04"

%test
# test that script is a success