# This is an optimized dockerfile but not for directly usage of developing when
# it involves onnruntime stuff.

ARG OS_VERSION=18.04
ARG PYTHON_VERSION=3.7.2
ARG PYTORCH_VERSION="1.8.1"
ARG CUDA_VERSION="11.1"
ARG CUDNN_VERSION="8"
# TensorRT version fixed by downloaded .tar package

FROM pytorch/pytorch:${PYTORCH_VERSION}-cuda${CUDA_VERSION}-cudnn${CUDNN_VERSION}-devel as torch
ARG ONNXRUNTIME_VERSION=1.5.1
ARG ONNX_VERSION=1.9.0
ARG ONNX_SIMP_VERSION=0.3.6
RUN --mount=type=cache,id=apt-dev,target=/var/cache/apt \
    apt-get update && apt-get install -y --no-install-recommends \
        build-essential \
        ca-certificates \
        ccache \
        cmake \
        wget \
        git \
    && rm -rf /var/lib/apt/lists/*
WORKDIR /workspace
COPY TensorRT-7.2.3.4.Ubuntu-18.04.x86_64-gnu.cuda-11.1.cudnn8.1.tar.gz /workspace/TensorRT-7.2.3.4.Ubuntu-18.04.x86_64-gnu.cuda-11.1.cudnn8.1.tar.gz
ENV TENSORRT_DIR=/workspace/TensorRT-7.2.3.4
ENV ONNXRUNTIME_DIR=/workspace/onnxruntime-linux-x64-${ONNXRUNTIME_VERSION}
ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$ONNXRUNTIME_DIR/lib:$TENSORRT_DIR/lib
RUN tar -xvzf TensorRT-7.2.3.4.Ubuntu-18.04.x86_64-gnu.cuda-11.1.cudnn8.1.tar.gz \
    && pip install $TENSORRT_DIR/python/tensorrt-7.2.3.4-cp38-none-linux_x86_64.whl \
    && pip install $TENSORRT_DIR/onnx_graphsurgeon/onnx_graphsurgeon-0.2.6-py2.py3-none-any.whl \
    && pip install $TENSORRT_DIR/graphsurgeon/graphsurgeon-0.4.5-py2.py3-none-any.whl \
    && wget "https://github.com/microsoft/onnxruntime/releases/download/v${ONNXRUNTIME_VERSION}/onnxruntime-linux-x64-${ONNXRUNTIME_VERSION}.tgz" \
    && tar -zxvf onnxruntime-linux-x64-${ONNXRUNTIME_VERSION}.tgz \
    && rm -rf TensorRT-7.2.3.4.Ubuntu-18.04.x86_64-gnu.cuda-11.1.cudnn8.1.tar.gz onnxruntime-linux-x64-${ONNXRUNTIME_VERSION}.tgz \
    && pip --no-cache-dir install onnxruntime==${ONNXRUNTIME_VERSION} \
    && git clone https://github.com/open-mmlab/mmcv/ /workspace/mmcv \
    && cd mmcv \
    && MMCV_WITH_OPS=1 MMCV_WITH_TRT=1 MMCV_WITH_ORT=1 pip install -e . \
    && pip --no-cache-dir install onnx-simplifier==${ONNX_SIMP_VERSION} onnx==${ONNX_VERSION} pycocotools terminaltables

FROM ubuntu:${OS_VERSION} as official
ARG PYTORCH_VERSION="1.8.1"
RUN --mount=type=cache,id=apt-dev,target=/var/cache/apt \
    apt-get update && apt-get install -y --no-install-recommends \
        ca-certificates \
        libjpeg-dev \
        libpng-dev \
        ccache \
        cmake \
        gcc \
        git \
        # MMDet Requirements
        ffmpeg libsm6 libxext6 git ninja-build libglib2.0-0 libsm6 libxrender-dev libxext6 \
    && rm -rf /var/lib/apt/lists/*
ENV PATH /opt/conda/bin:$PATH
ENV NVIDIA_VISIBLE_DEVICES all
ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
ENV LD_LIBRARY_PATH /usr/local/nvidia/lib:/usr/local/nvidia/lib64
ENV PYTORCH_VERSION ${PYTORCH_VERSION}
WORKDIR /workspace

FROM official as final
COPY --from=torch /workspace /workspace
COPY --from=torch /opt/conda /opt/conda
RUN /opt/conda/bin/conda clean -ya
WORKDIR /workspace
CMD /bin/bash