# ==============================================================
# flash_attn 휠 빌드 전용 이미지
#
# 결과 이미지: /wheels/flash_attn-{version}-cp312-cp312-linux_x86_64.whl 포함
# Gitea 레지스트리 태그: 10.89.0.8:3000/gm/flash-attn-builder:cu128-torch211
#
# 빌드: podman build -t 10.89.0.8:3000/gm/flash-attn-builder:cu128-torch211 .
# 재빌드 필요 조건: torch 버전 변경, flash_attn 버전 변경, CUDA 버전 변경
# ==============================================================

# ── Stage 1: 컴파일 ──────────────────────────────────────────
FROM docker.io/nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04 AS builder

ENV DEBIAN_FRONTEND=noninteractive
ENV PIP_NO_CACHE_DIR=1
ENV PATH=/opt/venv/bin:$PATH

RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.12 python3.12-venv python3.12-dev python3-pip \
    build-essential ninja-build git \
    && rm -rf /var/lib/apt/lists/*

RUN python3.12 -m venv /opt/venv && pip install --upgrade pip setuptools wheel

ARG TORCH_VERSION=2.11.0+cu128
ARG TORCH_INDEX=https://download.pytorch.org/whl/cu128

# torch만 설치 (flash_attn 컴파일에 torch 헤더 필요, GPU 불필요)
RUN pip install torch==${TORCH_VERSION} --extra-index-url ${TORCH_INDEX}

ARG FLASH_ATTN_VERSION=2.8.3
ARG MAX_JOBS=4

RUN mkdir /wheels && \
    MAX_JOBS=${MAX_JOBS} pip wheel flash-attn==${FLASH_ATTN_VERSION} \
    --no-build-isolation --no-deps -w /wheels

# ── Stage 2: 휠만 담은 최소 이미지 ──────────────────────────
FROM alpine:3
COPY --from=builder /wheels/ /wheels/
CMD ["ls", "/wheels/"]