gpu/single-gpu/home_benjamin_hdbasin/code/scripts/run_data_mixture.sh

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

MODE="${1:-smoke}"
OUT_ROOT="${2:-runs/data_mixture_${MODE}_$(date -u +%Y%m%d_%H%M%S)}"
PYTHON="${PYTHON:-python3}"
mkdir -p "${OUT_ROOT}" results figures reports
RUN_LOG="${OUT_ROOT}/run.log"
STATUS="${OUT_ROOT}/stage_status.tsv"
GPU_LOG="${OUT_ROOT}/gpu_telemetry.csv"
: > "${RUN_LOG}"
echo -e "stage\tstatus\tstarted_at\tfinished_at" > "${STATUS}"
echo "timestamp,name,utilization_gpu,memory_used_mib,power_draw_w" > "${GPU_LOG}"

monitor_gpu() {
  while true; do
    if command -v nvidia-smi >/dev/null 2>&1; then
      nvidia-smi --query-gpu=timestamp,name,utilization.gpu,memory.used,power.draw --format=csv,noheader,nounits >> "${GPU_LOG}" 2>/dev/null || true
    else
      date -u "+%Y-%m-%d %H:%M:%S,no_gpu,0,0,0" >> "${GPU_LOG}"
    fi
    sleep 30
  done
}

run_stage() {
  local name="$1"; shift
  local start finish
  start="$(date -u +%Y-%m-%dT%H:%M:%SZ)"
  echo "stage_start=${name} at=${start}" | tee -a "${RUN_LOG}"
  if "$@" 2>&1 | tee -a "${RUN_LOG}"; then
    finish="$(date -u +%Y-%m-%dT%H:%M:%SZ)"
    echo -e "${name}\tpass\t${start}\t${finish}" >> "${STATUS}"
    return 0
  fi
  finish="$(date -u +%Y-%m-%dT%H:%M:%SZ)"
  echo -e "${name}\tfail\t${start}\t${finish}" >> "${STATUS}"
  return 1
}

monitor_gpu &
MONITOR_PID=$!
trap 'kill "${MONITOR_PID}" 2>/dev/null || true' EXIT

{
  echo "started_at=$(date -u +%Y-%m-%dT%H:%M:%SZ)"
  echo "mode=${MODE}"
  echo "out_root=${OUT_ROOT}"
  echo "gpu=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null || echo none)"
  if [ "${MODE}" = "smoke" ]; then
    run_stage data_mixture_smoke "${PYTHON}" scripts/run_data_mixture_benchmark.py \
      --outdir "${OUT_ROOT}/data_mixture_smoke" \
      --model-id "${MIX_SMOKE_MODEL_ID:-Qwen/Qwen2.5-0.5B-Instruct}" \
      --seeds "${MIX_SMOKE_SEEDS:-1}" \
      --budget "${MIX_SMOKE_BUDGET:-2}" \
      --train-size "${MIX_SMOKE_TRAIN_SIZE:-48}" \
      --val-size "${MIX_SMOKE_VAL_SIZE:-24}" \
      --max-train-steps "${MIX_SMOKE_MAX_TRAIN_STEPS:-8}" \
      --eval-batches "${MIX_SMOKE_EVAL_BATCHES:-6}" \
      --methods "${MIX_SMOKE_METHODS:-hdbasinflow}" \
      --no-save-adapters
  else
    adapter_args=()
    if [ "${MIX_SAVE_ADAPTERS:-1}" = "0" ]; then
      adapter_args=(--no-save-adapters)
    fi
    run_stage data_mixture_full "${PYTHON}" scripts/run_data_mixture_benchmark.py \
      --outdir "${OUT_ROOT}/data_mixture_full" \
      --model-id "${MIX_MODEL_ID:-Qwen/Qwen2.5-1.5B-Instruct}" \
      --seeds "${MIX_SEEDS:-5}" \
      --budget "${MIX_BUDGET:-16}" \
      --train-size "${MIX_TRAIN_SIZE:-1024}" \
      --val-size "${MIX_VAL_SIZE:-256}" \
      --max-train-steps "${MIX_MAX_TRAIN_STEPS:-120}" \
      --eval-batches "${MIX_EVAL_BATCHES:-0}" \
      --methods "${MIX_METHODS:-hdbasinflow}" \
      "${adapter_args[@]}"
  fi
  if [ "${MODE}" = "smoke" ]; then
    SRC="${OUT_ROOT}/data_mixture_smoke"
  else
    SRC="${OUT_ROOT}/data_mixture_full"
  fi
  if [ -d "${SRC}" ]; then
    mkdir -p results figures reports
    cp -R "${SRC}/figures/." figures/ 2>/dev/null || true
    cp -R "${SRC}/reports/." reports/ 2>/dev/null || true
  fi
  echo "finished_at=$(date -u +%Y-%m-%dT%H:%M:%SZ)"
} 2>&1 | tee -a "${RUN_LOG}"

kill "${MONITOR_PID}" 2>/dev/null || true
trap - EXIT