# Docker Image for PaddlePaddle Ascend NPU

FROM registry.baidubce.com/device/paddle-cpu:ubuntu20-npu-base-x86_64-gcc84
LABEL maintainer="PaddlePaddle Authors <paddle-dev@baidu.com>"

ARG CANN_VERSION=8.0.T113
ARG SYSTEM=x86_64
ARG NPU_VERSION=910b

# HwHiAiUser
RUN groupadd -g 1000 HwHiAiUser && \
    useradd -u 1000 -g 1000 -m -d /home/HwHiAiUser HwHiAiUser

RUN mkdir -p /usr/local/Ascend/driver
WORKDIR /usr/local/Ascend

# install CANN requirement
# https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/700alpha003/softwareinstall/instg/instg_0026.html
RUN apt-get update -y && apt-get install -y zlib1g zlib1g-dev libsqlite3-dev openssl libssl-dev libffi-dev libbz2-dev \
    libxslt1-dev unzip pciutils net-tools libblas-dev gfortran libblas3 liblapack-dev liblapack3 libopenblas-dev zstd

RUN pip3.10 install --upgrade pip setuptools wheel

RUN pip3.10 install 'numpy>=1.19.2' 'decorator>=4.4.0' 'sympy>=1.5.1' 'cffi>=1.12.3' 'protobuf>=3.13.0'

RUN pip3.10 install attrs pyyaml pathlib2 scipy requests psutil absl-py

# update envs for driver
ENV LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64:$LD_LIBRARY_PATH
ENV LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64/common:$LD_LIBRARY_PATH
ENV LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64/driver:$LD_LIBRARY_PATH

# Install Ascend toolkit
RUN wget -q --no-proxy https://paddle-ascend.bj.bcebos.com/cvmodel/ascend-materials/Ascend-cann-toolkit_${CANN_VERSION}_linux-${SYSTEM}.run --no-check-certificate && \
    chmod +x Ascend-cann-toolkit_${CANN_VERSION}_linux-*.run && \
    ./Ascend-cann-toolkit_${CANN_VERSION}_linux-*.run --install --quiet && \
    echo "source /usr/local/Ascend/ascend-toolkit/set_env.sh" >>~/.bashrc && \
    rm -rf Ascend-cann-toolkit_${CANN_VERSION}_linux-*.run

# Install Ascend Kernels
ARG ASCEND_KLS=Ascend-cann-kernels-${NPU_VERSION}_${CANN_VERSION}_linux-${SYSTEM}.run
RUN wget -q --no-proxy https://paddle-ascend.bj.bcebos.com/cvmodel/ascend-materials/${ASCEND_KLS} --no-check-certificate && \
    chmod +x ${ASCEND_KLS} && ./${ASCEND_KLS} --install --quiet && rm -rf ${ASCEND_KLS}

# Install Ascend nnal
RUN wget -q --no-proxy https://paddle-ascend.bj.bcebos.com/cvmodel/ascend-materials/Ascend-cann-nnal_${CANN_VERSION}_linux-${SYSTEM}.run --no-check-certificate && \
    chmod +x Ascend-cann-nnal_${CANN_VERSION}_linux-${SYSTEM}.run && \
    . /usr/local/Ascend/ascend-toolkit/set_env.sh && \
    ./Ascend-cann-nnal_${CANN_VERSION}_linux-${SYSTEM}.run --install --quiet &&\
    rm -rf Ascend-cann-nnal_${CANN_VERSION}_linux-${SYSTEM}.run && \
    echo "source /usr/local/Ascend/nnal/atb/set_env.sh" >> /root/.bashrc

# install post process ops
RUN wget -q --no-proxy https://paddle-ascend.bj.bcebos.com/code-share-master.zip --no-check-certificate && \
    . /usr/local/Ascend/ascend-toolkit/set_env.sh && \
    unzip code-share-master.zip && \
    cd code-share-master/build && bash build_ops.sh > /dev/null && \
    chmod +x aie_ops.run && ./aie_ops.run --extract=/usr/local/Ascend/ && \
    rm -rf /usr/local/Ascend/code-share-master

# update env for ascendc
ENV ASCEND_CUSTOM_OPP_PATH=/usr/local/Ascend/vendors/aie_ascendc

# DEV image should open error level log
# 0 debug; 1 info; 2 warning; 3 error; 4 null
ENV ASCEND_GLOBAL_LOG_LEVEL=3

# environment for HCCL
ENV HCCL_CONNECT_TIMEOUT=7200
ENV HCCL_WHITELIST_DISABLE=1
ENV HCCL_SECURITY_MODE=1
ENV HCCL_BUFFSIZE=120

# environment for PaddlePaddle
ENV FLAGS_npu_storage_format=0
ENV FLAGS_use_stride_kernel=0
ENV FLAGS_allocator_strategy=naive_best_fit
ENV PADDLE_XCCL_BACKEND=npu

# map this folder in docker run
RUN rm -rf /usr/local/Ascend/driver

# Clean
RUN apt-get clean -y
RUN pip cache purge

EXPOSE 22