Dockerfile

ARG BASE_NAME=cpu

###############################################################################
# NVIDIA BASE IMAGE
FROM nvcr.io/nvidia/pytorch:26.01-py3 AS nvidia

RUN apt-get update -y && apt-get install -y python3-venv slurm-wlm libslurm-dev

ENV VIRTUAL_ENV=/app/.venv
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
RUN python -m venv $VIRTUAL_ENV --system-site-packages && \
    . $VIRTUAL_ENV/bin/activate

# Put HPC-X MPI in the PATH, i.e. mpirun
ENV PATH=$PATH:/opt/hpcx/ompi/bin
ARG LD_LIBRARY_PATH=""
ENV LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:/opt/hpcx/ompi/lib:/usr/local/lib/python3.12/dist-packages/torch/lib

ARG TORCH_VERSION="2.9.1"
ARG TORCH_CUDA_ARCH_LIST="7.5"

RUN pip install uv && \
    uv pip install ninja && \
    pip install --upgrade "protobuf>=6.30.0"

ENV PIP_CONSTRAINT=""

ARG INSTALL_ROOT=/app/cray
WORKDIR ${INSTALL_ROOT}

ENV BASE_NAME=nvidia

ENV TORCHINDUCTOR_MAX_AUTOTUNE=0
ENV TORCHINDUCTOR_COORDINATE_DESCENT_TUNING=0
ENV TORCH_COMPILE_DISABLE=1

###############################################################################
# NVIDIA DGX SPARK BASE IMAGE
# aarch64 Grace CPU + Blackwell GPU (SM 12.0). The NGC PyTorch image is
# multi-arch, so pulling this tag on linux/arm64 yields the aarch64 variant.
FROM nvidia AS spark

ENV BASE_NAME=spark

###############################################################################
# CPU BASE IMAGE
FROM ubuntu:24.04 AS cpu

RUN --mount=type=cache,target=/var/cache/apt \
    apt-get update -y \
    && apt-get install -y python3 python3-pip python3-venv \
    openmpi-bin libopenmpi-dev libpmix-dev slurm-wlm libslurm-dev \
    cmake

ENV VIRTUAL_ENV=/app/.venv
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
RUN python3 -m venv $VIRTUAL_ENV && \
    . $VIRTUAL_ENV/bin/activate

ARG TORCH_VERSION="2.10.0"

RUN pip install uv && \
    uv pip install torch==${TORCH_VERSION}+cpu --index-url https://download.pytorch.org/whl/cpu && \
    uv pip install ninja

# Put torch on the LD_LIBRARY_PATH
ENV LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:/app/.venv/lib64/python3.12/site-packages/torch/lib

ARG INSTALL_ROOT=/app/cray
WORKDIR ${INSTALL_ROOT}

ENV BASE_NAME=cpu

###############################################################################
# AMD BASE IMAGE
FROM gdiamos/rocm-base-mi300:v0.9926 AS amd

ENV BASE_NAME=amd

#RUN pip install pyhip>=1.1.0
ENV HIP_FORCE_DEV_KERNARG=1

ARG INSTALL_ROOT=/app/cray
WORKDIR ${INSTALL_ROOT}

ENV LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:/app/venv/lib/python3.12/site-packages/torch/lib:/usr/local/rdma-lib

###############################################################################
# ScalarLM React UI build stage
#
# Builds the first-party React SPA from ./ui into /build/dist and hands it off
# to the final stage as /app/ui-bundle. No Node.js runtime is carried forward;
# the final image only contains the static bundle.
###############################################################################
FROM node:24.2.0 AS scalarlm_ui_builder

WORKDIR /build

# Copy package metadata first so Docker can cache `npm ci` across source edits.
COPY ui/package.json ui/package-lock.json* ./

# Lockfile may not exist on a fresh checkout; fall back to npm install.
RUN if [ -f package-lock.json ]; then npm ci; else npm install; fi

COPY ui/ .
RUN npm run build

###############################################################################
# VLLM BUILD STAGE
FROM ${BASE_NAME} AS vllm

RUN --mount=type=cache,target=/var/cache/apt \
    apt-get update -y \
    && apt-get install -y curl git ccache vim numactl gcc-12 g++-12 libomp-dev libnuma-dev \
    && apt-get install -y ffmpeg libsm6 libxext6 libgl1 libdnnl-dev \
    && update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 10 --slave /usr/bin/g++ g++ /usr/bin/g++-12

ARG INSTALL_ROOT=/app/cray

WORKDIR ${INSTALL_ROOT}

# Install build dependencies FIRST
RUN pip install setuptools-scm

# Configure vLLM source - can use either local directory or remote repo.
#
# VLLM_BRANCH defaults to `main` on vllm-fork. The fixes that previously
# required pinning to scalarlm-on-v0.19.0 at a specific SHA (TorchAllocator
# crash, Triton scratch-allocator memleak, torch 2.10 ABI) are now merged
# into vllm-fork's main branch, so the branch tip is sufficient.
#
# VLLM_COMMIT remains available as an opt-in pin if a deployment needs
# reproducibility across time or wants to roll back to a specific SHA;
# leave it empty to use BRANCH tip (the default).
ARG VLLM_SOURCE=remote
ARG VLLM_BRANCH=main
ARG VLLM_COMMIT=
ARG VLLM_REPO=https://github.com/supermassive-intelligence/vllm-fork.git

# Handle vLLM source - support both local and remote modes.
# build-copy-vllm.sh and apply_patches.py are copied in a single COPY
# step so the image stays under Docker's 127-layer overlay-fs stacking
# cap. The script sources apply_patches.py from its own SCRIPT_DIR.
COPY scripts/build-copy-vllm.sh scripts/vllm_patches/apply_patches.py ${INSTALL_ROOT}/

