FROM nvcr.io/nvidia/cuda-dl-base:25.03-cuda12.8-devel-ubuntu24.04

WORKDIR /workspace

# Init
RUN apt update

# cv2 dependencies
RUN apt-get install ffmpeg libsm6 libxext6 -y

# Util
RUN apt-get install screen -y

# Oh-my-zsh
RUN sh -c "$(wget -O- https://github.com/deluan/zsh-in-docker/releases/download/v1.2.1/zsh-in-docker.sh)" -- \
    -t robbyrussell \
    -p git \
    -p https://github.com/zsh-users/zsh-autosuggestions \
    -p https://github.com/zsh-users/zsh-syntax-highlighting \
    -p history \
    -p fzf
RUN git clone --depth 1 https://github.com/junegunn/fzf.git --branch v0.61.0 --single-branch ~/.fzf
RUN ~/.fzf/install
RUN chsh -s /bin/zsh

# Python environment setup
# Miniconda
RUN mkdir -p /root/miniconda3 && \
    wget https://repo.anaconda.com/miniconda/Miniconda3-py310_25.1.1-2-Linux-x86_64.sh && \
    bash Miniconda3-py310_25.1.1-2-Linux-x86_64.sh -b -u -p /root/miniconda3 && \
    rm Miniconda3-py310_25.1.1-2-Linux-x86_64.sh && \
    source /root/miniconda3/bin/activate && \
    conda init --all && \
    conda install -c conda-forge libstdcxx-ng=14 -y

# HF Stuff
RUN source /root/miniconda3/bin/activate && \
    pip install hf_transfer

RUN mkdir -p /root/wheels

COPY wheels/flash_attn-2.7.4.post1-cp310-cp310-linux_x86_64.whl /root/wheels
COPY wheels/triton-3.3.0+git95326d9f-cp310-cp310-linux_x86_64.whl /root/wheels
COPY wheels/vllm-0.8.4.dev0+g296c657.d20250410.cu128-cp310-cp310-linux_x86_64.whl /root/wheels
COPY wheels/xformers-0.0.30+4fa0149.d20250410-cp310-cp310-linux_x86_64.whl /root/wheels

RUN /root/miniconda3/bin/pip install --force-reinstall torch torchvision torchaudio --pre --index-url https://download.pytorch.org/whl/nightly/cu128
RUN /root/miniconda3/bin/pip install /root/wheels/xformers-0.0.30+4fa0149.d20250410-cp310-cp310-linux_x86_64.whl
RUN /root/miniconda3/bin/pip install /root/wheels/flash_attn-2.7.4.post1-cp310-cp310-linux_x86_64.whl
RUN /root/miniconda3/bin/pip install /root/wheels/vllm-0.8.4.dev0+g296c657.d20250410.cu128-cp310-cp310-linux_x86_64.whl
RUN /root/miniconda3/bin/pip install /root/wheels/triton-3.3.0+git95326d9f-cp310-cp310-linux_x86_64.whl

# Clean up
RUN rm -rf /root/wheels
RUN /root/miniconda3/bin/pip cache purge

# SSH
RUN apt-get install openssh-server -y && \
    mkdir -p /root/.ssh && \
    chmod 700 /root/.ssh && \
    sed -i 's/#Port 22/Port 1022/' /etc/ssh/sshd_config

RUN echo "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAILLRIFSAs+bubZQULKHVe7HG7oXhP2E7eqg2+Qn5sKsY zhaomh1998@outlook.com" >> /root/.ssh/authorized_keys
RUN chmod 600 /root/.ssh/authorized_keys

COPY entrypoint.sh /entrypoint.sh
RUN chmod +x /entrypoint.sh
ENTRYPOINT ["/entrypoint.sh"]