wehub-resource-sync
88 lines
2.7 KiB
Bash
88 lines
2.7 KiB
Bash
#!/usr/bin/env bash
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Download (or just verify the local cache of) a HuggingFace model
|
|
# repo, with bounded retry + exponential backoff so flaky HF mirrors
|
|
# don't fail the whole CI run.
|
|
#
|
|
# Usage:
|
|
# download_model.sh <repo_id> [<repo_id> ...]
|
|
# MODEL_ID=facebook/opt-125m download_model.sh
|
|
#
|
|
# Environment:
|
|
# MODEL_ID fallback when no positional args given
|
|
# HF_TOKEN optional HuggingFace token for authenticated
|
|
# requests (avoids anonymous rate limits)
|
|
# HF_DOWNLOAD_MAX_RETRIES default 5
|
|
# HF_DOWNLOAD_RETRY_DELAY default 30 (seconds, doubled per retry)
|
|
# HF_DOWNLOAD_FAIL_ON_ERROR default 0 (1 -> exit non-zero on failure)
|
|
|
|
set -euo pipefail
|
|
|
|
if [ "$#" -eq 0 ]; then
|
|
if [ -z "${MODEL_ID:-}" ]; then
|
|
echo "!! download_model.sh: no model id provided"
|
|
echo " pass repo ids as args or set MODEL_ID=..."
|
|
exit 2
|
|
fi
|
|
set -- "${MODEL_ID}"
|
|
fi
|
|
|
|
MAX_RETRIES="${HF_DOWNLOAD_MAX_RETRIES:-5}"
|
|
RETRY_DELAY="${HF_DOWNLOAD_RETRY_DELAY:-30}"
|
|
FAIL_ON_ERROR="${HF_DOWNLOAD_FAIL_ON_ERROR:-0}"
|
|
|
|
MAX_RETRIES="${MAX_RETRIES}" RETRY_DELAY="${RETRY_DELAY}" \
|
|
FAIL_ON_ERROR="${FAIL_ON_ERROR}" HF_TOKEN="${HF_TOKEN:-}" \
|
|
python3 - "$@" <<'PY'
|
|
import os
|
|
import random
|
|
import sys
|
|
import time
|
|
|
|
from huggingface_hub import snapshot_download
|
|
|
|
max_retries = int(os.environ["MAX_RETRIES"])
|
|
base_delay = int(os.environ["RETRY_DELAY"])
|
|
fail_on_error = os.environ["FAIL_ON_ERROR"] == "1"
|
|
hf_token = os.environ.get("HF_TOKEN") or None
|
|
|
|
repos = sys.argv[1:]
|
|
failures = []
|
|
|
|
for repo in repos:
|
|
# Try local cache first to avoid unnecessary HF API calls
|
|
# (which can 429 on busy CI runners even when the model is cached).
|
|
try:
|
|
snapshot_download(repo, local_files_only=True, token=hf_token)
|
|
print(f"CACHED: {repo} (local, no network)")
|
|
continue
|
|
except Exception:
|
|
pass
|
|
|
|
delay = base_delay
|
|
ok = False
|
|
for attempt in range(max_retries):
|
|
try:
|
|
print(f"Attempt {attempt + 1}/{max_retries}: {repo}")
|
|
snapshot_download(repo, token=hf_token)
|
|
print(f"OK: {repo}")
|
|
ok = True
|
|
break
|
|
except Exception as exc:
|
|
print(f"Attempt {attempt + 1} failed for {repo}: {exc}")
|
|
if attempt < max_retries - 1:
|
|
# Add ±25% jitter to avoid thundering herd
|
|
jitter = delay * (0.75 + random.random() * 0.5)
|
|
print(
|
|
"Waiting %.1fs before retry..." % jitter)
|
|
time.sleep(jitter)
|
|
delay *= 2
|
|
if not ok:
|
|
failures.append(repo)
|
|
|
|
if failures:
|
|
print("All retry attempts failed for: %s" % ", ".join(failures))
|
|
sys.exit(1 if fail_on_error else 0)
|
|
PY
|