# use version with cudnn 8.5 to match torch==1.13.1 that uses 8.5.0.96
# has Python 3.8.10
FROM nvcr.io/nvidia/tensorrt:22.08-py3
CMD nvidia-smi

# Ignore interactive questions during `docker build`
ENV DEBIAN_FRONTEND noninteractive

# Install and update tools to minimize security vulnerabilities
RUN apt-get update
RUN apt-get install -y software-properties-common wget apt-utils patchelf git libprotobuf-dev protobuf-compiler cmake \
    bzip2 ca-certificates libglib2.0-0 libxext6 libsm6 libxrender1 mercurial subversion libopenmpi-dev && \
    apt-get clean
RUN unattended-upgrade
RUN apt-get autoremove -y

RUN python -m pip install -U pip

RUN pip install torch scipy datasets evaluate diffusers

RUN pip install transformers
RUN pip install onnxruntime-gpu

# Install Optimum
COPY . /workspace/optimum
RUN pip install /workspace/optimum[onnxruntime-gpu,tests,exporters-gpu]

ENV TEST_LEVEL=1
ENV RUN_SLOW=1
CMD pytest exporters --durations=0 -s -vvvvv -m gpu_test
