[SPARK-27024] Executor interface for cluster managers to support GPU and other resources

core/src/main/scala/org/apache/spark/ResourceDiscoverer.scala

@@ -0,0 +1,89 @@
+/*
+ * 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
+
+import java.io.File
+
+import com.fasterxml.jackson.core.JsonParseException
+import org.json4s.{DefaultFormats, MappingException}
+import org.json4s.jackson.JsonMethods._
+
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.config._
+import org.apache.spark.util.Utils.executeAndGetOutput
+
+/**
+ * Discovers resources (GPUs/FPGAs/etc).
+ * This class find resources by running and parses the output of the user specified script
+ * from the config spark.{driver/executor}.resource.{resourceType}.discoveryScript.
+ * The output of the script it runs is expected to be JSON in the format of the
+ * ResourceInformation class, with addresses being optional.
+ *
+ * For example:  {"name": "gpu","count":2, "units":"", "addresses": ["0","1"]}
+ */
+private[spark] object ResourceDiscoverer extends Logging {
+
+  private implicit val formats = DefaultFormats
+
+  def findResources(sparkconf: SparkConf, isDriver: Boolean): Map[String, ResourceInformation] = {
+    val prefix = if (isDriver) {
+      SPARK_DRIVER_RESOURCE_PREFIX
+    } else {
+      SPARK_EXECUTOR_RESOURCE_PREFIX
+    }
+    // get unique resource types by grabbing first part config with multiple periods,
+    // ie resourceType.count, grab resourceType part
+    val resourceTypes = sparkconf.getAllWithPrefix(prefix).map { case (k, _) =>
+      k.split('.').head
+    }.toSet
+    resourceTypes.map { rtype => {
+      val rInfo = getResourceInfoForType(sparkconf, prefix, rtype)
+      (rtype -> rInfo)
+    }}.toMap
+  }
+
+  private def getResourceInfoForType(
+      sparkconf: SparkConf,
+      prefix: String,
+      resourceType: String): ResourceInformation = {
+    val discoveryConf = prefix + resourceType + SPARK_RESOURCE_DISCOVERY_SCRIPT_POSTFIX
+    val script = sparkconf.getOption(discoveryConf)
+    val result = if (script.nonEmpty) {
+      val scriptFile = new File(script.get)
+      // check that script exists and try to execute
+      if (scriptFile.exists()) {
+        try {
+          val output = executeAndGetOutput(Seq(script.get), new File("."))
+          val parsedJson = parse(output)
+          parsedJson.extract[ResourceInformation]
+        } catch {
+          case e @ (_: SparkException | _: MappingException | _: JsonParseException) =>
+            throw new SparkException(s"Error running the resource discovery script: $scriptFile" +
+              s" for $resourceType", e)
+        }
+      } else {
+        throw new SparkException(s"Resource script: $scriptFile to discover $resourceType" +
+          s" doesn't exist!")
+      }
+    } else {
+      throw new SparkException(s"User is expecting to use $resourceType resources but " +
+        s"didn't specify a script via conf: $discoveryConf, to find them!")
+    }
+    result
+  }
+}

core/src/main/scala/org/apache/spark/ResourceInformation.scala

@@ -0,0 +1,41 @@
+/*
+ * 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
+
+import org.apache.spark.annotation.Evolving
+
+/**
+ * Class to hold information about a type of Resource. A resource could be a GPU, FPGA, etc.
+ * The units are resource specific and could be something like MB or GB for memory.
+ * The array of addresses are resource specific and its up to the user to interpret the address.
+ * The units and addresses could be empty if they don't apply to that resource.
+ *
+ * One example is GPUs, where the addresses would be the indices of the GPUs, the count would be the
+ * number of GPUs and the units would be an empty string.
+ *
+ * @param name the name of the resource
+ * @param units the units of the resources, can be an empty string if units don't apply
+ * @param count the number of resources available
+ * @param addresses an optional array of strings describing the addresses of the resource
+ */
+@Evolving
+case class ResourceInformation(
+    val name: String,
+    val units: String,
+    val count: Long,
+    val addresses: Array[String] = Array.empty)

core/src/main/scala/org/apache/spark/executor/CoarseGrainedExecutorBackend.scala

