FROM whatcanyousee/verl:ngc-th2.6.0-cu124-vllm0.8.3

# Define environments
ENV MAX_JOBS=64

RUN apt-get update && \
    apt-get install -y aria2 libfreeimage3 libfreeimage-dev zlib1g

# 4. Install Sglang
RUN pip install --no-deps "sglang[all]>=0.4.5.post3"

# 5. Install cudnn
RUN aria2c --max-tries=9999 https://developer.download.nvidia.com/compute/cudnn/9.8.0/local_installers/cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb && \
    dpkg -i cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb && \
    cp /var/cudnn-local-repo-ubuntu2204-9.8.0/cudnn-*-keyring.gpg /usr/share/keyrings/ && \
    apt-get update && \
    apt-get -y install cudnn-cuda-12 && \
    rm cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb

RUN pip install -i https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple nvidia-cudnn-cu12==9.8.0.87

# 6. Install Apex
RUN git clone https://github.com/NVIDIA/apex.git && \
    cd apex && \
    pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./

# 7. Install TransformerEngine
RUN export NVTE_FRAMEWORK=pytorch && pip3 install --no-deps git+https://github.com/NVIDIA/TransformerEngine.git@v2.2

# 8. Install Megatron-LM
RUN pip3 install --no-deps git+https://github.com/NVIDIA/Megatron-LM.git@core_r0.12.0

# 9. Fix opencv
RUN pip install opencv-python

RUN pip install opencv-fixer && \
    python -c "from opencv_fixer import AutoFix; AutoFix()"