timlawrenz/gnn-ruby-code-study

1#!/usr/bin/env bash2# Runner script for GNN complexity prediction experiments.3# Outputs METRICS:{json} for Ratiocinator fleet parsing.4#5# Environment variables (set by Ratiocinator fleet):6#   CONV_TYPE     - GNN convolution type: GCN, SAGE, GAT, GIN, GraphConv (default: SAGE)7#   HIDDEN_DIM    - Hidden dimension (default: 64)8#   NUM_LAYERS    - Number of GNN layers (default: 3)9#   DROPOUT       - Dropout rate (default: 0.1)10#   LEARNING_RATE - Learning rate (default: 0.001)11#   EPOCHS        - Training epochs (default: 50)12#   BATCH_SIZE    - Batch size (default: 32)13#   DATASET_PATH  - Path to dataset dir (default: dataset/)14 15set -uo pipefail16 17CONV_TYPE="${CONV_TYPE:-SAGE}"18HIDDEN_DIM="${HIDDEN_DIM:-64}"19NUM_LAYERS="${NUM_LAYERS:-3}"20DROPOUT="${DROPOUT:-0.1}"21LEARNING_RATE="${LEARNING_RATE:-0.001}"22EPOCHS="${EPOCHS:-50}"23BATCH_SIZE="${BATCH_SIZE:-32}"24DATASET_PATH="${DATASET_PATH:-dataset/}"25OUTPUT_PATH="models/experiment_model.pt"26 27echo "=== GNN Complexity Arm ==="28echo "CONV_TYPE=$CONV_TYPE HIDDEN_DIM=$HIDDEN_DIM NUM_LAYERS=$NUM_LAYERS"29echo "DROPOUT=$DROPOUT LR=$LEARNING_RATE EPOCHS=$EPOCHS BATCH=$BATCH_SIZE"30 31# Pull LFS files if they are pointers (e.g., after shallow clone)32if command -v git-lfs &>/dev/null || git lfs version &>/dev/null 2>&1; then33    echo "Pulling LFS files..."34    git lfs pull 2>&1 || echo "LFS pull returned non-zero (may be OK if files exist)"35elif [ -f "${DATASET_PATH}/validation.jsonl" ] && head -1 "${DATASET_PATH}/validation.jsonl" | grep -q "^version https://git-lfs"; then36    echo "ERROR: LFS pointer files detected but git-lfs not installed"37    echo "Install with: apt-get install -y git-lfs && git lfs pull"38    exit 139fi40 41# Ensure train/val split exists42if [ ! -f "${DATASET_PATH}/train.jsonl" ]; then43    echo "Creating train/val split..."44    python scripts/split_complexity_data.py \45        --input "${DATASET_PATH}/validation.jsonl" \46        --output-dir "${DATASET_PATH}"47fi48 49# Symlink val.jsonl as validation.jsonl if train.py expects it50if [ -f "${DATASET_PATH}/val.jsonl" ] && [ ! -f "${DATASET_PATH}/validation_split.jsonl" ]; then51    cp "${DATASET_PATH}/val.jsonl" "${DATASET_PATH}/validation_split.jsonl"52fi53 54mkdir -p models55 56# Symlink validation.jsonl → val.jsonl for compatibility57if [ -f "${DATASET_PATH}/val.jsonl" ]; then58    ORIG_VAL="${DATASET_PATH}/validation.jsonl"59    if [ -f "$ORIG_VAL" ] && ! [ -L "$ORIG_VAL" ]; then60        mv "$ORIG_VAL" "${DATASET_PATH}/validation_full.jsonl"61    fi62    ln -sf val.jsonl "${DATASET_PATH}/validation.jsonl"63fi64 65# Run training — stream output directly (no capturing)66TRAIN_LOG="/tmp/train_output_$$.log"67python train.py \68    --dataset_path "$DATASET_PATH" \69    --epochs "$EPOCHS" \70    --output_path "$OUTPUT_PATH" \71    --batch_size "$BATCH_SIZE" \72    --learning_rate "$LEARNING_RATE" \73    --hidden_dim "$HIDDEN_DIM" \74    --num_layers "$NUM_LAYERS" \75    --conv_type "$CONV_TYPE" \76    --dropout "$DROPOUT" \77    --num_workers 0 \78    2>&1 | tee "$TRAIN_LOG"79 80TRAIN_RC=${PIPESTATUS[0]}81if [ "$TRAIN_RC" -ne 0 ]; then82    echo "ERROR: train.py exited with code $TRAIN_RC"83    echo "METRICS:{\"error\": \"training_failed\", \"exit_code\": $TRAIN_RC}"84    exit 185fi86 87# Extract best validation loss from training output88BEST_VAL_LOSS=$(grep "Best validation loss" "$TRAIN_LOG" | grep -oP '[\d.]+' | tail -1)89 90# Run evaluation to get MAE on the validation set91python -c "92import sys, os, json, torch93sys.path.insert(0, os.path.join(os.path.dirname('.'), 'src'))94from data_processing import create_data_loaders95from models import RubyComplexityGNN96import numpy as np97 98device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')99 100checkpoint = torch.load('$OUTPUT_PATH', map_location=device, weights_only=False)101config = checkpoint['model_config']102model = RubyComplexityGNN(103    input_dim=config.get('input_dim', 74),104    hidden_dim=config.get('hidden_dim', 64),105    num_layers=config.get('num_layers', 3),106    conv_type=config.get('conv_type', 'SAGE'),107    dropout=config.get('dropout', 0.1)108).to(device)109model.load_state_dict(checkpoint['model_state_dict'])110model.eval()111 112val_path = os.path.join('${DATASET_PATH}', 'val.jsonl')113if not os.path.exists(val_path):114    val_path = os.path.join('${DATASET_PATH}', 'validation.jsonl')115_, val_loader = create_data_loaders(val_path, val_path, batch_size=64, shuffle=False, num_workers=0)116 117all_preds, all_targets = [], []118with torch.no_grad():119    for batch in val_loader:120        batch = batch.to(device)121        preds = model(batch).squeeze()122        all_preds.extend(preds.cpu().numpy().tolist())123        all_targets.extend(batch.y.cpu().numpy().tolist())124 125preds = np.array(all_preds)126targets = np.array(all_targets)127mae = float(np.mean(np.abs(preds - targets)))128mse = float(np.mean((preds - targets) ** 2))129r2 = float(1 - np.sum((targets - preds)**2) / np.sum((targets - np.mean(targets))**2))130 131print('METRICS:' + json.dumps({132    'val_mae': round(mae, 4),133    'val_mse': round(mse, 4),134    'val_r2': round(r2, 4),135    'best_val_loss': round(float('${BEST_VAL_LOSS:-0}'), 4),136    'conv_type': '$CONV_TYPE',137    'hidden_dim': $HIDDEN_DIM,138    'num_layers': $NUM_LAYERS,139    'dropout': $DROPOUT,140    'learning_rate': $LEARNING_RATE,141    'epochs': $EPOCHS142}))143" 2>&1144 145rm -f "$TRAIN_LOG"146