Support specifying models in GKE experiment runs (#122)

DonggeLiu · web-flow · commit 5b1bd3c0990c · 2024-02-23T15:33:36.000+11:00
Support specifying models in GKE experiment runs 1. Add `model` argument to `docker_run.sh`, `upload_report.sh`, and `web.py`. 2. Display the model name in experiment reports 3. Some TODOs to rewrite code to adapt to increasing complexity. Preview: <img width="1425" alt="image" src="https://github.com/google/oss-fuzz-gen/assets/20501961/96b12bb2-d945-43cb-b3eb-77f52f453297"> We can beautify the layout later.
diff --git a/report/docker_run.sh b/report/docker_run.sh
@@ -26,10 +26,13 @@
 ##   run_timeout: Timeout in seconds for each fuzzing target.
 ##     Default: 300
 
+# TODO(dongge): Re-write this script in Python as it gets more complex.
+
 BENCHMARK_SET=$1
 FREQUENCY_LABEL=$2
 RUN_TIMEOUT=$3
 SUB_DIR=$4
+MODEL=$5
 # Uses python3 by default and /venv/bin/python3 for Docker containers.
 PYTHON="$( [[ -x "/venv/bin/python3" ]] && echo "/venv/bin/python3" || echo "python3" )"
 export PYTHON
@@ -67,6 +70,12 @@ then
   echo "Sub-directory was not specified as the fourth argument. Defaulting to ${SUB_DIR:?}. Please consider using sub-directory to classify your experiment."
 fi
 
+# The LLM used to generate and fix fuzz targets.
+if [[ $MODEL = '' ]]
+then
+  MODEL='vertex_ai_code-bison-32k'
+  echo "LLM was not specified as the fifth argument. Defaulting to ${MODEL:?}."
+fi
 DATE=$(date '+%Y-%m-%d')
 LOCAL_RESULTS_DIR='results'
 # Experiment name is used to label the Cloud Builds and as part of the
@@ -79,7 +88,7 @@ EXPERIMENT_NAME="${DATE:?}-${FREQUENCY_LABEL:?}-${BENCHMARK_SET:?}"
 GCS_REPORT_DIR="${SUB_DIR:?}/${EXPERIMENT_NAME:?}"
 
 # Generate a report and upload it to GCS
-bash report/upload_report.sh "${LOCAL_RESULTS_DIR:?}" "${GCS_REPORT_DIR:?}" "${BENCHMARK_SET:?}" &
+bash report/upload_report.sh "${LOCAL_RESULTS_DIR:?}" "${GCS_REPORT_DIR:?}" "${BENCHMARK_SET:?}" "${MODEL:?}" &
 pid_report=$!
 
 # Run the experiment
@@ -91,7 +100,7 @@ $PYTHON run_all_experiments.py \
   --template-directory 'prompts/template_xml' \
   --work-dir ${LOCAL_RESULTS_DIR:?} \
   --num-samples 10 \
-  --mode 'vertex_ai_code-bison-32k'
+  --model "$MODEL"
 
 export ret_val=$?
 
diff --git a/report/templates/base.html b/report/templates/base.html
@@ -22,5 +22,6 @@
 }
 </style>
 <body>
+LLM: {{ model }}
  {% block content %}{% endblock %}
 </body>
diff --git a/report/upload_report.sh b/report/upload_report.sh
@@ -21,9 +21,12 @@
 ##   gcs_dir is the name of the directory for the report in gs://oss-fuzz-gcb-experiment-run-logs/Result-reports/.
 ##     Defaults to '$(whoami)-%YY-%MM-%DD'.
 
+# TODO(dongge): Re-write this script in Python as it gets more complex.
+
 RESULTS_DIR=$1
 GCS_DIR=$2
 BENCHMARK_SET=$3
+MODEL=$4
 WEB_PORT=8080
 DATE=$(date '+%Y-%m-%d')
 
@@ -38,15 +41,22 @@ fi
 
 if [[ $GCS_DIR = '' ]]
 then
-  GCS_DIR="$(whoami)-${DATE:?}"
-  echo "GCS directory was not specified as the second argument. Defaulting to ${GCS_DIR:?}."
+  echo "This script needs to take gcloud Bucket directory as the second argument. Consider using $(whoami)-${DATE:?}."
+  exit 1
+fi
+
+# The LLM used to generate and fix fuzz targets.
+if [[ $MODEL = '' ]]
+then
+  echo "This script needs to take LLM as the third argument."
+  exit 1
 fi
 
 mkdir results-report
 
 while true; do
   # Spin up the web server generating the report (and bg the process).
-  $PYTHON -m report.web "${RESULTS_DIR:?}" "${WEB_PORT:?}" "${BENCHMARK_SET:?}" &
+  $PYTHON -m report.web "${RESULTS_DIR:?}" "${WEB_PORT:?}" "${BENCHMARK_SET:?}" "$MODEL" &
   pid_web=$!
 
   cd results-report || exit 1
diff --git a/report/web.py b/report/web.py
@@ -266,18 +266,23 @@ def get_targets(benchmark: str, sample: str) -> list[Target]:
 
 @app.route('/')
 def index():
-  return render_template('index.html', benchmarks=list_benchmarks())
+  return render_template('index.html',
+                         benchmarks=list_benchmarks(),
+                         model=model)
 
 
 @app.route('/json')
 def index_json():
-  return render_template('index.json', benchmarks=list_benchmarks())
+  return render_template('index.json',
+                         benchmarks=list_benchmarks(),
+                         model=model)
 
 
 @app.route('/sort')
 def index_sort():
   return render_template('index.html',
-                         benchmarks=sort_benchmarks(list_benchmarks()))
+                         benchmarks=sort_benchmarks(list_benchmarks()),
+                         model=model)
 
 
 @app.route('/benchmark/<benchmark>')
@@ -334,8 +339,10 @@ def serve(directory: str, port: int, benchmark_set: str):
 
 
 if __name__ == '__main__':
+  # TODO(Dongge): Use argparser as this script gets more complex.
   results_dir = sys.argv[1]
   server_port = int(sys.argv[2])
   benchmark_dir = sys.argv[3] if len(sys.argv) > 3 else ''
+  model = sys.argv[4] if len(sys.argv) > 4 else ''
 
   serve(results_dir, server_port, benchmark_dir)

Original file line number	Diff line number	Diff line change
`@@ -22,5 +22,6 @@`
`22`	`22`	`}`
`23`	`23`	`</style>`
`24`	`24`	`<body>`
	`25`	`+LLM: {{ model }}`
`25`	`26`	`{% block content %}{% endblock %}`
`26`	`27`	`</body>`