[SPARK-20327][CORE][YARN] Add CLI support for YARN custom resources, like GPUs

docs/running-on-yarn.md

@@ -121,6 +121,43 @@ To use a custom metrics.properties for the application master and executors, upd
     Use lower-case suffixes, e.g. <code>k</code>, <code>m</code>, <code>g</code>, <code>t</code>, and <code>p</code>, for kibi-, mebi-, gibi-, tebi-, and pebibytes, respectively.
   </td>
 </tr>
+<tr>
+  <td><code>spark.yarn.am.resource.{resource-type}</code></td>
+  <td><code>(none)</code></td>
+  <td>
+    Amount of resource to use for the YARN Application Master in client mode.
+    In cluster mode, use <code>spark.yarn.driver.resource.&lt;resource-type&gt;</code> instead.
+    Please note that this feature can be used only with YARN 3.0+
+    For reference, see YARN Resource Model documentation: https://hadoop.apache.org/docs/r3.0.1/hadoop-yarn/hadoop-yarn-site/ResourceModel.html
+    <p/>
+    Example: 
+    To request GPU resources from YARN, use: <code>spark.yarn.am.resource.yarn.io/gpu</code>
+  </td>
+</tr>
+<tr>
+  <td><code>spark.yarn.driver.resource.{resource-type}</code></td>
+  <td><code>(none)</code></td>
+  <td>
+    Amount of resource to use for the YARN Application Master in cluster mode.
+    Please note that this feature can be used only with YARN 3.0+
+    For reference, see YARN Resource Model documentation: https://hadoop.apache.org/docs/r3.0.1/hadoop-yarn/hadoop-yarn-site/ResourceModel.html
+    <p/>
+    Example: 
+    To request GPU resources from YARN, use: <code>spark.yarn.driver.resource.yarn.io/gpu</code>
+  </td>
+</tr>
+<tr>
+  <td><code>spark.yarn.executor.resource.{resource-type}</code></td>
+  <td><code>(none)</code></td>
+ <td>
+     Amount of resource to use per executor process.
+     Please note that this feature can be used only with YARN 3.0+
+     For reference, see YARN Resource Model documentation: https://hadoop.apache.org/docs/r3.0.1/hadoop-yarn/hadoop-yarn-site/ResourceModel.html
+     <p/>
+     Example: 
+     To request GPU resources from YARN, use: <code>spark.yarn.executor.resource.yarn.io/gpu</code>
+ </td>
+</tr>
 <tr>
   <td><code>spark.yarn.am.cores</code></td>
   <td><code>1</code></td>

resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala

