#!/usr/bin/env bash
# Build the unsloth-blackwell image on any Linux host with Docker. The build host's
# GPU is NOT used: nvcc cross-compiles.
#
# Usage:
#   ./build.sh                 # builds unsloth-blackwell:latest pinned to unsloth main
#   ./build.sh --rocm          # builds unsloth-rocm:latest for AMD GPUs
#   TAG=2026.05.1 ./build.sh   # custom tag
#   UNSLOTH_REF=v2026.5.6 UNSLOTH_ZOO_REF=v2026.5.4 ./build.sh   # pin git refs
#
# ROCm: the default is the ROCm 7.2.4 base with the pytorch.org rocm7.2 wheels
# (RDNA2, RDNA3, RDNA4, CDNA). Strix APUs and RDNA4 cards get AMD's per-arch
# wheels with --gfx (or ROCM_GFX=...), which switches the index to
# repo.amd.com/rocm/whl/<family>/ and the torch 2.11 line, as install.sh does:
#   ./build.sh --rocm --gfx gfx1151        # Strix Halo
#   ROCM_GFX=gfx1201 ./build.sh --rocm     # RX 9070 XT
# gfx906 (Radeon VII / MI50) needs the last base that carries it, and --gfx gfx906
# leaves out bitsandbytes (no prebuilt gfx906 kernels):
#   ROCM_VERSION=6.3.4 ./build.sh --rocm --gfx gfx906
set -euo pipefail

cd "$(dirname "$0")"

ROCM=0
ROCM_GFX="${ROCM_GFX:-}"
while [[ $# -gt 0 ]]; do
    case "$1" in
        --rocm) ROCM=1 ;;
        --gfx)  [[ $# -ge 2 ]] || { echo "--gfx needs an arch (e.g. gfx1151)" >&2; exit 2; }
                ROCM_GFX="$2"; shift ;;
        --gfx=*) ROCM_GFX="${1#--gfx=}" ;;
        *) echo "unknown option: $1 (build.sh takes --rocm [--gfx <arch>])" >&2; exit 2 ;;
    esac
    shift
done
if [[ -n "$ROCM_GFX" && $ROCM -eq 0 ]]; then
    echo "--gfx / ROCM_GFX only applies to --rocm builds" >&2
    exit 2
fi

TAG="${TAG:-latest}"
PYTHON_VERSION="${PYTHON_VERSION:-3.12}"
UNSLOTH_REF="${UNSLOTH_REF:-main}"
UNSLOTH_ZOO_REF="${UNSLOTH_ZOO_REF:-main}"
UNSLOTH_NOTEBOOKS_REF="${UNSLOTH_NOTEBOOKS_REF:-main}"
if [[ $ROCM -eq 1 ]]; then
    IMAGE_NAME="${IMAGE_NAME:-unsloth-rocm}"
    # 7.2 is the floor: rocm6.4 tops out below what unsloth-zoo needs, and RDNA4 needs 7.x.
    ROCM_VERSION="${ROCM_VERSION:-7.2.4}"
    # the index follows the base unless named (7.2.4 -> rocm7.2)
    TORCH_INDEX_URL="${TORCH_INDEX_URL:-https://download.pytorch.org/whl/rocm${ROCM_VERSION%.*}}"
else
    IMAGE_NAME="${IMAGE_NAME:-unsloth-blackwell}"
    CUDA_VERSION="${CUDA_VERSION:-12.8.1}"
    UBUNTU_VERSION="${UBUNTU_VERSION:-24.04}"
fi

# Resolve refs to a commit: docker caches RUN layers by command string, so "main" would
# never rebuild.
resolve_git_ref() {
    local repo="$1" ref="$2" ls_out sha
    printf '%s' "$ref" | grep -Eq '^[0-9a-f]{40}$' && { printf '%s' "$ref"; return 0; }
    command -v git >/dev/null 2>&1 || { printf '%s' "$ref"; return 0; }
    # ls-remote fails only when the remote is unreachable; warn and pass the name through.
    # Also ask for ${ref}^{}: an annotated tag's own sha is not a commit, so uv refuses it.
    if ! ls_out="$(git ls-remote "$repo" "$ref" "${ref}^{}" 2>/dev/null)"; then
        echo "warning: ${repo} unreachable; passing mutable ref '${ref}' (docker may reuse a cached layer)" >&2
        printf '%s' "$ref"
        return 0
    fi
    sha="$(printf '%s\n' "$ls_out" | awk 'NR==1{s=$1; r=$2} $2==r"^{}"{s=$1} END{print s}')"
    [ -n "$sha" ] || sha="$ref"     # not a branch/tag: a short sha or already-gone ref
    printf '%s' "$sha"
}
UNSLOTH_REF="$(resolve_git_ref https://github.com/unslothai/unsloth "$UNSLOTH_REF")"
UNSLOTH_ZOO_REF="$(resolve_git_ref https://github.com/unslothai/unsloth-zoo "$UNSLOTH_ZOO_REF")"
# Freeze the notebooks ref for the same layer-cache reason. Dockerfile.rocm carries neither
# the notebooks nor the llama.cpp prebuilt.
if [[ $ROCM -eq 0 ]]; then
UNSLOTH_NOTEBOOKS_REF="$(resolve_git_ref https://github.com/unslothai/notebooks "$UNSLOTH_NOTEBOOKS_REF")"

