Dockerfile.cu128

# Surogate Docker Image
# High-performance LLM pre-training/fine-tuning framework
#
# Build Args:
#   PACKAGE_VERSION: Surogate package version (e.g., 0.1.1)
#
# Build:
#   docker build -f Dockerfile.cu128 \
#     --build-arg PACKAGE_VERSION=0.1.1 \
#     -t ghcr.io/invergent-ai/surogate:0.1.1-cu128 .
#
# Run:
#   docker run --gpus all ghcr.io/invergent-ai/surogate:0.1.1-cu128 --help
#   docker run --gpus all \
#     -v /path/to/config.yaml:/config.yaml \
#     -v /path/to/output:/output \
#     ghcr.io/invergent-ai/surogate:0.1.1-cu128 sft /config.yaml

FROM ubuntu:noble-20260210.1
LABEL org.opencontainers.image.source=https://github.com/invergent-ai/surogate
ARG PACKAGE_VERSION=0.1.1
ARG CUDA_TAG=cu128
ARG WHEEL_URL=https://github.com/invergent-ai/surogate/releases/download/v${PACKAGE_VERSION}/surogate-${PACKAGE_VERSION}+${CUDA_TAG}-cp312-abi3-manylinux_2_39_x86_64.whl
ARG DSTACK_RUNNER_VERSION=0.20.18

ENV DEBIAN_FRONTEND=noninteractive
ENV CUDA_LIB_ROOT="/root/.venv/lib/python3.12/site-packages/nvidia"
ENV UV_NO_CACHE=1
# torch 2.11 pulls the full CUDA wheel set (cudnn/cublas are 350-400MB each);
# uv's 30s default request timeout loses that race on a slow link.
ENV UV_HTTP_TIMEOUT=600

RUN apt-get update \
    && apt-get install -y --no-install-recommends --allow-change-held-packages ca-certificates curl libdw-dev python3-dev build-essential rsync wget netcat-openbsd gcc patch pciutils fuse3 openssh-server sudo openssh-server git \
    && rm -rf /var/lib/apt/lists/*

RUN echo "${CUDA_LIB_ROOT}/cuda_runtime/lib" >> /etc/ld.so.conf.d/cuda-12.conf \
    && echo "${CUDA_LIB_ROOT}/cuda_nvrtc/lib" >> /etc/ld.so.conf.d/cuda-12.conf \
    && echo "${CUDA_LIB_ROOT}/cublas/lib" >> /etc/ld.so.conf.d/cuda-12.conf \
    && echo "${CUDA_LIB_ROOT}/cufile/lib" >> /etc/ld.so.conf.d/cuda-12.conf \
    && echo "${CUDA_LIB_ROOT}/cudnn/lib" >> /etc/ld.so.conf.d/cuda-12.conf \
    && echo "${CUDA_LIB_ROOT}/nccl/lib" >> /etc/ld.so.conf.d/cuda-12.conf \
    && ldconfig

# Install uv
RUN curl -LsSf https://astral.sh/uv/install.sh | UV_INSTALL_DIR="$HOME/.local/bin" sh

WORKDIR /root

# Create virtual environment for Surogate
RUN ~/.local/bin/uv venv /root/.venv --python=3.12

ENV PATH="/root/.venv/bin:$PATH" \
    VIRTUAL_ENV="/root/.venv"

# Install wheel from URL (single RUN so nvidia version overrides don't leave duplicates in lower layers).
# ldconfig runs last: the cache built above was made before these libs existed, so without
# refreshing it the loader cannot find libcudart and the entrypoint fails to import _surogate.
RUN ~/.local/bin/uv pip install "torch==2.11.0+${CUDA_TAG}" "torchvision==0.26.0+${CUDA_TAG}" "torchaudio==2.11.0+${CUDA_TAG}" --index-url https://download.pytorch.org/whl/${CUDA_TAG} \
 && ~/.local/bin/uv pip install "vllm==0.25.1" \
 && ~/.local/bin/uv pip install ${WHEEL_URL} \
 && ~/.local/bin/uv pip install "nvidia-cuda-runtime-cu12==12.8.90" "nvidia-nccl-cu12==2.29.3" "nvidia-cufile-cu12==1.14.1.1" "nvidia-cuda-nvrtc-cu12==12.8.93" "nvidia-cudnn-cu12==9.19.0.56" \
 && ldconfig

# Add dstack launcher
RUN curl --connect-timeout 60 --max-time 240 --retry 1 --output /usr/local/bin/dstack-runner https://dstack-runner-downloads.s3.eu-west-1.amazonaws.com/${DSTACK_RUNNER_VERSION}/binaries/dstack-runner-linux-amd64 \
    && chmod +x /usr/local/bin/dstack-runner

# Add the jackalope dashboard (standalone binary, no Node needed). Non-fatal if
# the release isn't published yet — `surogate jackalope` fetches it on first run.
ARG JACKALOPE_TAG=jackalope-latest
RUN curl -fsSL "https://github.com/invergent-ai/surogate/releases/download/${JACKALOPE_TAG}/jackalope-linux-x64" \
      -o /root/.venv/bin/jackalope && chmod +x /root/.venv/bin/jackalope \
    || echo "jackalope binary not available yet — skipping"

# Default entrypoint
ENTRYPOINT ["/root/.venv/bin/surogate"]
CMD ["--help"]