Upload onstart_v2.sh with huggingface_hub

onstart_v2.sh

@@ -1,12 +1,12 @@
 #!/bin/bash
-# Khudi AI v2 — On-start script (FIXED VERSION)
-# All issues from v1 addressed
+# Khudi AI v2 — On-start script (v3 - ALL FIXES APPLIED)
+# All issues from v1 + v2.1 addressed
 
-set -e
-exec > >(tee -a /workspace/train.log) 2>&1
+# Don't use 'set -e' - we want to handle errors gracefully
+# instead of exiting on first failure
 
 echo "============================================"
-echo "🇵🇰 Khudi AI v2 — SAFE TRAINING (v2.1)"
+echo "🇵🇰 Khudi AI v2 — SAFE TRAINING (v3)"
 echo "============================================"
 echo "Start: $(date)"
 
@@ -18,11 +18,11 @@ IDLE_COUNT=0
 # ============================================================
 # SAFETY: 14-hour auto-shutdown
 # ============================================================
-( sleep 50400 && echo "⏰ 14 HOUR TIMEOUT" && pkill -9 -f python ) &
+( sleep 50400 && echo "⏰ 14 HOUR TIMEOUT" && pkill -9 -f "python train_v2" ) &
 TIMEOUT_PID=$!
 
 # ============================================================
