# syntax=docker/dockerfile:1.7
# =============================================================================
# FastVideo CUDA development image (unified, parameterized)
# =============================================================================
# Replaces the former per-Python Dockerfiles
# (Dockerfile.python3.10 / 3.11 / 3.12 / 3.12.cuda12.9.1): pick a variant with
# --build-arg instead of maintaining a separate near-identical file.
#
#   # Default image (Python 3.12, CUDA 13.0.0, Ubuntu 22.04, Hopper sm_90a):
#   docker build -f docker/Dockerfile -t fastvideo-dev:py3.12 .
#
#   # Python 3.10 / 3.11:
#   docker build -f docker/Dockerfile --build-arg PYTHON_VERSION=3.11 \
#       -t fastvideo-dev:py3.11 .
#
#   # CUDA 13.0.0: match the torch backend + flash-attn wheel to the base CUDA.
#   docker build -f docker/Dockerfile \
#       --build-arg CUDA_VERSION=13.0.0 \
#       --build-arg UV_TORCH_BACKEND=cu130 \
#       --build-arg FLASH_ATTN_WHEEL_TAG=cu130torch2.12 \
#       -t fastvideo-dev:py3.12-cuda13.0.0 .
#
#   # DGX Spark / GB10 (aarch64, Blackwell sm_121): the matching arm64
#   # prebuilt flash-attn wheel is selected automatically from TARGETARCH.
#   docker build --platform linux/arm64 -f docker/Dockerfile \
#       --build-arg TORCH_CUDA_ARCH_LIST=12.1 -t fastvideo-dev:spark .
#
# The ARG defaults below are the source of truth for the default image.
# ROCm builds use docker/rocm7_1.Dockerfile (different base + flow).
# =============================================================================

# --- Base image selection (overridable for private mirrors / new CUDA) -------
# Build on the -devel- image: nvcc + the full toolkit are needed to compile
# fastvideo-kernel. CUDA_VERSION/UBUNTU_VERSION feed the default tag, or override
# BUILD_BASE_IMAGE wholesale for a mirror.
ARG CUDA_VERSION=13.0.0
ARG UBUNTU_VERSION=22.04
ARG BUILD_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu${UBUNTU_VERSION}

FROM ${BUILD_BASE_IMAGE} AS base

# Redeclare the ARGs we need inside the stage (FROM consumes the global ones).
ARG CUDA_VERSION
ARG PYTHON_VERSION=3.12

# CUDA arch fastvideo-kernel is compiled for. Build machines have no GPU, so
# build.sh uses this instead of probing a live device. Default Hopper (sm_90a),
# matching the released kernel wheel; override e.g. TORCH_CUDA_ARCH_LIST=12.1
# for GB10/Blackwell (DGX Spark).
ARG TORCH_CUDA_ARCH_LIST=9.0a

# Cap CMake/Ninja concurrency to fit standard 16 GB build runners. This default
# matches the CUDA 13 kernel-release build; the CUDA 12.6 image workflow uses 4.
ARG CMAKE_BUILD_PARALLEL_LEVEL=1

# flash-attn: when FLASH_ATTN_WHEEL_TAG is set, install the matching prebuilt wheel
# for TARGETARCH. The prebuilt-wheel project publishes x86_64 and arm64 wheels in
# different releases, so each architecture has its own release URL. Set the wheel
# tag to empty only when an explicit source build is desired. The wheel's
# cuXXXtorchY.Y tag must match UV_TORCH_BACKEND (below): cu130torch2.12 for the
# default 13.0 base, cu126torch2.12 for 12.6 (both backends are on torch 2.12).
ARG FLASH_ATTN_VERSION=2.8.3
ARG FLASH_ATTN_WHEEL_TAG=cu130torch2.12
ARG FLASH_ATTN_WHEEL_RELEASE=https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.17
ARG FLASH_ATTN_WHEEL_RELEASE_ARM64=https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.22

# The prebuilt flash-attn wheel ships an FA4 `flash_attn.cute` built against the
# cutlass-4.4 `cute.core.ThrMma` API, which crashes on the cutlass-dsl 4.5 that
# flashinfer/quack pull in. After the wheel install we overlay this cutlass-4.5-safe
# upstream cute (flash-attn-4) so the image runs FA4 instead of the FA2 fallback.
ARG FA4_CUTE_REF=14c377950125c70b7a9dabf9c561fca53715ac7d

# Provided automatically by BuildKit/buildx (e.g. "amd64" / "arm64") and used to
# select the prebuilt flash-attn wheel. Empty under a plain `docker build` without
# BuildKit, which we treat as x86_64 for compatibility with the historical build.
ARG TARGETARCH

