#!/usr/bin/env bash
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=lib.sh
source "${SCRIPT_DIR}/lib.sh"

DATA_DIR="/mnt/shared/datasets/gsm8k"
MODEL_ID="Qwen/Qwen2.5-0.5B-Instruct"

echo "[head] prepare dataset (idempotent) -> ${DATA_DIR}"
dexec "${HEAD_CONTAINER}" bash -lc "mkdir -p ${DATA_DIR} && if [[ -f ${DATA_DIR}/train.parquet && -f ${DATA_DIR}/test.parquet ]]; then echo 'dataset_exists: skip'; else python3 /workspace/verl/examples/data_preprocess/gsm8k.py --local_save_dir ${DATA_DIR}; fi"

echo "[head] ensure model cached to persistent HF_HOME (idempotent) -> ${MODEL_ID}"
PY_CODE="$(cat <<'PY'
import os

model_id = os.environ["MODEL_ID"]

hf_home = os.environ.get("HF_HOME", "/mnt/shared/hf")
os.environ.setdefault("HF_HOME", hf_home)
os.environ.setdefault("HUGGINGFACE_HUB_CACHE", os.path.join(hf_home, "hub"))
os.environ.setdefault("TRANSFORMERS_CACHE", os.path.join(hf_home, "transformers"))

from huggingface_hub import snapshot_download

try:
    snapshot_download(repo_id=model_id, local_files_only=True)
    print("model_cache_exists: skip", model_id)
except Exception:
    print("model_cache_missing: downloading", model_id)
    snapshot_download(repo_id=model_id)
    print("model_cached_ok:", model_id)
PY
)"

printf "%s\n" "${PY_CODE}" | dexec "${HEAD_CONTAINER}" bash -lc "MODEL_ID='${MODEL_ID}' python3 -"
