#!/usr/bin/env bash
# =============================================================================
# runpod-gpu-train-models-sync.sh — baja al hf-cache del volume los modelos base de ai-toolkit
# de los grupos pedidos, solo si entran
#
# Se corre DENTRO del pod, después de train-setup.sh (usa el python de venv-train):
#   curl -fsSL https://dl.ramsesvii.com/runpod-gpu-train-models-sync.sh | TRAIN_MODELS=cuentos bash
# Manifiesto: runpod-gpu-train-models.txt  (grupo | repo_hf | archivo o *)
# Mide lo que falta contra el espacio libre del volume y, si no entra con TRAIN_MARGIN_GB de
# margen, no baja nada y sale con 2: el volume es compartido y llenarlo rompe a los demás.
# El libre sale de df; si el total de df supera la cuota del volume (ve el clúster de red), se usa
# GPU_VOLUME_GB (env del pod) menos du. Nunca borra. Log en /workspace/train-models-sync.log.
# [agent:server-audit][date:2026-09-28][action:created]
# =============================================================================
set -u
W="${TRAIN_WORKSPACE:-/workspace}"
DL="${GPU_MANIFEST_BASE:-https://dl.ramsesvii.com}"
PY="${TRAIN_PYTHON:-$W/venv-train/bin/python}"
LOG=$W/train-models-sync.log
export HF_HOME="${HF_HOME:-$W/hf-cache}"
say() { echo "[$(date -u +%FT%TZ)] $*" | tee -a "$LOG"; }

[ -n "${TRAIN_MODELS:-}" ] || { say "TRAIN_MODELS vacío (p. ej. TRAIN_MODELS=cuentos): nada que bajar"; exit 0; }
[ -x "$PY" ] || { say "ERROR: falta $PY; correr antes runpod-gpu-train-setup.sh"; exit 1; }
MANIFEST="${TRAIN_MANIFEST:-}"
[ -n "$MANIFEST" ] || MANIFEST=$(curl -fsSL "$DL/runpod-gpu-train-models.txt" 2>/dev/null || true)
[ -n "$MANIFEST" ] || { say "ERROR: no pude leer el manifiesto"; exit 1; }

# Libre en bytes: df, salvo que su total supere la cuota del volume (df ve el clúster de red).
read -r DF_TOTAL DF_AVAIL < <(df -B1 --output=size,avail "$W" | tail -1)
FREE=$DF_AVAIL; SRC=df
if [ -n "${GPU_VOLUME_GB:-}" ] && [ "$DF_TOTAL" -gt $(( GPU_VOLUME_GB * 1000000000 * 3 / 2 )) ]; then
  say "df ve ${DF_TOTAL} B (más que la cuota de ${GPU_VOLUME_GB} GB): midiendo con du, tarda en el volume de red"
  USED=$(du -s -B1 "$W" 2>/dev/null | cut -f1)
  FREE=$(( GPU_VOLUME_GB * 1000000000 - USED )); SRC="cuota ${GPU_VOLUME_GB} GB - du"
fi

TRAIN_MODELS="$TRAIN_MODELS" MANIFEST="$MANIFEST" FREE="$FREE" SRC="$SRC" \
MARGIN_GB="${TRAIN_MARGIN_GB:-5}" DRY="${TRAIN_DRY_RUN:-0}" "$PY" - <<'PY' 2>&1 | tee -a "$LOG"
import os, sys, fnmatch
from huggingface_hub import HfApi, hf_hub_download, snapshot_download, try_to_load_from_cache
G = lambda b: "%.2f GB" % (b / 1e9)
groups = {g.strip() for g in os.environ["TRAIN_MODELS"].split(",") if g.strip()}
free, margin = int(os.environ["FREE"]), int(float(os.environ["MARGIN_GB"]) * 1e9)
api, todo, need = HfApi(), [], 0
for line in os.environ["MANIFEST"].splitlines():
    line = line.split("#", 1)[0].strip()
    if not line: continue
    grp, repo, pat = (x.strip() for x in line.split("|"))
    if grp not in groups: continue
    info = api.model_info(repo, files_metadata=True)
    files = [s for s in info.siblings if fnmatch.fnmatch(s.rfilename, pat)]
    if not files: sys.exit("ERROR: %s no tiene %s" % (repo, pat))
    miss = [s for s in files if not isinstance(try_to_load_from_cache(repo, s.rfilename, revision=info.sha), str)]
    b = sum((s.lfs.size if s.lfs else s.size) or 0 for s in miss)
    print("  %-40s %-34s faltan %d/%d archivos, %s" % (repo, pat, len(miss), len(files), G(b)))
    need += b
    if miss: todo.append((repo, pat))
print("grupos %s: faltan %s | libre %s (%s) | margen %s" % (",".join(sorted(groups)), G(need), G(free), os.environ["SRC"], G(margin)))
if not todo: print("todo en caché"); sys.exit(0)
if need + margin > free:
    print("NO ENTRA: faltan %s y hay %s libres con %s de margen. No bajo nada." % (G(need), G(free), G(margin))); sys.exit(2)
if os.environ["DRY"] == "1": print("dry-run: entra, no bajo nada"); sys.exit(0)
# sin revision: así queda refs/main y ai-toolkit (que pide main) lo resuelve desde el caché
for repo, pat in todo:
    print("bajando %s %s" % (repo, pat), flush=True)
    if pat == "*": snapshot_download(repo)
    else: hf_hub_download(repo, pat)
print("listo: %s en %s" % (G(need), os.environ["HF_HOME"]))
PY
rc=${PIPESTATUS[0]}
exit "$rc"