dougvk/Unlimited-OCR-RDNA4

1#!/usr/bin/env bash2set -euo pipefail3 4ROOT=$(CDPATH= cd -- "$(dirname -- "$0")/.." && pwd)5PYTHON_BIN=python3.126VENV_DIR="$ROOT/.venv"7DRY_RUN=08DEPENDENCIES_ONLY=09LOCK_FILE="$ROOT/requirements/bootstrap.lock"10 11usage() {12  printf '%s\n' \13    "Usage: scripts/bootstrap-rocm.sh [--python PATH] [--venv PATH] [--dependencies-only] [--dry-run]" \14    "" \15    "Creates an isolated Python 3.12 environment with AMD's verified ROCm 7.2.1 wheels." \16    "It never modifies the system Python or the host ROCm installation."17}18 19while (($#)); do20  case "$1" in21    --python)22      PYTHON_BIN=${2:?--python requires a path}23      shift 224      ;;25    --venv)26      VENV_DIR=${2:?--venv requires a path}27      shift 228      ;;29    --dependencies-only)30      DEPENDENCIES_ONLY=131      shift32      ;;33    -n|--dry-run)34      DRY_RUN=135      shift36      ;;37    -h|--help)38      usage39      exit 040      ;;41    *)42      printf 'error: unknown argument: %s\n' "$1" >&243      usage >&244      exit 245      ;;46  esac47done48 49ROCM_RELEASE=7.2.150CACHE_BASE=${XDG_CACHE_HOME:-${HOME}/.cache}/unlimited-ocr-rdna451WHEEL_DIR="$CACHE_BASE/wheels/rocm-$ROCM_RELEASE"52TASK_TMP="$CACHE_BASE/tmp"53 54TORCH_FILE='torch-2.9.1+rocm7.2.1.lw.gitff65f5bc-cp312-cp312-linux_x86_64.whl'55TORCH_URL='https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1/torch-2.9.1%2Brocm7.2.1.lw.gitff65f5bc-cp312-cp312-linux_x86_64.whl'56TORCH_SHA='fb45ace0a27e9f0d0e3c4c6efd8932162743f8376f2aa4752a4d31ef5a1bd3d7'57 58VISION_FILE='torchvision-0.24.0+rocm7.2.1.gitb919bd0c-cp312-cp312-linux_x86_64.whl'59VISION_URL='https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1/torchvision-0.24.0%2Brocm7.2.1.gitb919bd0c-cp312-cp312-linux_x86_64.whl'60VISION_SHA='d5fca8cda173235a3b7434baeebe04c3ebffec3c6fc191e79aa8aa300633f2c9'61 62TRITON_FILE='triton-3.5.1+rocm7.2.1.gita272dfa8-cp312-cp312-linux_x86_64.whl'63TRITON_URL='https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1/triton-3.5.1%2Brocm7.2.1.gita272dfa8-cp312-cp312-linux_x86_64.whl'64TRITON_SHA='07787af1d28c273852f897bfeaa7bca29f2fa4a13ca0f28f535832b240ce7016'65 66if ((DRY_RUN)); then67  printf 'python=%s\nvenv=%s\nrocm_release=%s\nwheel_cache=%s\ndependency_lock=%s\n' \68    "$PYTHON_BIN" "$VENV_DIR" "$ROCM_RELEASE" "$WHEEL_DIR" "$LOCK_FILE"69  printf 'dependencies_only=%s\n' "$DEPENDENCIES_ONLY"70  printf 'download=%s\ndownload=%s\ndownload=%s\n' "$TORCH_URL" "$VISION_URL" "$TRITON_URL"71  printf 'next=%s/bin/unlimited-ocr-rdna4 doctor --device 0\n' "$VENV_DIR"72  exit 073fi74 75case "$(uname -s)-$(uname -m)" in76  Linux-x86_64) ;;77  *)78    printf 'error: the verified bootstrap supports Linux x86_64 only\n' >&279    exit 380    ;;81esac82 83if [[ ! -r /opt/rocm/.info/version ]]; then84  printf 'error: /opt/rocm/.info/version is missing; install a supported host ROCm stack first\n' >&285  exit 386fi87HOST_ROCM=$(tr -d '\r\n' < /opt/rocm/.info/version)88if [[ $HOST_ROCM != "$ROCM_RELEASE" ]]; then89  printf 'error: this bootstrap is verified for host ROCm 7.2.1, found %s\n' "$HOST_ROCM" >&290  exit 391fi92 93PYTHON_VERSION=$("$PYTHON_BIN" -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")')94if [[ $PYTHON_VERSION != 3.12 ]]; then95  printf 'error: AMD wheel set requires Python 3.12, found %s\n' "$PYTHON_VERSION" >&296  exit 397fi98command -v curl >/dev/null || { printf 'error: curl is required\n' >&2; exit 3; }99command -v sha256sum >/dev/null || { printf 'error: sha256sum is required\n' >&2; exit 3; }100command -v flock >/dev/null || { printf 'error: flock is required\n' >&2; exit 3; }101[[ -f $LOCK_FILE && ! -L $LOCK_FILE ]] || { printf 'error: dependency lock is missing or symlinked: %s\n' "$LOCK_FILE" >&2; exit 3; }102 103mkdir -p "$WHEEL_DIR" "$TASK_TMP"104export TMPDIR="$TASK_TMP"105export PIP_CACHE_DIR="$CACHE_BASE/pip"106exec 9>"$WHEEL_DIR/.bootstrap.lock"107flock 9108 109download_wheel() {110  local filename=$1111  local url=$2112  local expected=$3113  local destination="$WHEEL_DIR/$filename"114  local partial="$destination.part"115 116  if [[ -L $destination || -L $partial ]]; then117    printf 'error: refusing symlinked wheel cache entry: %s\n' "$destination" >&2118    exit 4119  fi120  if [[ -f $destination ]]; then121    if printf '%s  %s\n' "$expected" "$destination" | sha256sum --check --status; then122      printf 'wheel=reused path=%s\n' "$destination"123      return124    fi125    printf 'wheel=discard-corrupt path=%s\n' "$destination" >&2126    rm -f -- "$destination"127  fi128 129  for attempt in 1 2; do130    printf 'downloading=%s attempt=%s\n' "$filename" "$attempt"131    if ! curl --fail --location --proto '=https' --tlsv1.2 --continue-at - --output "$partial" "$url"; then132      printf 'wheel=discard-failed-partial path=%s\n' "$partial" >&2133      rm -f -- "$partial"134      continue135    fi136    if printf '%s  %s\n' "$expected" "$partial" | sha256sum --check --status; then137      mv "$partial" "$destination"138      return139    fi140    printf 'wheel=discard-corrupt-partial path=%s\n' "$partial" >&2141    rm -f -- "$partial"142  done143  printf 'error: downloaded wheel checksum mismatch after clean retry: %s\n' "$filename" >&2144  exit 4145}146 147download_wheel "$TORCH_FILE" "$TORCH_URL" "$TORCH_SHA"148download_wheel "$VISION_FILE" "$VISION_URL" "$VISION_SHA"149download_wheel "$TRITON_FILE" "$TRITON_URL" "$TRITON_SHA"150 151validate_venv() {152  if [[ -L $VENV_DIR || ! -d $VENV_DIR || -L $VENV_DIR/pyvenv.cfg || ! -f $VENV_DIR/pyvenv.cfg ]]; then153    printf 'error: existing --venv path is not a regular Python virtual environment: %s\n' "$VENV_DIR" >&2154    exit 3155  fi156  if [[ ! -x $VENV_DIR/bin/python ]]; then157    printf 'error: virtual environment has no executable bin/python: %s\n' "$VENV_DIR" >&2158    exit 3159  fi160  "$VENV_DIR/bin/python" - "$VENV_DIR" <<'PY'161import os162import sys163 164expected = os.path.realpath(sys.argv[1])165if os.path.realpath(sys.prefix) != expected or sys.prefix == sys.base_prefix:166    raise SystemExit("virtual environment prefix validation failed")167if sys.version_info[:2] != (3, 12):168    raise SystemExit(f"virtual environment requires Python 3.12, found {sys.version_info.major}.{sys.version_info.minor}")169PY170}171 172if [[ -L $VENV_DIR ]]; then173  printf 'error: --venv path must not be a symlink: %s\n' "$VENV_DIR" >&2174  exit 3175fi176if [[ -e $VENV_DIR && ! -d $VENV_DIR ]]; then177  printf 'error: --venv path exists and is not a directory: %s\n' "$VENV_DIR" >&2178  exit 3179fi180if [[ ! -e $VENV_DIR ]]; then181  "$PYTHON_BIN" -m venv "$VENV_DIR"182fi183validate_venv184"$VENV_DIR/bin/python" -m pip install --require-hashes --only-binary=:all: -r "$LOCK_FILE"185"$VENV_DIR/bin/python" -m pip install \186  --no-deps \187  "$WHEEL_DIR/$TORCH_FILE" \188  "$WHEEL_DIR/$VISION_FILE" \189  "$WHEEL_DIR/$TRITON_FILE"190if ((DEPENDENCIES_ONLY)); then191  "$VENV_DIR/bin/python" -m pip check192  printf 'bootstrap=PASS mode=dependencies-only venv=%s\n' "$VENV_DIR"193  exit 0194fi195BUILD_DIR=$(mktemp -d "$TASK_TMP/package-build.XXXXXX")196trap 'rm -rf -- "$BUILD_DIR"' EXIT197"$VENV_DIR/bin/python" -m pip wheel --no-build-isolation --no-deps --wheel-dir "$BUILD_DIR" "$ROOT"198PACKAGE_WHEEL=$(find "$BUILD_DIR" -maxdepth 1 -type f -name 'unlimited_ocr_rdna4-*.whl' -print -quit)199[[ -n $PACKAGE_WHEEL ]] || { printf 'error: local package wheel was not produced\n' >&2; exit 4; }200"$VENV_DIR/bin/python" -m pip install --force-reinstall --no-deps "$PACKAGE_WHEEL"201"$VENV_DIR/bin/python" -m pip check202"$VENV_DIR/bin/unlimited-ocr-rdna4" doctor --device "${UNLIMITED_OCR_DEVICE:-0}"203printf 'bootstrap=PASS venv=%s\n' "$VENV_DIR"204