Allow list of excluded tasks from benchmarks

.github/workflows/all-benchmark.yml

@@ -46,6 +46,24 @@ 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: ''
+      vpc_id:
+        description: 'AWS VPC id'
+        required: true
+        type: string
+      vpc_cidr:
+        description: 'AWS VPC CIDR'
+        required: true
+        type: string
 
 env:
   # nightly run default parameters
@@ -90,6 +108,10 @@ 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 }}
+      vpc_id: ${{ inputs.vpc_id }}
+      vpc_cidr: ${{ inputs.vpc_cidr }}
     secrets:
       SNAPSHOT_AWS_ACCESS_KEY_ID: ${{ secrets.SNAPSHOT_AWS_ACCESS_KEY_ID }}
       SNAPSHOT_AWS_SECRET_ACCESS_KEY: ${{ secrets.SNAPSHOT_AWS_SECRET_ACCESS_KEY }}

.github/workflows/all-ingest.yml

@@ -47,6 +47,14 @@ on:
         required: false
         type: string
         default: eu-north-1
+      vpc_id:
+        description: 'AWS VPC id'
+        required: true
+        type: string
+      vpc_cidr:
+        description: 'AWS VPC CIDR'
+        required: true
+        type: string
 
 env:
   WORKLOADS: ${{ inputs.workloads }}
@@ -96,6 +104,8 @@ jobs:
       snapshot_version: ${{ inputs.snapshot_version }}
       # assumes the only possible event is workflow_dispatch
       ci_tag: 'manual'
+      vpc_id: ${{ inputs.vpc_id }}
+      vpc_cidr: ${{ inputs.vpc_cidr }}
     secrets:
       SNAPSHOT_AWS_ACCESS_KEY_ID: ${{ secrets.SNAPSHOT_AWS_ACCESS_KEY_ID }}
       SNAPSHOT_AWS_SECRET_ACCESS_KEY: ${{ secrets.SNAPSHOT_AWS_SECRET_ACCESS_KEY }}

.github/workflows/benchmark.yml

@@ -71,6 +71,19 @@ on:
         required: false
         type: string
         default: 'latest'
+      include_tasks:
+        description: 'Comma-separated list of tasks to include'
+        required: false
+        type: string
+        default: ''
+      vpc_id:
+        description: 'AWS VPC id'
+        required: true
+        type: string
+      vpc_cidr:
+        description: 'AWS VPC CIDR'
+        required: true
+        type: string
     secrets:
       AWS_ACCESS_KEY_ID:
         required: true
@@ -108,6 +121,8 @@ env:
   TF_VAR_os_version: ${{ inputs.os_version }}
   TF_VAR_distribution_version: ${{ inputs.os_version }}
   TF_VAR_snapshot_version: ${{ inputs.snapshot_version }}
+  TF_VAR_vpc_id: ${{ inputs.vpc_id }}
+  TF_VAR_vpc_cidr: ${{ inputs.vpc_cidr }}
   AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
   AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
   SSH_OPTIONS: -o StrictHostKeyChecking=no -o ServerAliveInterval=60 -o ServerAliveCountMax=10
@@ -324,11 +339,12 @@ jobs:
         run: |
             ssh -i $(terraform output -raw ssh_private_key_file) $SSH_OPTIONS \
               -tt ubuntu@$(terraform output -raw load-generation-ip) -- \
-              "EXTRA_CLIENT_OPTIONS=timeout:240 RUN_GROUP_ID=$RUN_GROUP_ID CI_TAG=$CI_TAG bash -i /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;\""
         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]

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)

infra/main.tf

@@ -57,26 +57,24 @@ resource "aws_key_pair" "ssh_key" {
   public_key = tls_private_key.ssh_key.public_key_openssh
 }
 
