From 553123f86ae5115ef50723d83cf5979cf9451de4 Mon Sep 17 00:00:00 2001
From: Nathan Lambert <nathanl@allenai.org>
Date: Thu, 24 Oct 2024 19:02:21 +0000
Subject: [PATCH] clean scripts

---
 scripts/eval/oe-eval.sh     | 7 ++++++-
 scripts/submit_eval_jobs.py | 6 ++++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/scripts/eval/oe-eval.sh b/scripts/eval/oe-eval.sh
index e250faaef..dc68df993 100755
--- a/scripts/eval/oe-eval.sh
+++ b/scripts/eval/oe-eval.sh
@@ -36,7 +36,7 @@ set -ex
 
 # Function to print usage
 usage() {
-    echo "Usage: $0 --model-name MODEL_NAME --model-location MODEL_LOCATION [--hf-upload] [--revision REVISION] [--max-length <max_length>]"
+    echo "Usage: $0 --model-name MODEL_NAME --model-location MODEL_LOCATION [--num_gpus GPUS] [--hf-upload] [--revision REVISION] [--max-length <max_length>]"
     exit 1
 }
 
@@ -45,13 +45,18 @@ while [[ "$#" -gt 0 ]]; do
     case $1 in
         --model-name) MODEL_NAME="$2"; shift ;;
         --model-location) MODEL_LOCATION="$2"; shift ;;
+        --num_gpus) NUM_GPUS="$2"; shift ;;
         --hf-upload) HF_UPLOAD="true" ;;
         --revision) REVISION="$2"; shift ;;
+        --max-length) MAX_LENGTH="$2"; shift ;;
         *) echo "Unknown parameter passed: $1"; usage ;;
     esac
     shift
 done
 
+# Optional: Default number of GPUs if not specified
+NUM_GPUS="${NUM_GPUS:-1}"
+
 # Check required arguments
 if [[ -z "$MODEL_NAME" || -z "$MODEL_LOCATION" ]]; then
     echo "Error: --model-name and --model-location are required."
diff --git a/scripts/submit_eval_jobs.py b/scripts/submit_eval_jobs.py
index 49e01bda7..d6753651a 100755
--- a/scripts/submit_eval_jobs.py
+++ b/scripts/submit_eval_jobs.py
@@ -101,6 +101,7 @@ def adjust_gpus(task_spec, experiment_group, model_name, gpu_multiplier):
 parser.add_argument("--run_oe_eval_experiments", action="store_true", help="Run the OE eval tool and experiments too.")
 parser.add_argument("--run_safety_evaluations", action="store_true", help="Run the OE safety evaluations too.")
 parser.add_argument("--skip_oi_evals", action="store_true", help="Don't run open instruct evals.")
+parser.add_argument("--oe_eval_max_length", type=int, default=4096, help="Max length for OE eval.")
 args = parser.parse_args()
 
 
@@ -585,6 +586,11 @@ def adjust_gpus(task_spec, experiment_group, model_name, gpu_multiplier):
         oe_eval_cmd += f" --model-location beaker://{model_info[1]}"
     if args.hf_revision:
         oe_eval_cmd += f" --revision {args.hf_revision}"
+    # add string with number of gpus
+    oe_eval_cmd += f" --num_gpus {task_spec['resources']['gpuCount']}"
+    if args.oe_eval_max_length:
+        oe_eval_cmd += f" --max-length {args.oe_eval_max_length}"
+    print(f"Running OE eval with command: {oe_eval_cmd}")
     subprocess.Popen(oe_eval_cmd, shell=True)
 
 # create an experiment that runs the safety eval tasks