# Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
#
# 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.


# -------------------------------------------------- #
# This is a Docker image dedicated to develop
# FasterTransformer.
# -------------------------------------------------- #

ARG DOCKER_VERSION=22.09
ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:${DOCKER_VERSION}-py3
FROM ${BASE_IMAGE}

RUN apt-get update && \
    apt-get install -y --no-install-recommends bc git-lfs&& \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

# backend build
WORKDIR /workspace/FasterTransformer
ADD . /workspace/FasterTransformer

RUN git submodule update --init --recursive && \
    git clone https://github.com/NVIDIA/NeMo /workspace/FasterTransformer/3rdparty/NeMo && \
    cd /workspace/FasterTransformer/3rdparty/NeMo && \
    git checkout e01ab855422126f2fbef3aeecba95cb1a756a05d && \
    cd -

# Originally, we need to re-install the apex package for NeMo.
# However, we don't really need apex in tests about NeMo because we 
# only use NeMo to do tokenization, dataset loading and model conversion.
# So, remove the re-installation because it is time-consuming. 
# RUN pip uninstall -y apex && \
#     pip install git+https://github.com/NVIDIA/apex --global-option="--cpp_ext" --global-option="--cuda_ext" --global-option="--fast_layer_norm" && \
#     pip install 3rdparty/NeMo[nlp]
RUN pip install --upgrade --no-cache-dir pip \
 && pip install --no-cache-dir 3rdparty/NeMo[nlp]

ENV NCCL_LAUNCH_MODE=GROUP
ARG SM=80
ARG BUILD_MIXED_GEMM=OFF
ARG FORCE_BACKEND_REBUILD=0
ARG SPARSITY_SUPPORT=OFF
ARG BUILD_MULTI_GPU=ON
ARG ENABLE_FP8=OFF
RUN mkdir /var/run/sshd -p && \
    mkdir build -p && cd build && \
    wget https://developer.download.nvidia.com/compute/libcusparse-lt/0.1.0/local_installers/libcusparse_lt-linux-x86_64-0.1.0.2.tar.gz && \
    tar -xzvf libcusparse_lt-linux-x86_64-0.1.0.2.tar.gz && \
    cmake -DSM=${SM} -DCMAKE_BUILD_TYPE=Release -DBUILD_PYT=ON -DSPARSITY_SUPPORT=${SPARSITY_SUPPORT} -DMEASURE_BUILD_TIME=ON \
      -DBUILD_CUTLASS_MIXED_GEMM=${BUILD_MIXED_GEMM} -DCUSPARSELT_PATH=/workspace/FasterTransformer/build/libcusparse_lt/ \
      -DBUILD_MULTI_GPU=${BUILD_MULTI_GPU} -DBUILD_TRT=ON -DENABLE_FP8=${ENABLE_FP8} .. && \
    make -j"$(grep -c ^processor /proc/cpuinfo)"
