[Spark-15155][Mesos] Optionally ignore default role resources

hellertime · hellertime · commit d64b26171c01 · 2016-10-13T08:08:31.000-04:00
Add new boolean property: `spark.mesos.ignoreDefaultRoleResources`.
When this property is set Spark will only accept resources from the role passed in the `spark.mesos.role` property.
If `spark.mesos.role` has not been set, `spark.mesos.ignoreDefaultRoleResources` has no effect.

Additional unit tests added to `MesosSchedulerBackendSuite`, extending the original multi-role test suite.
diff --git a/docs/running-on-mesos.md b/docs/running-on-mesos.md
@@ -428,6 +428,13 @@ See the [configuration page](configuration.html) for information on Spark config
     and resource weight sharing.
   </td>
 </tr>
+<tr>
+  <td><code>spark.mesos.ignoreDefaultRoleResources</code></td>
+  <td>false</td>
+  <td>
+    Only if `spark.mesos.role` has been set, ignore mesos resources with the role `*`.
+  </td>
+</tr>
 <tr>
   <td><code>spark.mesos.constraints</code></td>
   <td>(none)</td>
diff --git a/mesos/src/main/scala/org/apache/spark/deploy/mesos/MesosDriverDescription.scala b/mesos/src/main/scala/org/apache/spark/deploy/mesos/MesosDriverDescription.scala
@@ -41,7 +41,7 @@ private[spark] class MesosDriverDescription(
     val cores: Double,
     val supervise: Boolean,
     val command: Command,
-    schedulerProperties: Map[String, String],
+    val schedulerProperties: Map[String, String],
     val submissionId: String,
     val submissionDate: Date,
     val retryState: Option[MesosClusterRetryState] = None)