-# SAFETY: Health check (only starts AFTER training begins)
+# SAFETY: Health check (only after training starts)
 # ============================================================
 start_health_check() {
     while true; do
@@ -33,12 +33,12 @@ start_health_check() {
             IDLE_COUNT=$((IDLE_COUNT + 1))
             if [ "$IDLE_COUNT" -gt 12 ]; then
                 echo "[HEALTH] GPU idle 60+ min - EMERGENCY EXIT"
-                pkill -f "python.*train" 2>/dev/null
                 # Save what we have
                 if [ -d /workspace/khudi-v2-output/checkpoints ]; then
-                    LATEST=$(ls /workspace/khudi-v2-output/checkpoints/ | grep checkpoint- | sort -V | tail -1)
+                    LATEST=$(ls /workspace/khudi-v2-output/checkpoints/ 2>/dev/null | grep checkpoint- | sort -V | tail -1)
                     if [ -n "$LATEST" ]; then
-                        python -c "
+                        echo "Saving $LATEST to HF..."
+                        HF_TOKEN="$HF_TOKEN_VALUE" python -c "
 from huggingface_hub import HfApi
 import os
 api = HfApi()
@@ -52,6 +52,7 @@ print('Saved $LATEST to HF')
 " 2>&1
                     fi
                 fi
+                pkill -9 -f "python train_v2" 2>/dev/null
                 kill -9 $TIMEOUT_PID 2>/dev/null
                 exit 2
             fi
@@ -73,7 +74,7 @@ pip install -q --no-cache-dir \
     "accelerate==0.34.0" \
     "bitsandbytes==0.43.3" \
     "huggingface_hub==0.25.0" \
-    "sentencepiece" "protobuf" 2>&1 | tail -3
+    "sentencepiece" "protobuf" 2>&1 | tail -3 || echo "  ⚠️ Some packages may have failed"
 
 # ============================================================
 # Step 2: Setup directories
@@ -84,31 +85,42 @@ cd /workspace
 mkdir -p khudi-v2-output/checkpoints khudi-v2-output/hf_cache
 
 # ============================================================
-# Step 3: Download train script from HF (FIX: was missing)
+# Step 3: Download train script from HF
 # ============================================================
 echo "[3/8] Downloading training scripts from HF..."
 HF_TOKEN_VALUE="HF_TOKEN_PLACEHOLDER"
 
-wget -q --header="Authorization: Bearer $HF_TOKEN_VALUE" \
-     -O train_v2.py \
-     "https://huggingface.co/ZaoKing/khudi-ai-v2/resolve/main/train_v2.py"
+# FIX: Add retry logic
+download_with_retry() {
+    local url="$1"
+    local output="$2"
+    local max_retries=3
+    local retry=0
 
-if [ ! -s train_v2.py ]; then
+    while [ $retry -lt $max_retries ]; do
+        if wget -q --header="Authorization: Bearer $HF_TOKEN_VALUE" -O "$output" "$url" && [ -s "$output" ]; then
+            return 0
+        fi
+        retry=$((retry + 1))
+        echo "  Retry $retry/$max_retries for $output..."
+        sleep 5
+    done
+    return 1
+}
+
+if ! download_with_retry "https://huggingface.co/ZaoKing/khudi-ai-v2/resolve/main/train_v2.py" "train_v2.py"; then
     echo "❌ FATAL: Could not download train_v2.py"
     exit 1
 fi
 echo "  ✅ train_v2.py downloaded: $(wc -l < train_v2.py) lines"
 
 # ============================================================
-# Step 4: Download v2 dataset (FIX: was 1 URL, need 4 chunks)
+# Step 4: Download v2 dataset (4 chunks)
 # ============================================================
 echo "[4/8] Downloading v2 dataset (4 chunks)..."
 for i in 0 1 2 3; do
     echo "  Downloading chunk $i..."
-    wget -q --header="Authorization: Bearer $HF_TOKEN_VALUE" \
-         -O "v2_chunk_$i.jsonl" \
-         "https://huggingface.co/ZaoKing/khudi-ai-v2/resolve/main/data/v2_chunk_$i.jsonl"
-    if [ ! -s "v2_chunk_$i.jsonl" ]; then
+    if ! download_with_retry "https://huggingface.co/ZaoKing/khudi-ai-v2/resolve/main/data/v2_chunk_$i.jsonl" "v2_chunk_$i.jsonl"; then
         echo "  ❌ FATAL: Could not download chunk $i"
         exit 1
     fi
@@ -116,13 +128,15 @@ done
 
 # Combine chunks
 cat v2_chunk_0.jsonl v2_chunk_1.jsonl v2_chunk_2.jsonl v2_chunk_3.jsonl > v2_data.jsonl
-echo "  ✅ Combined: $(wc -l < v2_data.jsonl) samples ($(ls -lh v2_data.jsonl | awk '{print $5}'))"
+TOTAL=$(wc -l < v2_data.jsonl)
+SIZE=$(ls -lh v2_data.jsonl | awk '{print $5}')
+echo "  ✅ Combined: $TOTAL samples ($SIZE)"
 
 # ============================================================
 # Step 5: Quick data validation
 # ============================================================
 echo "[5/8] Validating data..."
-python -c "
+HF_TOKEN="$HF_TOKEN_VALUE" python -c "
 import json
 count = 0
 errors = 0
@@ -137,7 +151,10 @@ with open('/workspace/v2_data.jsonl') as f:
         except:
             errors += 1
 print(f'Valid: {count}, Errors: {errors}')
-" 2>&1
+if count < 1000:
+    print('FATAL: Not enough valid samples!')
+    exit(1)
+" || { echo "❌ Data validation failed"; exit 1; }
 
 # ============================================================
 # Step 6: Check for resume checkpoint
@@ -153,29 +170,31 @@ else
 fi
 
 # ============================================================
-# Step 7: Start training (health check only NOW)
+# Step 7: Start training
 # ============================================================
 echo "[7/8] Starting health check + training..."
 start_health_check &
 HEALTH_PID=$!
 echo "  Health check PID: $HEALTH_PID"
 
-# Export token for the training script
+# Export ALL required env vars
 export HF_TOKEN="$HF_TOKEN_VALUE"
 export HUGGINGFACE_HUB_TOKEN="$HF_TOKEN_VALUE"
 export DATA_PATH="/workspace/v2_data.jsonl"
 export OUTPUT_DIR="/workspace/khudi-v2-output"
 export MODEL_NAME="Qwen/Qwen3.5-9B"
 
-python train_v2.py $RESUME_FLAG 2>&1
+# FIX: Pass --epochs 1 explicitly
+cd /workspace
+python train_v2.py $RESUME_FLAG --epochs 1 2>&1
 
 TRAIN_EXIT=$?
 echo "Training exit: $TRAIN_EXIT"
 echo "End: $(date)"
 
 # Stop health check
-kill $HEALTH_PID 2>/dev/null
-kill $TIMEOUT_PID 2>/dev/null
+kill $HEALTH_PID 2>/dev/null || true
+kill $TIMEOUT_PID 2>/dev/null || true
 
 # ============================================================
 # Step 8: Final status
@@ -188,7 +207,7 @@ else
 fi
 
 # Upload final log
-python -c "
+HF_TOKEN="$HF_TOKEN_VALUE" python -c "
 from huggingface_hub import HfApi
 api = HfApi()
 api.upload_file(
@@ -199,11 +218,11 @@ api.upload_file(
     commit_message='Training log v2'
 )
 print('Log uploaded')
-" 2>&1
+" 2>&1 || echo "  Log upload failed"
 
 # Stop instance
 echo "Stopping instance..."
-pkill -9 -f python 2>/dev/null
+pkill -9 -f "python train_v2" 2>/dev/null || true
 sleep 2
-sudo shutdown -h now 2>/dev/null || shutdown -h now 2>/dev/null
+sudo shutdown -h now 2>/dev/null || shutdown -h now 2>/dev/null || true
 echo "Script complete"