File size: 7,365 Bytes
f340984
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
#!/usr/bin/env bash
set -euo pipefail

ROOT=$(CDPATH= cd -- "$(dirname -- "$0")/.." && pwd)
PYTHON_BIN=python3.12
VENV_DIR="$ROOT/.venv"
DRY_RUN=0
DEPENDENCIES_ONLY=0
LOCK_FILE="$ROOT/requirements/bootstrap.lock"

usage() {
  printf '%s\n' \
    "Usage: scripts/bootstrap-rocm.sh [--python PATH] [--venv PATH] [--dependencies-only] [--dry-run]" \
    "" \
    "Creates an isolated Python 3.12 environment with AMD's verified ROCm 7.2.1 wheels." \
    "It never modifies the system Python or the host ROCm installation."
}

while (($#)); do
  case "$1" in
    --python)
      PYTHON_BIN=${2:?--python requires a path}
      shift 2
      ;;
    --venv)
      VENV_DIR=${2:?--venv requires a path}
      shift 2
      ;;
    --dependencies-only)
      DEPENDENCIES_ONLY=1
      shift
      ;;
    -n|--dry-run)
      DRY_RUN=1
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      printf 'error: unknown argument: %s\n' "$1" >&2
      usage >&2
      exit 2
      ;;
  esac
done

ROCM_RELEASE=7.2.1
CACHE_BASE=${XDG_CACHE_HOME:-${HOME}/.cache}/unlimited-ocr-rdna4
WHEEL_DIR="$CACHE_BASE/wheels/rocm-$ROCM_RELEASE"
TASK_TMP="$CACHE_BASE/tmp"

TORCH_FILE='torch-2.9.1+rocm7.2.1.lw.gitff65f5bc-cp312-cp312-linux_x86_64.whl'
TORCH_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'
TORCH_SHA='fb45ace0a27e9f0d0e3c4c6efd8932162743f8376f2aa4752a4d31ef5a1bd3d7'

VISION_FILE='torchvision-0.24.0+rocm7.2.1.gitb919bd0c-cp312-cp312-linux_x86_64.whl'
VISION_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'
VISION_SHA='d5fca8cda173235a3b7434baeebe04c3ebffec3c6fc191e79aa8aa300633f2c9'

TRITON_FILE='triton-3.5.1+rocm7.2.1.gita272dfa8-cp312-cp312-linux_x86_64.whl'
TRITON_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'
TRITON_SHA='07787af1d28c273852f897bfeaa7bca29f2fa4a13ca0f28f535832b240ce7016'

if ((DRY_RUN)); then
  printf 'python=%s\nvenv=%s\nrocm_release=%s\nwheel_cache=%s\ndependency_lock=%s\n' \
    "$PYTHON_BIN" "$VENV_DIR" "$ROCM_RELEASE" "$WHEEL_DIR" "$LOCK_FILE"
  printf 'dependencies_only=%s\n' "$DEPENDENCIES_ONLY"
  printf 'download=%s\ndownload=%s\ndownload=%s\n' "$TORCH_URL" "$VISION_URL" "$TRITON_URL"
  printf 'next=%s/bin/unlimited-ocr-rdna4 doctor --device 0\n' "$VENV_DIR"
  exit 0
fi

case "$(uname -s)-$(uname -m)" in
  Linux-x86_64) ;;
  *)
    printf 'error: the verified bootstrap supports Linux x86_64 only\n' >&2
    exit 3
    ;;
esac

if [[ ! -r /opt/rocm/.info/version ]]; then
  printf 'error: /opt/rocm/.info/version is missing; install a supported host ROCm stack first\n' >&2
  exit 3
fi
HOST_ROCM=$(tr -d '\r\n' < /opt/rocm/.info/version)
if [[ $HOST_ROCM != "$ROCM_RELEASE" ]]; then
  printf 'error: this bootstrap is verified for host ROCm 7.2.1, found %s\n' "$HOST_ROCM" >&2
  exit 3
fi

