[WIP] Load tests with BacklogGenerator

pranavbhandari24 · Jun 6, 2023 · 7e17d7b · 7e17d7b
1 parent 812bc7a
commit 7e17d7b
Show file tree

Hide file tree

Showing 17 changed files with 1,322 additions and 166 deletions.
diff --git a/.github/workflows/load-tests.yml b/.github/workflows/load-tests.yml
@@ -0,0 +1,41 @@
+name: Load Tests
+
+on:
+  schedule:
+  # at 00:00 every Monday.
+  - cron: '0 0 * * 1'
+  workflow_dispatch:
+    inputs:
+      testPattern:
+        description: 'Tests to run'
+        type: string
+        required: false
+        # run all load tests by default
+        default: '*LT'
+
+permissions: read-all
+
+jobs:
+  load_tests:
+    name: Dataflow Templates Load tests
+    timeout-minutes: 2880 # 2 days
+    # Run on any runner that matches all the specified runs-on values.
+    runs-on: [ self-hosted, perf ]
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v3
+      with:
+        fetch-depth: 0
+    - name: Setup Environment
+      id: setup-java-env
+      uses: ./.github/actions/setup-java-env
+    - name: Get Host IP
+      id: variables
+      run: echo "hostIP=$(gcloud compute instances list | grep $(hostname) | awk '{print $4}' )" >> $GITHUB_OUTPUT
+    - name: Run load tests with Maven
+      run: mvn test -Dtest=${TEST_PATTERN} -Dproject=apache-beam-testing -DartifactBucket=gs://apache-beam-testing-pranavbhandari -DhostIp=${HOST_IP} -DexportProject=cloud-teleport-testing -DexportDataset=performance_tests -DexportTable=template_performance_metrics -Djib.skip=true -DskipUnitTests -DfailIfNoTests=false -DtrimStackTrace=false -fae
+      env:
+        TEST_PATTERN: ${{ inputs.testPattern }}
+        HOST_IP: ${{ steps.variables.outputs.hostIP }}
+    - name: Cleanup Java Environment
+      uses: ./.github/actions/cleanup-java-env
diff --git a/it/common/src/main/java/com/google/cloud/teleport/it/common/AbstractPipelineLauncher.java b/it/common/src/main/java/com/google/cloud/teleport/it/common/AbstractPipelineLauncher.java
@@ -281,6 +281,7 @@ public JobState waitUntilActive(String project, String region, String jobId) thr
 
   @Override
   public void cleanupAll() throws IOException {
+    LOG.info("Cleaning up dataflow jobs...");
     for (String jobId : launchedJobs) {
       try {
         JobState state = getJobStatus(TestProperties.project(), TestProperties.region(), jobId);
@@ -291,5 +292,6 @@ public void cleanupAll() throws IOException {
         LOG.warn("Unable to cancel {}. Encountered error.", jobId, e);
       }
     }
+    LOG.info("Dataflow jobs successfully cleaned up.");
   }
 }
diff --git a/it/google-cloud-platform/pom.xml b/it/google-cloud-platform/pom.xml
@@ -107,7 +107,6 @@
         <dependency>
             <groupId>org.apache.beam</groupId>
             <artifactId>beam-sdks-java-io-synthetic</artifactId>
-            <scope>test</scope>
         </dependency>
         <dependency>
             <groupId>org.apache.beam</groupId>
@@ -119,5 +118,34 @@
             <artifactId>beam-sdks-java-test-utils</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.apache.beam</groupId>
+            <artifactId>beam-sdks-java-io-parquet</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.hadoop</groupId>
+            <artifactId>hadoop-client</artifactId>
+            <version>3.3.2</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.beam</groupId>
+            <artifactId>beam-sdks-java-io-jdbc</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.beam</groupId>
+            <artifactId>beam-sdks-java-extensions-google-cloud-platform-core</artifactId>
+        </dependency>
+        <!-- Jdbc dependencies -->
+        <dependency>
+            <groupId>org.postgresql</groupId>
+            <artifactId>postgresql</artifactId>
+            <version>${postgresql.version}</version>
+        </dependency>
+        <!--TODO: remove this -->
+        <dependency>
+            <groupId>com.google.cloud.teleport</groupId>
+            <artifactId>it-jdbc</artifactId>
+            <version>${project.version}</version>
+        </dependency>
     </dependencies>
 </project>
