Additional clean-up and fixes on top of existing fix

Author: pwendell

core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/MesosSchedulerBackend.scala

@@ -208,22 +208,25 @@ private[spark] class MesosSchedulerBackend(
    */
   override def resourceOffers(d: SchedulerDriver, offers: JList[Offer]) {
     inClassLoader() {
-      val (acceptedOffers, declinedOffers) = offers.partition { o =>
+      // Fail-fast on offers we know will be rejected
+      val (usableOffers, unUsableOffers) = offers.partition { o =>
         val mem = getResource(o.getResourcesList, "mem")
         val cpus = getResource(o.getResourcesList, "cpus")
         val slaveId = o.getSlaveId.getValue
+        // TODO(pwendell): Should below be 1 + scheduler.CPUS_PER_TASK?
         (mem >= MemoryUtils.calculateTotalMemory(sc) &&
           // need at least 1 for executor, 1 for task
           cpus >= 2 * scheduler.CPUS_PER_TASK) ||
           (slaveIdsWithExecutors.contains(slaveId) &&
             cpus >= scheduler.CPUS_PER_TASK)
       }
 
-      val offerableWorkers = acceptedOffers.map { o =>
+      val workerOffers = usableOffers.map { o =>
         val cpus = if (slaveIdsWithExecutors.contains(o.getSlaveId.getValue)) {
           getResource(o.getResourcesList, "cpus").toInt
         } else {
           // If the executor doesn't exist yet, subtract CPU for executor
+          // TODO(pwendell): Should below just subtract "1"?
           getResource(o.getResourcesList, "cpus").toInt -
             scheduler.CPUS_PER_TASK
         }
@@ -233,17 +236,21 @@ private[spark] class MesosSchedulerBackend(
           cpus)
       }
 
-      val slaveIdToOffer = acceptedOffers.map(o => o.getSlaveId.getValue -> o).toMap
+      val slaveIdToOffer = usableOffers.map(o => o.getSlaveId.getValue -> o).toMap
 
       val mesosTasks = new HashMap[String, JArrayList[MesosTaskInfo]]
 
+
+      val slavesIdsOfAcceptedOffers = HashSet[String]()
+
       // Call into the TaskSchedulerImpl
-      scheduler.resourceOffers(offerableWorkers)
-        .filter(!_.isEmpty)
+      val acceptedOffers = scheduler.resourceOffers(workerOffers).filter(!_.isEmpty)
+      acceptedOffers
         .foreach { offer =>
           offer.foreach { taskDesc =>
             val slaveId = taskDesc.executorId
             slaveIdsWithExecutors += slaveId
+            slavesIdsOfAcceptedOffers += slaveId
             taskIdToSlaveId(taskDesc.taskId) = slaveId
             mesosTasks.getOrElseUpdate(slaveId, new JArrayList[MesosTaskInfo])
               .add(createMesosTask(taskDesc, slaveId))
@@ -257,11 +264,14 @@ private[spark] class MesosSchedulerBackend(
         d.launchTasks(Collections.singleton(slaveIdToOffer(slaveId).getId), tasks, filters)
       }
 
-      for (o <- acceptedOffers if !slaveIdsWithExecutors.contains(o.getSlaveId.getValue)) {
+      // Decline offers that weren't used
+      // NOTE: This logic assumes that we only get a single offer for each host in a given batch
+      for (o <- usableOffers if !slavesIdsOfAcceptedOffers.contains(o.getSlaveId.getValue)) {
         d.declineOffer(o.getId)
       }
 
-      declinedOffers.foreach(o => d.declineOffer(o.getId))
+      // Decline offers we ruled out immediately
+      unUsableOffers.foreach(o => d.declineOffer(o.getId))
     }
   }
 

core/src/test/scala/org/apache/spark/scheduler/mesos/MesosSchedulerBackendSuite.scala

@@ -34,7 +34,8 @@ import scala.collection.mutable.ArrayBuffer
 import scala.util.Random
 
 class MesosSchedulerBackendSuite extends FunSuite with LocalSparkContext with EasyMockSugar {
-  test("mesos resource offer is launching tasks") {
+
+  test("mesos resource offers result in launching tasks") {
     def createOffer(id: Int, mem: Int, cpu: Int) = {
       val builder = Offer.newBuilder()
       builder.addResourcesBuilder()
@@ -53,34 +54,38 @@ class MesosSchedulerBackendSuite extends FunSuite with LocalSparkContext with Ea
     val taskScheduler = EasyMock.createMock(classOf[TaskSchedulerImpl])
 
     val sc = EasyMock.createMock(classOf[SparkContext])
-
     EasyMock.expect(sc.executorMemory).andReturn(100).anyTimes()
     EasyMock.expect(sc.getSparkHome()).andReturn(Option("/path")).anyTimes()
     EasyMock.expect(sc.executorEnvs).andReturn(new mutable.HashMap).anyTimes()
     EasyMock.expect(sc.conf).andReturn(new SparkConf).anyTimes()
     EasyMock.replay(sc)
+
     val minMem = MemoryUtils.calculateTotalMemory(sc).toInt
     val minCpu = 4
+
     val offers = new java.util.ArrayList[Offer]
     offers.add(createOffer(1, minMem, minCpu))
     offers.add(createOffer(2, minMem - 1, minCpu))
     offers.add(createOffer(3, minMem, minCpu))
+
     val backend = new MesosSchedulerBackend(taskScheduler, sc, "master")
-    val workerOffers = new ArrayBuffer[WorkerOffer](2)
-    workerOffers.append(new WorkerOffer(
+
+    val expectedWorkerOffers = new ArrayBuffer[WorkerOffer](2)
+    expectedWorkerOffers.append(new WorkerOffer(
       offers.get(0).getSlaveId.getValue,
       offers.get(0).getHostname,
       2
     ))
-    workerOffers.append(new WorkerOffer(
+    expectedWorkerOffers.append(new WorkerOffer(
       offers.get(2).getSlaveId.getValue,
       offers.get(2).getHostname,
       2
     ))
     val taskDesc = new TaskDescription(1L, "s1", "n1", 0, ByteBuffer.wrap(new Array[Byte](0)))
-    EasyMock.expect(taskScheduler.resourceOffers(EasyMock.eq(workerOffers))).andReturn(Seq(Seq(taskDesc)))
+    EasyMock.expect(taskScheduler.resourceOffers(EasyMock.eq(expectedWorkerOffers))).andReturn(Seq(Seq(taskDesc)))
     EasyMock.expect(taskScheduler.CPUS_PER_TASK).andReturn(2).anyTimes()
     EasyMock.replay(taskScheduler)
+
     val capture = new Capture[util.Collection[TaskInfo]]
     EasyMock.expect(
       driver.launchTasks(
@@ -95,8 +100,10 @@ class MesosSchedulerBackendSuite extends FunSuite with LocalSparkContext with Ea
     backend.resourceOffers(driver, offers)
     EasyMock.verify(driver)
     assert(capture.getValue.size() == 1)
+
     val taskInfo = capture.getValue.iterator().next()
     assert(taskInfo.getName.equals("n1"))
+
     val cpus = taskInfo.getResourcesList.get(0)
     assert(cpus.getName.equals("cpus"))
     assert(cpus.getScalar.getValue.equals(2.0))