ARG VERSION=pytorch2.7.1_dtk26.04_v2.0_Hygon
ARG BASE_IMAGE=aicompute.tencentcloudcr.com/hyperdrive/hygon:${VERSION}
FROM ${BASE_IMAGE}

ARG COS_BASE_URL="https://haihub-model-1251001002.cos.ap-shanghai.myqcloud.com/HyperDrive/haiguang/hygon-bench/train"
ARG PIP_INDEX_URL="https://pypi.tuna.tsinghua.edu.cn/simple"
ARG COS_WHL_BASE="https://hygon-thpcbench-1251001002.cos.ap-shanghai.myqcloud.com/hygon/whl"
ARG LLVM_RUN_URL="${COS_WHL_BASE}/dtk_llvm.run"
ARG TORCH_WHL_URL="${COS_WHL_BASE}/torch-2.7.1+das.opt1.dtk2604.2605091817.g39c19f-cp310-cp310-manylinux_2_28_x86_64.whl"
ARG TRITON_WHL_URL="${COS_WHL_BASE}/triton-3.2.0+das.opt1.dtk2604.torch271.2605122039.gc13c4e-cp310-cp310-manylinux_2_28_x86_64.whl"
ARG MMCV_WHL_URL="${COS_WHL_BASE}/mmcv_full-1.6.1+das1.1.gite18c519.abi1.dtk26.torch2.5-cp310-cp310-manylinux_2_28_x86_64.whl"
ARG MMCV_2_WHL_URL="${COS_WHL_BASE}/mmcv-2.2.0+das.opt1.dtk2604.20260204.gb9d4e636-cp310-cp310-manylinux_2_28_x86_64.whl"
ARG LIGHTOP_WHL_URL="${COS_WHL_BASE}/lightop-0.7.0+das.dtk2604.torch271-cp310-cp310-manylinux_2_28_x86_64.whl"
ARG HIPDNN_WHL_URL="${COS_WHL_BASE}/hipdnn-1.0.1+dtk2604-cp310-cp310-linux_x86_64.whl"
ARG FASTPT_WHL_URL="${COS_WHL_BASE}/fastpt-2.2.0+das.dtk2604.torch271.2605091101.gc8b32c.noglog-cp310-cp310-linux_x86_64.whl"


SHELL ["/bin/bash", "-c"]

RUN apt-get update && \
    apt-get install -y --no-install-recommends zstd aria2 pv rsync python3-tk && \
    rm -rf /var/lib/apt/lists/*

# Base image provides DTK. Replace the Python stack with the torch 2.7.1
# packages from the internal package mirrors.
RUN pip uninstall -y torch triton && \
    curl -fL -o /tmp/dtk_llvm.run "${LLVM_RUN_URL}" && \
    chmod +x /tmp/dtk_llvm.run && \
    /tmp/dtk_llvm.run && \
    rm -f /tmp/dtk_llvm.run && \
    pip install "${TORCH_WHL_URL}" && \
    pip install "${TRITON_WHL_URL}" && \
    pip install "${MMCV_WHL_URL}" && \
    pip install "${FASTPT_WHL_URL}" && \
    if [ -n "${HIPDNN_WHL_URL}" ]; then \
        pip install "${HIPDNN_WHL_URL}"; \
    else \
        pip install hipdnn -i https://pypi.sourcefind.cn/release/dtk2604/; \
    fi

# Download Hyhal for compile if dtk_llvm did not already provide it.
RUN if [ -e /opt/hyhal/lib ]; then \
        echo "Use existing /opt/hyhal"; \
    else \
        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; \
    fi

# Install common deps
RUN pip install -i "${PIP_INDEX_URL}" \
    absl-py==2.3.1 antlr4-python3-runtime==4.9.3 asttokens==3.0.1 black==25.12.0 \
    decorator==5.2.1 descartes==1.1.0 executing==2.2.1 fire==0.7.1 flake8==7.3.0 \
    fvcore==0.1.5.post20221221 iopath==0.1.10 ipython==8.37.0 jedi==0.19.2 \
    joblib==1.5.3 lyft-dataset-sdk==0.0.8 Markdown==3.10 matplotlib-inline==0.2.1 \
    mccabe==0.7.0 mmdet==2.25.1 mmengine==0.10.7 mmsegmentation==0.25.0 \
    mypy_extensions==1.1.0 narwhals==2.14.0 nuscenes-devkit==1.1.11 \
    scikit-learn==1.7.2 termcolor==3.2.0 terminaltables==3.1.10 trimesh==2.35.39 \
    omegaconf==2.3.0 seaborn==0.13.2 yarl==1.20.0 yapf==0.40.1 \
    numba==0.56.4 numpy==1.23.5 future tensorboard av2 scikit-image==0.19.3 plyfile==1.1 && \
    sed 's/1.6.0/1.7.0/' -i /usr/local/lib/python3.10/dist-packages/mmseg/__init__.py

ENV TORCH_CUDA_ARCH_LIST="9.0"

# Build detectron2 with fastpt, matching the mmdet3d compile flow. Hyhal must
# stay available until both source-built packages are complete.
COPY 3rdparty/detectron2 /workspace/detectron2
RUN pip uninstall -y detectron2 && \
    cd /workspace/detectron2 && /bin/bash -c \
    "source /usr/local/bin/fastpt -C && \
    FORCE_CUDA=1 python3 setup.py bdist_wheel && \
    pip install --no-deps dist/detectron2-*.whl && \
    rm -rf /workspace/detectron2 && \
    (source /usr/local/bin/fastpt -E || true)"

# Install modified mmdet3d
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 bdist_wheel && \
    pip install --no-deps dist/mmdet3d-*.whl && \
    rm -rf /workspace/mmdetection3d && \
    (source /usr/local/bin/fastpt -E || true) && \
    rm -rf /opt/hyhal"

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

WORKDIR /workspace
