Skip to content

Commit 875cfd1

Browse files
Amar3ttoAmar3tto
committed
Refactor DefaultPipelineLauncher
1 parent 22b2f38 commit 875cfd1

25 files changed

Lines changed: 65 additions & 85 deletions

File tree

‎it/common/src/main/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncher.java‎

Lines changed: 28 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818
package org.apache.beam.it.common.dataflow;
1919

2020
import static org.apache.beam.it.common.logging.LogStrings.formatForLogging;
21-
import static org.apache.beam.sdk.testing.TestPipeline.PROPERTY_BEAM_TEST_PIPELINE_OPTIONS;
2221
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
2322

2423
import com.fasterxml.jackson.databind.ObjectMapper;
@@ -41,7 +40,6 @@
4140
import java.util.stream.StreamSupport;
4241
import org.apache.beam.it.common.PipelineLauncher;
4342
import org.apache.beam.it.common.utils.PipelineUtils;
44-
import org.apache.beam.it.common.IOLoadTestBase;
4543
import org.apache.beam.runners.dataflow.DataflowPipelineJob;
4644
import org.apache.beam.sdk.PipelineResult;
4745
import org.apache.beam.sdk.metrics.DistributionResult;
@@ -73,6 +71,8 @@ public class DefaultPipelineLauncher extends AbstractPipelineLauncher {
7371
private static final String READ_PIPELINE_NAME_OVERWRITE = "readPipelineNameOverride";
7472
private static final String WRITE_PIPELINE_NAME_OVERWRITE = "writePipelineNameOverride";
7573
private static final Pattern JOB_ID_PATTERN = Pattern.compile("Submitted job: (\\S+)");
74+
/** Namespace for Beam provided pipeline metrics (set up by Metrics transform). */
75+
public static final String BEAM_METRICS_NAMESPACE = "BEAM_METRICS";
7676

7777
// For unsupported runners (other than dataflow), implement launcher methods by operating with
7878
// PipelineResult.
@@ -101,6 +101,18 @@ public class DefaultPipelineLauncher extends AbstractPipelineLauncher {
101101
.put(PipelineResult.State.UNRECOGNIZED, JobState.UNKNOWN)
102102
.build();
103103

104+
// To make PipelineLauncher.getMetric work in a unified way for both runner provided metrics and
105+
// pipeline defined
106+
// metrics, here we wrap Beam provided metrics as a pre-defined metrics name
107+
// [name_space:metric_type:metric_name
108+
// which will be recognized by getMetric method
109+
public enum PipelineMetricsType {
110+
COUNTER,
111+
STARTTIME,
112+
ENDTIME,
113+
RUNTIME,
114+
}
115+
104116
private DefaultPipelineLauncher(Builder builder) {
105117
super(
106118
new Dataflow(
@@ -169,7 +181,7 @@ private static <T> void checkIfMetricResultIsUnique(
169181
resultCount <= 1,
170182
"More than one metric result matches name: %s in namespace %s. Metric results count: %s",
171183
name,
172-
IOLoadTestBase.BEAM_METRICS_NAMESPACE,
184+
BEAM_METRICS_NAMESPACE,
173185
resultCount);
174186
}
175187

@@ -181,14 +193,14 @@ private static Iterable<MetricResult<DistributionResult>> getDistributions(
181193
.queryMetrics(
182194
MetricsFilter.builder()
183195
.addNameFilter(
184-
MetricNameFilter.named(IOLoadTestBase.BEAM_METRICS_NAMESPACE, metricName))
196+
MetricNameFilter.named(BEAM_METRICS_NAMESPACE, metricName))
185197
.build());
186198
return metrics.getDistributions();
187199
}
188200

189201
/** Pull Beam pipeline defined metrics given the jobId. */
190202
public Long getBeamMetric(
191-
String jobId, IOLoadTestBase.PipelineMetricsType metricType, String metricName) {
203+
String jobId, PipelineMetricsType metricType, String metricName) {
192204
PipelineResult pipelineResult =
193205
MANAGED_JOBS.getOrDefault(jobId, UNMANAGED_JOBS.getOrDefault(jobId, null));
194206
if (pipelineResult != null) {
@@ -198,7 +210,7 @@ public Long getBeamMetric(
198210
.queryMetrics(
199211
MetricsFilter.builder()
200212
.addNameFilter(
201-
MetricNameFilter.named(IOLoadTestBase.BEAM_METRICS_NAMESPACE, metricName))
213+
MetricNameFilter.named(BEAM_METRICS_NAMESPACE, metricName))
202214
.build());
203215

204216
switch (metricType) {
@@ -212,7 +224,7 @@ public Long getBeamMetric(
212224
LOG.error(
213225
"Failed to get metric {}, from namespace {}",
214226
metricName,
215-
IOLoadTestBase.BEAM_METRICS_NAMESPACE);
227+
BEAM_METRICS_NAMESPACE);
216228
}
217229
return UNKNOWN_METRIC_VALUE;
218230
case STARTTIME:
@@ -230,9 +242,9 @@ public Long getBeamMetric(
230242
.map(element -> Objects.requireNonNull(element.getAttempted()).getMax())
231243
.max(Long::compareTo)
232244
.orElse(UNKNOWN_METRIC_VALUE);
233-
if (metricType == IOLoadTestBase.PipelineMetricsType.STARTTIME) {
245+
if (metricType == PipelineMetricsType.STARTTIME) {
234246
return lowestMin;
235-
} else if (metricType == IOLoadTestBase.PipelineMetricsType.ENDTIME) {
247+
} else if (metricType == PipelineMetricsType.ENDTIME) {
236248
return greatestMax;
237249
} else {
238250
if (lowestMin != UNKNOWN_METRIC_VALUE && greatestMax != UNKNOWN_METRIC_VALUE) {
@@ -254,15 +266,15 @@ public Long getBeamMetric(
254266
@Override
255267
public Double getMetric(String project, String region, String jobId, String metricName)
256268
throws IOException {
257-
if (metricName.startsWith(IOLoadTestBase.BEAM_METRICS_NAMESPACE)) {
269+
if (metricName.startsWith(BEAM_METRICS_NAMESPACE)) {
258270
String[] nameSpacedMetrics = metricName.split(":", 3);
259271
Preconditions.checkState(
260272
nameSpacedMetrics.length == 3,
261273
String.format(
262274
"Invalid Beam metrics name: %s, expected: '%s:metric_type:metric_name'",
263-
metricName, IOLoadTestBase.BEAM_METRICS_NAMESPACE));
264-
IOLoadTestBase.PipelineMetricsType metricType =
265-
IOLoadTestBase.PipelineMetricsType.valueOf(nameSpacedMetrics[1]);
275+
metricName, BEAM_METRICS_NAMESPACE));
276+
PipelineMetricsType metricType =
277+
PipelineMetricsType.valueOf(nameSpacedMetrics[1]);
266278

267279
// Pipeline defined metrics are long values. Have to cast to double that is what the base
268280
// class defined.
@@ -437,16 +449,15 @@ private List<String> extractOptions(String project, String region, LaunchConfig
437449
// add pipeline options from beamTestPipelineOptions system property to preserve the
438450
// pipeline options already set in TestPipeline.
439451
@Nullable
440-
String beamTestPipelineOptions = System.getProperty(PROPERTY_BEAM_TEST_PIPELINE_OPTIONS);
452+
String beamTestPipelineOptions = System.getProperty(org.apache.beam.sdk.testing.TestPipeline.PROPERTY_BEAM_TEST_PIPELINE_OPTIONS);
441453
if (!Strings.isNullOrEmpty(beamTestPipelineOptions)) {
442454
try {
443455
additionalOptions.addAll(MAPPER.readValue(beamTestPipelineOptions, List.class));
444456
} catch (IOException e) {
445457
throw new RuntimeException(
446458
"Unable to instantiate test options from system property "
447-
+ PROPERTY_BEAM_TEST_PIPELINE_OPTIONS
448-
+ ":"
449-
+ System.getProperty(PROPERTY_BEAM_TEST_PIPELINE_OPTIONS),
459+
+ org.apache.beam.sdk.testing.TestPipeline.PROPERTY_BEAM_TEST_PIPELINE_OPTIONS
460+
+ ":" + beamTestPipelineOptions,
450461
e);
451462
}
452463
}

it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerTest.java renamed to it/common/src/test/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerTest.java

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
* See the License for the specific language governing permissions and
1616
* limitations under the License.
1717
*/
18-
package org.apache.beam.it.gcp.bigquery;
18+
package org.apache.beam.it.common.bigquery;
1919

2020
import static com.google.common.truth.Truth.assertThat;
2121
import static org.junit.Assert.assertThrows;
@@ -41,8 +41,6 @@
4141
import com.google.cloud.bigquery.TableId;
4242
import com.google.cloud.bigquery.TableInfo;
4343
import com.google.cloud.bigquery.TimePartitioning;
44-
import org.apache.beam.it.common.bigquery.BigQueryResourceManager;
45-
import org.apache.beam.it.common.bigquery.BigQueryResourceManagerException;
4644
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
4745
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
4846
import org.junit.Before;

it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerUtilsTest.java renamed to it/common/src/test/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerUtilsTest.java

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,14 +15,13 @@
1515
* See the License for the specific language governing permissions and
1616
* limitations under the License.
1717
*/
18-
package org.apache.beam.it.gcp.bigquery;
18+
package org.apache.beam.it.common.bigquery;
1919

2020
import static org.apache.beam.it.common.bigquery.BigQueryResourceManagerUtils.checkValidTableId;
2121
import static org.junit.Assert.assertThrows;
2222

2323
import java.util.Arrays;
2424

25-
import org.apache.beam.it.common.bigquery.BigQueryResourceManagerUtils;
2625
import org.junit.Test;
2726

2827
/** Unit tests for {@link BigQueryResourceManagerUtils}. */

it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/AbstractPipelineLauncherTest.java renamed to it/common/src/test/java/org/apache/beam/it/common/dataflow/AbstractPipelineLauncherTest.java

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
* See the License for the specific language governing permissions and
1616
* limitations under the License.
1717
*/
18-
package org.apache.beam.it.gcp.dataflow;
18+
package org.apache.beam.it.common.dataflow;
1919

2020
import static com.google.common.truth.Truth.assertThat;
2121
import static org.junit.Assert.assertThrows;
@@ -34,7 +34,6 @@
3434
import java.io.IOException;
3535
import java.net.SocketTimeoutException;
3636
import org.apache.beam.it.common.PipelineLauncher.JobState;
37-
import org.apache.beam.it.common.dataflow.AbstractPipelineLauncher;
3837
import org.junit.Rule;
3938
import org.junit.Test;
4039
import org.junit.runner.RunWith;

it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/DefaultPipelineLauncherTest.java renamed to it/common/src/test/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncherTest.java

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,17 +15,16 @@
1515
* See the License for the specific language governing permissions and
1616
* limitations under the License.
1717
*/
18-
package org.apache.beam.it.gcp.dataflow;
18+
package org.apache.beam.it.common.dataflow;
1919

20+
import static org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.BEAM_METRICS_NAMESPACE;
2021
import static org.junit.Assert.assertEquals;
2122
import static org.junit.Assert.assertTrue;
2223

2324
import java.io.IOException;
2425
import java.time.Instant;
26+
import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType;
2527
import org.apache.beam.it.common.PipelineLauncher;
26-
import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher;
27-
import org.apache.beam.it.gcp.IOLoadTestBase;
28-
import org.apache.beam.it.gcp.IOLoadTestBase.PipelineMetricsType;
2928
import org.apache.beam.sdk.io.GenerateSequence;
3029
import org.apache.beam.sdk.testing.TestPipeline;
3130
import org.apache.beam.sdk.testutils.metrics.TimeMonitor;
@@ -49,7 +48,7 @@ public void testPipelineMetrics() throws IOException {
4948

5049
pipeline
5150
.apply(GenerateSequence.from(0).to(numElements))
52-
.apply(ParDo.of(new TimeMonitor<>(IOLoadTestBase.BEAM_METRICS_NAMESPACE, timeMetrics)))
51+
.apply(ParDo.of(new TimeMonitor<>(BEAM_METRICS_NAMESPACE, timeMetrics)))
5352
.apply(ParDo.of(new IOLoadTestBase.CountingFn<>(counterMetrics)));
5453

5554
PipelineLauncher.LaunchConfig options =

it/common/src/test/java/org/apache/beam/it/common/IOLoadTestBase.java renamed to it/common/src/test/java/org/apache/beam/it/common/dataflow/IOLoadTestBase.java

Lines changed: 6 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
* See the License for the specific language governing permissions and
1616
* limitations under the License.
1717
*/
18-
package org.apache.beam.it.common;
18+
package org.apache.beam.it.common.dataflow;
1919

2020
import com.google.cloud.Timestamp;
2121
import java.io.IOException;
@@ -24,7 +24,9 @@
2424
import java.util.Map;
2525
import java.util.UUID;
2626

27-
import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher;
27+
import org.apache.beam.it.common.PipelineLauncher;
28+
import org.apache.beam.it.common.TestProperties;
29+
import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType;
2830
import org.apache.beam.sdk.metrics.Counter;
2931
import org.apache.beam.sdk.metrics.Metrics;
3032
import org.apache.beam.sdk.testutils.NamedTestResult;
@@ -38,6 +40,8 @@
3840
import org.slf4j.Logger;
3941
import org.slf4j.LoggerFactory;
4042

43+
import static org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.BEAM_METRICS_NAMESPACE;
44+
4145
/** Base class for IO Load tests. */
4246
@RunWith(JUnit4.class)
4347
@SuppressWarnings({
@@ -89,21 +93,6 @@ public void processElement(ProcessContext ctx) {
8993
}
9094
}
9195

92-
// To make PipelineLauncher.getMetric work in a unified way for both runner provided metrics and
93-
// pipeline defined
94-
// metrics, here we wrap Beam provided metrics as a pre-defined metrics name
95-
// [name_space:metric_type:metric_name
96-
// which will be recognized by getMetric method
97-
public enum PipelineMetricsType {
98-
COUNTER,
99-
STARTTIME,
100-
ENDTIME,
101-
RUNTIME,
102-
}
103-
104-
/** Namespace for Beam provided pipeline metrics (set up by Metrics transform). */
105-
public static final String BEAM_METRICS_NAMESPACE = "BEAM_METRICS";
106-
10796
/** Given a metrics name, return Beam metrics name. */
10897
public static String getBeamMetricsName(PipelineMetricsType metricstype, String metricsName) {
10998
return BEAM_METRICS_NAMESPACE + ":" + metricstype + ":" + metricsName;

it/common/src/test/java/org/apache/beam/it/common/IOStressTestBase.java renamed to it/common/src/test/java/org/apache/beam/it/common/dataflow/IOStressTestBase.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
* See the License for the specific language governing permissions and
1616
* limitations under the License.
1717
*/
18-
package org.apache.beam.it.common;
18+
package org.apache.beam.it.common.dataflow;
1919

2020
import java.io.Serializable;
2121
import java.time.Duration;

it/common/src/test/java/org/apache/beam/it/common/LoadTestBase.java renamed to it/common/src/test/java/org/apache/beam/it/common/dataflow/LoadTestBase.java

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
* See the License for the specific language governing permissions and
1616
* limitations under the License.
1717
*/
18-
package org.apache.beam.it.common;
18+
package org.apache.beam.it.common.dataflow;
1919

2020
import static org.apache.beam.it.common.logging.LogStrings.formatForLogging;
2121
import static org.apache.beam.it.common.dataflow.AbstractPipelineLauncher.RUNNER_V2;
@@ -41,9 +41,11 @@
4141
import java.util.Map.Entry;
4242
import java.util.regex.Pattern;
4343

44+
import org.apache.beam.it.common.PipelineLauncher;
4445
import org.apache.beam.it.common.PipelineLauncher.LaunchInfo;
46+
import org.apache.beam.it.common.PipelineOperator;
47+
import org.apache.beam.it.common.TestProperties;
4548
import org.apache.beam.it.common.bigquery.BigQueryResourceManager;
46-
import org.apache.beam.it.common.dataflow.AbstractPipelineLauncher;
4749
import org.apache.beam.it.common.monitoring.MonitoringClient;
4850
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
4951
import org.checkerframework.checker.nullness.qual.Nullable;

‎it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/datagenerator/package-info.java‎

Lines changed: 0 additions & 20 deletions
This file was deleted.

‎it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/WordCountIT.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@
2525
import java.time.Duration;
2626
import java.util.Arrays;
2727
import java.util.List;
28+
29+
import org.apache.beam.it.common.dataflow.IOLoadTestBase;
2830
import org.apache.beam.it.common.PipelineLauncher.LaunchConfig;
2931
import org.apache.beam.it.common.PipelineLauncher.LaunchInfo;
3032
import org.apache.beam.it.common.PipelineLauncher.Sdk;

0 commit comments

Comments
 (0)