@@ -154,6 +154,8 @@ private[spark] class Client(
    * available in the alpha API.
    */
   def submitApplication(): ApplicationId = {
+    ResourceRequestHelper.validateResources(sparkConf)
+
     var appId: ApplicationId = null
     try {
       launcherBackend.connect()
@@ -234,6 +236,13 @@ private[spark] class Client(
   def createApplicationSubmissionContext(
       newApp: YarnClientApplication,
       containerContext: ContainerLaunchContext): ApplicationSubmissionContext = {
+    val amResources =
+      if (isClusterMode) {
+        sparkConf.getAllWithPrefix(config.YARN_DRIVER_RESOURCE_TYPES_PREFIX).toMap
+      } else {
+        sparkConf.getAllWithPrefix(config.YARN_AM_RESOURCE_TYPES_PREFIX).toMap
+      }
+    logDebug(s"AM resources: $amResources")
     val appContext = newApp.getApplicationSubmissionContext
     appContext.setApplicationName(sparkConf.get("spark.app.name", "Spark"))
     appContext.setQueue(sparkConf.get(QUEUE_NAME))
@@ -256,6 +265,10 @@ private[spark] class Client(
     val capability = Records.newRecord(classOf[Resource])
     capability.setMemory(amMemory + amMemoryOverhead)
     capability.setVirtualCores(amCores)
+    if (amResources.nonEmpty) {
+      ResourceRequestHelper.setResourceRequests(amResources, capability)
+    }
+    logDebug(s"Created resource capability for AM request: $capability")
 
     sparkConf.get(AM_NODE_LABEL_EXPRESSION) match {
       case Some(expr) =>

resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/ResourceRequestHelper.scala

@@ -0,0 +1,155 @@
+/*
+ * 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.deploy.yarn
+
+import java.lang.{Long => JLong}
+import java.lang.reflect.InvocationTargetException
+
+import scala.collection.mutable
+import scala.util.Try
+
+import org.apache.hadoop.yarn.api.records.Resource
+
+import org.apache.spark.{SparkConf, SparkException}
+import org.apache.spark.deploy.yarn.config._
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.config._
+import org.apache.spark.util.Utils
+
+/**
+ * This helper class uses some of Hadoop 3 methods from the YARN API,
+ * so we need to use reflection to avoid compile error when building against Hadoop 2.x
+ */
+private object ResourceRequestHelper extends Logging {
+  private val AMOUNT_AND_UNIT_REGEX = "([0-9]+)([A-Za-z]*)".r
+  private val RESOURCE_INFO_CLASS = "org.apache.hadoop.yarn.api.records.ResourceInformation"
+
+  /**
+   * Validates sparkConf and throws a SparkException if any of standard resources (memory or cores)
+   * is defined with the property spark.yarn.x.resource.y
+   * Need to reject all combinations of AM / Driver / Executor and memory / CPU cores resources, as
+   * Spark has its own names for them (memory, cores),
+   * but YARN have its names too: (memory, memory-mb, mb) and (cores, vcores, cpu-vcores).
+   * We need to disable every possible way YARN could receive the resource definitions above.
+   */
+  def validateResources(sparkConf: SparkConf): Unit = {
+    val resourceDefinitions = Seq[(String, String)](
+      (AM_MEMORY.key, YARN_AM_RESOURCE_TYPES_PREFIX + "memory"),
+      (DRIVER_MEMORY.key, YARN_DRIVER_RESOURCE_TYPES_PREFIX + "memory"),
+      (EXECUTOR_MEMORY.key, YARN_EXECUTOR_RESOURCE_TYPES_PREFIX + "memory"),
+      (AM_MEMORY.key, YARN_AM_RESOURCE_TYPES_PREFIX + "mb"),
+      (DRIVER_MEMORY.key, YARN_DRIVER_RESOURCE_TYPES_PREFIX + "mb"),
+      (EXECUTOR_MEMORY.key, YARN_EXECUTOR_RESOURCE_TYPES_PREFIX + "mb"),
+      (AM_MEMORY.key, YARN_AM_RESOURCE_TYPES_PREFIX + "memory-mb"),
+      (DRIVER_MEMORY.key, YARN_DRIVER_RESOURCE_TYPES_PREFIX + "memory-mb"),
+      (EXECUTOR_MEMORY.key, YARN_EXECUTOR_RESOURCE_TYPES_PREFIX + "memory-mb"),
+      (AM_CORES.key, YARN_AM_RESOURCE_TYPES_PREFIX + "cores"),
+      (DRIVER_CORES.key, YARN_DRIVER_RESOURCE_TYPES_PREFIX + "cores"),
+      (EXECUTOR_CORES.key, YARN_EXECUTOR_RESOURCE_TYPES_PREFIX + "cores"),
+      (AM_CORES.key, YARN_AM_RESOURCE_TYPES_PREFIX + "vcores"),
+      (DRIVER_CORES.key, YARN_DRIVER_RESOURCE_TYPES_PREFIX + "vcores"),
+      (EXECUTOR_CORES.key, YARN_EXECUTOR_RESOURCE_TYPES_PREFIX + "vcores"),
+      (AM_CORES.key, YARN_AM_RESOURCE_TYPES_PREFIX + "cpu-vcores"),
+      (DRIVER_CORES.key, YARN_DRIVER_RESOURCE_TYPES_PREFIX + "cpu-vcores"),
+      (EXECUTOR_CORES.key, YARN_EXECUTOR_RESOURCE_TYPES_PREFIX + "cpu-vcores"))
+    val errorMessage = new mutable.StringBuilder()
+
+    resourceDefinitions.foreach { case (sparkName, resourceRequest) =>
+      if (sparkConf.contains(resourceRequest)) {
+        errorMessage.append(s"Error: Do not use $resourceRequest, " +
+            s"please use $sparkName instead!\n")
+      }
+    }
+
+    if (errorMessage.nonEmpty) {
+      throw new SparkException(errorMessage.toString())
+    }
+  }
+
+  /**
+   * Sets resource amount with the corresponding unit to the passed resource object.
+   * @param resources resource values to set
+   * @param resource resource object to update
+   */
+  def setResourceRequests(
+      resources: Map[String, String],
+      resource: Resource): Unit = {
+    require(resource != null, "Resource parameter should not be null!")
+
+    logDebug(s"Custom resources requested: $resources")
+    if (!isYarnResourceTypesAvailable()) {
+      if (resources.nonEmpty) {
+        logWarning("Ignoring custom resource requests because " +
+            "the version of YARN does not support it!")
+      }
+      return
+    }
+
+    val resInfoClass = Utils.classForName(RESOURCE_INFO_CLASS)
+    val setResourceInformationMethod =
+      resource.getClass.getMethod("setResourceInformation", classOf[String], resInfoClass)
+    resources.foreach { case (name, rawAmount) =>
+      try {
+        val AMOUNT_AND_UNIT_REGEX(amountPart, unitPart) = rawAmount
+        val amount = amountPart.toLong
+        val unit = unitPart match {
+          case "g" => "G"
+          case "t" => "T"
+          case "p" => "P"
+          case _ => unitPart
+        }
+        logDebug(s"Registering resource with name: $name, amount: $amount, unit: $unit")
+        val resourceInformation = createResourceInformation(name, amount, unit, resInfoClass)
+        setResourceInformationMethod.invoke(
+          resource, name, resourceInformation.asInstanceOf[AnyRef])
+      } catch {
+        case _: MatchError =>
+          throw new IllegalArgumentException(s"Resource request for '$name' ('$rawAmount') " +
+              s"does not match pattern $AMOUNT_AND_UNIT_REGEX.")
+        case e: InvocationTargetException if e.getCause != null => throw e.getCause
+      }
+    }
+  }
+
+  private def createResourceInformation(
+      resourceName: String,
+      amount: Long,
+      unit: String,
+      resInfoClass: Class[_]): Any = {
+    val resourceInformation =
+      if (unit.nonEmpty) {
+        val resInfoNewInstanceMethod = resInfoClass.getMethod("newInstance",
+          classOf[String], classOf[String], JLong.TYPE)
+        resInfoNewInstanceMethod.invoke(null, resourceName, unit, amount.asInstanceOf[JLong])
+      } else {
+        val resInfoNewInstanceMethod = resInfoClass.getMethod("newInstance",
+          classOf[String], JLong.TYPE)
+        resInfoNewInstanceMethod.invoke(null, resourceName, amount.asInstanceOf[JLong])
+      }
+    resourceInformation
+  }
+
+  /**
+   * Checks whether Hadoop 2.x or 3 is used as a dependency.
+   * In case of Hadoop 3 and later, the ResourceInformation class
+   * should be available on the classpath.
+   */
+  def isYarnResourceTypesAvailable(): Boolean = {
+    Try(Utils.classForName(RESOURCE_INFO_CLASS)).isSuccess
+  }
+}

resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala

@@ -140,10 +140,18 @@ private[yarn] class YarnAllocator(
   }
   // Number of cores per executor.
   protected val executorCores = sparkConf.get(EXECUTOR_CORES)
-  // Resource capability requested for each executors
-  private[yarn] val resource = Resource.newInstance(
-    executorMemory + memoryOverhead + pysparkWorkerMemory,
-    executorCores)
+
+  private val executorResourceRequests =
+    sparkConf.getAllWithPrefix(config.YARN_EXECUTOR_RESOURCE_TYPES_PREFIX).toMap
+
+  // Resource capability requested for each executor
+  private[yarn] val resource: Resource = {
+    val resource = Resource.newInstance(
+      executorMemory + memoryOverhead + pysparkWorkerMemory, executorCores)
+    ResourceRequestHelper.setResourceRequests(executorResourceRequests, resource)
+    logDebug(s"Created resource capability: $resource")
+    resource
+  }
 
   private val launcherPool = ThreadUtils.newDaemonCachedThreadPool(
     "ContainerLauncher", sparkConf.get(CONTAINER_LAUNCH_MAX_THREADS))
@@ -288,9 +296,16 @@ private[yarn] class YarnAllocator(
       s"executorsStarting: ${numExecutorsStarting.get}")
 
     if (missing > 0) {
-      logInfo(s"Will request $missing executor container(s), each with " +
-        s"${resource.getVirtualCores} core(s) and " +
-        s"${resource.getMemory} MB memory (including $memoryOverhead MB of overhead)")
+      if (log.isInfoEnabled()) {
+        var requestContainerMessage = s"Will request $missing executor container(s), each with " +
+            s"${resource.getVirtualCores} core(s) and " +
+            s"${resource.getMemory} MB memory (including $memoryOverhead MB of overhead)"
+        if (ResourceRequestHelper.isYarnResourceTypesAvailable() &&
+            executorResourceRequests.nonEmpty) {
+          requestContainerMessage ++= s" with custom resources: " + resource.toString
+        }
+        logInfo(requestContainerMessage)
+      }
 
       // Split the pending container request into three groups: locality matched list, locality
       // unmatched list and non-locality list. Take the locality matched container request into
@@ -456,13 +471,20 @@ private[yarn] class YarnAllocator(
     // memory, but use the asked vcore count for matching, effectively disabling matching on vcore
     // count.
     val matchingResource = Resource.newInstance(allocatedContainer.getResource.getMemory,
-          resource.getVirtualCores)
+      resource.getVirtualCores)
+
+    ResourceRequestHelper.setResourceRequests(executorResourceRequests, matchingResource)
+
+    logDebug(s"Calling amClient.getMatchingRequests with parameters: " +
+        s"priority: ${allocatedContainer.getPriority}, " +
+        s"location: $location, resource: $matchingResource")
     val matchingRequests = amClient.getMatchingRequests(allocatedContainer.getPriority, location,
       matchingResource)
 
     // Match the allocation to a request
     if (!matchingRequests.isEmpty) {
       val containerRequest = matchingRequests.get(0).iterator.next
+      logDebug(s"Removing container request via AM client: $containerRequest")
       amClient.removeContainerRequest(containerRequest)
       containersToUse += allocatedContainer
     } else {

resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/config.scala

@@ -345,4 +345,8 @@ package object config {
       .booleanConf
       .createWithDefault(false)
 
+  private[yarn] val YARN_EXECUTOR_RESOURCE_TYPES_PREFIX = "spark.yarn.executor.resource."
+  private[yarn] val YARN_DRIVER_RESOURCE_TYPES_PREFIX = "spark.yarn.driver.resource."
+  private[yarn] val YARN_AM_RESOURCE_TYPES_PREFIX = "spark.yarn.am.resource."
+
 }

resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/ClientSuite.scala

@@ -23,6 +23,7 @@ import java.util.Properties
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable.{HashMap => MutableHashMap}
+import scala.util.control.NonFatal
 
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
@@ -199,6 +200,20 @@ class ClientSuite extends SparkFunSuite with Matchers {
     appContext.getMaxAppAttempts should be (42)
   }
 
+  test("resource request (client mode)") {
+    val sparkConf = new SparkConf().set("spark.submit.deployMode", "client")
+      .set(YARN_AM_RESOURCE_TYPES_PREFIX + "fpga", "2")
+      .set(YARN_AM_RESOURCE_TYPES_PREFIX + "gpu", "3")
+    testResourceRequest(sparkConf, List("gpu", "fpga"), Seq(("fpga", 2), ("gpu", 3)))
+  }
+
+  test("resource request (cluster mode)") {
+    val sparkConf = new SparkConf().set("spark.submit.deployMode", "cluster")
+      .set(YARN_DRIVER_RESOURCE_TYPES_PREFIX + "fpga", "4")
+      .set(YARN_DRIVER_RESOURCE_TYPES_PREFIX + "gpu", "5")
+    testResourceRequest(sparkConf, List("gpu", "fpga"), Seq(("fpga", 4), ("gpu", 5)))
+  }
+
   test("spark.yarn.jars with multiple paths and globs") {
     val libs = Utils.createTempDir()
     val single = Utils.createTempDir()
@@ -433,4 +448,30 @@ class ClientSuite extends SparkFunSuite with Matchers {
     classpath(env)
   }
 
+  private def testResourceRequest(
+      sparkConf: SparkConf,
+      resources: List[String],
+      expectedResources: Seq[(String, Long)]): Unit = {
+    assume(ResourceRequestHelper.isYarnResourceTypesAvailable())
+    ResourceRequestTestHelper.initializeResourceTypes(resources)
+
+    val args = new ClientArguments(Array())
+
+    val appContext = Records.newRecord(classOf[ApplicationSubmissionContext])
+    val getNewApplicationResponse = Records.newRecord(classOf[GetNewApplicationResponse])
+    val containerLaunchContext = Records.newRecord(classOf[ContainerLaunchContext])
+
+    val client = new Client(args, sparkConf)
+    client.createApplicationSubmissionContext(
+      new YarnClientApplication(getNewApplicationResponse, appContext),
+      containerLaunchContext)
+
+    appContext.getAMContainerSpec should be (containerLaunchContext)
+    appContext.getApplicationType should be ("SPARK")
+
+    expectedResources.foreach { case (name, value) =>
+      ResourceRequestTestHelper.getResourceTypeValue(appContext.getResource, name) should be (value)
+    }
+  }
+
 }