[BEAM-1672] Use Accumulable MetricsContainers in Spark runner.

runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineResult.java

@@ -18,13 +18,15 @@
 
 package org.apache.beam.runners.spark;
 
+import static org.apache.beam.sdk.metrics.MetricsContainerStepMap.asAttemptedOnlyMetricResults;
+
 import java.io.IOException;
 import java.util.Objects;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
-import org.apache.beam.runners.spark.metrics.SparkMetricResults;
+import org.apache.beam.runners.spark.metrics.MetricsAccumulator;
 import org.apache.beam.runners.spark.translation.SparkContextFactory;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.PipelineResult;
@@ -41,7 +43,6 @@ public abstract class SparkPipelineResult implements PipelineResult {
  protected final Future pipelineExecution;
  protected JavaSparkContext javaSparkContext;
  protected PipelineResult.State state;
- private final SparkMetricResults metricResults = new SparkMetricResults();
 
  SparkPipelineResult(final Future<?> pipelineExecution, final JavaSparkContext javaSparkContext) {
  this.pipelineExecution = pipelineExecution;
@@ -106,7 +107,8 @@ public State waitUntilFinish(final Duration duration) {
 
  @Override
  public MetricResults metrics() {
- return metricResults;
+ return asAttemptedOnlyMetricResults(
+ MetricsAccumulator.getInstance().value());
  }
 
  @Override

runners/spark/src/main/java/org/apache/beam/runners/spark/io/SourceRDD.java

@@ -26,12 +26,12 @@
 import java.util.Iterator;
 import java.util.List;
 import org.apache.beam.runners.spark.metrics.MetricsAccumulator;
-import org.apache.beam.runners.spark.metrics.SparkMetricsContainer;
 import org.apache.beam.runners.spark.translation.SparkRuntimeContext;
 import org.apache.beam.sdk.io.BoundedSource;
 import org.apache.beam.sdk.io.Source;
 import org.apache.beam.sdk.io.UnboundedSource;
 import org.apache.beam.sdk.metrics.MetricsContainer;
+import org.apache.beam.sdk.metrics.MetricsContainerStepMap;
 import org.apache.beam.sdk.metrics.MetricsEnvironment;
 import org.apache.beam.sdk.util.WindowedValue;
 import org.apache.spark.Accumulator;
@@ -65,7 +65,7 @@ public static class Bounded<T> extends RDD<WindowedValue<T>> {
  private final SparkRuntimeContext runtimeContext;
  private final int numPartitions;
  private final String stepName;
- private final Accumulator<SparkMetricsContainer> metricsAccum;
+ private final Accumulator<MetricsContainerStepMap> metricsAccum;
 
  // to satisfy Scala API.
  private static final scala.collection.immutable.Seq<Dependency<?>> NIL =

runners/spark/src/main/java/org/apache/beam/runners/spark/io/SparkUnboundedSource.java

@@ -25,7 +25,6 @@
 import org.apache.beam.runners.spark.SparkPipelineOptions;
 import org.apache.beam.runners.spark.coders.CoderHelpers;
 import org.apache.beam.runners.spark.metrics.MetricsAccumulator;
-import org.apache.beam.runners.spark.metrics.SparkMetricsContainer;
 import org.apache.beam.runners.spark.stateful.StateSpecFunctions;
 import org.apache.beam.runners.spark.translation.SparkRuntimeContext;
 import org.apache.beam.runners.spark.translation.streaming.UnboundedDataset;
@@ -37,6 +36,7 @@
 import org.apache.beam.sdk.metrics.Gauge;
 import org.apache.beam.sdk.metrics.Metrics;
 import org.apache.beam.sdk.metrics.MetricsContainer;
+import org.apache.beam.sdk.metrics.MetricsContainerStepMap;
 import org.apache.beam.sdk.metrics.MetricsEnvironment;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
@@ -190,7 +190,7 @@ public scala.collection.immutable.List<DStream<?>> dependencies() {
  public scala.Option<RDD<BoxedUnit>> compute(Time validTime) {
  // compute parent.
  scala.Option<RDD<Metadata>> parentRDDOpt = parent.getOrCompute(validTime);
- final Accumulator<SparkMetricsContainer> metricsAccum = MetricsAccumulator.getInstance();
+ final Accumulator<MetricsContainerStepMap> metricsAccum = MetricsAccumulator.getInstance();
  long count = 0;
  SparkWatermarks sparkWatermark = null;
  Instant globalLowWatermarkForBatch = BoundedWindow.TIMESTAMP_MIN_VALUE;
@@ -211,7 +211,7 @@ public scala.Option<RDD<BoxedUnit>> compute(Time validTime) {
  ? partitionHighWatermark : globalHighWatermarkForBatch;
  // Update metrics reported in the read
  final Gauge gauge = Metrics.gauge(NAMESPACE, READ_DURATION_MILLIS);
- final MetricsContainer container = metadata.getMetricsContainer().getContainer(stepName);
+ final MetricsContainer container = metadata.getMetricsContainers().getContainer(stepName);
  try (Closeable ignored = MetricsEnvironment.scopedMetricsContainer(container)) {
  final long readDurationMillis = metadata.getReadDurationMillis();
  if (readDurationMillis > maxReadDuration) {
@@ -220,7 +220,7 @@ public scala.Option<RDD<BoxedUnit>> compute(Time validTime) {
  } catch (IOException e) {
  throw new RuntimeException(e);
  }
- metricsAccum.value().update(metadata.getMetricsContainer());
+ metricsAccum.value().updateAll(metadata.getMetricsContainers());
  }
 
  sparkWatermark =
@@ -260,20 +260,19 @@ public static class Metadata implements Serializable {
  private final Instant lowWatermark;
  private final Instant highWatermark;
  private final long readDurationMillis;
- private final SparkMetricsContainer metricsContainer;
+ private final MetricsContainerStepMap metricsContainers;
 
  public Metadata(
  long numRecords,
  Instant lowWatermark,
  Instant highWatermark,
  final long readDurationMillis,
- SparkMetricsContainer metricsContainer) {
+ MetricsContainerStepMap metricsContainer) {
  this.numRecords = numRecords;
  this.readDurationMillis = readDurationMillis;
- this.metricsContainer = metricsContainer;
+ this.metricsContainers = metricsContainer;
  this.lowWatermark = lowWatermark;
  this.highWatermark = highWatermark;
- metricsContainer.materialize();
  }
 
  long getNumRecords() {
@@ -292,8 +291,8 @@ public long getReadDurationMillis() {
  return readDurationMillis;
  }
 
- SparkMetricsContainer getMetricsContainer() {
- return metricsContainer;
+ MetricsContainerStepMap getMetricsContainers() {
+ return metricsContainers;
  }
  }
 

runners/spark/src/main/java/org/apache/beam/runners/spark/metrics/MetricsAccumulator.java

@@ -24,6 +24,7 @@
 import org.apache.beam.runners.spark.SparkPipelineOptions;
 import org.apache.beam.runners.spark.translation.streaming.Checkpoint;
 import org.apache.beam.runners.spark.translation.streaming.Checkpoint.CheckpointDir;
+import org.apache.beam.sdk.metrics.MetricsContainerStepMap;
 import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.fs.Path;
 import org.apache.spark.Accumulator;
@@ -44,7 +45,7 @@ public class MetricsAccumulator {
  private static final String ACCUMULATOR_NAME = "Beam.Metrics";
  private static final String ACCUMULATOR_CHECKPOINT_FILENAME = "metrics";
 
- private static volatile Accumulator<SparkMetricsContainer> instance = null;
+ private static volatile Accumulator<MetricsContainerStepMap> instance = null;
  private static volatile FileSystem fileSystem;
  private static volatile Path checkpointFilePath;
 
@@ -58,11 +59,13 @@ public static void init(SparkPipelineOptions opts, JavaSparkContext jsc) {
  Optional<CheckpointDir> maybeCheckpointDir =
  opts.isStreaming() ? Optional.of(new CheckpointDir(opts.getCheckpointDir()))
  : Optional.<CheckpointDir>absent();
- Accumulator<SparkMetricsContainer> accumulator =
- jsc.sc().accumulator(new SparkMetricsContainer(), ACCUMULATOR_NAME,
+ Accumulator<MetricsContainerStepMap> accumulator =
+ jsc.sc().accumulator(
+ new MetricsContainerStepMap(),
+ ACCUMULATOR_NAME,
  new MetricsAccumulatorParam());
  if (maybeCheckpointDir.isPresent()) {
- Optional<SparkMetricsContainer> maybeRecoveredValue =
+ Optional<MetricsContainerStepMap> maybeRecoveredValue =
  recoverValueFromCheckpoint(jsc, maybeCheckpointDir.get());
  if (maybeRecoveredValue.isPresent()) {
  accumulator.setValue(maybeRecoveredValue.get());
@@ -75,22 +78,23 @@ public static void init(SparkPipelineOptions opts, JavaSparkContext jsc) {
  }
  }
 
- public static Accumulator<SparkMetricsContainer> getInstance() {
+ public static Accumulator<MetricsContainerStepMap> getInstance() {
  if (instance == null) {
  throw new IllegalStateException("Metrics accumulator has not been instantiated");
  } else {
  return instance;
  }
  }
 
- private static Optional<SparkMetricsContainer> recoverValueFromCheckpoint(
+ private static Optional<MetricsContainerStepMap> recoverValueFromCheckpoint(
  JavaSparkContext jsc,
  CheckpointDir checkpointDir) {
  try {
  Path beamCheckpointPath = checkpointDir.getBeamCheckpointDir();
  checkpointFilePath = new Path(beamCheckpointPath, ACCUMULATOR_CHECKPOINT_FILENAME);
  fileSystem = checkpointFilePath.getFileSystem(jsc.hadoopConfiguration());
- SparkMetricsContainer recoveredValue = Checkpoint.readObject(fileSystem, checkpointFilePath);
+ MetricsContainerStepMap recoveredValue =
+ Checkpoint.readObject(fileSystem, checkpointFilePath);
  if (recoveredValue != null) {
  LOG.info("Recovered metrics from checkpoint.");
  return Optional.of(recoveredValue);
@@ -117,7 +121,7 @@ private static void checkpoint() throws IOException {
  }
 
  /**
- * Spark Listener which checkpoints {@link SparkMetricsContainer} values for fault-tolerance.
+ * Spark Listener which checkpoints {@link MetricsContainerStepMap} values for fault-tolerance.
  */
  public static class AccumulatorCheckpointingSparkListener extends JavaStreamingListener {
  @Override

runners/spark/src/main/java/org/apache/beam/runners/spark/metrics/MetricsAccumulatorParam.java

@@ -18,25 +18,31 @@
 
 package org.apache.beam.runners.spark.metrics;
 
+import org.apache.beam.sdk.metrics.MetricsContainerStepMap;
 import org.apache.spark.AccumulatorParam;
 
 
 /**
  * Metrics accumulator param.
  */
-class MetricsAccumulatorParam implements AccumulatorParam<SparkMetricsContainer> {
+class MetricsAccumulatorParam implements AccumulatorParam<MetricsContainerStepMap> {
  @Override
- public SparkMetricsContainer addAccumulator(SparkMetricsContainer c1, SparkMetricsContainer c2) {
- return c1.update(c2);
+ public MetricsContainerStepMap addAccumulator(
+ MetricsContainerStepMap c1,
+ MetricsContainerStepMap c2) {
+ return addInPlace(c1, c2);
  }
 
  @Override
- public SparkMetricsContainer addInPlace(SparkMetricsContainer c1, SparkMetricsContainer c2) {
- return c1.update(c2);
+ public MetricsContainerStepMap addInPlace(
+ MetricsContainerStepMap c1,
+ MetricsContainerStepMap c2) {
+ c1.updateAll(c2);
+ return c1;
  }
 
  @Override
- public SparkMetricsContainer zero(SparkMetricsContainer initialValue) {
- return new SparkMetricsContainer();
+ public MetricsContainerStepMap zero(MetricsContainerStepMap initialValue) {
+ return new MetricsContainerStepMap();
  }
 }

runners/spark/src/main/java/org/apache/beam/runners/spark/metrics/SparkBeamMetric.java

@@ -18,6 +18,8 @@
 
 package org.apache.beam.runners.spark.metrics;
 
+import static org.apache.beam.sdk.metrics.MetricsContainerStepMap.asAttemptedOnlyMetricResults;
+
 import com.codahale.metrics.Metric;
 import com.google.common.annotations.VisibleForTesting;
 import java.util.HashMap;
@@ -27,20 +29,23 @@
 import org.apache.beam.sdk.metrics.MetricName;
 import org.apache.beam.sdk.metrics.MetricQueryResults;
 import org.apache.beam.sdk.metrics.MetricResult;
+import org.apache.beam.sdk.metrics.MetricResults;
+import org.apache.beam.sdk.metrics.MetricsContainerStepMap;
 import org.apache.beam.sdk.metrics.MetricsFilter;
 
 
 /**
- * An adapter between the {@link SparkMetricsContainer} and Codahale's {@link Metric} interface.
+ * An adapter between the {@link MetricsContainerStepMap} and Codahale's {@link Metric} interface.
  */
 class SparkBeamMetric implements Metric {
  private static final String ILLEGAL_CHARACTERS = "[^A-Za-z0-9\\._-]";
  private static final String ILLEGAL_CHARACTERS_AND_PERIOD = "[^A-Za-z0-9_-]";
 
- private final SparkMetricResults metricResults = new SparkMetricResults();
-
  Map<String, ?> renderAll() {
  Map<String, Object> metrics = new HashMap<>();
+ MetricResults metricResults =
+ asAttemptedOnlyMetricResults(
+ MetricsAccumulator.getInstance().value());
  MetricQueryResults metricQueryResults =
  metricResults.queryMetrics(MetricsFilter.builder().build());
  for (MetricResult<Long> metricResult : metricQueryResults.counters()) {

runners/spark/src/main/java/org/apache/beam/runners/spark/metrics/SparkBeamMetricSource.java

@@ -24,7 +24,7 @@
 
 /**
  * A Spark {@link Source} that is tailored to expose a {@link SparkBeamMetric},
- * wrapping an underlying {@link SparkMetricsContainer} instance.
+ * wrapping an underlying {@link org.apache.beam.sdk.metrics.MetricResults} instance.
  */
 public class SparkBeamMetricSource implements Source {
  private static final String METRIC_NAME = "Metrics";