ARG VERSION=pytorch_dtk26.04_v1.0_HygonBW1101
FROM aicompute.tencentcloudcr.com/hyperdrive/hygon:${VERSION}

ARG COS_BASE_URL="https://haihub-model-1251001002.cos.ap-shanghai.myqcloud.com/HyperDrive/haiguang/hygon-bench/train"

# MMCV 已内置在基础镜像

# Download Hyhal for compile
RUN aria2c --enable-color=false --show-console-readout=false -x 8 -s 8 -k 1M -d /tmp -o hyhal.tar.gz ${COS_BASE_URL}/common/hyhal.tar.gz && \
    mkdir -p /opt/hyhal && \
    tar -zxf /tmp/hyhal.tar.gz --strip-components=2 -C /opt/hyhal && \
    rm -rf /tmp/hyhal.tar.gz

# Install mmcv1.6
RUN pip uninstall -y mmcv && \
    pip install ${COS_BASE_URL}/autodrive/mmcv_full-1.6.1%2Bdas1.1.gite18c519.abi1.dtk26.torch2.5-cp310-cp310-manylinux_2_28_x86_64.whl

# Install fastpt
RUN pip install ${COS_BASE_URL}/autodrive/fastpt-2.1.1%2Bdas.dtk25042-cp310-cp310-linux_x86_64.whl

# Install detectron
COPY 3rdparty/detectron2 /workspace/detectron2
RUN pip uninstall -y detectron2 && \
    cd /workspace/detectron2 && \
    /bin/bash -c \
    "source /usr/local/bin/fastpt -C && \
    python3 setup.py install && \
    rm -rf /workspace/detectron2 && source /usr/local/bin/fastpt -E"

# Install common deps
RUN pip install numpy==1.23.5 numba==0.56.4 yapf==0.40.1 setuptools==61.0.0 \
    IPython fvcore omegaconf seaborn tensorboard lyft-dataset-sdk nuscenes-devkit trimesh \
    mmdet==2.26.0 mmsegmentation==0.29.1 av2 scikit-image==0.19.3 plyfile==1.1

# Install modified mmdet3d
ENV TORCH_CUDA_ARCH_LIST="9.0"
COPY 3rdparty/mmdetection3d /workspace/mmdetection3d
RUN pip uninstall -y mmdet3d
RUN cd /workspace/mmdetection3d && /bin/bash -c \
    "source /usr/local/bin/fastpt -C && \
    FORCE_CUDA=1 python3 setup.py install && \
    rm -rf /workspace/mmdetection3d && source /usr/local/bin/fastpt -E && rm -rf /opt/hyhal"

# Install hipdnn
RUN wget --progress=dot:giga -O /tmp/hipDNN.tar.gz ${COS_BASE_URL}/autodrive/hipDNN.tar.gz && \
    mkdir -p /opt/dtk/share/hipdnn && \
    tar -zxf /tmp/hipDNN.tar.gz --strip-components=2 -C /opt/dtk/share/hipdnn && \
    rm -f /tmp/hipDNN.tar.gz && \
    export LD_LIBRARY_PATH="/opt/dtk/share/hipdnn/lib:$LD_LIBRARY_PATH" && \
    pip install /opt/dtk/share/hipdnn/hipdnn-1.0.1+dtk250404-cp310-cp310-linux_x86_64.whl && \
    printf "export LD_LIBRARY_PATH=\"/opt/dtk/share/hipdnn/lib:\$LD_LIBRARY_PATH\"\n\n" | tee -a /root/.bashrc /root/.hygon_base_env

# Remove hyhal for runtime
RUN rm -rf /opt/hyhal

WORKDIR /workspace
