#!/bin/bash
# vast.ai entrypoint for the MM-Jev pipeline (stage-2 multimodal). Logs: /content/pipe/{setup,pipeline}.log (+ HF work repo)
mkdir -p /content/pipe && cd /content
L=/content/pipe/setup.log
export DEBIAN_FRONTEND=noninteractive PIP_BREAK_SYSTEM_PACKAGES=1
step() { echo "=== $1 $(date -u +%T)" | tee -a $L; shift; for t in 1 2 3; do "$@" >> $L 2>&1 && return 0; echo "retry $t" >> $L; sleep 10; done; echo "STEP FAILED: $*" | tee -a $L; return 1; }
step apt bash -c "apt-get update -qq && apt-get install -y -qq ffmpeg git"
step pip pip install -q transformers==5.16.1 timm==1.0.29 datasets==4.8.5 accelerate==1.14.0 "huggingface_hub>=1.4" \
     librosa==0.11.0 soundfile opencv-python-headless matplotlib safetensors py-spy torchcodec
step code python -c "from huggingface_hub import snapshot_download as s; s('$MMJEV_WORK_REPO', allow_patterns=['code/*'], local_dir='/content/_w')"
cp /content/_w/code/* /content/
{ nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv; df -h /content | tail -1; free -g | sed -n 2p; nproc
  python -c 'import torch,transformers;print("torch",torch.__version__,"tf",transformers.__version__, torch.cuda.is_available())'; } 2>&1 | tee -a $L
while true; do
  python /content/pipeline.py >> /content/pipe/stdout.log 2>&1
  code=$?
  echo "$(date -u +%H:%M:%S) pipeline exited with $code" | tee -a /content/pipe/pipeline.log
  [ $code -eq 3 ] || break
  sleep 5
done
python - <<PY
from huggingface_hub import HfApi
for loc, rem in (("/content/pipe/setup.log", "$MMJEV_RUN_PREFIX/setup.log"), ("/content/pipe/pipeline.log", "$MMJEV_RUN_PREFIX/pipeline.log")):
    try: HfApi().upload_file(path_or_fileobj=loc, path_in_repo=rem, repo_id="$MMJEV_WORK_REPO", commit_message="final log")
    except Exception as e: print("push failed", e)
PY
# stop the instance so billing ends (results are on the Hub)
[ -n "$CONTAINER_ID" ] && [ -n "$CONTAINER_API_KEY" ] && curl -s -X PUT "https://console.vast.ai/api/v0/instances/$CONTAINER_ID/" -H "Authorization: Bearer $CONTAINER_API_KEY" -H "Content-Type: application/json" -d '{"state": "stopped"}'