-resource "aws_vpc" "vpc" {
-  cidr_block           = "10.0.0.0/16"
-  enable_dns_hostnames = true
-  enable_dns_support   = true
+data "aws_vpc" "vpc" {
+  id = var.vpc_id
 }
 
 resource "aws_subnet" "subnet" {
-  cidr_block        = cidrsubnet(aws_vpc.vpc.cidr_block, 3, 1)
-  vpc_id            = aws_vpc.vpc.id
+  cidr_block        = [var.vpc_cidr]
+  vpc_id            = data.aws_vpc.vpc.id
   availability_zone = var.aws_subnet_zone
 }
 
 resource "aws_internet_gateway" "gtw" {
-  vpc_id = aws_vpc.vpc.id
+  vpc_id = data.aws_vpc.vpc.id
 }
 
 resource "aws_security_group" "allow_osb" {
   name        = "${terraform.workspace}-allow-osb"
   description = "Allow ES/OS/OSB inbound traffic and all outbound traffic"
-  vpc_id      = aws_vpc.vpc.id
+  vpc_id      = data.aws_vpc.vpc.id
 
   lifecycle {
     create_before_destroy = true
@@ -114,7 +112,7 @@ resource "aws_vpc_security_group_egress_rule" "allow_all_traffic_ipv4" {
 }
 
 resource "aws_route_table" "route-table-test-env" {
-  vpc_id = aws_vpc.vpc.id
+  vpc_id = data.aws_vpc.vpc.id
   route {
     cidr_block = "0.0.0.0/0"
     gateway_id = aws_internet_gateway.gtw.id

infra/modules/elasticsearch/es-cluster.yaml

@@ -11,11 +11,13 @@ write_files:
     permissions: '0644'
   - path: /es_cluster.sh
     encoding: gz+b64
-    content: ${es_cluster_script}
+    content: |
+      ${es_cluster_script}
     owner: root:root
     permissions: '0755'
   - path: /ssh_pub_key
-    content: ${authorized_ssh_key}
+    content: |
+      ${authorized_ssh_key}
     owner: ubuntu:ubuntu
     permissions: '0644'
 
@@ -38,4 +40,4 @@ runcmd:
   - [ sysctl, -p, /etc/sysctl.d/99-custom.conf ]
   - [ chown, -R, ubuntu:ubuntu, /mnt ]
   - [ mv, /ssh_pub_key, /home/ubuntu/.ssh/authorized_keys ]
-  - [ sudo, -u, ubuntu, /es_cluster.sh, ${es_password}, ${es_version}, ${es_arch}, ${es_snapshot_access_key}, ${es_snapshot_secret_key} ]
+  - [ sudo, -u, ubuntu, /es_cluster.sh, "${es_password}", "${es_version}", "${es_arch}", "${es_snapshot_access_key}", "${es_snapshot_secret_key}" ]

infra/modules/elasticsearch/es-load-generation.yaml

@@ -11,12 +11,14 @@ write_files:
     permissions: '0644'
   - path: /load_generation.sh
     encoding: gz+b64
-    content: ${load_script}
+    content: |
+      ${load_script}
     owner: root:root
     permissions: '0755'
   - path: /fix_files.sh
     encoding: gz+b64
-    content: ${fix_files_script}
+    content: |
+      ${fix_files_script}
     owner: root:root
     permissions: '0755'
   - path: /es_indexes/big5/es_index_8.json
@@ -45,7 +47,8 @@ write_files:
     permissions: '0644'
   - path: /id_rsa
     encoding: gz+b64
-    content: ${ssh_private_key}
+    content: |
+      ${ssh_private_key}
     owner: ubuntu:ubuntu
     permissions: '0600'
 
@@ -78,8 +81,8 @@ runcmd:
   - [ chown, -R, ubuntu:ubuntu, /mnt ]
   - [ mv, /id_rsa, /home/ubuntu/.ssh/id_rsa ]
   - [ chown, ubuntu:ubuntu, /home/ubuntu/.ssh/id_rsa ]
-  - [ sudo, -u, ubuntu, /load_generation.sh, ${es_cluster}, elastic, ${es_password}, ${distribution_version}, ${es_version}, "ES", ${instance_type}, ${cluster_instance_id}]
-  - [ sudo, -u, ubuntu, /fix_files.sh, https://${es_cluster}:9200, elastic, ${es_password} ]
+  - [ sudo, -u, ubuntu, /load_generation.sh, "${es_cluster}", elastic, "${es_password}", "${distribution_version}", "${es_version}", "ES", "${instance_type}", "${cluster_instance_id}"]
+  - [ sudo, -u, ubuntu, /fix_files.sh, "https://${es_cluster}:9200", elastic, "${es_password}" ]
   - [sed, -i, '/^env.name =/c\env.name = ${benchmark_environment}', /mnt/.benchmark/benchmark.ini]
   - [sed, -i, '/^datastore\./s/^/# /',  /mnt/.benchmark/benchmark.ini]
   - |

infra/modules/elasticsearch/main.tf

@@ -91,7 +91,7 @@ resource "aws_instance" "load-generation" {
           workload = var.workload,
         }
       ))),
-      ssh_private_key = base64gzip(var.ssh_priv_key)
+      ssh_private_key = yamlencode(base64gzip(var.ssh_priv_key))
     }
   )
   user_data_replace_on_change = true

infra/modules/elasticsearch/variables.tf

@@ -130,10 +130,6 @@ variable "workload" {
 variable "snapshot_version" {
   description = "Specific version of the snapshot to restore"
   type        = string
-  validation {
-    condition     = can(regex("^\\d{4}-\\d{2}-\\d{2}_\\d{2}-\\d{2}-\\d{2}$", var.snapshot_version))
-    error_message = "Snapshot version must be: YYYY-mm-dd_HH-MM-ss"
-  }
 }
 
 variable "osb_version" {

infra/modules/opensearch/main.tf

@@ -36,7 +36,7 @@ resource "aws_instance" "target-cluster" {
       os_arch                = local.cluster_arch,
       os_snapshot_access_key = var.snapshot_user_aws_access_key_id,
       os_snapshot_secret_key = var.snapshot_user_aws_secret_access_key,
-      authorized_ssh_key     = var.ssh_pub_key,
+      authorized_ssh_key     = yamlencode(base64gzip(var.ssh_pub_key)),
     }
   )
   user_data_replace_on_change = true
@@ -84,7 +84,7 @@ resource "aws_instance" "load-generation" {
       instance_type           = var.cluster_instance_type
       cluster_instance_id     = aws_instance.target-cluster.id
       fix_files_script        = yamlencode(base64gzip(file("${path.module}/fix_files.sh")))
-      ssh_private_key         = base64gzip(var.ssh_priv_key)
+      ssh_private_key         = yamlencode(base64gzip(var.ssh_priv_key))
     }
   )
   user_data_replace_on_change = true

infra/modules/opensearch/os-cluster.yaml

@@ -11,11 +11,14 @@ write_files:
     permissions: '0644'
   - path: /os_cluster.sh
     encoding: gz+b64
-    content: ${os_cluster_script}
+    content: |
+      ${os_cluster_script}
     owner: root:root
     permissions: '0755'
   - path: /ssh_pub_key
-    content: ${authorized_ssh_key}
+    encoding: gz+b64
+    content: |
+      ${authorized_ssh_key}
     owner: ubuntu:ubuntu
     permissions: '0644'
 
@@ -38,4 +41,4 @@ runcmd:
   - [ sysctl, -p, /etc/sysctl.d/99-custom.conf ]
   - [ chown, -R, ubuntu:ubuntu, /mnt ]
   - [ mv, /ssh_pub_key, /home/ubuntu/.ssh/authorized_keys ]
-  - [ sudo, -u, ubuntu, /os_cluster.sh, ${os_password}, ${os_version}, ${os_arch}, ${os_snapshot_access_key}, ${os_snapshot_secret_key}  ]
+  - [ sudo, -u, ubuntu, /os_cluster.sh, "${os_password}", "${os_version}", "${os_arch}", "${os_snapshot_access_key}", "${os_snapshot_secret_key}" ]

infra/modules/opensearch/os-load-generation.yaml

@@ -33,7 +33,8 @@ write_files:
     permissions: '0644'
   - path: /id_rsa
     encoding: gz+b64
-    content: ${ssh_private_key}
+    content: |
+      ${ssh_private_key}
     owner: ubuntu:ubuntu
     permissions: '0600'
 
@@ -66,7 +67,7 @@ runcmd:
   - [ chown, -R, ubuntu:ubuntu, /mnt ]
   - [ mv, /id_rsa, /home/ubuntu/.ssh/id_rsa ]
   - [ chown, ubuntu:ubuntu, /home/ubuntu/.ssh/id_rsa ]
-  - [ sudo, -u, ubuntu, /load_generation.sh, ${os_cluster}, admin, ${os_password}, ${distribution_version}, ${os_version}, "OS", ${instance_type}, ${cluster_instance_id}]
+  - [ sudo, -u, ubuntu, /load_generation.sh, "${os_cluster}", "admin", "${os_password}", "${distribution_version}", "${os_version}", "OS", "${instance_type}", "${cluster_instance_id}"]
   - [ sudo, -u, ubuntu, /fix_files.sh ]
   - [sed, -i, '/^env.name =/c\env.name = ${benchmark_environment}', /mnt/.benchmark/benchmark.ini]
   - [sed, -i, '/^datastore\./s/^/# /',  /mnt/.benchmark/benchmark.ini]

infra/modules/opensearch/variables.tf

@@ -130,10 +130,6 @@ variable "workload" {
 variable "snapshot_version" {
   description = "Specific version of the snapshot to restore"
   type        = string
-  validation {
-    condition     = can(regex("^\\d{4}-\\d{2}-\\d{2}_\\d{2}-\\d{2}-\\d{2}$", var.snapshot_version))
-    error_message = "Snapshot version must be: YYYY-mm-dd_HH-MM-ss"
-  }
 }
 
 variable "osb_version" {

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

infra/variables.tf

@@ -111,14 +111,25 @@ variable "snapshot_version" {
   description = "Version of the snapshot to deploy (latest, new, or a specific version)"
   type        = string
   default     = "latest"
-  validation {
-    condition     = can(regex("^(latest|new)$", var.snapshot_version)) || can(regex("^\\d{4}-\\d{2}-\\d{2}_\\d{2}-\\d{2}-\\d{2}$", var.snapshot_version))
-    error_message = "Snapshot version must be one of: latest, new, or a specific version (YYYY-mm-dd_HH-MM-ss)"
-  }
+# TODO(Evan): Address this after finishing parallel tests
+# validation {
+#   condition     = can(regex("^(latest|new)$", var.snapshot_version)) || can(regex("^\\d{4}-\\d{2}-\\d{2}_\\d{2}-\\d{2}-\\d{2}$", var.snapshot_version))
+#   error_message = "Snapshot version must be one of: latest, new, or a specific version (YYYY-mm-dd_HH-MM-ss)"
+# }
 }
 
 variable "osb_version" {
   description = "OpenSearch Benchmark version"
   type        = string
   default     = "1.11.0"
 }
+
+variable "vpc_id" {
+  description = "VPC id to use"
+  type        = string
+}
+
+variable "vpc_cidr" {
+  description = "VPC subnet cidr to use"
+  type        = string
+}