quick_start.sh

#!/bin/bash
#
# Quick Start Script for Large-Scale CyberBattle GPU Training
#
# Usage:
#   ./quick_start.sh small      # 1000 nodes, 5k episodes (~5 min)
#   ./quick_start.sh medium     # 2000 nodes, 10k episodes (~10 min)
#   ./quick_start.sh large      # 5000 nodes, 20k episodes (~30 min)
#   ./quick_start.sh test       # 100 nodes, 100 episodes (quick test)

set -e

# Colors
GREEN='\033[0;32m'
BLUE='\033[0;34m'
YELLOW='\033[1;33m'
NC='\033[0m'

print_header() {
    echo -e "\n${BLUE}═══════════════════════════════════════════════════════════════════${NC}"
    echo -e "${BLUE}  $1${NC}"
    echo -e "${BLUE}═══════════════════════════════════════════════════════════════════${NC}\n"
}

print_info() {
    echo -e "${GREEN}➜${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}⚠${NC} $1"
}

# Check arguments
if [ $# -eq 0 ]; then
    echo "GPU-Accelerated DQN Training for Large-Scale CyberBattle"
    echo ""
    echo "Usage: $0 <preset>"
    echo ""
    echo "Presets:"
    echo "  test      100 nodes, 100 episodes    (~30 seconds - quick functionality test)"
    echo "  small     1000 nodes, 5k episodes    (~5 minutes)"
    echo "  medium    2000 nodes, 10k episodes   (~10 minutes)"
    echo "  large     5000 nodes, 20k episodes   (~30 minutes)"
    echo "  xlarge    10000 nodes, 50k episodes  (~2 hours)"
    echo ""
    echo "Examples:"
    echo "  $0 test      # Quick test"
    echo "  $0 small     # Recommended starting point"
    echo "  $0 medium    # Production training"
    echo ""
    exit 1
fi

PRESET=$1

# Set configuration based on preset
case $PRESET in
    test)
        NODES=100
        EPISODES=100
        MAX_STEPS=500
        BATCH_SIZE=32
        DESCRIPTION="Quick Test (100 nodes, 100 episodes)"
        ;;
    small)
        NODES=1000
        EPISODES=5000
        MAX_STEPS=1500
        BATCH_SIZE=128
        DESCRIPTION="Small Network (1000 nodes, 5k episodes)"
        ;;
    medium)
        NODES=2000
        EPISODES=10000
        MAX_STEPS=2000
        BATCH_SIZE=256
        DESCRIPTION="Medium Network (2000 nodes, 10k episodes)"
        ;;
    large)
        NODES=5000
        EPISODES=20000
        MAX_STEPS=3000
        BATCH_SIZE=512
        DESCRIPTION="Large Network (5000 nodes, 20k episodes)"
        ;;
    xlarge)
        NODES=10000
        EPISODES=50000
        MAX_STEPS=5000
        BATCH_SIZE=512
        DESCRIPTION="Extra Large Network (10000 nodes, 50k episodes)"
        ;;
    *)
        echo "Unknown preset: $PRESET"
        echo "Available presets: test, small, medium, large, xlarge"
        exit 1
        ;;
esac

print_header "Large-Scale CyberBattle GPU Training"

print_info "Configuration: $DESCRIPTION"
echo "  Network size: $NODES nodes"
echo "  Training episodes: $EPISODES"
echo "  Max steps per episode: $MAX_STEPS"
echo "  Batch size: $BATCH_SIZE"
echo ""

# Check if we're in the right directory
if [ ! -f "python/train_large_scale.py" ]; then
    print_warning "Training script not found. Are you in the project root?"
    echo "Expected: python/train_large_scale.py"
    exit 1
fi

# Check if build directory exists (for C++ trainer)
if [ ! -d "build" ]; then
    print_warning "Build directory not found. Building project..."
    mkdir -p build
    cd build
    cmake -DCMAKE_BUILD_TYPE=Release .. && make -j$(nproc)
    cd ..
    print_info "Build complete"
fi

# Check GPU availability
print_info "Checking GPU availability..."
if command -v nvidia-smi &> /dev/null; then
    GPU_INFO=$(nvidia-smi --query-gpu=name,memory.total --format=csv,noheader | head -1)
    echo "  GPU detected: $GPU_INFO"
else
    print_warning "No GPU detected. Training will use CPU (much slower)"
    read -p "Continue anyway? (y/n) " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        exit 1
    fi
fi

# Check Python environment
print_info "Checking Python environment..."
if ! python3 -c "import numpy" 2>/dev/null; then
    print_warning "NumPy not found. Please install: pip3 install numpy"
    exit 1
fi

# Check CyberBattleSim (optional - we have fallback)
if python3 -c "import cyberbattle" 2>/dev/null; then
    echo "  CyberBattleSim: installed ✓"
else
    print_warning "CyberBattleSim not found (using fallback mode)"
    echo "  Install with: pip3 install git+https://github.com/microsoft/CyberBattleSim.git"
fi

echo ""

# Confirm before starting
if [ "$PRESET" != "test" ]; then
    print_warning "This will start training. Results will be saved to results/large_scale/"
    read -p "Continue? (y/n) " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        exit 1
    fi
fi

print_header "Starting Training"

# Create timestamp for results
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
OUTPUT_DIR="results/large_scale/run_${NODES}nodes_${TIMESTAMP}"

print_info "Results will be saved to: $OUTPUT_DIR"
echo ""

# Run training
python3 python/train_large_scale.py \
    --nodes $NODES \
    --episodes $EPISODES \
    --max-steps $MAX_STEPS \
    --batch-size $BATCH_SIZE \
    --output-dir "$OUTPUT_DIR" \
    --save-freq 500 \
    --eval-freq 500

EXIT_CODE=$?

echo ""

if [ $EXIT_CODE -eq 0 ]; then
    print_header "Training Complete!"
    
    print_info "Results saved to: $OUTPUT_DIR"
    echo ""
    echo "View results:"
    echo "  📊 Training curves: $OUTPUT_DIR/plots/training_curves.png"
    echo "  📈 Metrics: $OUTPUT_DIR/training_metrics.json"
    echo "  ⚙️  Configuration: $OUTPUT_DIR/config.json"
    echo ""
    
    # Show quick summary
    if [ -f "$OUTPUT_DIR/training_metrics.json" ]; then
        print_info "Training Summary:"
        python3 << EOF
import json
with open('$OUTPUT_DIR/training_metrics.json') as f:
    metrics = json.load(f)
    rewards = metrics['episode_rewards']
    successes = metrics.get('episode_successes', [])
    
    print(f"  Episodes completed: {len(rewards)}")
    print(f"  Mean reward: {sum(rewards)/len(rewards):.2f}")
    print(f"  Final 100 episodes: {sum(rewards[-100:])/100:.2f}")
    if successes:
        print(f"  Success rate: {sum(successes)/len(successes)*100:.1f}%")
        print(f"  Final 100 episodes: {sum(successes[-100:])/len(successes[-100:])*100:.1f}%")
EOF
    fi
    
    echo ""
    print_info "Next steps:"
    echo "  • View training curves: open $OUTPUT_DIR/plots/training_curves.png"
    echo "  • Analyze metrics: cat $OUTPUT_DIR/training_metrics.json | python3 -m json.tool"
    echo "  • Train larger network: ./quick_start.sh large"
    echo ""
    
else
    print_warning "Training failed with exit code $EXIT_CODE"
    echo "Check logs in: $OUTPUT_DIR/logs/"
    exit $EXIT_CODE
fi