PYTHON_VERSION=$("$PYTHON_BIN" -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")')
if [[ $PYTHON_VERSION != 3.12 ]]; then
  printf 'error: AMD wheel set requires Python 3.12, found %s\n' "$PYTHON_VERSION" >&2
  exit 3
fi
command -v curl >/dev/null || { printf 'error: curl is required\n' >&2; exit 3; }
command -v sha256sum >/dev/null || { printf 'error: sha256sum is required\n' >&2; exit 3; }
command -v flock >/dev/null || { printf 'error: flock is required\n' >&2; exit 3; }
[[ -f $LOCK_FILE && ! -L $LOCK_FILE ]] || { printf 'error: dependency lock is missing or symlinked: %s\n' "$LOCK_FILE" >&2; exit 3; }

mkdir -p "$WHEEL_DIR" "$TASK_TMP"
export TMPDIR="$TASK_TMP"
export PIP_CACHE_DIR="$CACHE_BASE/pip"
exec 9>"$WHEEL_DIR/.bootstrap.lock"
flock 9

download_wheel() {
  local filename=$1
  local url=$2
  local expected=$3
  local destination="$WHEEL_DIR/$filename"
  local partial="$destination.part"

  if [[ -L $destination || -L $partial ]]; then
    printf 'error: refusing symlinked wheel cache entry: %s\n' "$destination" >&2
    exit 4
  fi
  if [[ -f $destination ]]; then
    if printf '%s  %s\n' "$expected" "$destination" | sha256sum --check --status; then
      printf 'wheel=reused path=%s\n' "$destination"
      return
    fi
    printf 'wheel=discard-corrupt path=%s\n' "$destination" >&2
    rm -f -- "$destination"
  fi

  for attempt in 1 2; do
    printf 'downloading=%s attempt=%s\n' "$filename" "$attempt"
    if ! curl --fail --location --proto '=https' --tlsv1.2 --continue-at - --output "$partial" "$url"; then
      printf 'wheel=discard-failed-partial path=%s\n' "$partial" >&2
      rm -f -- "$partial"
      continue
    fi
    if printf '%s  %s\n' "$expected" "$partial" | sha256sum --check --status; then
      mv "$partial" "$destination"
      return
    fi
    printf 'wheel=discard-corrupt-partial path=%s\n' "$partial" >&2
    rm -f -- "$partial"
  done
  printf 'error: downloaded wheel checksum mismatch after clean retry: %s\n' "$filename" >&2
  exit 4
}

download_wheel "$TORCH_FILE" "$TORCH_URL" "$TORCH_SHA"
download_wheel "$VISION_FILE" "$VISION_URL" "$VISION_SHA"
download_wheel "$TRITON_FILE" "$TRITON_URL" "$TRITON_SHA"

validate_venv() {
  if [[ -L $VENV_DIR || ! -d $VENV_DIR || -L $VENV_DIR/pyvenv.cfg || ! -f $VENV_DIR/pyvenv.cfg ]]; then
    printf 'error: existing --venv path is not a regular Python virtual environment: %s\n' "$VENV_DIR" >&2
    exit 3
  fi
  if [[ ! -x $VENV_DIR/bin/python ]]; then
    printf 'error: virtual environment has no executable bin/python: %s\n' "$VENV_DIR" >&2
    exit 3
  fi
  "$VENV_DIR/bin/python" - "$VENV_DIR" <<'PY'
import os
import sys

expected = os.path.realpath(sys.argv[1])
if os.path.realpath(sys.prefix) != expected or sys.prefix == sys.base_prefix:
    raise SystemExit("virtual environment prefix validation failed")
if sys.version_info[:2] != (3, 12):
    raise SystemExit(f"virtual environment requires Python 3.12, found {sys.version_info.major}.{sys.version_info.minor}")
PY
}

if [[ -L $VENV_DIR ]]; then
  printf 'error: --venv path must not be a symlink: %s\n' "$VENV_DIR" >&2
  exit 3
fi
if [[ -e $VENV_DIR && ! -d $VENV_DIR ]]; then
  printf 'error: --venv path exists and is not a directory: %s\n' "$VENV_DIR" >&2
  exit 3
fi
if [[ ! -e $VENV_DIR ]]; then
  "$PYTHON_BIN" -m venv "$VENV_DIR"
fi
validate_venv
"$VENV_DIR/bin/python" -m pip install --require-hashes --only-binary=:all: -r "$LOCK_FILE"
"$VENV_DIR/bin/python" -m pip install \
  --no-deps \
  "$WHEEL_DIR/$TORCH_FILE" \
  "$WHEEL_DIR/$VISION_FILE" \
  "$WHEEL_DIR/$TRITON_FILE"
if ((DEPENDENCIES_ONLY)); then
  "$VENV_DIR/bin/python" -m pip check
  printf 'bootstrap=PASS mode=dependencies-only venv=%s\n' "$VENV_DIR"
  exit 0
fi
BUILD_DIR=$(mktemp -d "$TASK_TMP/package-build.XXXXXX")
trap 'rm -rf -- "$BUILD_DIR"' EXIT
"$VENV_DIR/bin/python" -m pip wheel --no-build-isolation --no-deps --wheel-dir "$BUILD_DIR" "$ROOT"
PACKAGE_WHEEL=$(find "$BUILD_DIR" -maxdepth 1 -type f -name 'unlimited_ocr_rdna4-*.whl' -print -quit)
[[ -n $PACKAGE_WHEEL ]] || { printf 'error: local package wheel was not produced\n' >&2; exit 4; }
"$VENV_DIR/bin/python" -m pip install --force-reinstall --no-deps "$PACKAGE_WHEEL"
"$VENV_DIR/bin/python" -m pip check
"$VENV_DIR/bin/unlimited-ocr-rdna4" doctor --device "${UNLIMITED_OCR_DEVICE:-0}"
printf 'bootstrap=PASS venv=%s\n' "$VENV_DIR"