ENV DEBIAN_FRONTEND=noninteractive

SHELL ["/bin/bash", "-c"]

WORKDIR /FastVideo

RUN apt-get update && apt-get install -y --no-install-recommends \
    wget \
    git \
    ca-certificates \
    openssh-server \
    zsh \
    vim \
    curl \
    ffmpeg \
    libgl1 \
    libglib2.0-0 \
    libx11-dev \
    gcc-11 \
    g++-11 \
    clang-11 \
    cmake \
    pkg-config \
    build-essential \
    libssl-dev \
    && rm -rf /var/lib/apt/lists/*

# Rust toolchain: some dependencies only ship sdists on aarch64 and need cargo
# to build. The dormant legacy Modal image layers the identical apt set +
# rustup on top of this image (fastvideo/tests/modal/pr_test.py); baking both
# here keeps its manual rollback path reproducible without changing the Slurm
# runner's package surface.
RUN set -o pipefail && \
    curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable --profile minimal && \
    /root/.cargo/bin/cargo --version && /root/.cargo/bin/rustc --version
ENV PATH=/root/.cargo/bin:${PATH}

# Set up C++20 compilers for ThunderKittens
RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11

# Set CUDA environment variables. Use the version-agnostic /usr/local/cuda
# symlink the nvidia images provide so this works for any CUDA_VERSION.
ENV CUDA_HOME=/usr/local/cuda
ENV PATH=${CUDA_HOME}/bin:${PATH}
ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH

# The uv package cache lives in a BuildKit cache mount (see RUNs below), not in an
# image layer, so we can drop the old --no-cache-dir to speed up rebuilds without
# bloating the image. Copy link mode avoids hardlink errors across the mount.
ENV UV_CACHE_DIR=/opt/uv/cache
ENV UV_LINK_MODE=copy

# pyproject no longer pins a PyTorch index; select the build at resolve time.
# Build hosts have no GPU driver (auto-detect would fall back to CPU), so pin
# explicitly, matching the base CUDA: cu130 (default, 13.0) / cu126 (12.6), both
# on torch 2.12. Keep FLASH_ATTN_WHEEL_TAG in sync.
ARG UV_TORCH_BACKEND=cu130
ENV UV_TORCH_BACKEND=${UV_TORCH_BACKEND}

# Install uv and source its environment
RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
    echo 'source $HOME/.local/bin/env' >> /root/.bashrc

# Copy only dependency metadata first to leverage Docker cache. The kernel is
# deliberately excluded from this layer: arm64 resolves it to the in-tree source,
# which needs the complete checkout and CUDA architecture settings below.
COPY pyproject.toml ./
COPY docker/uv-excludes ./docker/uv-excludes

# Create a dummy README to satisfy the installation
RUN echo "# Placeholder" > README.md

# Create the venv and install the (slow-changing) dependency tree + flash-attn
# from pyproject.toml alone, before the source is copied in, so that editing
# source below does not bust this expensive layer.
RUN --mount=type=cache,target=/opt/uv/cache \
    source $HOME/.local/bin/env && \
    uv venv --python ${PYTHON_VERSION} --seed /opt/venv && \
    source /opt/venv/bin/activate && \
    uv pip install --upgrade pip && \
    uv pip install --excludes docker/uv-excludes ".[dev]" && \
    python -c "import cv2; print('OpenCV', cv2.__version__)" && \
    PYTAG=cp$(echo "${PYTHON_VERSION}" | tr -d .) && \
    case "${TARGETARCH:-amd64}" in \
        amd64) \
            FLASH_ATTN_WHEEL_PLATFORM=linux_x86_64; \
            FLASH_ATTN_TARGET_RELEASE=${FLASH_ATTN_WHEEL_RELEASE} \
            ;; \
        arm64) \
            FLASH_ATTN_WHEEL_PLATFORM=linux_aarch64; \
            FLASH_ATTN_TARGET_RELEASE=${FLASH_ATTN_WHEEL_RELEASE_ARM64} \
            ;; \
        *) \
            echo "Unsupported TARGETARCH for flash-attn: ${TARGETARCH}" >&2; \
            exit 1 \
            ;; \
    esac && \
    if [ -n "${FLASH_ATTN_WHEEL_TAG}" ]; then \
        echo "Installing prebuilt flash-attn wheel (${FLASH_ATTN_WHEEL_TAG}, ${PYTAG}, ${FLASH_ATTN_WHEEL_PLATFORM})" && \
        uv pip install "${FLASH_ATTN_TARGET_RELEASE}/flash_attn-${FLASH_ATTN_VERSION}+${FLASH_ATTN_WHEEL_TAG}-${PYTAG}-${PYTAG}-${FLASH_ATTN_WHEEL_PLATFORM}.whl"; \
    else \
        echo "Building flash-attn ${FLASH_ATTN_VERSION} from source" && \
        uv pip install flash-attn==${FLASH_ATTN_VERSION} --no-build-isolation; \
    fi

# Overlay the CuTe-DSL-4.6-compatible upstream FA4 cute (FA4_CUTE_REF) over the
# wheel/source one so the image runs FA4, not the FA2 fallback. This pulls the FA4
# runtime stack (cutlass-dsl, quack-kernels, apache-tvm-ffi, torch-c-dlpack-ext) --
# the same deps the [dreamverse] extra already installs in CI; the installed torch
# is left untouched. The cute package only ships flash_attn/cute (never
# flash_attn/__init__.py), so FA2/varlen/bert_padding stay from the install above;
# rmtree clears the wheel's stale cute files first to avoid an install conflict.
# Then verify both survive so a broken overlay fails the build instead of shipping
# an FA2-less image. The pinned stack is validated on ARM64 GB200 (sm_100) as well
# as x86; FA4 remains opt-in through FASTVIDEO_FA4=1 so lanes with FA2 baselines
# keep their existing numerics.
RUN --mount=type=cache,target=/opt/uv/cache \
    source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
    python -c "import glob, shutil; [shutil.rmtree(d, ignore_errors=True) for d in glob.glob('/opt/venv/lib/python*/site-packages/flash_attn/cute')]" && \
    uv pip install "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@${FA4_CUTE_REF}#subdirectory=flash_attn/cute" && \
    python -c "import flash_attn; assert hasattr(flash_attn, 'flash_attn_func'), 'FA2 was clobbered by the cute overlay'; import flash_attn.cute; print('FA2 + FA4 cute OK')"

