DockerfileCuda.jax

FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04

USER root
## add a conda for python 3 environment
# Install miniconda - this is from 
# https://github.com/ContinuumIO/docker-images/blob/master/miniconda3/debian/Dockerfile
# =================================
# hadolint ignore=DL3008
# xterm for x11 forwarding

RUN apt-get update -q && \
    apt-get install -q -y --no-install-recommends \
    bzip2 \
    ca-certificates \
    git \
    libglib2.0-0 \
    libsm6 \
    libxext6 \
    libxrender1 \
    mercurial \
    subversion \
    wget \
    xterm \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*

ENV PATH /opt/conda/bin:$PATH

# Leave these args here to better use the Docker build cache
ARG CONDA_VERSION=py38_4.9.2
ARG CONDA_MD5=122c8c9beb51e124ab32a0fa6426c656

RUN wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-${CONDA_VERSION}-Linux-x86_64.sh -O miniconda.sh && \
    echo "${CONDA_MD5}  miniconda.sh" > miniconda.md5 && \
    if ! md5sum --status -c miniconda.md5; then exit 1; fi && \
    mkdir -p /opt && \
    sh miniconda.sh -b -p /opt/conda && \
    rm miniconda.sh miniconda.md5 && \
    ln -s /opt/conda/etc/profile.d/conda.sh /etc/profile.d/conda.sh && \
    echo ". /opt/conda/etc/profile.d/conda.sh" >> ~/.bashrc && \
    #echo "conda activate base" >> ~/.bashrc && \
    find /opt/conda/ -follow -type f -name '*.a' -delete && \
    find /opt/conda/ -follow -type f -name '*.js.map' -delete && \
    /opt/conda/bin/conda clean -afy

# Install basic libraries
RUN conda update conda && \
    conda config --add channels conda-forge && \
    conda install pip \
    matplotlib \
    scipy \
    numpy \
    pandas \
    ipykernel && \
    conda clean -ya

# Install jax ecosystem in cuda
RUN conda install -c nvidia cuda-nvcc 
RUN pip install --upgrade pip && \
    pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html && \
    pip install diffrax optax flax equinox tensorflow

# install convex optimization package
RUN pip install --upgrade pip && \
    pip install cvxpy==1.3.0 diffcp cvxpylayers gymnasium==0.27.1

# install probalistic programming package
RUN conda update conda && \
    conda config --add channels conda-forge && \
    conda install numpyro && \
    conda clean -ya

# create user and render using X11 forwarding 
RUN export uid=1000 gid=1000 && \
    mkdir -p /home/developer && \
    mkdir -p /etc/sudoers.d && \
    echo "developer:x:${uid}:${gid}:Developer,,,:/home/developer:/bin/bash" >> /etc/passwd && \
    echo "developer:x:${uid}:" >> /etc/group && \
    echo "developer ALL=(ALL) NOPASSWD: ALL" > /etc/sudoers.d/developer && \
    chmod 0440 /etc/sudoers.d/developer && \
    chown ${uid}:${gid} -R /home/developer && \
    mkdir -m 1777 /tmp/.X11-unix

USER developer
ENV HOME /home/developer

# activate conda environment
RUN conda init bash && \
    . ~/.bashrc
SHELL ["conda", "run", "-n", "base", "/bin/bash", "-c"]