diff --git a/mesos/src/main/scala/org/apache/spark/scheduler/cluster/mesos/MesosClusterScheduler.scala b/mesos/src/main/scala/org/apache/spark/scheduler/cluster/mesos/MesosClusterScheduler.scala
@@ -551,12 +551,16 @@ private[spark] class MesosClusterScheduler(
       currentOffers: List[ResourceOffer],
       tasks: mutable.HashMap[OfferID, ArrayBuffer[TaskInfo]]): Unit = {
     for (submission <- candidates) {
+      val acceptedResourceRoles = getAcceptedResourceRoles(submission.schedulerProperties)
       val driverCpu = submission.cores
       val driverMem = submission.mem
       logTrace(s"Finding offer to launch driver with cpu: $driverCpu, mem: $driverMem")
       val offerOption = currentOffers.find { o =>
-        getResource(o.resources, "cpus") >= driverCpu &&
-        getResource(o.resources, "mem") >= driverMem
+        val acceptableResources = o.resources.asScala
+          .filter((r: Resource) => acceptedResourceRoles(r.getRole))
+          .asJava
+        getResource(acceptableResources, "cpus") >= driverCpu &&
+        getResource(acceptableResources, "mem") >= driverMem
       }
       if (offerOption.isEmpty) {
         logDebug(s"Unable to find offer to launch driver id: ${submission.submissionId}, " +
diff --git a/mesos/src/main/scala/org/apache/spark/scheduler/cluster/mesos/MesosCoarseGrainedSchedulerBackend.scala b/mesos/src/main/scala/org/apache/spark/scheduler/cluster/mesos/MesosCoarseGrainedSchedulerBackend.scala
@@ -382,6 +382,8 @@ private[spark] class MesosCoarseGrainedSchedulerBackend(
     val remainingResources = mutable.Map(offers.map(offer =>
       (offer.getId.getValue, offer.getResourcesList)): _*)
 
+    val acceptedResourceRoles = getAcceptedResourceRoles(sc.conf)
+
     var launchTasks = true
 
     // TODO(mgummelt): combine offers for a single slave
@@ -393,15 +395,19 @@ private[spark] class MesosCoarseGrainedSchedulerBackend(
       for (offer <- offers) {
         val slaveId = offer.getSlaveId.getValue
         val offerId = offer.getId.getValue
-        val resources = remainingResources(offerId)
+        val resources =
+          remainingResources(offerId).asScala
+            .filter((r: Resource) => acceptedResourceRoles(r.getRole))
+            .asJava
 
-        if (canLaunchTask(slaveId, resources)) {
+        if (canLaunchTask(slaveId, resources, acceptedResourceRoles)) {
           // Create a task
           launchTasks = true
           val taskId = newMesosTaskId()
           val offerCPUs = getResource(resources, "cpus").toInt
           val taskGPUs = Math.min(
-            Math.max(0, maxGpus - totalGpusAcquired), getResource(resources, "gpus").toInt)
+            Math.max(0, maxGpus - totalGpusAcquired),
+            getResource(resources, "gpus").toInt)
 
           val taskCPUs = executorCores(offerCPUs)
           val taskMemory = executorMemory(sc)
@@ -466,7 +472,10 @@ private[spark] class MesosCoarseGrainedSchedulerBackend(
       cpuResourcesToUse ++ memResourcesToUse ++ portResourcesToUse ++ gpuResourcesToUse)
   }
 
-  private def canLaunchTask(slaveId: String, resources: JList[Resource]): Boolean = {
+  private def canLaunchTask(
+      slaveId: String,
+      resources: JList[Resource],
+      acceptedResourceRoles: Set[String]): Boolean = {
     val offerMem = getResource(resources, "mem")
     val offerCPUs = getResource(resources, "cpus").toInt
     val cpus = executorCores(offerCPUs)
diff --git a/mesos/src/main/scala/org/apache/spark/scheduler/cluster/mesos/MesosFineGrainedSchedulerBackend.scala b/mesos/src/main/scala/org/apache/spark/scheduler/cluster/mesos/MesosFineGrainedSchedulerBackend.scala
@@ -249,10 +249,15 @@ private[spark] class MesosFineGrainedSchedulerBackend(
           .setRefuseSeconds(rejectOfferDurationForUnmetConstraints).build())
       }
 
+      val acceptedResourceRoles = getAcceptedResourceRoles(sc.conf)
+
       // Of the matching constraints, see which ones give us enough memory and cores
       val (usableOffers, unUsableOffers) = offersMatchingConstraints.partition { o =>
-        val mem = getResource(o.getResourcesList, "mem")
-        val cpus = getResource(o.getResourcesList, "cpus")
+        val acceptableResources = o.getResourcesList.asScala
+          .filter((r: Resource) => acceptedResourceRoles(r.getRole))
+          .asJava
+        val mem = getResource(acceptableResources, "mem")
+        val cpus = getResource(acceptableResources, "cpus")
         val slaveId = o.getSlaveId.getValue
         val offerAttributes = toAttributeMap(o.getAttributesList)
 
diff --git a/mesos/src/main/scala/org/apache/spark/scheduler/cluster/mesos/MesosSchedulerUtils.scala b/mesos/src/main/scala/org/apache/spark/scheduler/cluster/mesos/MesosSchedulerUtils.scala
@@ -49,6 +49,44 @@ trait MesosSchedulerUtils extends Logging {
   // Driver for talking to Mesos
   protected var mesosDriver: SchedulerDriver = null
 
+  /**
+   * Returns the configured set of roles that an offer can be selected from
+   * @param conf Spark configuration
+   */
+  protected def getAcceptedResourceRoles(conf: SparkConf): Set[String] = {
+    getAcceptedResourceRoles(
+      conf.getBoolean("spark.mesos.ignoreDefaultRoleResources", false),
+      conf.getOption("spark.mesos.role"))
+  }
+  /**
+   * Returns the configured set of roles that an offer can be selected from
+   * @param props Mesos driver description schedulerProperties map
+   */
+  protected def getAcceptedResourceRoles(props: Map[String, String]): Set[String] = {
+    getAcceptedResourceRoles(
+      props.get("spark.mesos.ignoreDefaultRoleResources") match {
+        case Some(truth) => truth.toBoolean
+        case None => false
+      },
+      props.get("spark.mesos.role"))
+  }
+  /**
+   * Internal version of getAcceptedResourceRoles
+   * @param ignoreDefaultRoleResources user specified property
+   * @param role user specified property
+   */
+  private def getAcceptedResourceRoles(
+      ignoreDefaultRoleResources: Boolean,
+      role: Option[String]) = {
+    val roles = ignoreDefaultRoleResources match {
+      case true if role.isDefined => Set(role)
+      case _ => Set(Some("*"), role)
+    }
+    val acceptedRoles = roles.flatten
+    logDebug(s"Accepting resources from role(s): ${acceptedRoles.mkString(",")}")
+    acceptedRoles
+  }
+
   /**
    * Creates a new MesosSchedulerDriver that communicates to the Mesos master.
    *
diff --git a/mesos/src/test/scala/org/apache/spark/scheduler/cluster/mesos/MesosClusterSchedulerSuite.scala b/mesos/src/test/scala/org/apache/spark/scheduler/cluster/mesos/MesosClusterSchedulerSuite.scala
diff --git a/mesos/src/test/scala/org/apache/spark/scheduler/cluster/mesos/MesosFineGrainedSchedulerBackendSuite.scala b/mesos/src/test/scala/org/apache/spark/scheduler/cluster/mesos/MesosFineGrainedSchedulerBackendSuite.scala
