FROM nvidia/cuda:10.1-cudnn7-devel-ubuntu18.04

# Install some basic utilities
RUN apt-get update && apt-get install -y \
    curl \
    ca-certificates \
    sudo \
    git \
    bzip2 \
    libx11-6 \
    vim \
    nano \
    libglib2.0-0 \
    rsync\
    && rm -rf /var/lib/apt/lists/*

# Create a working directory
RUN mkdir /app
WORKDIR /app

# Create a non-root user and switch to it
RUN adduser --disabled-password --gecos '' --shell /bin/bash user \
    && chown -R user:user /app
RUN echo "user ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/90-user
USER user

# All users can use /home/user as their home directory
ENV HOME=/home/user
RUN chmod 777 /home/user

# Install Miniconda
RUN curl -Lso ~/miniconda.sh https://repo.continuum.io/miniconda/Miniconda3-4.7.12.1-Linux-x86_64.sh \
    && chmod +x ~/miniconda.sh \
    && ~/miniconda.sh -b -p ~/miniconda \
    && rm ~/miniconda.sh
ENV PATH=/home/user/miniconda/bin:$PATH
ENV CONDA_AUTO_UPDATE_CONDA=false

# Create a Python 3.7 environment
RUN /home/user/miniconda/bin/conda create -y --name py37 python=3.7 \
    && /home/user/miniconda/bin/conda clean -ya
ENV CONDA_DEFAULT_ENV=py37
ENV CONDA_PREFIX=/home/user/miniconda/envs/$CONDA_DEFAULT_ENV
ENV PATH=$CONDA_PREFIX/bin:$PATH

RUN conda install -y -c pytorch \
    cudatoolkit=10.1 \
    pytorch=1.7.0 \
    torchvision \
    torchaudio \
    tqdm \
    && conda clean -ya

RUN pip install scipy

RUN pip install --no-index torch-scatter -f https://pytorch-geometric.com/whl/torch-1.7.0+cu101.html
RUN pip install --no-index torch-sparse -f https://pytorch-geometric.com/whl/torch-1.7.0+cu101.html
RUN pip install --no-index torch-cluster -f https://pytorch-geometric.com/whl/torch-1.7.0+cu101.html
RUN pip install --no-index torch-spline-conv -f https://pytorch-geometric.com/whl/torch-1.7.0+cu101.html
RUN pip install torch-geometric
RUN pip install "pybind11[global]"

RUN sudo apt-get update && sudo apt-get install -y libopenblas-dev \
    && sudo rm -rf /var/lib/apt/lists/*

ENV TORCH_CUDA_ARCH_LIST="6.1 7.0+PTX 7.5+PTX"

RUN mkdir -p /app/third_party/exptune
ADD third_party/exptune /app/third_party/exptune
RUN pip install /app/third_party/exptune

RUN mkdir -p /app/main/experiments
ADD requirements.txt /app/main/
RUN pip install -r /app/main/requirements.txt
ADD run_pretrained.sh /app/main
ADD train_ablation.sh /app/main
ADD train_main_table.sh /app/main

RUN mkdir -p /app/kernels/cpu
RUN mkdir -p /app/kernels/gpu
ADD kernels/gpu /app/kernels/gpu
RUN pip install /app/kernels/gpu

ADD kernels/cpu /app/kernels/cpu
RUN pip install /app/kernels/cpu

ADD kernels/test.py /app/kernels/
ADD kernels/check_reference.py /app/kernels/
ADD kernels/prof.py /app/kernels/
ADD kernels/profiler.sh /app/kernels/

ADD main.py /app/main/
ADD experiments /app/main/experiments

ENV DATA_LOC=/app/datasets
RUN mkdir -p /app/datasets
ENV PYTHONPATH="/app/main:${PYTHONPATH}"

WORKDIR /app/main
