scripts/sync_training_outputs.py

#!/usr/bin/env python3
"""
Sync training outputs and checkpoints to Hugging Face.
Monitors training jobs and uploads results when complete.
"""
import time
import subprocess
import json
from pathlib import Path
from huggingface_hub import HfApi, upload_file, upload_folder

REPO_ID = "anhtld/vla"
TRAINING_JOBS = [14758888]  # h=16 evaluation job
CHECK_INTERVAL = 300  # 5 minutes
SCRATCH = Path("/scratch/knguy52/dovla/experiments")

def check_job_status(job_id):
    """Check if SLURM job completed"""
    try:
        result = subprocess.run(
            ["sacct", "-j", str(job_id), "--format=State", "--noheader"],
            capture_output=True,
            text=True,
            check=True
        )
        state = result.stdout.strip().split()[0]
        return state
    except:
        return "UNKNOWN"

def upload_checkpoint(checkpoint_path, seed):
    """Upload single checkpoint to HF"""
    try:
        print(f"📦 Uploading checkpoint: {checkpoint_path.name}")
        upload_file(
            path_or_fileobj=str(checkpoint_path),
            path_in_repo=f"checkpoints/h16_seed{seed}_{checkpoint_path.name}",
            repo_id=REPO_ID,
            commit_message=f"Add h=16 training checkpoint (seed {seed})"
        )
        print(f"✅ Uploaded: {checkpoint_path.name}")
        return True
    except Exception as e:
        print(f"❌ Upload failed: {e}")
        return False

def upload_training_results(run_dir, seed):
    """Upload training logs and results"""
    try:
        print(f"📊 Uploading training results (seed {seed})...")
        
        # Upload best checkpoint
        best_pt = run_dir / "best.pt"
        if best_pt.exists():
            upload_checkpoint(best_pt, seed)
        
        # Upload training log
        log_files = list(Path("logs").glob(f"train_h16_*_{seed}.out"))
        for log_file in log_files:
            if log_file.exists():
                upload_file(
                    path_or_fileobj=str(log_file),
                    path_in_repo=f"training_logs/{log_file.name}",
                    repo_id=REPO_ID,
                    commit_message=f"Add training log (seed {seed})"
                )
                print(f"✅ Uploaded log: {log_file.name}")
        
        # Upload results JSON if exists
        results_json = run_dir / "results.json"
        if results_json.exists():
            upload_file(
                path_or_fileobj=str(results_json),
                path_in_repo=f"results/h16_seed{seed}_results.json",
                repo_id=REPO_ID,
                commit_message=f"Add training results (seed {seed})"
            )
            print(f"✅ Uploaded: results.json")
        
        return True
    except Exception as e:
        print(f"❌ Upload failed: {e}")
        return False

def main():
    print("="*60)
    print("🔄 Training Output Monitor & Uploader")
    print(f"Monitoring jobs: {TRAINING_JOBS}")
    print(f"Check interval: {CHECK_INTERVAL}s")
    print("="*60)
    print()
    
    uploaded = set()
    
    while True:
        for job_id in TRAINING_JOBS:
            if job_id in uploaded:
                continue
            
            status = check_job_status(job_id)
            print(f"[{time.strftime('%H:%M:%S')}] Job {job_id}: {status}")
            
            if status == "COMPLETED":
                print(f"🎉 Job {job_id} completed! Uploading outputs...")
                
                # Upload for each seed (0, 1, 2)
                for seed in range(3):
                    run_dir = SCRATCH / f"h16_policy_runs/seed_{seed}"
                    if run_dir.exists():
                        upload_training_results(run_dir, seed)
                
                uploaded.add(job_id)
                print(f"✅ All outputs uploaded for job {job_id}")
            
            elif status == "FAILED":
                print(f"❌ Job {job_id} failed, skipping upload")
                uploaded.add(job_id)
        
        if len(uploaded) == len(TRAINING_JOBS):
            print("✅ All jobs processed, exiting")
            break
        
        time.sleep(CHECK_INTERVAL)

if __name__ == "__main__":
    main()