[tab] slurm/jobs/mihir/juelich/crowd-pilot/eval/eval_miles_run_batch_unit_test.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --mem=300GB
#SBATCH --partition=booster
#SBATCH --account=envcomp
#SBATCH --output=/p/project1/envcomp/mihir/logs/eval/eval_miles_run/%x_%j.log
#SBATCH --error=/p/project1/envcomp/mihir/logs/eval/eval_miles_run/%x_%j.log
#SBATCH --job-name=eval_miles_run_batch

# =============================================================================
# eval_miles_run_batch.sh
# Batch evaluation script - generates for all checkpoints first, then evaluates
# all in one go to avoid loading/unloading the judge model repeatedly.
# =============================================================================

source .venv/bin/activate

module load Stages/2026 CUDA/13

RUN_NAME=$1

# Extract RUN_ID (last underscore-separated field)
RUN_ID="${RUN_NAME##*_}"

# Extract model size from RUN_NAME and map to Qwen3 model name
# Pattern: qwen{size}_lora_... -> Qwen3-{size}
model_prefix="${RUN_NAME%%_*}"  # e.g., "qwen1.7B"
model_size="${model_prefix#qwen}"  # e.g., "1.7B"

# Map to official Qwen3 model names
case "$model_size" in
    0.6B|0.6b)   BASE_MODEL="Qwen3-0.6B" ;;
    1.7B|1.7b)   BASE_MODEL="Qwen3-1.7B" ;;
    4B|4b)       BASE_MODEL="Qwen3-4B" ;;
    8B|8b)       BASE_MODEL="Qwen3-8B" ;;
    14B|14b)     BASE_MODEL="Qwen3-14B" ;;
    32B|32b)     BASE_MODEL="Qwen3-32B" ;;
    30B-A3B|30b-a3b)     BASE_MODEL="Qwen3-30B-A3B" ;;
    235B-A22B|235b-a22b) BASE_MODEL="Qwen3-235B-A22B" ;;
    *)
        echo "Error: Unknown model size '$model_size' in RUN_NAME '$RUN_NAME'"
        exit 1
        ;;
esac

echo "Extracted from RUN_NAME: RUN_ID=$RUN_ID, BASE_MODEL=$BASE_MODEL"


CKPT_DIR="/p/scratch/envcomp/mihir/huggingface/shared_data/${RUN_ID}/checkpoints"
HF_DIR="/p/scratch/envcomp/mihir/huggingface"

USE_LORA=false
# Generation
generator_host="0.0.0.0"
generator_port=30000
context_length=15000
temperature=0.7
num_samples=10
top_p=0.8
top_k=20
min_p=0.0

test_cases_file="data/eval/handcrafted_test_cases/unit_test_cases.jsonl"
output_dir="data/eval/output/${RUN_ID}"

# Local logging configuration (for offline mode)
local_log_dir="data/eval/local_logs"

mkdir -p $output_dir
mkdir -p $local_log_dir

# Evaluation
judge_host="localhost"
judge_port=30000
judge_model_path="/p/scratch/envcomp/mihir/huggingface/Qwen3-30B-A3B"
judge_context_length=40960
judge_prompt_file="data/prompts/judge_prompt_v2.md"
wandb_name=$RUN_NAME
wandb_project="crowd-pilot-miles"
wandb_eval_type="unit-test-v1"


# Check if this is a lora run or full fine tune
get_latest_iter() {
    ls -1 "$CKPT_DIR" 2>/dev/null | grep -E '^iter_[0-9]+$' | sort -t_ -k2 -n | tail -1
}

latest_iter=$(get_latest_iter)
if [[ -f "$CKPT_DIR/$latest_iter/adapter_model.safetensors" ]]; then
    echo "LoRA adapters found. Evaluating accordingly."
    USE_LORA=true
else
    USE_LORA=false
    echo "No LoRA adapters found. Evaluating as full fine-tune."
fi


# =============================================================================
# PHASE 1: Generate for all checkpoints
# =============================================================================
echo ""
echo "=============================================="
echo "PHASE 1: Generating outputs for all checkpoints"
echo "=============================================="

# Arrays to collect files for batch evaluation
GENERATIONS_FILES=""
EVALUATIONS_FILES=""
EVAL_STEPS=""

# First, generate for base model (step 0)
step=0
generations_file="${output_dir}/${RUN_NAME}_${step}_generations.jsonl"
evaluations_file="${output_dir}/${RUN_NAME}_${step}_evaluations.jsonl"

