run_workflow.sh

#!/usr/bin/env bash
set -euo pipefail

usage() {
  cat <<'USAGE'
Usage: ./run_workflow.sh <model> [options]

Model must be one of: gpt-4o, gpt-5.1, gpt-5.1-reasoning, codex

Options:
  --parallel           Run GEPA build, memory build, and baseline in parallel where possible.
  --gepa-vanilla-prompt Start GEPA from the simple prompt "Translate this C project to Rust code."
                        instead of the full bullet_point prompt + rule reminder.
  --only STEPS         Run only the given step(s), for ablations. STEPS is a comma-separated
                        subset of: baseline, gepa, memory. Default: run all three.
                        Examples:
                          --only baseline
                          --only gepa
                          --only gepa --gepa-vanilla-prompt
                          --only baseline,gepa,memory

GEPA Options (for ablations):
  --training-size N    Number of training examples for GEPA (default: 2)
  --population-size N  GEPA population size (default: 2)
  --generations N      Number of GEPA generations (default: 2)
  --batch-size N       GEPA batch size (default: 2)
  --no-llm-feedback    Disable LLM feedback in GEPA (default: enabled)
USAGE
}

MODEL="${1:-}"
shift || true
PARALLEL=0
GEPA_VANILLA_PROMPT=0
ONLY_STEPS=""
GEPA_TRAINING_SIZE=2
GEPA_POPULATION_SIZE=2
GEPA_GENERATIONS=2
GEPA_BATCH_SIZE=2
GEPA_USE_LLM_FEEDBACK="True"
while [[ $# -gt 0 ]]; do
  case "$1" in
    --parallel) PARALLEL=1; shift ;;
    --gepa-vanilla-prompt) GEPA_VANILLA_PROMPT=1; shift ;;
    --only)
      shift || true
      ONLY_STEPS="${1:-}"
      if [[ -z "$ONLY_STEPS" ]]; then
        echo "ERROR: --only requires a value (e.g. baseline,gepa,memory)" >&2
        usage >&2
        exit 1
      fi
      shift || true
      ;;
    --training-size)
      shift || true
      GEPA_TRAINING_SIZE="${1:-}"
      if [[ -z "$GEPA_TRAINING_SIZE" ]]; then
        echo "ERROR: --training-size requires a value" >&2
        usage >&2
        exit 1
      fi
      shift || true
      ;;
    --population-size)
      shift || true
      GEPA_POPULATION_SIZE="${1:-}"
      if [[ -z "$GEPA_POPULATION_SIZE" ]]; then
        echo "ERROR: --population-size requires a value" >&2
        usage >&2
        exit 1
      fi
      shift || true
      ;;
    --generations)
      shift || true
      GEPA_GENERATIONS="${1:-}"
      if [[ -z "$GEPA_GENERATIONS" ]]; then
        echo "ERROR: --generations requires a value" >&2
        usage >&2
        exit 1
      fi
      shift || true
      ;;
    --batch-size)
      shift || true
      GEPA_BATCH_SIZE="${1:-}"
      if [[ -z "$GEPA_BATCH_SIZE" ]]; then
        echo "ERROR: --batch-size requires a value" >&2
        usage >&2
        exit 1
      fi
      shift || true
      ;;
    --no-llm-feedback) GEPA_USE_LLM_FEEDBACK="False"; shift ;;
    *) echo "Unknown option: $1" >&2; usage >&2; exit 1 ;;
  esac
done

# Parse --only into flags (default: run all)
RUN_BASELINE=1
RUN_GEPA=1
RUN_MEMORY=1
if [[ -n "$ONLY_STEPS" ]]; then
  RUN_BASELINE=0
  RUN_GEPA=0
  RUN_MEMORY=0
  for step in $(echo "$ONLY_STEPS" | tr ',' ' '); do
    case "$step" in
      baseline) RUN_BASELINE=1 ;;
      gepa) RUN_GEPA=1 ;;
      memory) RUN_MEMORY=1 ;;
      *) echo "ERROR: unknown step in --only: $step (use baseline,gepa,memory)" >&2; exit 1 ;;
    esac
  done
fi
if [[ -z "$MODEL" ]]; then
  usage
  exit 1
fi

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

