#!/bin/bash
# ============================================================================
# Daimon Training — Launch Script
# Liberation Labs
# ============================================================================
#
# This script runs pre-flight checks and then launches full-parameter SFT
# via DeepSpeed ZeRO-2 on a single H200 SXM GPU.
#
# Usage:
#   bash /workspace/runpod-template/launch.sh
#
# To resume from a checkpoint after pod restart:
#   bash /workspace/runpod-template/launch.sh
#   (Automatic — the training script finds the latest checkpoint)
#
# To use a custom config:
#   DAIMON_CONFIG=/path/to/config.yaml bash /workspace/runpod-template/launch.sh
#
# ╔════════════════════════════════════════════════════════════════╗
# ║  AGENT INTERFERENCE GUARD                                     ║
# ║  This is FULL SFT. If OOM occurs, adjust batch/seq/accum.    ║
# ║  DO NOT switch to LoRA/QLoRA. Report errors to Thomas.        ║
# ╚════════════════════════════════════════════════════════════════╝
# ============================================================================

set -e

SCRIPT_DIR=$(dirname "$(readlink -f "$0")")
CONFIG="${DAIMON_CONFIG:-$SCRIPT_DIR/train_daimon_config.yaml}"
DS_CONFIG="$SCRIPT_DIR/ds_config_zero2.json"

# ── Agent interference check ──────────────────────────────────────────────
if grep -qi "lora\|peft\|adapter" "$CONFIG" 2>/dev/null; then
    echo "╔══════════════════════════════════════════════════════════╗"
    echo "║  BLOCKED: LoRA/adapter config detected in YAML.         ║"
    echo "║  This is FULL SFT. Do not modify the training method.   ║"
    echo "║  If OOM: adjust batch_size, grad_accum, or seq_length.  ║"
    echo "║  Report errors to Thomas. Do not 'fix' by downgrading.  ║"
    echo "╚══════════════════════════════════════════════════════════╝"
    exit 1
fi
if pip list 2>/dev/null | grep -qi "^peft "; then
    echo "WARNING: peft is installed. This is a full SFT pipeline."
    echo "peft should not be in this environment. Continuing but flagging."
fi
LOG_DIR="/workspace/daimon-sft/logs"
mkdir -p "$LOG_DIR"

echo "============================================================"
echo "  DAIMON FULL-PARAMETER SFT — LAUNCH"
echo "  Liberation Labs"
echo "  $(date)"
echo "============================================================"

# ── Pre-flight checks ──────────────────────────────────────────────────────

echo ""
echo "=== Pre-flight Checks ==="

# Check GPU count (need exactly 1 for this config)
GPU_COUNT=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | wc -l)
if [ "$GPU_COUNT" -lt 1 ]; then
    echo "FATAL: No GPUs detected."
    exit 1
fi
echo "  GPUs: $GPU_COUNT"

# Check VRAM (need >= 140GB for bf16 model + activations + gradients)
VRAM_MIB=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits | head -1 | tr -d ' ')
if [ "$VRAM_MIB" -lt 140000 ]; then
    echo "FATAL: GPU has ${VRAM_MIB} MiB VRAM. Need >= 140,000 MiB (H200 SXM)."
    echo "Full SFT memory budget: ~70GB model + ~20GB activations = ~90GB GPU."
    exit 1
fi
echo "  VRAM: ${VRAM_MIB} MiB"

# Check system RAM (critical for CPU-offloaded optimizer)
TOTAL_RAM_KB=$(grep MemTotal /proc/meminfo | awk '{print $2}')
TOTAL_RAM_GB=$((TOTAL_RAM_KB / 1024 / 1024))
if [ "$TOTAL_RAM_GB" -lt 180 ]; then
    echo "FATAL: System RAM is ${TOTAL_RAM_GB}GB. Need >= 180GB."
    echo "Adafactor optimizer states (~35GB) + gradients (~70GB) are CPU-offloaded."
    echo "Total CPU memory needed: ~105GB, but need headroom for OS and data loading."
    exit 1
