[tab] slurm/jobs/mihir/juelich/crowd-pilot/2-wave-broken-resolutions/launch_fleet.sh (shellscript)
#!/usr/bin/env bash
#
# Fleet Launcher Script for Wave 2 (BROKEN_RESOLUTIONS)
#
# This script loops through all configuration combinations and submits jobs.
# 
# Configurations:
#   - Model sizes: 1.7b, 4b, 8b
#   - Batch sizes: 8, 16, 32
#   - LoRA: true, false
#
# Usage:
#   ./launch_fleet.sh           # Submit all jobs
#   ./launch_fleet.sh --dry-run # Preview jobs without submitting
#

set -euo pipefail

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
SBATCH_TEMPLATE="${SCRIPT_DIR}/sbatch_template.sh"

# Configuration arrays
MODEL_SIZES=("1.7b" "4b" "8b")
BATCH_SIZES=("8" "16" "32")
USE_LORA_OPTIONS=("true" "false")

# Parse arguments
DRY_RUN=false
if [[ "${1:-}" == "--dry-run" ]]; then
    DRY_RUN=true
    echo "=== DRY RUN MODE - No jobs will be submitted ==="
    echo ""
fi

# Create log directory if it doesn't exist
LOG_DIR="/p/project1/envcomp/mihir/logs/train/crowd-pilot-wave2-broken-resolutions"
if [[ "${DRY_RUN}" == "false" ]]; then
    mkdir -p "${LOG_DIR}"
fi

echo "============================================"
echo "Fleet Launcher for Wave 2 (BROKEN_RESOLUTIONS)"
echo "============================================"
echo ""
echo "Configurations:"
echo "  Model sizes: ${MODEL_SIZES[*]}"
echo "  Batch sizes: ${BATCH_SIZES[*]}"
echo "  LoRA options: ${USE_LORA_OPTIONS[*]}"
echo ""
echo "Total jobs: $((${#MODEL_SIZES[@]} * ${#BATCH_SIZES[@]} * ${#USE_LORA_OPTIONS[@]}))"
echo ""
echo "============================================"
echo ""

# Counter for submitted jobs
SUBMITTED=0
FAILED=0

# Loop through all configurations
for MODEL_SIZE in "${MODEL_SIZES[@]}"; do
    for BATCH_SIZE in "${BATCH_SIZES[@]}"; do
        for USE_LORA in "${USE_LORA_OPTIONS[@]}"; do
            
            # Create a descriptive job name
            if [[ "${USE_LORA}" == "true" ]]; then
                LORA_LABEL="lora"
            else
                LORA_LABEL="no_lora"
            fi
            JOB_NAME="qwen_${MODEL_SIZE}_${LORA_LABEL}_bs_${BATCH_SIZE}"
            
            # Calculate required nodes and CP size
            # Rule of thumb based on previous wave scripts:
            CONTEXT_PARALLEL_SIZE=1 
            if [[ "${MODEL_SIZE}" == "8b" ]]; then
                if [[ "${USE_LORA}" == "true" ]]; then
                    if [[ "${BATCH_SIZE}" == "32" ]]; then NUM_NODES=8;
                    elif [[ "${BATCH_SIZE}" == "16" ]]; then NUM_NODES=4;
                    else NUM_NODES=2; fi
                else
                    if [[ "${BATCH_SIZE}" == "32" ]]; then NUM_NODES=16; CONTEXT_PARALLEL_SIZE=8;
                    elif [[ "${BATCH_SIZE}" == "16" ]]; then NUM_NODES=8; CONTEXT_PARALLEL_SIZE=2;
                    else NUM_NODES=4; CONTEXT_PARALLEL_SIZE=2; fi
                fi
            elif [[ "${MODEL_SIZE}" == "4b" ]]; then
                if [[ "${BATCH_SIZE}" == "32" ]]; then NUM_NODES=8;
                elif [[ "${BATCH_SIZE}" == "16" ]]; then NUM_NODES=4;
                else NUM_NODES=2; fi
            else # 1.7b
                if [[ "${BATCH_SIZE}" == "32" ]]; then NUM_NODES=4;
                else NUM_NODES=2; fi
            fi

            echo "Job: ${JOB_NAME}"
            echo "  MODEL_SIZE: ${MODEL_SIZE}"
            echo "  BATCH_SIZE: ${BATCH_SIZE}"
            echo "  USE_LORA:   ${USE_LORA}"
            echo "  NUM_NODES:  ${NUM_NODES}"
            echo "  CP_SIZE:    ${CONTEXT_PARALLEL_SIZE}"
            
            if [[ "${DRY_RUN}" == "true" ]]; then
                echo "  [DRY RUN] Would submit: sbatch --job-name=${JOB_NAME} --nodes=${NUM_NODES} --export=MODEL_SIZE=${MODEL_SIZE},BATCH_SIZE=${BATCH_SIZE},USE_LORA=${USE_LORA},CONTEXT_PARALLEL_SIZE=${CONTEXT_PARALLEL_SIZE} ${SBATCH_TEMPLATE}"
            else
                # Submit the job with calculated nodes
                if sbatch \
                    --job-name="${JOB_NAME}" \
                    --nodes="${NUM_NODES}" \
                    --export="ALL,MODEL_SIZE=${MODEL_SIZE},BATCH_SIZE=${BATCH_SIZE},USE_LORA=${USE_LORA},CONTEXT_PARALLEL_SIZE=${CONTEXT_PARALLEL_SIZE}" \
                    "${SBATCH_TEMPLATE}"; then
                    ((SUBMITTED+=1))
                    echo "  [SUBMITTED]"
                else
                    ((FAILED+=1))
                    echo "  [FAILED TO SUBMIT]"
                fi
            fi
            echo ""
        done
    done
done

echo "============================================"
echo "Fleet Launch Summary"
echo "============================================"
if [[ "${DRY_RUN}" == "true" ]]; then
    echo "DRY RUN - No jobs were submitted"
    echo "Total jobs that would be submitted: $((${#MODEL_SIZES[@]} * ${#BATCH_SIZES[@]} * ${#USE_LORA_OPTIONS[@]}))"
else
    echo "Jobs submitted: ${SUBMITTED}"
    echo "Jobs failed to submit: ${FAILED}"
fi
echo "============================================"