@@ -17,6 +17,7 @@
 
 package org.apache.spark.executor
 
+import java.io.{BufferedInputStream, FileInputStream}
 import java.net.URL
 import java.nio.ByteBuffer
 import java.util.Locale
@@ -26,11 +27,18 @@ import scala.collection.mutable
 import scala.util.{Failure, Success}
 import scala.util.control.NonFatal
 
+import com.fasterxml.jackson.databind.exc.MismatchedInputException
+import org.json4s.DefaultFormats
+import org.json4s.JsonAST.JArray
+import org.json4s.MappingException
+import org.json4s.jackson.JsonMethods._
+
 import org.apache.spark._
 import org.apache.spark.TaskState.TaskState
 import org.apache.spark.deploy.SparkHadoopUtil
 import org.apache.spark.deploy.worker.WorkerWatcher
 import org.apache.spark.internal.Logging
+import org.apache.spark.internal.config._
 import org.apache.spark.rpc._
 import org.apache.spark.scheduler.{ExecutorLossReason, TaskDescription}
 import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages._
@@ -44,9 +52,12 @@ private[spark] class CoarseGrainedExecutorBackend(
     hostname: String,
     cores: Int,
     userClassPath: Seq[URL],
-    env: SparkEnv)
+    env: SparkEnv,
+    resourcesFile: Option[String])
   extends ThreadSafeRpcEndpoint with ExecutorBackend with Logging {
 
+  private implicit val formats = DefaultFormats
+
   private[this] val stopping = new AtomicBoolean(false)
   var executor: Executor = null
   @volatile var driver: Option[RpcEndpointRef] = None
@@ -61,7 +72,7 @@ private[spark] class CoarseGrainedExecutorBackend(
       // This is a very fast action so we can use "ThreadUtils.sameThread"
       driver = Some(ref)
       ref.ask[Boolean](RegisterExecutor(executorId, self, hostname, cores, extractLogUrls,
-        extractAttributes))
+        extractAttributes, parseResources(resourcesFile)))
     }(ThreadUtils.sameThread).onComplete {
       // This is a very fast action so we can use "ThreadUtils.sameThread"
       case Success(msg) =>
@@ -71,6 +82,159 @@ private[spark] class CoarseGrainedExecutorBackend(
     }(ThreadUtils.sameThread)
   }
 