case "$MODEL" in
  gpt-4o)
    ENDPOINT="gpt-4o"
    CONFIG_FILE="gpt_4o.json"
    ;;
  gpt-5.1)
    # GPT-5.1 without reasoning (gpt_5_1_no_reasoning.json has "reasoning": {"effort": "none"})
    ENDPOINT="gpt-5.1"
    CONFIG_FILE="gpt_5_1_no_reasoning.json"
    ;;
  gpt-5.1-reasoning)
    # GPT-5.1 with reasoning (gpt_5_1.json has "reasoning": {"effort": "high"})
    ENDPOINT="gpt-5.1"
    CONFIG_FILE="gpt_5_1.json"
    ;;
  codex)
    ENDPOINT="codex"
    CONFIG_FILE="codex.json"
    ;;
  *)
    usage
    exit 1
    ;;
esac

CONFIG_PATH="${REPO_ROOT}/src/endpoints/configs/${CONFIG_FILE}"
if [[ ! -f "$CONFIG_PATH" ]]; then
  echo "Config not found: $CONFIG_PATH" >&2
  exit 1
fi

RUN_ID="$(date +%Y%m%d_%H%M%S)_$$"
RESULTS_ROOT="${REPO_ROOT}/run_results/${MODEL}/${RUN_ID}"
RESULTS_JSON="${RESULTS_ROOT}/results.json"
LOG_DIR="${RESULTS_ROOT}/logs"

GEPA_DIR="${REPO_ROOT}/gepa_outputs/GEPA_${MODEL}/${RUN_ID}"
MEM_DIR="${REPO_ROOT}/memory_outputs/Mem_${MODEL}/${RUN_ID}"

mkdir -p "$LOG_DIR" "$GEPA_DIR" "$MEM_DIR"
printf "[]\n" > "$RESULTS_JSON"

append_result() {
  local strategy="$1"
  local log_file="$2"
  local out_file="$3"
  python - "$MODEL" "$strategy" "$log_file" "$out_file" <<'PY'
import json
import re
import sys
from pathlib import Path

model = sys.argv[1]
strategy = sys.argv[2]
log_path = Path(sys.argv[3])
out_file = Path(sys.argv[4])

text = log_path.read_text(errors="ignore")
match = re.search(r"Overall stats:\s*(\{[\s\S]*?\})", text)
if not match:
    raise SystemExit(f"ERROR: could not find 'Overall stats' in {log_path}")

stats = json.loads(match.group(1))
record = {
    "model": model,
    "strategy": strategy,
    "stats": {
        "All tests pass": stats.get("All tests pass", 0),
        "Compile": stats.get("Compile", 0),
    },
}
out_file.write_text(json.dumps(record, indent=2))
print(json.dumps(record, indent=2))
PY
}

run_gepa_build() {
  echo "==> Building GEPA prompt (${MODEL})"
  echo "[$(date '+%Y-%m-%d %H:%M:%S')] Starting GEPA build..."
  GEPA_LOG="${LOG_DIR}/gepa_build.log"
  (
    cd "${REPO_ROOT}/src"
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] Running gepa.py..."
    python gepa.py \
      --benchmark_dir ../datasets/CBench \
      --rust_dir ../datasets/RBench \
      --output_dir "${GEPA_DIR}" \
      --endpoint "${ENDPOINT}" \
      --config "${CONFIG_PATH}" \
      --training_size "${GEPA_TRAINING_SIZE}" \
      --population_size "${GEPA_POPULATION_SIZE}" \
      --generations "${GEPA_GENERATIONS}" \
      --batch_size "${GEPA_BATCH_SIZE}" \
      --use_llm_feedback "${GEPA_USE_LLM_FEEDBACK}" \
      --repair_iterations 3 \
      --repairer_prompt ./prompts/repair_prompts/bullet_point/bullet_point.prompt \
      --repairer_format bullet_point_with_system_instructions \
      --repairer_strategy all \
      $([[ "$GEPA_VANILLA_PROMPT" -eq 1 ]] && echo "--use_vanilla_prompt True")
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] gepa.py completed with exit code: $?"
  ) 2>&1 | tee "$GEPA_LOG"
  echo "[$(date '+%Y-%m-%d %H:%M:%S')] GEPA build finished. Checking for best_prompt.txt..."
  if [[ -f "${GEPA_DIR}/best_prompt.txt" ]]; then
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] ✓ best_prompt.txt found"
  else
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] ✗ ERROR: best_prompt.txt NOT FOUND"
  fi
}

