# Pull a model or dataset from HuggingFace into ana-ml2's shared HF cache
# (`/tank/aimodels/huggingface/`). Repo-type-agnostic — handles GGUFs,
# safetensors, datasets, or anything else HF hosts.
#
# Consumers (llama-swap, vLLM, training jobs) read from the same cache:
#   - llama-swap's config.yaml references entries via `-hf <user>/<repo>`
#   - vLLM stacks reference via `--model <user>/<repo>` with HF_HOME mounted
#   - datasets.load_dataset() reads from HF_HOME/datasets too
#
# This playbook does NOT edit any consumer config. Run it for the download,
# then hand-edit the relevant `stacks/<consumer>/conf/*.yaml` and deploy
# that consumer separately. Keeps the download mechanical and the
# per-asset params (ctx-size, quant choice, sampler defaults, split filters,
# etc.) where they belong — in human-curated config.
#
# Usage:
#
#   # Full repo (e.g. safetensors for later AWQ quant + vLLM serving):
#   scripts/elway ana-ml2 --playbook playbooks/pull-hf-repo.yaml \
#       --var hf_repo=Skywork/Skywork-Reward-V2-Llama-3.1-8B
#
#   # Single quant pattern (e.g. one GGUF for llama-swap):
#   scripts/elway ana-ml2 --playbook playbooks/pull-hf-repo.yaml \
#       --var hf_repo=mradermacher/Selene-1-Mini-Llama-3.1-8B-GGUF \
#       --var allow_patterns='*Q6_K*'
#
#   # Dataset (preference data, eval set, etc.):
#   scripts/elway ana-ml2 --playbook playbooks/pull-hf-repo.yaml \
#       --var hf_repo=Skywork/Skywork-Reward-Preference-80K-v0.2 \
#       --var repo_type=dataset
#
# allow_patterns is forwarded to `hf download --include` — set it to filter
# down to specific files (e.g. `*UD-Q6_K_XL*` or `train.parquet`). Omit to
# pull every file in the repo.
#
# repo_type — one of `model` (default), `dataset`, `space`. Affects both the
# `hf download --repo-type` flag AND the cache dir prefix (`models--…`,
# `datasets--…`, or `spaces--…`).

vars:
  hf_repo: ""             # REQUIRED, e.g. unsloth/Qwen3.6-35B-A3B-GGUF
  allow_patterns: ""      # OPTIONAL, e.g. "*Q6_K*"
  repo_type: "model"      # OPTIONAL, one of: model | dataset | space
  hf_cache_dir: /tank/aimodels/huggingface

steps:
  # ── Sanity: refuse to run with an empty hf_repo or an unknown repo_type. ──
  - name: Check hf_repo + repo_type are valid
    shell: |
      [ -n "{{ hf_repo }}" ] || {
        echo "elway: --var hf_repo=<user>/<repo> is required" >&2
        exit 2
      }
      case "{{ repo_type }}" in
        model|dataset|space) ;;
        *) echo "elway: repo_type must be one of: model, dataset, space (got '{{ repo_type }}')" >&2; exit 2 ;;
      esac
    changed_when: "false"

  # ── Tooling: install `hf` CLI via pipx once, inject hf_transfer for
  #            the fast multi-connection download path. Both no-ops if
  #            already installed. ──────────────────────────────────────
  - name: Install huggingface_hub via pipx (one-time)
    shell: pipx install --quiet huggingface_hub
    creates: ~/.local/bin/hf

  - name: Inject hf_transfer for faster downloads
    # pipx inject is idempotent — if hf_transfer is already in the
    # huggingface_hub venv, this is a fast no-op. Cheaper to run
    # unconditionally than to glob-check the venv path (elway's
    # creates: doesn't expand shell globs).
    shell: pipx inject --quiet huggingface_hub hf_transfer
    changed_when: "false"

  # ── Pull. hf CLI is itself idempotent — re-runs only fetch missing
  #         blobs, so safe to invoke unconditionally. ─────────────────
  - name: Pull {{ hf_repo }} ({{ repo_type }}) into {{ hf_cache_dir }}
    shell: |
      export PATH="$HOME/.local/bin:$PATH"
      pat="{{ allow_patterns }}"
      args="--repo-type {{ repo_type }}"
      if [ -n "$pat" ]; then
        args="$args --include $pat"
      fi
      HF_HOME={{ hf_cache_dir }} \
      HF_HUB_ENABLE_HF_TRANSFER=1 \
      hf download {{ hf_repo }} $args
    # changed_when:false — hf CLI doesn't expose pulled-vs-cached cheaply,
    # and a `du` diff would itself take seconds on a multi-GB repo. Accept
    # always-reporting-ok here; the verify step below confirms presence.
    changed_when: "false"

  - name: Verify cache hit for {{ hf_repo }} ({{ repo_type }})
    shell: |
      # hf normalizes user/repo to <kind>--user--repo in the hub cache.
      # Cache-dir prefix tracks repo_type: models--, datasets--, spaces--.
      prefix="{{ repo_type }}s--"
      slug=$(printf '%s' "{{ hf_repo }}" | tr '/' '~')
      slug=${slug//\~/--}
      d="{{ hf_cache_dir }}/hub/${prefix}${slug}"
      [ -d "$d" ] || { echo "elway: cache dir not found for {{ hf_repo }} at $d" >&2; exit 1; }
      n=$(find "$d/snapshots/" -mindepth 2 -maxdepth 2 \( -type f -o -type l \) 2>/dev/null | wc -l)
      [ "$n" -gt 0 ] || { echo "elway: snapshot dir exists but is empty under $d" >&2; exit 1; }
      echo "elway: $n files present under $d/snapshots/"
    changed_when: "false"

verify:
  - name: hf CLI installed
    shell: test -x ~/.local/bin/hf
    changed_when: "false"

  - name: HF cache dir writable by SSH user
    shell: test -w {{ hf_cache_dir }}
    changed_when: "false"

  - name: Repo present in cache (post-pull)
    shell: |
      prefix="{{ repo_type }}s--"
      slug=$(printf '%s' "{{ hf_repo }}" | tr '/' '~')
      slug=${slug//\~/--}
      d="{{ hf_cache_dir }}/hub/${prefix}${slug}"
      test -d "$d" && [ -n "$(find $d/snapshots/ -mindepth 2 -maxdepth 2 \( -type f -o -type l \) 2>/dev/null | head -1)" ]
    changed_when: "false"