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 PIP_INDEX_URL="https://pypi.tuna.tsinghua.edu.cn/simple"

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

# 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 MapTRv2 /workspace/MapTRv2
COPY build_bev_pool_ext.py /tmp/build_bev_pool_ext.py

RUN pip install -i "${PIP_INDEX_URL}" mmdet==2.14.0 mmsegmentation==0.14.1 \
    numba==0.56.4 numpy==1.23.5 && \
    cd /workspace/MapTRv2 && \
    pip install -r requirement.txt && \
    python3 -m pip install --no-cache-dir shapely==2.0.7 && \
    source /usr/local/bin/fastpt -C && \
    cd /workspace/MapTRv2/projects/mmdet3d_plugin/maptr/modules/ops/geometric_kernel_attn && \
    python setup.py build install && \
    pip uninstall -y mmdet3d && \
    cd /workspace/MapTRv2/mmdetection3d && \
    FORCE_CUDA=1 MAX_JOBS=4 python setup.py bdist_wheel && \
    pip install --no-deps dist/mmdet3d-*.whl && \
    FORCE_CUDA=1 MAX_JOBS=4 python /tmp/build_bev_pool_ext.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 && \
    sed 's/1.4.0/1.7.0/' -i /usr/local/lib/python3.10/dist-packages/mmseg/__init__.py && \
    sed 's/1.4.0/1.7.0/' -i /usr/local/lib/python3.10/dist-packages/mmdet/__init__.py && \
    (source /usr/local/bin/fastpt -E || true)

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

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

RUN rm -f /root/.bash_history

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

WORKDIR /workspace/MapTRv2
