FROM nvcr.io/nvidia/pytorch:24.04-py3

ENV DEBIAN_FRONTEND=noninteractive
WORKDIR /workspace

RUN sed -i 's|http://archive.ubuntu.com|http://mirrors.huaweicloud.com|g' /etc/apt/sources.list && \
    sed -i 's|http://security.ubuntu.com|http://mirrors.huaweicloud.com|g' /etc/apt/sources.list

RUN apt update && apt install -y \
    build-essential \
    git \
    clang \
    pkg-config \
    ca-certificates \
    curl \
    protobuf-compiler \
    && rm -rf /var/lib/apt/lists/*

ARG CA_PATH=

COPY ${CA_PATH} /usr/local/share/ca-certificates/

RUN if [ -n "${CA_PATH}" ]; then \
        update-ca-certificates; \
    fi

ENV RUSTUP_DIST_SERVER=https://mirrors.ustc.edu.cn/rust-static/
ENV RUSTUP_UPDATE_ROOT=https://mirrors.ustc.edu.cn/rust-static/rustup/

RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y

ENV PATH=/root/.cargo/bin:$PATH

COPY . /workspace/xmig
WORKDIR /workspace/xmig/GPU

RUN mkdir -p /root/.cargo && \
    printf '[source.crates-io]\nreplace-with = "ustc"\n\n[source.ustc]\nregistry = "sparse+https://mirrors.ustc.edu.cn/crates.io-index/"\n' \
    > /root/.cargo/config.toml

RUN cargo fetch

ENV CUDA_HOME=/usr/local/cuda
ENV NCCL_HOME=/workspace/xmig/GPU/target/nccl-2.27.7-1
ENV PATH=${CUDA_HOME}/bin:${PATH}
ENV LD_LIBRARY_PATH=${NCCL_HOME}/build/lib:${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}

WORKDIR /workspace

CMD ["/bin/bash"]