fi
echo "  System RAM: ${TOTAL_RAM_GB} GB"

# Check disk space (full checkpoints are ~70GB each, save_total_limit=3 = ~210GB)
AVAIL_GB=$(df -BG /workspace | tail -1 | awk '{print $4}' | tr -d 'G')
if [ "$AVAIL_GB" -lt 200 ]; then
    echo "FATAL: Only ${AVAIL_GB}GB free on /workspace. Need at least 200GB."
    echo "Full model checkpoints are ~70GB each. With save_total_limit=3, need ~210GB."
    exit 1
fi
echo "  Disk: ${AVAIL_GB}GB free"

# Check model exists (either local or will be pulled from HF)
MODEL_DIR="/workspace/models/Qwen3.6-35B-A3B"
if [ -d "$MODEL_DIR" ] && [ -f "$MODEL_DIR/config.json" ]; then
    echo "  Model: $MODEL_DIR (local)"
else
    echo "  Model: will use HuggingFace ID from config"
    if [ -z "$HF_TOKEN" ]; then
        echo "  WARNING: HF_TOKEN not set. May fail if model is gated."
    fi
fi

# Check config exists
if [ ! -f "$CONFIG" ]; then
    echo "FATAL: Config not found at $CONFIG"
    exit 1
fi
echo "  Config: $CONFIG"

# Check DeepSpeed config exists
if [ ! -f "$DS_CONFIG" ]; then
    echo "FATAL: DeepSpeed config not found at $DS_CONFIG"
    exit 1
fi
echo "  DeepSpeed: $DS_CONFIG"

# Check training script exists
TRAIN_SCRIPT="$SCRIPT_DIR/train_daimon.py"
if [ ! -f "$TRAIN_SCRIPT" ]; then
    echo "FATAL: Training script not found at $TRAIN_SCRIPT"
    exit 1
fi
echo "  Script: $TRAIN_SCRIPT"

# Check for existing checkpoints (resume support)
OUTPUT_DIR=$(python3 -c "import yaml; print(yaml.safe_load(open('$CONFIG'))['output_dir'])" 2>/dev/null || echo "/workspace/daimon-sft")
CKPT_COUNT=$(find "$OUTPUT_DIR" -maxdepth 1 -name "checkpoint-*" -type d 2>/dev/null | wc -l)
if [ "$CKPT_COUNT" -gt 0 ]; then
    LATEST_CKPT=$(ls -dt "$OUTPUT_DIR"/checkpoint-* 2>/dev/null | head -1)
    echo "  Resume: $CKPT_COUNT checkpoint(s) found, will resume from $LATEST_CKPT"
else
    echo "  Resume: No checkpoints found, starting fresh"
fi

echo ""
echo "Pre-flight: PASSED"

# ── CUDA settings ──────────────────────────────────────────────────────────
export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"
export CUDA_DEVICE_MAX_CONNECTIONS=1

# Disable tokenizer parallelism (avoids fork deadlocks with DataLoader)
export TOKENIZERS_PARALLELISM=false

# ── Launch training via DeepSpeed ─────────────────────────────────────────
echo ""
echo "=== Launching Training ==="
echo "  Command: deepspeed --num_gpus=1 $TRAIN_SCRIPT --deepspeed $DS_CONFIG --config $CONFIG"
echo "  Logging to: $LOG_DIR/"
echo ""

# DeepSpeed launcher handles process management and ZeRO initialization.
# Single GPU — no NCCL communication needed.
deepspeed --num_gpus=1 "$TRAIN_SCRIPT" \
    --deepspeed "$DS_CONFIG" \
    --config "$CONFIG" \
    2>&1 | tee "$LOG_DIR/training_$(date +%Y%m%d_%H%M%S).log"

echo ""
echo "============================================================"
echo "  TRAINING FINISHED"
echo "  $(date)"
echo "  Output: $OUTPUT_DIR"
echo "============================================================"