# Handle vLLM source - single RUN command with conditional mount
# For remote: clone from repository
# For local: mount and copy from ./vllm directory
RUN --mount=type=bind,source=./vllm,target=/workspace/vllm,rw \
    bash ${INSTALL_ROOT}/build-copy-vllm.sh ${VLLM_SOURCE} ${INSTALL_ROOT}/vllm \
    /workspace/vllm ${VLLM_REPO} ${VLLM_BRANCH} ${VLLM_COMMIT}

WORKDIR ${INSTALL_ROOT}/vllm

# Set build environment variables for CPU compilation
ARG TORCH_CUDA_ARCH_LIST="7.5"
ARG VLLM_TARGET_DEVICE=cpu

ENV TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}
ENV VLLM_TARGET_DEVICE=${VLLM_TARGET_DEVICE}
ENV CMAKE_BUILD_TYPE=Release

# vLLM dependencies
COPY ./infra/requirements-vllm.txt ${INSTALL_ROOT}/requirements-vllm.txt
RUN uv pip install --no-compile --no-cache-dir -r ${INSTALL_ROOT}/requirements-vllm.txt && \
    python ${INSTALL_ROOT}/vllm/use_existing_torch.py

RUN \
    --mount=type=cache,target=/root/.cache/pip \
    --mount=type=cache,target=/root/.cache/ccache \
    --mount=type=cache,target=/app/cray/vllm/.deps \
    NPROC=$(nproc) && \
    MEM_BASED=$(free -g | awk '/^Mem:/ {print int($2/4)}') && \
    COMPUTED=$(( NPROC < MEM_BASED ? NPROC : MEM_BASED )) && \
    export MAX_JOBS=$(( COMPUTED < 16 ? COMPUTED : 16 )) && \
    echo "MAX_JOBS=${MAX_JOBS} (nproc=${NPROC}, mem-based=${MEM_BASED}, capped at 16)" && \
    pip install --no-build-isolation -e . --verbose

WORKDIR ${INSTALL_ROOT}

###############################################################################
# MAIN IMAGE
FROM vllm AS infra

# Build GPU-aware MPI
COPY ./infra/cray_infra/training/gpu_aware_mpi ${INSTALL_ROOT}/infra/cray_infra/training/gpu_aware_mpi
RUN python3 ${INSTALL_ROOT}/infra/cray_infra/training/gpu_aware_mpi/setup.py bdist_wheel --dist-dir=dist && \
    pip install dist/*.whl

RUN apt-get update -y  \
    && apt-get install -y build-essential \
    less curl wget net-tools vim iputils-ping strace gdb python3-dbg python3-dev \
    dmidecode \
    && rm -rf /var/lib/apt/lists/*

# Setup python path
ENV PYTHONPATH="${INSTALL_ROOT}/infra"
ENV PYTHONPATH="${PYTHONPATH}:${INSTALL_ROOT}/sdk"
ENV PYTHONPATH="${PYTHONPATH}:${INSTALL_ROOT}/ml"
ENV PYTHONPATH="${PYTHONPATH}:${INSTALL_ROOT}/test"
ENV PYTHONPATH="${PYTHONPATH}:${INSTALL_ROOT}/vllm"

# Megatron dependencies (GPU only)
# note this has to happen after vllm because it overrides some packages installed by vllm
COPY ./infra/requirements-megatron.txt ${INSTALL_ROOT}/requirements-megatron.txt
COPY ./infra/requirements-megatron-cpu.txt ${INSTALL_ROOT}/requirements-megatron-cpu.txt
COPY ./requirements.txt ${INSTALL_ROOT}/requirements.txt

RUN \
    if [ "$VLLM_TARGET_DEVICE" != "cpu" ]; then \
        # --no-deps so these wheels don't drag in transitive deps that
        # clobber the versions vllm installed above (this layer runs after
        # vllm on purpose). All megatron deps are pure-Python wheels — no
        # source build, so no MAX_JOBS cap or --no-build-isolation needed
        # (those were only for flash-attn, which is no longer installed).
        uv pip install --no-deps --no-compile --no-cache-dir -r ${INSTALL_ROOT}/requirements-megatron.txt; \
    fi && \
    if [ "$VLLM_TARGET_DEVICE" != "cuda" ]; then \
        uv pip install --no-compile --no-cache-dir -r ${INSTALL_ROOT}/requirements-megatron-cpu.txt; \
    fi && \
    uv pip install --no-compile --no-cache-dir -r ${INSTALL_ROOT}/requirements.txt

RUN mkdir -p ${INSTALL_ROOT}/jobs ${INSTALL_ROOT}/nfs

# Copy the rest of the platform code
COPY ./infra ${INSTALL_ROOT}/infra
COPY ./sdk ${INSTALL_ROOT}/sdk
COPY ./test ${INSTALL_ROOT}/test
COPY ./ml ${INSTALL_ROOT}/ml
COPY ./scripts ${INSTALL_ROOT}/scripts

# ScalarLM React UI static bundle — served by FastAPI at /app/*.
# Path must match SCALARLM_UI_BUNDLE_DIR in infra/cray_infra/api/fastapi/setup_ui.py.
COPY --from=scalarlm_ui_builder /build/dist /app/ui-bundle

WORKDIR ${INSTALL_ROOT}

# Build SLURM plugin
RUN /app/cray/infra/slurm_src/compile.sh

ENV LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:${PYTHONPATH:-}:/usr/local/lib/slurm
ENV SLURM_CONF=${INSTALL_ROOT}/nfs/slurm.conf
ENV VLLM_CPU_MOE_PREPACK=0