From 966735e0851185c6c3c48d661000d16542e4d4ff Mon Sep 17 00:00:00 2001
From: Yu Chin Fabian Lim <flim@sg.ibm.com>
Date: Fri, 17 May 2024 06:14:19 +0000
Subject: [PATCH] fix benches and add verify configs

Signed-off-by: Yu Chin Fabian Lim <flim@sg.ibm.com>
---
 scripts/generate_sample_configurations.py  |  2 +-
 scripts/run_benchmarks.sh                  | 13 +++++++++++++
 scripts/verify_generated_configurations.sh | 22 ++++++++++++++++++++++
 tox.ini                                    | 14 +++++++++++++-
 4 files changed, 49 insertions(+), 2 deletions(-)
 create mode 100755 scripts/verify_generated_configurations.sh

diff --git a/scripts/generate_sample_configurations.py b/scripts/generate_sample_configurations.py
index cd354cbf..9f239041 100644
--- a/scripts/generate_sample_configurations.py
+++ b/scripts/generate_sample_configurations.py
@@ -157,7 +157,7 @@ def read_configuration(path: str) -> Dict:
 #   config.
 COMBINATIONS = [
     ("accelerated-peft-autogptq", (KEY_AUTO_GPTQ,)),
-    # ("accelerated-peft-bnb-nf4", (KEY_BNB_NF4,)),
+    ("accelerated-peft-bnb-nf4", (KEY_BNB_NF4,)),
 ]
 
 
diff --git a/scripts/run_benchmarks.sh b/scripts/run_benchmarks.sh
index a281cc53..d1d9f19e 100644
--- a/scripts/run_benchmarks.sh
+++ b/scripts/run_benchmarks.sh
@@ -30,6 +30,7 @@ DATA_CACHE=data/cache.json
 # env inputs
 DRY_RUN=${DRY_RUN:-"false"}
 NO_DATA_PROCESSING=${NO_DATA_PROCESSING:-"false"}
+NO_OVERWRITE=${NO_OVERWRITE:-"false"}
 
 # inputs
 NUM_GPUS_MATRIX=${1-"1 2"}
@@ -42,6 +43,18 @@ echo "RESULT_DIR: $RESULT_DIR"
 echo "SCENARIOS_CONFIG: $SCENARIOS_CONFIG"
 echo "SCENARIOS_FILTER: $SCENARIOS_FILTER"
 
+if [ -n "$RESULT_DIR" ]; then
+    echo "The results directory is not empty. "
+    if [ "$NO_OVERWRITE" = "true" ]; then 
+        echo "Results dir $RESULT_DIR is not empty, but NO_OVERWRITE=true"
+        echo "If intending to overwrite please delete the folder manually"
+        echo "or do not set NO_OVERWRITE"
+        exit 1
+    fi
+    echo "Deleting $RESULT_DIR"
+    rm -rf $RESULT_DIR
+fi
+
 # tag on the directories
 SCENARIOS_CONFIG=$WORKING_DIR/$SCENARIOS_CONFIG
 DEFAULTS_CONFIG=$WORKING_DIR/$DEFAULTS_CONFIG
diff --git a/scripts/verify_generated_configurations.sh b/scripts/verify_generated_configurations.sh
new file mode 100755
index 00000000..83344796
--- /dev/null
+++ b/scripts/verify_generated_configurations.sh
@@ -0,0 +1,22 @@
+#!/bin/bash
+
+OUTPUT_DIR=${1:-sample-configurations}
+
+GIT_DIFF=$(git diff HEAD -- $OUTPUT_DIR)
+echo "git diff of configurations with HEAD:"
+echo "$GIT_DIFF"
+
+function echoWarning() {
+  LIGHT_YELLOW='\033[1;33m'
+  NC='\033[0m' # No Color
+  echo -e "${LIGHT_YELLOW}${1}${NC}"
+}
+
+if [ ! -z "$GIT_DIFF" ]; then
+    echoWarning "At least one of the configs in the plugins should have changed."
+    echoWarning "Please run 'tox -e gen-configs' to ensure that the sample-configurations are correctly generated!"
+    echoWarning "After that commit the generated sample-configurations to remove this error."
+    exit 1
+fi
+
+echo "sample configurations up to date with configs in plugin directories"
diff --git a/tox.ini b/tox.ini
index b9f48607..70dfb7a8 100644
--- a/tox.ini
+++ b/tox.ini
@@ -9,15 +9,27 @@ skip_install = true
 commands = 
     python scripts/generate_sample_configurations.py {posargs:sample-configurations}
 
+[testenv:verify-configs]
+description = verify that sample configurations for all plugins are properly generated
+skip_install = true
+commands = 
+    bash scripts/verify_generated_configurations.sh {posargs:sample-configurations}
+allowlist_externals = bash
+
 # put this here first, consider moving it later
 [testenv:run-benches]
 description = run benchmarks 
 skip_install = true
+setenv = 
+    FHT_BRANCH = {env:FHT_BRANCH:main}
+    DRY_RUN = {env:DRY_RUN:false}
+    NO_DATA_PROCESSING = {env:NO_DATA_PROCESSING:false}
+    NO_OVERWRITE = {env:NO_OVERWRITE:false}
 commands = 
     # need a version of fms-hf-tuning that has integrated the framework
     # NOTE: have to install this first coz havnt merged
     # - this repo has a lot of pins, so we just install it first
-    pip install "fms-hf-tuning[flash-attn] @ git+https://github.com/fabianlim/fms-hf-tuning.git@acceleration-framework"
+    pip install "fms-hf-tuning[flash-attn] @ git+https://github.com/fabianlim/fms-hf-tuning.git@"{env:FHT_BRANCH}
 
     # some models need this for tokenizers
     pip install protobuf