Skip to content

Commit 6a0d43f

Browse files
khorshuhengkhorshuheng
committed
Add disk size as Dataflow Job Configuration
1 parent 7980e33 commit 6a0d43f

3 files changed

Lines changed: 46 additions & 0 deletions

File tree

‎core/src/main/java/feast/core/job/dataflow/DataflowRunnerConfig.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ public DataflowRunnerConfig(DataflowRunnerConfigOptions runnerConfigOptions) {
4242
this.tempLocation = runnerConfigOptions.getTempLocation();
4343
this.maxNumWorkers = runnerConfigOptions.getMaxNumWorkers();
4444
this.deadLetterTableSpec = runnerConfigOptions.getDeadLetterTableSpec();
45+
this.diskSizeGb = runnerConfigOptions.getDiskSizeGb();
4546
this.labels = runnerConfigOptions.getLabelsMap();
4647
validate();
4748
}
@@ -85,6 +86,9 @@ public DataflowRunnerConfig(DataflowRunnerConfigOptions runnerConfigOptions) {
8586
/* BigQuery table specification, e.g. PROJECT_ID:DATASET_ID.PROJECT_ID */
8687
public String deadLetterTableSpec;
8788

89+
/* Disk size to use on each remote Compute Engine worker instance */
90+
public Integer diskSizeGb;
91+
8892
public Map<String, String> labels;
8993

9094
/** Validates Dataflow runner configuration options */

‎core/src/test/java/feast/core/job/dataflow/DataflowRunnerConfigTest.java‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ public void shouldConvertToPipelineArgs() throws IllegalAccessException {
4242
.setUsePublicIps(false)
4343
.setWorkerMachineType("n1-standard-1")
4444
.setDeadLetterTableSpec("project_id:dataset_id.table_id")
45+
.setDiskSizeGb(100)
4546
.putLabels("key", "value")
4647
.build();
4748

@@ -60,9 +61,46 @@ public void shouldConvertToPipelineArgs() throws IllegalAccessException {
6061
"--usePublicIps=false",
6162
"--workerMachineType=n1-standard-1",
6263
"--deadLetterTableSpec=project_id:dataset_id.table_id",
64+
"--diskSizeGb=100",
6365
"--labels={\"key\":\"value\"}")
6466
.toArray(String[]::new);
6567
assertThat(args.size(), equalTo(expectedArgs.length));
6668
assertThat(args, containsInAnyOrder(expectedArgs));
6769
}
70+
71+
@Test
72+
public void shouldIgnoreOptionalArguments() throws IllegalAccessException {
73+
DataflowRunnerConfigOptions opts =
74+
DataflowRunnerConfigOptions.newBuilder()
75+
.setProject("my-project")
76+
.setRegion("asia-east1")
77+
.setZone("asia-east1-a")
78+
.setTempLocation("gs://bucket/tempLocation")
79+
.setNetwork("default")
80+
.setSubnetwork("regions/asia-east1/subnetworks/mysubnetwork")
81+
.setMaxNumWorkers(1)
82+
.setAutoscalingAlgorithm("THROUGHPUT_BASED")
83+
.setUsePublicIps(false)
84+
.setWorkerMachineType("n1-standard-1")
85+
.build();
86+
87+
DataflowRunnerConfig dataflowRunnerConfig = new DataflowRunnerConfig(opts);
88+
List<String> args = Lists.newArrayList(dataflowRunnerConfig.toArgs());
89+
String[] expectedArgs =
90+
Arrays.asList(
91+
"--project=my-project",
92+
"--region=asia-east1",
93+
"--zone=asia-east1-a",
94+
"--tempLocation=gs://bucket/tempLocation",
95+
"--network=default",
96+
"--subnetwork=regions/asia-east1/subnetworks/mysubnetwork",
97+
"--maxNumWorkers=1",
98+
"--autoscalingAlgorithm=THROUGHPUT_BASED",
99+
"--usePublicIps=false",
100+
"--workerMachineType=n1-standard-1",
101+
"--labels={}")
102+
.toArray(String[]::new);
103+
assertThat(args.size(), equalTo(expectedArgs.length));
104+
assertThat(args, containsInAnyOrder(expectedArgs));
105+
}
68106
}

‎protos/feast/core/Runner.proto‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,4 +77,8 @@ message DataflowRunnerConfigOptions {
7777

7878
/* Labels to apply to the dataflow job */
7979
map<string, string> labels = 13;
80+
81+
/* Disk size to use on each remote Compute Engine worker instance */
82+
int32 diskSizeGb = 14;
83+
8084
}

0 commit comments

Comments
 (0)