# BeeBuddy training image — pre-installs everything our HF Jobs scripts
# need so cloud-job cold start drops from ~10–15 min to ~1–2 min.
#
# Base: official pytorch image with cu12.8 (matches our PEP 723 torch pin).
# Pre-installed:
#   - All pure-Python ML deps (transformers, trl, peft, datasets, accelerate,
#     bitsandbytes, huggingface-hub, etc.) — populates UV's wheel cache so
#     venv creation is hardlinks (~5s) instead of downloads (~30s).
#   - CUDA-compiled deps (flash-linear-attention, causal-conv1d) — the
#     slow ones; saves 5–10 min/job since they don't recompile each time.
#
# Build & push:
#   bash training/docker/build.sh
#   bash training/docker/build.sh --push
#
# Use: set DEFAULT_IMAGE in scripts/submit_cloud.py to this image's URL.

FROM pytorch/pytorch:2.10.0-cuda12.8-cudnn9-devel

ENV PIP_NO_CACHE_DIR=1 \
    HF_HUB_ENABLE_HF_TRANSFER=1 \
    PYTHONUNBUFFERED=1 \
    DEBIAN_FRONTEND=noninteractive

# System deps for source builds + git installs.
RUN apt-get update && apt-get install -y --no-install-recommends \
        git build-essential ca-certificates curl && \
    rm -rf /var/lib/apt/lists/*

# uv (required by HF Jobs to execute UV scripts).
COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv

# Pre-install pure-Python deps with --system so they land in system
# site-packages AND populate uv's wheel cache. Subsequent `uv pip install`
# calls (from the PEP 723 metadata in our cloud scripts) will hardlink
# from the cache instead of downloading.
RUN uv pip install --system --break-system-packages \
        "transformers>=5.2.0,<5.3.0" \
        "trl>=0.29.0" \
        "datasets>=3.0.0" \
        "peft>=0.13.0" \
        "accelerate>=1.0.0" \
        "bitsandbytes>=0.45.0" \
        "huggingface-hub[hf_transfer]>=0.25.0" \
        "hf-xet>=1.0.0" \
        "trackio>=0.2.0" \
        "einops" \
        "rouge-score>=0.1.2" \
        "setuptools" \
        "ninja"

# CUDA-compiled deps — pre-building these is the whole point of this image.
# --no-build-isolation lets them link against the installed PyTorch/CUDA.
# Using uv (instead of plain pip) bypasses PEP 668's externally-managed
# refusal that newer Debian/Ubuntu images enforce, while keeping us on
# uv's fast resolver + wheel cache.
#
# NOTE: flash-attn is intentionally omitted — it takes 30-60 min to compile
# from source and we don't currently use it in the trainer (Gated DeltaNet
# routes through FLA, not flash-attn). Add it later if perf testing
# motivates it.
RUN uv pip install --system --break-system-packages --no-build-isolation \
        "flash-linear-attention>=0.4.0" \
        "causal-conv1d==1.6.0"

# Sanity check at build time (catches the kind of CUDA-vs-torch mismatch
# that bit us mid-job — fails the build instead of failing every job).
RUN python -c "import torch; assert torch.version.cuda, 'no CUDA'; print('torch:', torch.__version__, 'cuda:', torch.version.cuda)"
RUN python -c "import causal_conv1d; print('causal_conv1d OK')"
RUN python -c "import fla; print('fla OK')"

WORKDIR /workspace