run_gepa() {
  GEPA_PROMPT="${GEPA_DIR}/best_prompt.txt"
  if [[ ! -f "$GEPA_PROMPT" ]]; then
    echo "GEPA prompt not found: $GEPA_PROMPT" >&2
    exit 1
  fi
  echo "==> Running model with GEPA prompt (${MODEL})"
  GEPA_RUN_LOG="${LOG_DIR}/run_gepa.log"
  (
    cd "${REPO_ROOT}/src"
    python run.py \
      --mode "normal" \
      --benchmark_dir "../datasets/CBench" \
      --output_dir "../results/gepa_${MODEL}/${RUN_ID}/bullet_point_with_system_instructions" \
      --use_gepa_prompt \
      --gepa_prompt_path "../gepa_outputs/GEPA_${MODEL}/${RUN_ID}/best_prompt.txt" \
      --prompt_format bullet_point_with_system_instructions \
      --prompt_strategy all \
      --repairer_prompt ./prompts/repair_prompts/bullet_point/bullet_point.prompt \
      --repairer_format bullet_point_with_system_instructions \
      --repairer_strategy all \
      --iterations 3 \
      --config "${CONFIG_PATH}" \
      --endpoint "${ENDPOINT}" \
      --rust_dir "../datasets/RBench"
  ) 2>&1 | tee "$GEPA_RUN_LOG"
  append_result "gepa" "$GEPA_RUN_LOG" "${RESULTS_ROOT}/result_gepa.json"
}

run_memory_build() {
  echo "==> Building Memory prompt (${MODEL})"
  echo "[$(date '+%Y-%m-%d %H:%M:%S')] Starting Memory build..."
  MEMORY_LOG="${LOG_DIR}/memory_build.log"
  (
    cd "${REPO_ROOT}"
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] Running memory_system.py..."
    python src/memory_system.py \
      --benchmark_dir datasets/CBench \
      --rust_dir datasets/RBench \
      --num_programs 2 \
      --max_repair_attempts 2 \
      --use_llm_hints True \
      --output_dir "${MEM_DIR}" \
      --endpoint "${ENDPOINT}" \
      --config "${CONFIG_PATH}" \
      --hint_model "${ENDPOINT}"
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] memory_system.py completed with exit code: $?"
  ) 2>&1 | tee "$MEMORY_LOG"
  echo "[$(date '+%Y-%m-%d %H:%M:%S')] Memory build finished. Checking for memory_traces.txt..."
  if [[ -f "${MEM_DIR}/memory_traces.txt" ]]; then
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] ✓ memory_traces.txt found"
  else
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] ✗ ERROR: memory_traces.txt NOT FOUND"
  fi
}

run_memory() {
  MEMORY_TRACES="${MEM_DIR}/memory_traces.txt"
  if [[ ! -f "$MEMORY_TRACES" ]]; then
    echo "Memory traces not found: $MEMORY_TRACES" >&2
    exit 1
  fi
  echo "==> Running model with Memory prompt (${MODEL})"
  MEMORY_RUN_LOG="${LOG_DIR}/run_memory.log"
  (
    cd "${REPO_ROOT}/src"
    python run.py \
      --mode "normal" \
      --benchmark_dir "../datasets/CBench" \
      --output_dir "../results/memory_${MODEL}/${RUN_ID}/bullet_point_with_system_instructions" \
      --use_memory \
      --memory_traces_path "../memory_outputs/Mem_${MODEL}/${RUN_ID}/memory_traces.txt" \
      --prompt_format bullet_point_with_system_instructions \
      --prompt_strategy all \
      --repairer_prompt ./prompts/repair_prompts/bullet_point/bullet_point.prompt \
      --repairer_format bullet_point_with_system_instructions \
      --repairer_strategy all \
      --iterations 3 \
      --config "${CONFIG_PATH}" \
      --endpoint "${ENDPOINT}" \
      --rust_dir "../datasets/RBench"
  ) 2>&1 | tee "$MEMORY_RUN_LOG"
  append_result "memory" "$MEMORY_RUN_LOG" "${RESULTS_ROOT}/result_memory.json"
}

