bench.sh

#!/usr/bin/env bash
# A/B benchmark for ScriptEngine: compares a branch against master.
# Run from an EC-Earth4 experiment's runtime/ directory.
#
# Usage:
#   bash /path/to/bench.sh <branch> [iterations] [se-repo-url]
#
# Examples:
#   bash ~/scriptengine-bench/bench.sh perf/parallel-loops
#   bash ~/scriptengine-bench/bench.sh perf/parallel-loops 5
#   bash ~/scriptengine-bench/bench.sh my-fix 3 https://github.com/myfork/scriptengine.git

set -euo pipefail

BRANCH=${1:?'Usage: bench.sh <branch> [iterations] [se-repo-url]'}
N=${2:-3}
SE_REPO=${3:-https://github.com/uwefladrich/scriptengine.git}

CACHE_DIR=${SE_BENCH_CACHE:-$HOME/.cache/scriptengine-bench}
SE_CLONE="$CACHE_DIR/scriptengine"
SRC_A="$CACHE_DIR/src-branch"
SRC_B="$CACHE_DIR/src-master"

# --- Clone/update SE repo ---
echo "=== Preparing source trees ==="
mkdir -p "$CACHE_DIR"

if [ ! -d "$SE_CLONE/.git" ]; then
    echo "  Cloning $SE_REPO ..."
    git clone --quiet "$SE_REPO" "$SE_CLONE"
else
    echo "  Updating existing clone..."
    git -C "$SE_CLONE" fetch --all --quiet
fi

# Build source tree for branch
echo "  Extracting branch: $BRANCH"
rm -rf "$SRC_A"
mkdir -p "$SRC_A"
git -C "$SE_CLONE" archive "origin/$BRANCH" -- src/scriptengine 2>/dev/null \
    || git -C "$SE_CLONE" archive "$BRANCH" -- src/scriptengine \
    | tar -x -C "$SRC_A" --strip-components=1

# Build source tree for master
echo "  Extracting branch: master"
rm -rf "$SRC_B"
mkdir -p "$SRC_B"
git -C "$SE_CLONE" archive origin/master -- src/scriptengine \
    | tar -x -C "$SRC_B" --strip-components=1

LABEL_A=$(basename "$BRANCH")
LABEL_B="master"

# --- Detect experiment from current directory ---
RUNTIME_DIR=$(pwd)

EXPERIMENT_YML=$(ls *_experiment_*.yml 2>/dev/null | head -1 || true)
CONFIG_YML=$(ls *_*_*_config_*.yml 2>/dev/null | head -1 || true)
LEG_YML=$(ls *_*_*_*.yml 2>/dev/null | grep -v experiment | grep -v config | grep -v platform | grep -v bench | grep -v user | head -1 || true)
PLATFORM_YML=$(ls ../platforms/*.yml 2>/dev/null | head -1 || true)

if [ -z "$EXPERIMENT_YML" ] || [ -z "$CONFIG_YML" ]; then
    echo ""
    echo "ERROR: Cannot detect experiment YAML files in $(pwd)"
    echo "Expected files matching: *_experiment_*.yml and *_*_*_config_*.yml"
    echo "Run this script from an EC-Earth4 runtime/ directory."
    exit 1
fi

# Detect RUNDIR
if [ -z "${RUNDIR:-}" ]; then
    EXPID=$(echo "$EXPERIMENT_YML" | sed 's/_experiment.*//')
    DATE_DIR=$(echo "$CONFIG_YML" | grep -oP '\d{8}' | head -1 || true)

    if [ -n "$DATE_DIR" ]; then
        # Convention: run_dir parent is $SCRATCH/$EXPID/$DATE_DIR
        # Try to find it relative to current location
        SCRATCH_CANDIDATE=$(cd ../.. && pwd)
        RUNDIR_GUESS="$SCRATCH_CANDIDATE/$DATE_DIR"
        if [ ! -d "$(dirname "$RUNDIR_GUESS")" ]; then
            RUNDIR_GUESS=""
        fi
    fi

    echo ""
    echo "Detected experiment: $EXPID (date: ${DATE_DIR:-unknown})"
    echo "Runtime dir: $RUNTIME_DIR"

    if [ -n "${RUNDIR_GUESS:-}" ]; then
        echo "Guessed run directory: $RUNDIR_GUESS"
        read -p "Is this correct? [Y/n] " confirm
        if [[ "${confirm:-Y}" =~ ^[Nn] ]]; then
            read -p "Enter the run directory path (will be rm -rf'd between setup runs): " RUNDIR
        else
            RUNDIR="$RUNDIR_GUESS"
        fi
    else
        read -p "Enter the run directory path (will be rm -rf'd between setup runs): " RUNDIR
    fi
fi

# Build SE_ARGS from detected files
SE_ARGS="--loglevel info"
[ -n "$EXPERIMENT_YML" ] && SE_ARGS="$SE_ARGS $EXPERIMENT_YML"
[ -n "$CONFIG_YML" ] && SE_ARGS="$SE_ARGS $CONFIG_YML"
[ -n "${LEG_YML:-}" ] && SE_ARGS="$SE_ARGS $LEG_YML"
[ -n "${PLATFORM_YML:-}" ] && SE_ARGS="$SE_ARGS $PLATFORM_YML"

# Append bench.yml if it exists (skips submit/SLURM tasks)
if [ -f "bench.yml" ]; then
    SE_ARGS="$SE_ARGS bench.yml"
fi

echo ""
echo "=== SE A/B Benchmark (N=$N) ==="
echo "  A ($LABEL_A): $BRANCH"
echo "  B ($LABEL_B): master"
echo "  se: $(which se)"
echo "  args: $SE_ARGS"
echo "  rundir: $RUNDIR"
echo ""

# Run se and return wall-clock seconds (decimal)
run_seconds() {
    local t
    t=$( { time se $SE_ARGS > /dev/null 2>&1 || true; } 2>&1 )
    echo "$t" | grep real | sed 's/real\s*//' | awk -F'[ms]' '{print $1*60+$2}'
}

# Warmup: run both versions once to populate Lustre page cache
echo "=== WARMUP ==="
export PYTHONPATH=$SRC_A
rm -rf "$RUNDIR"
se $SE_ARGS > /dev/null 2>&1 || true
export PYTHONPATH=$SRC_B
rm -rf "$RUNDIR"
se $SE_ARGS > /dev/null 2>&1 || true
echo "  done"
echo ""

# Collect timings: alternating A/B to reduce drift bias
declare -a A_SETUP B_SETUP A_CONFIG B_CONFIG

for i in $(seq 1 $N); do
    echo "--- Iteration $i/$N ---"

    # A: setup
    export PYTHONPATH=$SRC_A
    rm -rf "$RUNDIR"
    t=$(run_seconds)
    A_SETUP+=($t)
    printf "  %-10s setup:  %6.1fs\n" "$LABEL_A" "$t"

    # A: config (run_dir exists, only config tasks run)
    t=$(run_seconds)
    A_CONFIG+=($t)
    printf "  %-10s config: %6.1fs\n" "$LABEL_A" "$t"

    # B: setup
    export PYTHONPATH=$SRC_B
    rm -rf "$RUNDIR"
    t=$(run_seconds)
    B_SETUP+=($t)
    printf "  %-10s setup:  %6.1fs\n" "$LABEL_B" "$t"

    # B: config
    t=$(run_seconds)
    B_CONFIG+=($t)
    printf "  %-10s config: %6.1fs\n" "$LABEL_B" "$t"

    echo ""
done

# Summary: median of N runs
median() {
    printf '%s\n' "$@" | sort -n | awk -v n=$# 'NR==int(n/2)+1{print}'
}

MA_SETUP=$(median "${A_SETUP[@]}")
MB_SETUP=$(median "${B_SETUP[@]}")
MA_CONFIG=$(median "${A_CONFIG[@]}")
MB_CONFIG=$(median "${B_CONFIG[@]}")

echo "=== SUMMARY (median of $N runs) ==="
printf "  %-10s setup:  %6.1fs    config: %6.1fs\n" "$LABEL_A" "$MA_SETUP" "$MA_CONFIG"
printf "  %-10s setup:  %6.1fs    config: %6.1fs\n" "$LABEL_B" "$MB_SETUP" "$MB_CONFIG"
DIFF_SETUP=$(awk "BEGIN{printf \"%.1f\", $MB_SETUP - $MA_SETUP}")
DIFF_CONFIG=$(awk "BEGIN{printf \"%.1f\", $MB_CONFIG - $MA_CONFIG}")
PCT_SETUP=$(awk "BEGIN{if($MB_SETUP>0) printf \"%.0f\", ($MB_SETUP-$MA_SETUP)/$MB_SETUP*100; else print 0}")
PCT_CONFIG=$(awk "BEGIN{if($MB_CONFIG>0) printf \"%.0f\", ($MB_CONFIG-$MA_CONFIG)/$MB_CONFIG*100; else print 0}")
echo ""
echo "  setup  diff: ${DIFF_SETUP}s (${PCT_SETUP}% faster)"
echo "  config diff: ${DIFF_CONFIG}s (${PCT_CONFIG}% faster)"
echo ""
echo "Done."