echo "Generating for base model (step 0)..."
python src/evaluation/sglang_generations.py \
    --test_cases_file $test_cases_file \
    --generations_file $generations_file \
    --system_prompt_file "data/prompts/swe_mini_v1.md" \
    --model_path $HF_DIR/$BASE_MODEL \
    --context_length $context_length \
    --server_host $generator_host \
    --server_port $generator_port \
    --temperature $temperature \
    --top_p $top_p \
    --top_k $top_k \
    --min_p $min_p \
    --num_samples $num_samples \
    --max_connections 32 \
    --concurrency 4 \
    --mem_fraction_static 0.85 \
    --tp_size 4

# Add to batch
GENERATIONS_FILES="${generations_file}"
EVALUATIONS_FILES="${evaluations_file}"
EVAL_STEPS="${step}"

# Now generate for each checkpoint
for dir in $(ls -d "${CKPT_DIR}"/iter_*/ | sort -V); do
    step_name=$(basename "$dir")
    step=$((10#${step_name#iter_}))
    generations_file="${output_dir}/${RUN_NAME}_${step}_generations.jsonl"
    evaluations_file="${output_dir}/${RUN_NAME}_${step}_evaluations.jsonl"
    
    if $USE_LORA ; then
        echo "LoRA adapter dir: $dir, Step: $step"
        GENERATION_ARGS=(
            --model_path $HF_DIR/$BASE_MODEL
            --lora_paths $dir
        )
    else
        cp $HF_DIR/$BASE_MODEL/*.json $dir
        MERGED_DIR=${CKPT_DIR}/merged_${step}
        if [ -d $MERGED_DIR ]; then
            echo "Merged checkpoint directory exists. Using for eval: $MERGED_DIR"
        else
            python ../tools/merge_fsdp_checkpoint.py --checkpoint-dir $dir --output-dir $MERGED_DIR --hf-source-dir $HF_DIR/$BASE_MODEL
        fi
        GENERATION_ARGS=(
            --model_path $MERGED_DIR
        )
        echo 'Not using LORA'
    fi

    # Generate
    echo "Generating for step $step..."
    python src/evaluation/sglang_generations.py \
        --test_cases_file $test_cases_file \
        --generations_file $generations_file \
        --system_prompt_file "data/prompts/swe_mini_v1.md" \
        --context_length $context_length \
        --server_host $generator_host \
        --server_port $generator_port \
        --temperature $temperature \
        --top_p $top_p \
        --top_k $top_k \
        --min_p $min_p \
        --num_samples $num_samples \
        --max_connections 32 \
        --concurrency 4 \
        --mem_fraction_static 0.85 \
        --tp_size 4 \
        ${GENERATION_ARGS[@]}

    # Append to batch lists (comma-separated)
    GENERATIONS_FILES="${GENERATIONS_FILES},${generations_file}"
    EVALUATIONS_FILES="${EVALUATIONS_FILES},${evaluations_file}"
    EVAL_STEPS="${EVAL_STEPS},${step}"
done

# =============================================================================
# PHASE 2: Batch evaluation with single judge model load
# =============================================================================
echo ""
echo "=============================================="
echo "PHASE 2: Batch evaluation (single judge load)"
echo "=============================================="

echo "Running batch evaluation with:"
echo "  Generations files: $GENERATIONS_FILES"
echo "  Evaluations files: $EVALUATIONS_FILES"
echo "  Eval steps: $EVAL_STEPS"

python src/evaluation/sglang_eval.py \
    --wandb_name $wandb_name \
    --wandb_id eval_unit_test_$RUN_ID \
    --generations_files "$GENERATIONS_FILES" \
    --evaluations_files "$EVALUATIONS_FILES" \
    --eval_steps "$EVAL_STEPS" \
    --judge_model_path $judge_model_path \
    --server_host $judge_host \
    --server_port $judge_port \
    --context_length $judge_context_length \
    --judge_prompt_file $judge_prompt_file \
    --temperature $temperature \
    --wandb_eval_type $wandb_eval_type \
    --tp_size 4 \
    --max_connections 32 \
    --concurrency 4 \
    --mem_fraction_static 0.85 \
    --use_local_logger \
    --local_log_dir $local_log_dir

echo ""
echo "=============================================="
echo "Evaluation complete. Local logs saved to:"
echo "  $local_log_dir/eval_$RUN_ID/"
echo ""
echo "To sync to wandb from the login node, run:"
echo "  python src/evaluation/sync_local_logs_to_wandb.py --log_dir $local_log_dir/eval_$RUN_ID"
echo "=============================================="