run_step5_only.py

#!/usr/bin/env python3
"""
Step 5 전용 실행 스크립트
기존 AZR 학습 데이터로 VeRL PPO 학습만 실행
"""

import os
import sys
import argparse
from pathlib import Path

# 경로 설정
sys.path.append('/home/ubuntu/RLVR/TestTime-RLVR-v2')
sys.path.append('/home/ubuntu/RLVR/TestTime-RLVR-v2/test')

from test.utils.iterative_trainer import IterativeTrainer

def main():
    parser = argparse.ArgumentParser(description='Run VeRL training (Step 5) only with existing data')
    parser.add_argument('--data_path', type=str, required=True,
                        help='Path to existing azr_training_data directory')
    parser.add_argument('--round', type=int, default=1,
                        help='Round number for logging (default: 1)')
    parser.add_argument('--experiment_name', type=str, default=None,
                        help='Custom experiment name')
    parser.add_argument('--config', type=str, 
                        default='/home/ubuntu/RLVR/TestTime-RLVR-v2/test/configs/ttrlvr_azr_ppo_4gpu.yaml',
                        help='VeRL config file path')
    
    args = parser.parse_args()
    
    # 데이터 경로 검증
    data_path = Path(args.data_path)
    if not data_path.exists():
        print(f"❌ Error: Data path does not exist: {data_path}")
        return 1
    
    # 필수 파일들 확인
    required_files = ['induction.parquet', 'deduction.parquet', 'abduction.parquet']
    missing_files = []
    for file_name in required_files:
        if not (data_path / file_name).exists():
            missing_files.append(file_name)
    
    if missing_files:
        print(f"❌ Error: Missing required files: {missing_files}")
        return 1
    
    print(f"✅ Found all required training data files in: {data_path}")
    
    # 파일 크기 정보 출력
    for file_name in required_files:
        file_path = data_path / file_name
        file_size = file_path.stat().st_size
        print(f"  📄 {file_name}: {file_size:,} bytes")
    
    # IterativeTrainer 초기화
    print(f"🚀 Initializing trainer with config: {args.config}")
    trainer = IterativeTrainer(config_path=args.config)
    
    # Step 5 전용 VeRL 학습 실행
    print(f"🎓 Starting VeRL training (Step 5 only)")
    print(f"📂 Data path: {data_path}")
    print(f"🔄 Round: {args.round}")
    
    try:
        result = trainer.run_verl_training_only(
            training_data_path=str(data_path),
            round_num=args.round,
            experiment_name=args.experiment_name
        )
        
        if result.get('success', False):
            print(f"✅ VeRL training completed successfully!")
            print(f"⏱️  Duration: {result.get('duration', 'N/A')} seconds")
            if 'model_path' in result:
                print(f"🤖 Updated model: {result['model_path']}")
        else:
            print(f"❌ VeRL training failed: {result.get('error', 'Unknown error')}")
            return 1
            
    except Exception as e:
        print(f"💥 Training failed with exception: {e}")
        import traceback
        traceback.print_exc()
        return 1
    
    print(f"🎉 Step 5 training completed!")
    return 0

if __name__ == "__main__":
    exit_code = main()
    sys.exit(exit_code)