From 3964a900a2e2f5cb88d84c74b7f8a5a084265c2d Mon Sep 17 00:00:00 2001
From: Dongge Liu <donggeliu@google.com>
Date: Thu, 29 Feb 2024 16:12:51 +1100
Subject: [PATCH] Allow delaying each experiment (#131)

Current large-scale experiments tend to exceed various quota limits
(e.g., cloud build).
This PR delays each benchmark's experiment by N seconds to amortize the
quota usage, at the cost of starting some experiments relatively later
than others.
For example, a 10-second delay will launch the last benchmark experiment
4+ hours later than the first.
This is acceptable because it is only used for large-scale experiments,
which tend to have hours of fuzzing timeout.

This is a temporary solution. Ideally, we want a universal controller
for safer and more efficient usage control.
---
 report/docker_run.sh   | 10 ++++++++++
 run_all_experiments.py | 18 +++++++++++++++---
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/report/docker_run.sh b/report/docker_run.sh
index 3bd0b3d138..7c203c5707 100755
--- a/report/docker_run.sh
+++ b/report/docker_run.sh
@@ -33,6 +33,7 @@ FREQUENCY_LABEL=$2
 RUN_TIMEOUT=$3
 SUB_DIR=$4
 MODEL=$5
+DELAY=$6
 # Uses python3 by default and /venv/bin/python3 for Docker containers.
 PYTHON="$( [[ -x "/venv/bin/python3" ]] && echo "/venv/bin/python3" || echo "python3" )"
 export PYTHON
@@ -76,6 +77,14 @@ then
   MODEL='vertex_ai_code-bison-32k'
   echo "LLM was not specified as the fifth argument. Defaulting to ${MODEL:?}."
 fi
+
+# The delay used to amortize quota usage.
+if [[ $DELAY = '' ]]
+then
+  DELAY='0'
+  echo "DELAY was not specified as the sixth argument. Defaulting to ${DELAY:?}."
+fi
+
 DATE=$(date '+%Y-%m-%d')
 LOCAL_RESULTS_DIR='results'
 # Experiment name is used to label the Cloud Builds and as part of the
@@ -100,6 +109,7 @@ $PYTHON run_all_experiments.py \
   --template-directory 'prompts/template_xml' \
   --work-dir ${LOCAL_RESULTS_DIR:?} \
   --num-samples 10 \
+  --delay "${DELAY:?}" \
   --model "$MODEL"
 
 export ret_val=$?
diff --git a/run_all_experiments.py b/run_all_experiments.py
index 01d1c68fa5..ea141bbc09 100755
--- a/run_all_experiments.py
+++ b/run_all_experiments.py
@@ -18,6 +18,7 @@
 import logging
 import os
 import sys
+import time
 import traceback
 from multiprocessing import Pool
 
@@ -156,6 +157,12 @@ def parse_args() -> argparse.Namespace:
                       action='store_true',
                       default=False,
                       help='Add context to function under test.')
+  parser.add_argument(
+      '--delay',
+      type=int,
+      default=0,
+      help=('Delay each experiment by certain seconds (e.g., 10s) to avoid '
+            'exceeding quota limit in large scale experiments.'))
 
   args = parser.parse_args()
   if args.num_samples:
@@ -219,10 +226,15 @@ def main():
       experiment_results.append(result)
       _print_experiment_result(result)
   else:
+    experiment_tasks = []
     with Pool(NUM_EXP) as p:
-      for result in p.starmap(run_experiments, experiment_configs):
-        experiment_results.append(result)
-        _print_experiment_result(result)
+      for config in experiment_configs:
+        experiment_task = p.apply_async(run_experiments,
+                                        config,
+                                        callback=_print_experiment_result)
+        experiment_tasks.append(experiment_task)
+        time.sleep(args.delay)
+      experiment_results = [task.get() for task in experiment_tasks]
 
   _print_experiment_results(experiment_results)