+  // we only support converting from units that are byte based, this will
+  // either convert it to bytes or just return the count if the units can't be
+  // parsed as bytes.
+  private def tryConvertUnitsToBytes(count: String, units: String): Long = {
+    try {
+        Utils.byteStringAsBytes(count + units)
+    } catch {
+      case e: NumberFormatException =>
+        // Ignore units not of byte types and just use count
+        logWarning("Illegal resource unit type, spark only " +
+          s"supports conversion of byte types, units: $units, " +
+          "ignoring the type and using the raw count.", e)
+        count.toLong
+    }
+  }
+
+  // Check that the executor resources at startup will satisfy the user specified task
+  // requirements (spark.task.resource.*) and that they match the executor configs
+  // specified by the user (spark.executor.resource.*) to catch mismatches between what
+  // the user requested and what resource manager gave or what the discovery script found.
+  private def checkExecResourcesMeetTaskRequirements(
+      taskResourceConfigs: Array[(String, String)],
+      actualExecResources: Map[String, ResourceInformation]): Unit = {
+
+    // get just the map of resource name to count
+    val taskResourcesAndCounts = taskResourceConfigs.
+      withFilter { case (k, v) => k.endsWith(SPARK_RESOURCE_COUNT_POSTFIX)}.
+      map { case (k, v) => (k.dropRight(SPARK_RESOURCE_COUNT_POSTFIX.size), v)}
+
+    case class ResourceRealCounts(execCount: Long, taskCount: Long)
+
+    // SPARK will only base it off the counts and known byte units, if
+    // user is trying to use something else we will have to add a plugin later
+    taskResourcesAndCounts.foreach { case (rName, taskCount) =>
+      if (actualExecResources.contains(rName)) {
+        val execResourceInfo = actualExecResources(rName)
+        val taskUnits = env.conf.getOption(
+          SPARK_TASK_RESOURCE_PREFIX + rName + SPARK_RESOURCE_UNITS_POSTFIX)
+        val userExecUnitsConfigName =
+          SPARK_EXECUTOR_RESOURCE_PREFIX + rName + SPARK_RESOURCE_UNITS_POSTFIX
+        val userExecConfigUnits = env.conf.getOption(userExecUnitsConfigName)
+        val realCounts = if (execResourceInfo.units.nonEmpty) {
+          if (taskUnits.nonEmpty && taskUnits.get.nonEmpty) {
+            if (userExecConfigUnits.isEmpty || userExecConfigUnits.get.isEmpty) {
+              throw new SparkException(s"Resource: $rName has units in task config " +
+                s"and executor startup config but the user specified executor resource " +
+                s"config is missing the units config - see ${userExecUnitsConfigName}.")
+            }
+            val execCountWithUnits =
+              tryConvertUnitsToBytes(execResourceInfo.count.toString, execResourceInfo.units)
+            val taskCountWithUnits =
+              tryConvertUnitsToBytes(taskCount, taskUnits.get)
+            ResourceRealCounts(execCountWithUnits, taskCountWithUnits)
+          } else {
+            throw new SparkException(
+              s"Resource: $rName has an executor units config: ${execResourceInfo.units}, but " +
+                s"the task units config is missing.")
+          }
+        } else {
+          if (taskUnits.nonEmpty && taskUnits.get.nonEmpty) {
+            throw new SparkException(
+              s"Resource: $rName has a task units config: ${taskUnits.get}, but the executor " +
+              s"units config is missing.")
+          }
+          ResourceRealCounts(execResourceInfo.count, taskCount.toLong)
+        }
+        if (realCounts.execCount < realCounts.taskCount) {
+          throw new SparkException(s"Executor resource: $rName, count: ${realCounts.execCount} " +
+            s"isn't large enough to meet task requirements of: ${realCounts.taskCount}")
+        }
+        // also make sure the executor resource count on start matches the
+        // spark.executor.resource configs specified by user
+        val userExecCountConfigName =
+          SPARK_EXECUTOR_RESOURCE_PREFIX + rName + SPARK_RESOURCE_COUNT_POSTFIX
+        val userExecConfigCount = env.conf.getOption(userExecCountConfigName).
+          getOrElse(throw new SparkException(s"Executor resource: $rName not specified " +
+            s"via config: $userExecCountConfigName, but required " +
+            s"by the task, please fix your configuration"))
+        val execConfigCountWithUnits = if (userExecConfigUnits.nonEmpty) {
+          tryConvertUnitsToBytes(userExecConfigCount, userExecConfigUnits.get)
+        } else {
+          userExecConfigCount.toLong
+        }
+        if (execConfigCountWithUnits != realCounts.execCount) {
+          throw new SparkException(s"Executor resource: $rName, count: ${realCounts.execCount} " +
+            s"doesn't match what user requests for executor count: $execConfigCountWithUnits, " +
+            s"via $userExecCountConfigName")
+        }
+      } else {
+        throw new SparkException(s"Executor resource config missing required task resource: $rName")
+      }
+    }
+  }
+
+  // visible for testing
+  def parseResources(resourcesFile: Option[String]): Map[String, ResourceInformation] = {
+    // only parse the resources if a task requires them
+    val taskResourceConfigs = env.conf.getAllWithPrefix(SPARK_TASK_RESOURCE_PREFIX)
+    val resourceInfo = if (taskResourceConfigs.nonEmpty) {
+      val execResources = resourcesFile.map { resourceFileStr => {
+        val source = new BufferedInputStream(new FileInputStream(resourceFileStr))
+        val resourceMap = try {
+          val parsedJson = parse(source).asInstanceOf[JArray].arr
+          parsedJson.map(_.extract[ResourceInformation]).map(x => (x.name -> x)).toMap
+        } catch {
+          case e @ (_: MappingException | _: MismatchedInputException | _: ClassCastException) =>
+            throw new SparkException(
+              s"Exception parsing the resources in $resourceFileStr", e)
+        } finally {
+          source.close()
+        }
+        resourceMap
+      }}.getOrElse(ResourceDiscoverer.findResources(env.conf, false))
+
+      if (execResources.isEmpty) {
+        throw new SparkException(s"User specified resources per task via: " +
+          s"$SPARK_TASK_RESOURCE_PREFIX, but can't find any resources available on the executor.")
+      }
+
+      // check that the executor has all the resources required by the application/task
+      checkExecResourcesMeetTaskRequirements(taskResourceConfigs, execResources)
+
+      // make sure the addresses make sense with count
+      execResources.foreach { case (rName, rInfo) =>
+        // check to make sure we have enough addresses when any specified, if we have
+        // more don't worry about it
+        if (rInfo.addresses.nonEmpty && rInfo.addresses.size < rInfo.count) {
+          throw new SparkException(s"The number of resource addresses is expected to either " +
+            s"be >= to the count or be empty if not applicable! Resource: $rName, " +
+            s"count: ${rInfo.count}, number of addresses: ${rInfo.addresses.size}")
+        }
+      }
+
+      logInfo(s"Executor ${executorId} using resources: ${execResources.keys}")
+      if (log.isDebugEnabled) {
+        logDebug("===============================================================================")
+        logDebug("Executor Resources:")
+        execResources.foreach{ case (k, v) =>
+          logDebug(s"$k -> [name: ${v.name}, units: ${v.units}, count: ${v.count}," +
+            s" addresses: ${v.addresses.deep}]")}
+        logDebug("===============================================================================")
+      }
+      execResources
+    } else {
+      if (resourcesFile.nonEmpty) {
+        logWarning(s"A resources file was specified but the application is not configured " +
+          s"to use any resources, see the configs with prefix: ${SPARK_TASK_RESOURCE_PREFIX}")
+      }
+      Map.empty[String, ResourceInformation]
+    }
+    resourceInfo
+  }
+
   def extractLogUrls: Map[String, String] = {
     val prefix = "SPARK_LOG_URL_"
     sys.env.filterKeys(_.startsWith(prefix))
@@ -188,13 +352,14 @@ private[spark] object CoarseGrainedExecutorBackend extends Logging {
       cores: Int,
       appId: String,
       workerUrl: Option[String],
-      userClassPath: mutable.ListBuffer[URL])
+      userClassPath: mutable.ListBuffer[URL],
+      resourcesFile: Option[String])
 
   def main(args: Array[String]): Unit = {
     val createFn: (RpcEnv, Arguments, SparkEnv) =>
       CoarseGrainedExecutorBackend = { case (rpcEnv, arguments, env) =>
       new CoarseGrainedExecutorBackend(rpcEnv, arguments.driverUrl, arguments.executorId,
-        arguments.hostname, arguments.cores, arguments.userClassPath, env)
+        arguments.hostname, arguments.cores, arguments.userClassPath, env, arguments.resourcesFile)
     }
     run(parseArguments(args, this.getClass.getCanonicalName.stripSuffix("$")), createFn)
     System.exit(0)