diff --git a/...e-cloud-platform/src/main/java/com/google/cloud/teleport/it/gcp/TemplateLoadTestBase.java b/...e-cloud-platform/src/main/java/com/google/cloud/teleport/it/gcp/TemplateLoadTestBase.java
@@ -15,18 +15,24 @@
  */
 package com.google.cloud.teleport.it.gcp;
 
+import com.google.auto.value.AutoValue;
 import com.google.cloud.teleport.it.common.PipelineLauncher;
+import com.google.cloud.teleport.it.common.PipelineLauncher.LaunchConfig;
 import com.google.cloud.teleport.it.gcp.dataflow.ClassicTemplateClient;
 import com.google.cloud.teleport.it.gcp.dataflow.FlexTemplateClient;
 import com.google.cloud.teleport.metadata.Template;
 import com.google.cloud.teleport.metadata.TemplateLoadTest;
+import java.util.Collections;
+import java.util.List;
 
 /** Base class for Templates Load Tests. */
 public class TemplateLoadTestBase extends LoadTestBase {
 
+  public static final List<String> MACHINE_TYPES =
+      List.of("n1-standard-1", "n1-standard-2", "n1-standard-4", "e2-standard-2", "e2-standard-4");
+
   PipelineLauncher launcher() {
     // If there is a TemplateLoadTest annotation, return appropriate dataflow template client
-    // Else, return null.
     TemplateLoadTest annotation = getClass().getAnnotation(TemplateLoadTest.class);
     if (annotation == null) {
       throw new RuntimeException(
@@ -48,4 +54,63 @@ PipelineLauncher launcher() {
       return ClassicTemplateClient.builder().setCredentials(CREDENTIALS).build();
     }
   }
+
+  protected LaunchConfig.Builder enableRunnerV2(LaunchConfig.Builder config) {
+    return config.addEnvironment(
+        "additionalExperiments", Collections.singletonList("use_runner_v2"));
+  }
+
+  protected LaunchConfig.Builder disableRunnerV2(LaunchConfig.Builder config) {
+    return config.addEnvironment(
+        "additionalExperiments", Collections.singletonList("disable_runner_v2"));
+  }
+
+  protected LaunchConfig.Builder enableStreamingEngine(LaunchConfig.Builder config) {
+    return config.addEnvironment("enableStreamingEngine", true);
+  }
+
+  /** Options for Backlog load tests. */
+  @AutoValue
+  public abstract static class BacklogConfiguration {
+
+    public abstract Integer getRowSize();
+
+    public abstract Long getNumRows();
+
+    public abstract Integer getGeneratorTimeout();
+
+    public abstract Integer getPipelineTimeout();
+
+    public static BacklogConfiguration of(long numRows, int generatorTimeout, int pipelineTimeout) {
+      return new AutoValue_TemplateLoadTestBase_BacklogConfiguration.Builder()
+          .setRowSize(1024)
+          .setNumRows(numRows)
+          .setGeneratorTimeout(generatorTimeout)
+          .setPipelineTimeout(pipelineTimeout)
+          .build();
+    }
+
+    public static BacklogConfiguration of(
+        int rowSize, long numRows, int generatorTimeout, int pipelineTimeout) {
+      return new AutoValue_TemplateLoadTestBase_BacklogConfiguration.Builder()
+          .setRowSize(rowSize)
+          .setNumRows(numRows)
+          .setGeneratorTimeout(generatorTimeout)
+          .setPipelineTimeout(pipelineTimeout)
+          .build();
+    }
+
+    @AutoValue.Builder
+    abstract static class Builder {
+      abstract Builder setRowSize(int rowSize);
+
+      abstract Builder setNumRows(long numRows);
+
+      abstract Builder setGeneratorTimeout(int timeOutMinutes);
+
+      abstract Builder setPipelineTimeout(int timeOutMinutes);
+
+      abstract BacklogConfiguration build();
+    }
+  }
 }