COPY . .

# Build immutable FastVideo kernel wheels for the published image. The requested
# architecture remains installed for normal image users; amd64 images also carry
# an SM89 artifact for L40S users and the dormant legacy rollback path.
ARG FASTVIDEO_KERNEL_PREBUILT_DIR=/opt/fastvideo-kernel-prebuilt
RUN --mount=type=cache,target=/opt/uv/cache \
    source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
    default_arch="${TORCH_CUDA_ARCH_LIST}" && \
    default_wheel_dir="${FASTVIDEO_KERNEL_PREBUILT_DIR}/${default_arch}" && \
    export TORCH_CUDA_ARCH_LIST="${default_arch}" && \
    cd fastvideo-kernel && \
    CMAKE_ARGS= CMAKE_BUILD_PARALLEL_LEVEL=${CMAKE_BUILD_PARALLEL_LEVEL} \
    ./build.sh --wheel-dir "${default_wheel_dir}" && \
    cd /FastVideo && \
    CMAKE_ARGS= python fastvideo/tests/modal/kernel_build_cache.py write-build-info \
        --wheel-dir "${default_wheel_dir}" \
        --output "${default_wheel_dir}/metadata.json" && \
    if [[ "${TARGETARCH:-amd64}" == "amd64" && "${default_arch}" != "8.9" ]]; then \
        export TORCH_CUDA_ARCH_LIST=8.9 && \
        l40s_wheel_dir="${FASTVIDEO_KERNEL_PREBUILT_DIR}/8.9" && \
        cd /FastVideo/fastvideo-kernel && \
        CMAKE_ARGS= CMAKE_BUILD_PARALLEL_LEVEL=${CMAKE_BUILD_PARALLEL_LEVEL} \
        ./build.sh --wheel-dir "${l40s_wheel_dir}" && \
        cd /FastVideo && \
        CMAKE_ARGS= python fastvideo/tests/modal/kernel_build_cache.py write-build-info \
            --wheel-dir "${l40s_wheel_dir}" \
            --output "${l40s_wheel_dir}/metadata.json"; \
    fi && \
    export TORCH_CUDA_ARCH_LIST="${default_arch}" && \
    default_wheel="$(find "${default_wheel_dir}" -maxdepth 1 -type f \
        \( -name 'fastvideo_kernel-*.whl' -o -name 'fastvideo-kernel-*.whl' \) \
        | sort | tail -n 1)" && \
    uv pip install "${default_wheel}" \
        --reinstall-package fastvideo-kernel --no-deps

# Install FastVideo itself (editable) now that the source is present, and set up
# shell configuration. Dependencies and the local kernel are already installed,
# so --no-deps prevents uv from resolving/building the ARM source override again.
RUN --mount=type=cache,target=/opt/uv/cache \
    source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
    uv pip install -e ".[dev]" --no-deps && \
    (git config --unset-all http.https://github.com/.extraheader || true) && \
    echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
    echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile

EXPOSE 22