# Defines a Singularity container with TensorFlow pre-installed
#

#
# Before bootstrapping this container, you must ensure that the following files
# are present in the current directory (alongside this definition file):
#
#   * cuda-linux64-rel-8.0.44-21122537.run  (* see below)
#   * NVIDIA-Linux-x86_64-375.20.run        (* see below)
#   * cudnn-8.0-linux-x64-v5.1.tgz          (https://developer.nvidia.com/cudnn)
#
# * The cuda-linux64 and NVIDIA-Linux files can be obtained by downloading the
# NVIDIA CUDA local runfile `cuda_8.0.44_linux.run` from:
#
#   https://developer.nvidia.com/cuda-downloads
#
# Then extract the necessary files by running:
#
#   sh cuda_8.0.44_linux.run --extract=<absolute/path/to/bootstrap/directory>
#
# IF YOUR HPC SYSTEM IS USING A DIFFERENT VERSION OF CUDA AND/OR NVIDIA DRIVERS
# YOU WILL NEED TO ADJUST THE ABOVE VERSION NUMBERS TO MATCH YOUR SYSTEM
#
# YOU WILL ALSO NEED TO DOWNLOAD THE APPROPRIATE DRIVER. For example,
# cuda_8.0.44_linux.run returns driver version 367.48.
#
# If you use this to create a container inside a virtual machine with no access to
# a GPU, comment out the final test.


BootStrap: docker
From: ubuntu:16.10


%runscript
    # When executed, the container will run Python with the TensorFlow module

    # Check the current environment
    chk_nvidia_uvm=$(grep nvidia_uvm /proc/modules)
    if [ -z "$chk_nvidia_uvm" ]; then
        echo "Problem detected on the host: the Linux kernel module nvidia_uvm is not loaded"
        exit 1
    fi

    exec /usr/bin/python "$@"


%setup
    # Runs from outside the container during Bootstrap

    NV_DRIVER_VERSION=375.20
    NV_CUDA_FILE=cuda-linux64-rel-8.0.44-21122537.run
    NV_CUDNN_FILE=cudnn-8.0-linux-x64-v5.1.tgz
    NV_DRIVER_FILE=NVIDIA-Linux-x86_64-${NV_DRIVER_VERSION}.run

    working_dir=$(pwd)

    echo "Unpacking NVIDIA driver into container..."
    cd ${SINGULARITY_ROOTFS}/usr/local/
    sh ${working_dir}/${NV_DRIVER_FILE} -x
    mv NVIDIA-Linux-x86_64-${NV_DRIVER_VERSION} NVIDIA-Linux-x86_64
    cd NVIDIA-Linux-x86_64/
    for n in *.$NV_DRIVER_VERSION; do
        ln -v -s $n ${n%.367.48}
    done
    ln -v -s libnvidia-ml.so.$NV_DRIVER_VERSION libnvidia-ml.so.1
    ln -v -s libcuda.so.$NV_DRIVER_VERSION libcuda.so.1
    cd $working_dir

    echo "Running NVIDIA CUDA installer..."
    sh $NV_CUDA_FILE -noprompt -nosymlink -prefix=${SINGULARITY_ROOTFS}/usr/local/cuda-8.0
    ln -r -s ${SINGULARITY_ROOTFS}/usr/local/cuda-8.0 ${SINGULARITY_ROOTFS}/usr/local/cuda

    echo "Unpacking cuDNN..."
    tar xvf $NV_CUDNN_FILE -C ${SINGULARITY_ROOTFS}/usr/local/

    echo "Adding NVIDIA PATHs to /environment..."
    NV_DRIVER_PATH=/usr/local/NVIDIA-Linux-x86_64
    echo "

LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:$NV_DRIVER_PATH:\$LD_LIBRARY_PATH
PATH=$NV_DRIVER_PATH:\$PATH
export PATH LD_LIBRARY_PATH
    
" >> $SINGULARITY_ROOTFS/environment


%post
    # Runs within the container during Bootstrap

    # Set up some required environment defaults
    export LC_ALL=C
    export PATH=/bin:/sbin:/usr/bin:/usr/sbin:$PATH

    # Install the necessary packages (from repo)
    apt-get update && apt-get install -y --no-install-recommends \
        build-essential \
        curl \
        git \
        libcurl4-openssl-dev \
        libfreetype6-dev \
        libpng-dev \
        libzmq3-dev \
        python-pip \
        pkg-config \
        python-dev \
        rsync \
        software-properties-common \
        unzip \
        zip \
        zlib1g-dev
    apt-get clean

    # Update to the latest pip (newer than repo)
    pip install --no-cache-dir --upgrade pip
    
    # Install other commonly-needed packages
    pip install --no-cache-dir --upgrade \
        future \
        matplotlib \
        scipy \
        sklearn

    # TensorFlow package versions as listed here:
    #   https://www.tensorflow.org/get_started/os_setup#test_the_tensorflow_installation
    #
    # Ubuntu/Linux 64-bit, GPU enabled, Python 2.7 (Requires CUDA toolkit 8.0 and CuDNN v5)
    export TF_BINARY_URL=https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-0.12.1-cp27-none-linux_x86_64.whl
    pip install --no-cache-dir --ignore-installed --upgrade $TF_BINARY_URL


%test
    # Sanity check that the container is operating
    export LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:/usr/local/NVIDIA-Linux-x86_64:$LD_LIBRARY_PATH

    # Ensure that TensorFlow can be imported
    /usr/bin/python -c "import tensorflow as tf"

    # Runs in less than 30 minutes on low-end CPU; in less than 2 minutes on GPU
    # Comment the following line if building the container inside a VM with no access to a GPU
    /usr/bin/python -m tensorflow.models.image.mnist.convolutional

