Dockerfile.gpu

ARG UBUNTU_VERSION
ARG CUDA_VERSION
####
## Stage 1 - Download and convert the notebook
####
FROM python:3.9-alpine3.19 AS notebook_autoconversion
RUN apk update && \
    apk add git

WORKDIR /home
RUN pip install --upgrade pip && \
    pip install nbformat==5.9.2

# Read all the arguments
ARG PATH_TO_NOTEBOOK
ARG PATH_TO_REQUIREMENTS
ARG SECTIONS_TO_REMOVE
ARG NOTEBOOK_NAME
ARG GPU_FLAG
ARG PYTHON_VERSION
ARG UBUNTU_VERSION
ARG CUDA_VERSION
ARG CUDNN_VERSION

# Custom cache invalidation
ARG CACHEBUST=1

# All the layers bellow this will not be cached
RUN echo "${CACHEBUST}"

# Download the notebook
ADD $PATH_TO_NOTEBOOK ./${NOTEBOOK_NAME}

# Autoconvert the notebook
RUN git clone --depth 1 https://github.com/HenriquesLab/DL4MicEverywhere.git 
RUN python DL4MicEverywhere/.tools/notebook_autoconversion/transform.py -p . -n ${NOTEBOOK_NAME} -s ${SECTIONS_TO_REMOVE} 
RUN mv colabless_${NOTEBOOK_NAME} ${NOTEBOOK_NAME}
RUN python DL4MicEverywhere/.tools/python_tools/create_docker_info.py "/home/docker_info.txt" \
        "${UBUNTU_VERSION}" "${CUDA_VERSION}" "${CUDNN_VERSION}" "${PATH_TO_NOTEBOOK}" "${PATH_TO_REQUIREMENTS}" \
        "${SECTIONS_TO_REMOVE}"  "${NOTEBOOK_NAME}" "${GPU_FLAG}" "${PYTHON_VERSION}"

####
## Stage 2 - Set the final image (install packages, python, python libraries)
####
FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu${UBUNTU_VERSION} AS final

# Read all the arguments
ARG PATH_TO_NOTEBOOK
ARG PATH_TO_REQUIREMENTS
ARG SECTIONS_TO_REMOVE
ARG NOTEBOOK_NAME
ARG GPU_FLAG
ARG PYTHON_VERSION
ARG UBUNTU_VERSION
ARG CUDA_VERSION
ARG CUDNN_VERSION

# Set the language
ENV LANG=C.UTF-8

# Set timezone for Python installation
ENV TZ=Europe/Lisbon
RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone

# Install an specific nodejs version (using NVM - Node Version Manager)
ENV NODE_VERSION=20.10.0
RUN apt-get update && \
    apt-get install -y curl && \
    curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.39.0/install.sh | bash
ENV NVM_DIR=/root/.nvm
RUN . "$NVM_DIR/nvm.sh" && nvm install ${NODE_VERSION} && \
    . "$NVM_DIR/nvm.sh" && nvm use v${NODE_VERSION} && \
    . "$NVM_DIR/nvm.sh" && nvm alias default v${NODE_VERSION}
ENV PATH="/root/.nvm/versions/node/v${NODE_VERSION}/bin/:${PATH}"

ENV LD_LIBRARY_PATH /usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH

RUN echo ${UBUNTU_VERSION%.*}${UBUNTU_VERSION#*.}

# Install common packages and cuda libraries 
RUN aux_cuda_version=${CUDA_VERSION%.*} && \
    apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${UBUNTU_VERSION%.*}${UBUNTU_VERSION#*.}/x86_64/3bf863cc.pub && \
    apt-get update && apt-get install -y --no-install-recommends \
        ca-certificates \
        cuda-command-line-tools-${CUDA_VERSION%.*.*}-${aux_cuda_version#*.} \
        cuda-nvrtc-${CUDA_VERSION%.*.*}-${aux_cuda_version#*.} \
        libcublas-${CUDA_VERSION%.*.*}-${aux_cuda_version#*.} \
        libcublas-dev-${CUDA_VERSION%.*.*}-${aux_cuda_version#*.}  \
        libcufft-${CUDA_VERSION%.*.*}-${aux_cuda_version#*.}  \
        libcurand-${CUDA_VERSION%.*.*}-${aux_cuda_version#*.}  \
        libcusolver-${CUDA_VERSION%.*.*}-${aux_cuda_version#*.}  \
        libcusparse-${CUDA_VERSION%.*.*}-${aux_cuda_version#*.}  \
        libcudnn8=${CUDNN_VERSION}-1+cuda${CUDA_VERSION%.*} \
        libgomp1 \
        build-essential \
        curl \
        libfreetype6-dev \
        pkg-config \
        software-properties-common \
        unzip \
        wget \
        git \
        libopenmpi-dev \
        libxext6 \
        libsm6 \
        ffmpeg  && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

# Instal Python and nvidia-cudnn-cuXX for Python
RUN add-apt-repository ppa:deadsnakes/ppa && \
    apt-get update && apt-get install -y \
    python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python3-pip python${PYTHON_VERSION}-venv && \
    rm -rf /var/lib/apt/lists/* && \
    python${PYTHON_VERSION} -m pip install pip --upgrade && \
    update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 0 && \
    update-alternatives --install /usr/bin/python python /usr/bin/python${PYTHON_VERSION} 0 

# Export the environment variable LD_LIBRARY_PATH
ENV LD_LIBRARY_PATH lib/:/usr/local/lib/python${PYTHON_VERSION}/dist-packages/nvidia/cudnn/lib:$LD_LIBRARY_PATH

# Set the working directory
WORKDIR /home 

# Install nbformat and jupyterlab
RUN pip install --upgrade pip && \
    if [  "$(printf '%s\n' "3.8" "${PYTHON_VERSION}" | sort -V | head -n1)" = "3.8" ] ; then \
        # For Python between 3.8 and 3.11
        pip install nbformat==5.9.2 ; \ 
    else \
        # For Python 3.7 or lower
        pip install nbformat==5.0.2 ; \
    fi && \
    if [  "$(printf '%s\n' "3.7" "${PYTHON_VERSION}" | sort -V | head -n1)" = "3.7" ] ; then \
        # For Python between 3.7 and 3.11
        pip install ipywidgets==8.1.0 && \ 
        pip install jupyterlab==3.6.0 && \ 
        pip install matplotlib==3.6.3 ; \  
    elif [  "$(printf '%s\n' "3.5" "${PYTHON_VERSION}" | sort -V | head -n1)" = "3.5" ] ; then \
        # For Python 3.4 or lower
        pip install ipywidgets==7.8.1 && \ 
        pip install jupyterlab==0.23.0 && \ 
        pip install matplotlib==2.1.0 ; \ 
    else \
        # For Python 3.5 and 3.6
        pip install ipywidgets==8.0.0a0 && \ 
        pip install jupyterlab==2.3.2 && \ 
        pip install matplotlib==3.0.3 ; \ 
    fi

# Custom cache invalidation
ARG CACHEBUST=1
# All the layers bellow this will not be cached
RUN echo "${CACHEBUST}"

# Download the notebook and requirements if they are not provided
ADD $PATH_TO_REQUIREMENTS ./requirements.txt

# Install the requirements and convert the notebook
RUN pip install --upgrade pip && \
    pip install -r requirements.txt && \
    rm requirements.txt
    
# Copy the converted notebook from stage 1
COPY --from=notebook_autoconversion /home/${NOTEBOOK_NAME} /home/docker_info.txt /home

ENV XLA_FLAGS=--xla_gpu_cuda_data_dir=/usr/local/cuda

CMD bash