[SPARK-18364][YARN] Expose metrics for YarnShuffleService

common/network-yarn/src/main/java/org/apache/spark/network/yarn/YarnShuffleService.java

@@ -35,6 +35,8 @@
 import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.fs.permission.FsPermission;
+import org.apache.hadoop.metrics2.impl.MetricsSystemImpl;
+import org.apache.hadoop.metrics2.lib.DefaultMetricsSystem;
 import org.apache.hadoop.yarn.api.records.ContainerId;
 import org.apache.hadoop.yarn.server.api.*;
 import org.apache.spark.network.util.LevelDBProvider;
@@ -168,6 +170,15 @@ protected void serviceInit(Configuration conf) throws Exception {
       TransportConf transportConf = new TransportConf("shuffle", new HadoopConfigProvider(conf));
       blockHandler = new ExternalShuffleBlockHandler(transportConf, registeredExecutorFile);
 
+      // register metrics on the block handler into the Node Manager's metrics system.
+      YarnShuffleServiceMetrics serviceMetrics =
+        new YarnShuffleServiceMetrics(blockHandler.getAllMetrics());
+
+      MetricsSystemImpl metricsSystem = (MetricsSystemImpl) DefaultMetricsSystem.instance();
+      metricsSystem.register(
+        "shuffleService", "Metrics on the Spark Shuffle Service", serviceMetrics);
+      logger.info("Registered metrics with Hadoop's DefaultMetricsSystem");
+
       // If authentication is enabled, set up the shuffle server to use a
       // special RPC handler that filters out unauthenticated fetch requests
       List<TransportServerBootstrap> bootstraps = Lists.newArrayList();

common/network-yarn/src/main/java/org/apache/spark/network/yarn/YarnShuffleServiceMetrics.java

@@ -0,0 +1,133 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.network.yarn;
+
+import com.codahale.metrics.*;
+import org.apache.hadoop.metrics2.MetricsCollector;
+import org.apache.hadoop.metrics2.MetricsInfo;
+import org.apache.hadoop.metrics2.MetricsRecordBuilder;
+import org.apache.hadoop.metrics2.MetricsSource;
+import org.apache.spark.network.shuffle.ExternalShuffleBlockHandler;
+
+import java.util.Map;
+
+/**
+ * Forward {@link ExternalShuffleBlockHandler.ShuffleMetrics} to hadoop metrics system.
+ * NodeManager by default exposes JMX endpoint where can be collected.
+ */
+class YarnShuffleServiceMetrics implements MetricsSource {
+
+  private final MetricSet metricSet;
+
+  YarnShuffleServiceMetrics(MetricSet metricSet) {
+    this.metricSet = metricSet;
+  }
+
+  /**
+   * Get metrics from the source
+   *
+   * @param collector to contain the resulting metrics snapshot
+   * @param all       if true, return all metrics even if unchanged.
+   */
+  @Override
+  public void getMetrics(MetricsCollector collector, boolean all) {
+    MetricsRecordBuilder metricsRecordBuilder = collector.addRecord("shuffleService");
+
+    for (Map.Entry<String, Metric> entry : metricSet.getMetrics().entrySet()) {
+      collectMetric(metricsRecordBuilder, entry.getKey(), entry.getValue());
+    }
+  }
+
+  /**
+   * The metric types used in {@link ExternalShuffleBlockHandler.ShuffleMetrics}
+   */
+  public static void collectMetric(MetricsRecordBuilder metricsRecordBuilder, String name, Metric metric) {
+    if (metric instanceof Timer) {
+      Timer t = (Timer) metric;
+      metricsRecordBuilder
+        .addCounter(new ShuffleServiceMetricsInfo(name + "_count", "Count of timer " + name),
+          t.getCount())
+        .addGauge(
+          new ShuffleServiceMetricsInfo(name + "_rate15", "15 minute rate of timer " + name),
+          t.getFifteenMinuteRate())
+        .addGauge(
+          new ShuffleServiceMetricsInfo(name + "_rate5", "5 minute rate of timer " + name),
+          t.getFiveMinuteRate())
+        .addGauge(
+          new ShuffleServiceMetricsInfo(name + "_rate1", "1 minute rate of timer " + name),
+          t.getOneMinuteRate())
+        .addGauge(new ShuffleServiceMetricsInfo(name + "_rateMean", "Mean rate of timer " + name),
+          t.getMeanRate());
+    } else if (metric instanceof Meter) {
+      Meter m = (Meter) metric;
+      metricsRecordBuilder
+        .addCounter(new ShuffleServiceMetricsInfo(name + "_count", "Count of meter " + name),
+          m.getCount())
+        .addGauge(
+          new ShuffleServiceMetricsInfo(name + "_rate15", "15 minute rate of meter " + name),
+          m.getFifteenMinuteRate())
+        .addGauge(
+          new ShuffleServiceMetricsInfo(name + "_rate5", "5 minute rate of meter " + name),
+          m.getFiveMinuteRate())
+        .addGauge(
+          new ShuffleServiceMetricsInfo(name + "_rate1", "1 minute rate of meter " + name),
+          m.getOneMinuteRate())
+        .addGauge(new ShuffleServiceMetricsInfo(name + "_rateMean", "Mean rate of meter " + name),
+          m.getMeanRate());
+    } else if (metric instanceof Gauge) {
+      final Object gaugeValue = ((Gauge) metric).getValue();
+      if (gaugeValue instanceof Integer) {
+        metricsRecordBuilder.addGauge(getShuffleServiceMetricsInfo(name), (Integer) gaugeValue);
+      } else if (gaugeValue instanceof Long) {
+        metricsRecordBuilder.addGauge(getShuffleServiceMetricsInfo(name), (Long) gaugeValue);
+      } else if (gaugeValue instanceof Float) {
+        metricsRecordBuilder.addGauge(getShuffleServiceMetricsInfo(name), (Float) gaugeValue);
+      } else if (gaugeValue instanceof Double) {
+        metricsRecordBuilder.addGauge(getShuffleServiceMetricsInfo(name), (Double) gaugeValue);
+      } else {
+        throw new IllegalStateException(
+                "Not supported class type of metric[" + name + "] for value " + gaugeValue);
+      }
+    }
+  }
+
+  private static MetricsInfo getShuffleServiceMetricsInfo(String name) {
+    return new ShuffleServiceMetricsInfo(name, "Value of gauge " + name);
+  }
+
+  private static class ShuffleServiceMetricsInfo implements MetricsInfo {
+
+    private final String name;
+    private final String description;
+
+    ShuffleServiceMetricsInfo(String name, String description) {
+      this.name = name;
+      this.description = description;
+    }
+
+    @Override
+    public String name() {
+      return name;
+    }
+
+    @Override
+    public String description() {
+      return description;
+    }
+  }
+}

resource-managers/yarn/src/test/scala/org/apache/spark/network/yarn/YarnShuffleServiceMetricsSuite.scala

@@ -0,0 +1,73 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.network.yarn
+
+import scala.collection.JavaConverters._
+
+import org.apache.hadoop.metrics2.MetricsRecordBuilder
+import org.mockito.Matchers._
+import org.mockito.Mockito.{mock, times, verify, when}
+import org.scalatest.Matchers
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.network.server.OneForOneStreamManager
+import org.apache.spark.network.shuffle.{ExternalShuffleBlockHandler, ExternalShuffleBlockResolver}
+
+class YarnShuffleServiceMetricsSuite extends SparkFunSuite with Matchers {
+
+  val streamManager = mock(classOf[OneForOneStreamManager])
+  val blockResolver = mock(classOf[ExternalShuffleBlockResolver])
+  when(blockResolver.getRegisteredExecutorsSize).thenReturn(42)
+
+  val metrics = new ExternalShuffleBlockHandler(streamManager, blockResolver).getAllMetrics
+
+  test("metrics named as expected") {
+    val allMetrics = Set(
+      "openBlockRequestLatencyMillis", "registerExecutorRequestLatencyMillis",
+      "blockTransferRateBytes", "registeredExecutorsSize")
+
+    metrics.getMetrics.keySet().asScala should be (allMetrics)
+  }
+
+  // these three metrics have the same effect on the collector
+  for (testname <- Seq("openBlockRequestLatencyMillis",
+      "registerExecutorRequestLatencyMillis",
+      "blockTransferRateBytes")) {
+    test(s"$testname - collector receives correct types") {
+      val builder = mock(classOf[MetricsRecordBuilder])
+      when(builder.addCounter(any(), anyLong())).thenReturn(builder)
+      when(builder.addGauge(any(), anyDouble())).thenReturn(builder)
+
+      YarnShuffleServiceMetrics.collectMetric(builder, testname,
+        metrics.getMetrics.get(testname))
+
+      verify(builder).addCounter(anyObject(), anyLong())
+      verify(builder, times(4)).addGauge(anyObject(), anyDouble())
+    }
+  }
+
+  // this metric writes only one gauge to the collector
+  test("registeredExecutorsSize - collector receives correct types") {
+    val builder = mock(classOf[MetricsRecordBuilder])
+
+    YarnShuffleServiceMetrics.collectMetric(builder, "registeredExecutorsSize",
+      metrics.getMetrics.get("registeredExecutorsSize"))
+
+    // only one
+    verify(builder).addGauge(anyObject(), anyInt())
+  }
+}