Allow list of excluded tasks from benchmarks

trailofbits · Nov 19, 2024 · b22cc7c · b22cc7c
1 parent 25b62d9
commit b22cc7c
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 11 deletions.
diff --git a/.github/workflows/all-benchmark.yml b/.github/workflows/all-benchmark.yml
@@ -46,6 +46,16 @@ on:
         required: false
         type: string
         default: eu-north-1
+      snapshot_version:
+        description: 'Snapshot version (`new` or `latest`)'
+        required: false
+        type: string
+        default: 'latest'
+      include_tasks:
+        description: 'Comma-separated list of tasks to include'
+        required: false
+        type: string
+        default: ''
 
 env:
   # nightly run default parameters
@@ -90,6 +100,8 @@ jobs:
       aws_region: ${{ inputs.aws_region || 'ca-central-1' }}
       # assumes the possible events are schedule or workflow_dispatch
       ci_tag: ${{ github.event_name == 'schedule' && 'scheduled' || 'manual' }}
+      snapshot_version: ${{ inputs.snapshot_version }}
+      include_tasks: ${{ inputs.include_tasks }}
     secrets:
       SNAPSHOT_AWS_ACCESS_KEY_ID: ${{ secrets.SNAPSHOT_AWS_ACCESS_KEY_ID }}
       SNAPSHOT_AWS_SECRET_ACCESS_KEY: ${{ secrets.SNAPSHOT_AWS_SECRET_ACCESS_KEY }}

diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml
@@ -71,6 +71,11 @@ on:
         required: false
         type: string
         default: 'latest'
+      include_tasks:
+        description: 'Comma-separated list of tasks to include'
+        required: false
+        type: string
+        default: ''
     secrets:
       AWS_ACCESS_KEY_ID:
         required: true
@@ -289,11 +294,12 @@ jobs:
         run: |
             ssh -i $(terraform output -raw ssh_private_key_file) $SSH_OPTIONS \
               -t ubuntu@$(terraform output -raw load-generation-ip) -- \
-              "set -x; bash -ixc \"EXTRA_CLIENT_OPTIONS=timeout:240 RUN_GROUP_ID=$RUN_GROUP_ID CI_TAG=$CI_TAG bash -x /mnt/benchmark.sh $BENCHMARK_TYPE;\""
+              "set -x; bash -ixc \"EXTRA_CLIENT_OPTIONS=timeout:240 RUN_GROUP_ID=$RUN_GROUP_ID CI_TAG=$CI_TAG bash -x /mnt/benchmark.sh $BENCHMARK_TYPE "$INCLUDE_TASKS";\""
         env:
             BENCHMARK_TYPE: ${{ inputs.benchmark_type }}
             RUN_GROUP_ID: ${{ needs.prepare_workspace.outputs.now_run_group_id }}
             CI_TAG: ${{ inputs.ci_tag }}
+            INCLUDE_TASKS: ${{ inputs.include_tasks }}
 
   destroy_workspace:
     needs: [prepare_workspace, run-osb, ingest-data]

diff --git a/infra/get_latest_snapshot_version.py b/infra/get_latest_snapshot_version.py
@@ -68,11 +68,7 @@ def is_version_format(version: str) -> bool:
 elif input_map["snapshot_version"] == "new":
     latest_version = new_snapshot_version()
 else:
-    subdirs = [x for x in subdirs if x == input_map["snapshot_version"]]
     latest_version = input_map["snapshot_version"]
-    if len(subdirs) != 1:
-        print(f"Snapshot version {latest_version} not found", file=sys.stderr)
-        sys.exit(1)
 
 output = {"latest_version": latest_version}
 print(f"Latest version: {latest_version}", file=sys.stderr)

diff --git a/infra/scripts/benchmark.sh b/infra/scripts/benchmark.sh
@@ -2,13 +2,14 @@
 
 source /mnt/utils.sh
 
-if [ $# -ne 1 ]; then
-    echo "Usage: bash benchmark.sh <run-type>"
+if [ $# -lt 1 ]; then
+    echo "Usage: bash benchmark.sh <run-type> [include-tasks]"
     echo "  where <run-type> is 'official' or 'dev'"
     exit 1
 fi
 
 RUN_TYPE=$1
+INCLUDE_TASKS="$2"
 
 if [ "$RUN_TYPE" != "official" ] && [ "$RUN_TYPE" != "dev" ]; then
     echo "Error: <run-type> must be 'official' or 'dev'"
@@ -86,10 +87,12 @@ set -x
 EXECUTION_DIR="/mnt/test_executions"
 mkdir -p "$EXECUTION_DIR"
 
-INCLUDE_TASKS="type:search,prod-queries"
-if [ "$ENGINE_TYPE" == "OS" ]; then
-    # ElasticSearch doesn't support the warmup-indices operation
-    INCLUDE_TASKS+=",warmup-indices"
+if [ -z "${INCLUDE_TASKS}" ]; then
+    INCLUDE_TASKS="type:search,prod-queries"
+    if [ "$ENGINE_TYPE" == "OS" ]; then
+        # ElasticSearch doesn't support the warmup-indices operation
+        INCLUDE_TASKS+=",warmup-indices"
+    fi
 fi
 
 # Queries only