Add disk size as Dataflow Job Configuration (#841)

Co-authored-by: Khor Shu Heng <[email protected]>
feast-dev · Jun 30, 2020 · cf5ed2d · cf5ed2d
1 parent 7980e33
commit cf5ed2d
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 0 deletions.
diff --git a/core/src/main/java/feast/core/job/dataflow/DataflowRunnerConfig.java b/core/src/main/java/feast/core/job/dataflow/DataflowRunnerConfig.java
@@ -42,6 +42,7 @@ public DataflowRunnerConfig(DataflowRunnerConfigOptions runnerConfigOptions) {
     this.tempLocation = runnerConfigOptions.getTempLocation();
     this.maxNumWorkers = runnerConfigOptions.getMaxNumWorkers();
     this.deadLetterTableSpec = runnerConfigOptions.getDeadLetterTableSpec();
+    this.diskSizeGb = runnerConfigOptions.getDiskSizeGb();
     this.labels = runnerConfigOptions.getLabelsMap();
     validate();
   }
@@ -85,6 +86,9 @@ public DataflowRunnerConfig(DataflowRunnerConfigOptions runnerConfigOptions) {
   /* BigQuery table specification, e.g. PROJECT_ID:DATASET_ID.PROJECT_ID */
   public String deadLetterTableSpec;
 
+  /* Disk size to use on each remote Compute Engine worker instance */
+  public Integer diskSizeGb;
+
   public Map<String, String> labels;
 
   /** Validates Dataflow runner configuration options */

diff --git a/core/src/test/java/feast/core/job/dataflow/DataflowRunnerConfigTest.java b/core/src/test/java/feast/core/job/dataflow/DataflowRunnerConfigTest.java
@@ -42,6 +42,7 @@ public void shouldConvertToPipelineArgs() throws IllegalAccessException {
             .setUsePublicIps(false)
             .setWorkerMachineType("n1-standard-1")
             .setDeadLetterTableSpec("project_id:dataset_id.table_id")
+            .setDiskSizeGb(100)
             .putLabels("key", "value")
             .build();
 
@@ -60,9 +61,46 @@ public void shouldConvertToPipelineArgs() throws IllegalAccessException {
                 "--usePublicIps=false",
                 "--workerMachineType=n1-standard-1",
                 "--deadLetterTableSpec=project_id:dataset_id.table_id",
+                "--diskSizeGb=100",
                 "--labels={\"key\":\"value\"}")
             .toArray(String[]::new);
     assertThat(args.size(), equalTo(expectedArgs.length));
     assertThat(args, containsInAnyOrder(expectedArgs));
   }
+
+  @Test
+  public void shouldIgnoreOptionalArguments() throws IllegalAccessException {
+    DataflowRunnerConfigOptions opts =
+        DataflowRunnerConfigOptions.newBuilder()
+            .setProject("my-project")
+            .setRegion("asia-east1")
+            .setZone("asia-east1-a")
+            .setTempLocation("gs://bucket/tempLocation")
+            .setNetwork("default")
+            .setSubnetwork("regions/asia-east1/subnetworks/mysubnetwork")
+            .setMaxNumWorkers(1)
+            .setAutoscalingAlgorithm("THROUGHPUT_BASED")
+            .setUsePublicIps(false)
+            .setWorkerMachineType("n1-standard-1")
+            .build();
+
+    DataflowRunnerConfig dataflowRunnerConfig = new DataflowRunnerConfig(opts);
+    List<String> args = Lists.newArrayList(dataflowRunnerConfig.toArgs());
+    String[] expectedArgs =
+        Arrays.asList(
+                "--project=my-project",
+                "--region=asia-east1",
+                "--zone=asia-east1-a",
+                "--tempLocation=gs://bucket/tempLocation",
+                "--network=default",
+                "--subnetwork=regions/asia-east1/subnetworks/mysubnetwork",
+                "--maxNumWorkers=1",
+                "--autoscalingAlgorithm=THROUGHPUT_BASED",
+                "--usePublicIps=false",
+                "--workerMachineType=n1-standard-1",
+                "--labels={}")
+            .toArray(String[]::new);
+    assertThat(args.size(), equalTo(expectedArgs.length));
+    assertThat(args, containsInAnyOrder(expectedArgs));
+  }
 }
diff --git a/protos/feast/core/Runner.proto b/protos/feast/core/Runner.proto
@@ -77,4 +77,8 @@ message DataflowRunnerConfigOptions {
 
     /* Labels to apply to the dataflow job */
     map<string, string> labels = 13;
+
+    /* Disk size to use on each remote Compute Engine worker instance */
+    int32 diskSizeGb = 14;
+
 }