# Resolved to a concrete tag so layer caching is correct; pin with LLAMA_PREBUILT_TAG.
resolve_latest_llama_tag() {
    curl -fsSL -o /dev/null -w '%{url_effective}' \
        "https://github.com/unslothai/llama.cpp/releases/latest" 2>/dev/null \
        | sed -n 's#.*/releases/tag/##p'
}
if [ -z "${LLAMA_PREBUILT_TAG:-}" ]; then
    LLAMA_PREBUILT_TAG="$(resolve_latest_llama_tag || true)"
    if [ -n "$LLAMA_PREBUILT_TAG" ]; then
        echo "Resolved latest llama.cpp release: ${LLAMA_PREBUILT_TAG}"
    else
        LLAMA_PREBUILT_TAG="latest"
        echo "Could not resolve latest llama.cpp tag here; passing 'latest' (resolved inside the build)"
    fi
fi
fi

if [[ $ROCM -eq 1 ]]; then
    echo "Building ${IMAGE_NAME}:${TAG}  [AMD ROCm]"
    echo "  ROCm           ${ROCM_VERSION}  Python ${PYTHON_VERSION}"
    if [[ "$ROCM_GFX" == gfx906 ]]; then
        echo "  torch index    ${TORCH_INDEX_URL} (gfx906: no bitsandbytes)"
    elif [[ -n "$ROCM_GFX" ]]; then
        echo "  torch index    AMD per-arch wheels for ${ROCM_GFX} (repo.amd.com/rocm/whl)"
    else
        echo "  torch index    ${TORCH_INDEX_URL}"
    fi
    echo "  unsloth        @${UNSLOTH_REF}"
    echo "  unsloth-zoo    @${UNSLOTH_ZOO_REF}"
    echo

    DOCKER_BUILDKIT=1 docker build \
        --progress=plain \
        -f Dockerfile.rocm \
        --build-arg ROCM_VERSION="${ROCM_VERSION}" \
        --build-arg PYTHON_VERSION="${PYTHON_VERSION}" \
        --build-arg TORCH_INDEX_URL="${TORCH_INDEX_URL}" \
        --build-arg ROCM_GFX="${ROCM_GFX}" \
        --build-arg UNSLOTH_REF="${UNSLOTH_REF}" \
        --build-arg UNSLOTH_ZOO_REF="${UNSLOTH_ZOO_REF}" \
        -t "${IMAGE_NAME}:${TAG}" \
        .

    echo
    echo "Built ${IMAGE_NAME}:${TAG}"
    echo
    echo "Smoke test on an AMD host (run.sh defaults to the published image, so name this one):"
    echo "  UNSLOTH_IMAGE=${IMAGE_NAME}:${TAG} UNSLOTH_WORKSPACE_VOLUME= bash run.sh --rocm python /workspace/smoke_test_rocm.py"
    exit 0
fi

echo "Building ${IMAGE_NAME}:${TAG}"
echo "  CUDA           ${CUDA_VERSION}  Ubuntu ${UBUNTU_VERSION}  Python ${PYTHON_VERSION}"
echo "  unsloth        @${UNSLOTH_REF}"
echo "  unsloth-zoo    @${UNSLOTH_ZOO_REF}"
echo "  llama.cpp      ${LLAMA_PREBUILT_TAG}"
echo "  notebooks      @${UNSLOTH_NOTEBOOKS_REF}"
# Read the arch list from the Dockerfile rather than repeating it.
ARCH_LIST="$(sed -n 's/^[[:space:]]*TORCH_CUDA_ARCH_LIST="\([^"]*\)".*/\1/p' \
             Dockerfile | head -n1)"
echo "  arch list      ${ARCH_LIST:-unknown}"
echo

DOCKER_BUILDKIT=1 docker build \
    --progress=plain \
    --build-arg CUDA_VERSION="${CUDA_VERSION}" \
    --build-arg UBUNTU_VERSION="${UBUNTU_VERSION}" \
    --build-arg PYTHON_VERSION="${PYTHON_VERSION}" \
    --build-arg UNSLOTH_REF="${UNSLOTH_REF}" \
    --build-arg UNSLOTH_ZOO_REF="${UNSLOTH_ZOO_REF}" \
    --build-arg LLAMA_PREBUILT_TAG="${LLAMA_PREBUILT_TAG}" \
    --build-arg UNSLOTH_NOTEBOOKS_REF="${UNSLOTH_NOTEBOOKS_REF}" \
    -t "${IMAGE_NAME}:${TAG}" \
    .

echo
echo "Built ${IMAGE_NAME}:${TAG}"
echo
# A failing nvidia-smi prints its error on stdout.
HOST_GPU=""
if HOST_GPU_QUERY="$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null)"; then
    HOST_GPU="$(printf '%s\n' "$HOST_GPU_QUERY" | head -1)"
fi
echo "Smoke test on this host${HOST_GPU:+ (${HOST_GPU})}:"
echo "  docker run --rm --gpus all ${IMAGE_NAME}:${TAG} python /workspace/smoke_test.py"
echo
echo "On another host, the same command after:"
echo "  docker pull ${IMAGE_NAME}:${TAG}   # or load .tar"