@@ -255,6 +420,7 @@ private[spark] object CoarseGrainedExecutorBackend extends Logging {
     var executorId: String = null
     var hostname: String = null
     var cores: Int = 0
+    var resourcesFile: Option[String] = None
     var appId: String = null
     var workerUrl: Option[String] = None
     val userClassPath = new mutable.ListBuffer[URL]()
@@ -274,6 +440,9 @@ private[spark] object CoarseGrainedExecutorBackend extends Logging {
         case ("--cores") :: value :: tail =>
           cores = value.toInt
           argv = tail
+        case ("--resourcesFile") :: value :: tail =>
+          resourcesFile = Some(value)
+          argv = tail
         case ("--app-id") :: value :: tail =>
           appId = value
           argv = tail
@@ -299,7 +468,7 @@ private[spark] object CoarseGrainedExecutorBackend extends Logging {
     }
 
     Arguments(driverUrl, executorId, hostname, cores, appId, workerUrl,
-      userClassPath)
+      userClassPath, resourcesFile)
   }
 
   private def printUsageAndExit(classNameForEntry: String): Unit = {
@@ -313,6 +482,7 @@ private[spark] object CoarseGrainedExecutorBackend extends Logging {
       |   --executor-id <executorId>
       |   --hostname <hostname>
       |   --cores <cores>
+      |   --resourcesFile <fileWithJSONResourceInformation>
       |   --app-id <appid>
       |   --worker-url <workerUrl>
       |   --user-class-path <url>