run_baseline() {
  echo "==> Running model with baseline prompt (${MODEL})"
  echo "[$(date '+%Y-%m-%d %H:%M:%S')] Starting baseline run..."
  BASELINE_RUN_LOG="${LOG_DIR}/run_baseline.log"
  (
    cd "${REPO_ROOT}/src"
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] Running run.py (baseline)..."
    python run.py \
      --mode "normal" \
      --benchmark_dir "../datasets/CBench" \
      --output_dir "../results/baseline_${MODEL}/${RUN_ID}/bullet_point_with_system_instructions" \
      --prompt ./prompts/transpilation_prompts/bullet_point/bullet_point_interface.prompt \
      --prompt_format bullet_point_with_system_instructions \
      --prompt_strategy all \
      --repairer_prompt ./prompts/repair_prompts/bullet_point/bullet_point.prompt \
      --repairer_format bullet_point_with_system_instructions \
      --repairer_strategy all \
      --iterations 3 \
      --config "${CONFIG_PATH}" \
      --endpoint "${ENDPOINT}" \
      --rust_dir "../datasets/RBench"
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] run.py (baseline) completed with exit code: $?"
  ) 2>&1 | tee "$BASELINE_RUN_LOG"
  echo "[$(date '+%Y-%m-%d %H:%M:%S')] Baseline run finished."
  append_result "baseline" "$BASELINE_RUN_LOG" "${RESULTS_ROOT}/result_baseline.json"
}

if [[ "$PARALLEL" -eq 1 ]]; then
  echo "[$(date '+%Y-%m-%d %H:%M:%S')] Running in PARALLEL mode..."
  GEPA_BUILD_PID=; MEM_BUILD_PID=; BASELINE_PID=
  [[ "$RUN_GEPA" -eq 1 ]] && { run_gepa_build & GEPA_BUILD_PID=$!; }
  [[ "$RUN_MEMORY" -eq 1 ]] && { run_memory_build & MEM_BUILD_PID=$!; }
  [[ "$RUN_BASELINE" -eq 1 ]] && { run_baseline & BASELINE_PID=$!; }
  [[ -n "${GEPA_BUILD_PID:-}" ]] && wait "$GEPA_BUILD_PID"
  [[ -n "${MEM_BUILD_PID:-}" ]] && wait "$MEM_BUILD_PID"
  [[ -n "${BASELINE_PID:-}" ]] && wait "$BASELINE_PID"
  echo "[$(date '+%Y-%m-%d %H:%M:%S')] Build/baseline processes completed."

  GEPA_PID=; MEM_PID=
  [[ "$RUN_GEPA" -eq 1 ]] && { run_gepa & GEPA_PID=$!; }
  [[ "$RUN_MEMORY" -eq 1 ]] && { run_memory & MEM_PID=$!; }
  [[ -n "${GEPA_PID:-}" ]] && wait "$GEPA_PID"
  [[ -n "${MEM_PID:-}" ]] && wait "$MEM_PID"
  echo "[$(date '+%Y-%m-%d %H:%M:%S')] All run processes completed."
else
  echo "[$(date '+%Y-%m-%d %H:%M:%S')] Running in SEQUENTIAL mode..."
  [[ "$RUN_GEPA" -eq 1 ]] && { run_gepa_build; run_gepa; }
  [[ "$RUN_MEMORY" -eq 1 ]] && { run_memory_build; run_memory; }
  [[ "$RUN_BASELINE" -eq 1 ]] && run_baseline
  echo "[$(date '+%Y-%m-%d %H:%M:%S')] All requested steps completed."
fi

echo "==> Final results (saved to ${RESULTS_JSON})"
python - "$RESULTS_ROOT" "$RESULTS_JSON" <<'PY'
import json
import sys
from pathlib import Path

root = Path(sys.argv[1])
out_path = Path(sys.argv[2])
records = []
for path in sorted(root.glob("result_*.json")):
    try:
        records.append(json.loads(path.read_text()))
    except json.JSONDecodeError:
        pass
out_path.write_text(json.dumps(records, indent=2))
print(json.dumps(records, indent=2))
PY