ARG VERSION=mmcv1.6_dtk26.04_v2.0_Hygon 
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"
ARG LIGHTOP_WHL_URL="https://download.sourcefind.cn:65024/file/4/lightop/DAS1.8/lightop-0.7.0+das.dtk2604.torch271-cp310-cp310-manylinux_2_28_x86_64.whl"

SHELL ["/bin/bash", "-c"]
ENV GPU_ARCHS=gfx938 \
    LIGHTOP_GPU_CUS=64

# Download Hyhal for compile if the base image does 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

COPY Sparse4D /workspace/Sparse4D

RUN pip install "${LIGHTOP_WHL_URL}" && \
    sed -i 's/^CUs = get_CUs()/CUs = os.getenv("LIGHTOP_GPU_CUS") or get_CUs() or "64"/' /usr/local/lib/python3.10/dist-packages/lightop/utils.py && \
    sed -i 's/torch.load(filename, map_location=map_location)/torch.load(filename, map_location=map_location, weights_only=False)/' /usr/local/lib/python3.10/dist-packages/mmcv/runner/checkpoint.py && \
    source /usr/local/bin/fastpt -C && \
    pip uninstall -y deformable_aggregation_ext && \
    cd /workspace/Sparse4D/projects/mmdet3d_plugin/ops && \
    FORCE_CUDA=1 python setup.py develop && \
    (source /usr/local/bin/fastpt -E || true)

# 兼容 THPC 老模板
COPY build.sh /opt/build_sparse4d.sh

COPY one_step_data.sh /workspace/Sparse4D/HyperDrive/tools/one_step_data.sh
COPY start_mmdet3d.sh /workspace/Sparse4D/tools/start_mmdet3d.sh

RUN rm -f /root/.bash_history

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

WORKDIR /workspace/Sparse4D
