#  Copyright (c) Meta Platforms, Inc. and affiliates.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
# CUDA Docker Image for AITemplate

FROM nvidia/cuda:11.6.2-devel-ubuntu20.04

# Base scripts
RUN apt-get update --fix-missing
RUN apt install -y python3 python3-dev python3-pip

# Environment variables
ENV PATH=/usr/local/nvidia/bin:${PATH}
ENV PATH=/usr/local/cuda/bin:${PATH}
ENV LIBRARY_PATH=/usr/local/cuda/lib64:${LIBRARY_PATH}
ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH}

ADD ./docker/install/ /Install
# necessary package
RUN bash /Install/install_basic_dep.sh

# for test
RUN bash /Install/install_test_dep.sh

# for docs
RUN bash /Install/install_doc_dep.sh


# install Pytorch
RUN pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113

# install NVIDIA cuda-python
RUN pip3 install 'cuda-python<12.0.0'

# for detection
RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get -y install tzdata
RUN bash /Install/install_detection_deps.sh

# Copy AITemplate to Docker
RUN mkdir /AITemplate
ADD ./COMMIT_INFO /AITemplate/COMMIT_INFO
ADD ./python /AITemplate/python
ADD ./3rdparty /AITemplate/3rdparty
ADD ./examples /AITemplate/examples
ADD ./tests /AITemplate/tests
ADD ./docs /AITemplate/docs
ADD ./static /AITemplate/static
ADD ./licenses /AITemplate/licenses
ADD ./docker/install/install_ait.sh /AITemplate/
RUN bash /AITemplate/install_ait.sh