[SPARK-13704][CORE][YARN] Re-implement RackResolver to reduce resolving time

core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala

@@ -375,9 +375,10 @@ private[spark] class TaskSchedulerImpl(
         executorIdToRunningTaskIds(o.executorId) = HashSet[Long]()
         newExecAvail = true
       }
-      for (rack <- getRackForHost(o.host)) {
-        hostsByRack.getOrElseUpdate(rack, new HashSet[String]()) += o.host
-      }
+    }
+    val hosts = offers.map(_.host)
+    for ((host, rack) <- hosts.zip(getRacksForHosts(hosts))) {
+      hostsByRack.getOrElseUpdate(rack, new HashSet[String]()) += host
     }
 
     // Before making any offers, remove any nodes from the blacklist whose blacklist has expired. Do
@@ -764,7 +765,8 @@ private[spark] class TaskSchedulerImpl(
     execs -= executorId
     if (execs.isEmpty) {
       hostToExecutors -= host
-      for (rack <- getRackForHost(host); hosts <- hostsByRack.get(rack)) {
+      val rack = getRackForHost(host)
+      for (hosts <- hostsByRack.get(rack)) {
         hosts -= host
         if (hosts.isEmpty) {
           hostsByRack -= rack
@@ -811,8 +813,39 @@ private[spark] class TaskSchedulerImpl(
     blacklistTrackerOpt.map(_.nodeBlacklist()).getOrElse(scala.collection.immutable.Set())
   }
 
-  // By default, rack is unknown
-  def getRackForHost(value: String): Option[String] = None
+  // Add a on-off switch to save time for rack resolving
+  private def skipRackResolving: Boolean = sc.conf.get(LOCALITY_WAIT_RACK) == 0L
+
+  /**
+   * Rack is `unknown` by default.
+   * It can be override in different TaskScheduler, like Yarn.
+   */
+  protected val defaultRackValue: String = "unknown"
+
+  /**
+   * Get racks info for a host list. This is the internal method of [[getRacksForHosts]].
+   * It should be override in different TaskScheduler. Return [[Nil]] by default.
+   */
+  protected def doGetRacksForHosts(hosts: Seq[String]): Seq[String] = Nil
+
+  def getRackForHost(hosts: String): String = {
+    if (skipRackResolving) {
+      defaultRackValue
+    } else {
+      doGetRacksForHosts(Seq(hosts)).head
+    }
+  }
+
+  /**
+   * null in return sequences will be replaced to [[defaultRackValue]].
+   */
+  def getRacksForHosts(hosts: Seq[String]): Seq[String] = {
+    if (skipRackResolving) {
+      hosts.map(_ => defaultRackValue)
+    } else {
+      doGetRacksForHosts(hosts).map(Option(_).getOrElse(defaultRackValue))
+    }
+  }
 
   private def waitBackendReady(): Unit = {
     if (backend.isReady) {

core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala

@@ -186,8 +186,23 @@ private[spark] class TaskSetManager(
 
   // Add all our tasks to the pending lists. We do this in reverse order
   // of task index so that tasks with low indices get launched first.
-  for (i <- (0 until numTasks).reverse) {
-    addPendingTask(i)
+  addPendingTasks()
+
+  private def addPendingTasks(): Unit = {
+    val (_, duration) = Utils.timeTakenMs {
+      val hostToIndices = new HashMap[String, ArrayBuffer[Int]]()
+      for (i <- (0 until numTasks).reverse) {
+        addPendingTask(i, Some(hostToIndices))
+      }
+      // Resolve the rack for each host. This can be somehow slow, so de-dupe the list of hosts,
+      // and assign the rack to all relevant task indices.
+      for (
+        (rack, indices) <- sched.getRacksForHosts(hostToIndices.keySet.toSeq)
+          .zip(hostToIndices.values)) {
+        pendingTasksForRack.getOrElseUpdate(rack, new ArrayBuffer) ++= indices
+      }
+    }
+    logDebug(s"Adding pending tasks take $duration ms")
   }
 
   /**
@@ -214,7 +229,9 @@ private[spark] class TaskSetManager(
   private[scheduler] var emittedTaskSizeWarning = false
 
   /** Add a task to all the pending-task lists that it should be on. */
-  private[spark] def addPendingTask(index: Int) {
+  private[spark] def addPendingTask(
+      index: Int,
+      initializingHostToIndices: Option[HashMap[String, ArrayBuffer[Int]]] = None) {
     for (loc <- tasks(index).preferredLocations) {
       loc match {
         case e: ExecutorCacheTaskLocation =>
@@ -234,8 +251,14 @@ private[spark] class TaskSetManager(
         case _ =>
       }
       pendingTasksForHost.getOrElseUpdate(loc.host, new ArrayBuffer) += index
-      for (rack <- sched.getRackForHost(loc.host)) {
-        pendingTasksForRack.getOrElseUpdate(rack, new ArrayBuffer) += index
+
+      initializingHostToIndices match {
+        case Some(hostToIndices) =>
+          // when TaskSetManager initializing, preferredLocation -> task indices
+          hostToIndices.getOrElseUpdate(loc.host, new ArrayBuffer) += index
+        case None =>
+          pendingTasksForRack.getOrElseUpdate(
+            sched.getRackForHost(loc.host), new ArrayBuffer) += index
       }
     }
 
@@ -249,24 +272,27 @@ private[spark] class TaskSetManager(
   /**
    * Return the pending tasks list for a given executor ID, or an empty list if
    * there is no map entry for that host
+   * This is visible for testing.
    */
-  private def getPendingTasksForExecutor(executorId: String): ArrayBuffer[Int] = {
+  private[scheduler] def getPendingTasksForExecutor(executorId: String): ArrayBuffer[Int] = {
     pendingTasksForExecutor.getOrElse(executorId, ArrayBuffer())
   }
 
   /**
    * Return the pending tasks list for a given host, or an empty list if
    * there is no map entry for that host
+   * This is visible for testing.
    */
-  private def getPendingTasksForHost(host: String): ArrayBuffer[Int] = {
+  private[scheduler] def getPendingTasksForHost(host: String): ArrayBuffer[Int] = {
     pendingTasksForHost.getOrElse(host, ArrayBuffer())
   }
 
   /**
    * Return the pending rack-local task list for a given rack, or an empty list if
    * there is no map entry for that rack
+   * This is visible for testing.
    */
-  private def getPendingTasksForRack(rack: String): ArrayBuffer[Int] = {
+  private[scheduler] def getPendingTasksForRack(rack: String): ArrayBuffer[Int] = {
     pendingTasksForRack.getOrElse(rack, ArrayBuffer())
   }
 
@@ -362,13 +388,12 @@ private[spark] class TaskSetManager(
 
       // Check for rack-local tasks
       if (TaskLocality.isAllowed(locality, TaskLocality.RACK_LOCAL)) {
-        for (rack <- sched.getRackForHost(host)) {
-          for (index <- speculatableTasks if canRunOnHost(index)) {
-            val racks = tasks(index).preferredLocations.map(_.host).flatMap(sched.getRackForHost)
-            if (racks.contains(rack)) {
-              speculatableTasks -= index
-              return Some((index, TaskLocality.RACK_LOCAL))
-            }
+        val rack = sched.getRackForHost(host)
+        for (index <- speculatableTasks if canRunOnHost(index)) {
+          val racks = tasks(index).preferredLocations.map(_.host).map(sched.getRackForHost)
+          if (racks.contains(rack)) {
+            speculatableTasks -= index
+            return Some((index, TaskLocality.RACK_LOCAL))
           }
         }
       }
@@ -412,10 +437,8 @@ private[spark] class TaskSetManager(
     }
 
     if (TaskLocality.isAllowed(maxLocality, TaskLocality.RACK_LOCAL)) {
-      for {
-        rack <- sched.getRackForHost(host)
-        index <- dequeueTaskFromList(execId, host, getPendingTasksForRack(rack))
-      } {
+      val rack = sched.getRackForHost(host)
+      for (index <- dequeueTaskFromList(execId, host, getPendingTasksForRack(rack))) {
         return Some((index, TaskLocality.RACK_LOCAL, false))
       }
     }

core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala

@@ -23,7 +23,7 @@ import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
 
 import org.mockito.ArgumentMatchers.{any, anyInt, anyString}
-import org.mockito.Mockito.{mock, never, spy, times, verify, when}
+import org.mockito.Mockito._
 import org.mockito.invocation.InvocationOnMock
 import org.mockito.stubbing.Answer
 
@@ -69,17 +69,23 @@ class FakeDAGScheduler(sc: SparkContext, taskScheduler: FakeTaskScheduler)
 // Get the rack for a given host
 object FakeRackUtil {
   private val hostToRack = new mutable.HashMap[String, String]()
+  var numBatchInvocation = 0
 
   def cleanUp() {
     hostToRack.clear()
+    numBatchInvocation = 0
   }
 
   def assignHostToRack(host: String, rack: String) {
     hostToRack(host) = rack
   }
 
-  def getRackForHost(host: String): Option[String] = {
-    hostToRack.get(host)
+  def getRacksForHosts(hosts: Seq[String]): Seq[String] = {
+    assert(hosts.toSet.size == hosts.size) // no dups in hosts
+    if (hosts.nonEmpty && hosts.length != 1) {
+      numBatchInvocation += 1
+    }
+    hosts.map(hostToRack.getOrElse(_, null))
   }
 }
 
@@ -99,12 +105,13 @@ class FakeTaskScheduler(sc: SparkContext, liveExecutors: (String, String)* /* ex
   val speculativeTasks = new ArrayBuffer[Int]
 
   val executors = new mutable.HashMap[String, String]
+
   for ((execId, host) <- liveExecutors) {
     addExecutor(execId, host)
   }
 
-  for ((execId, host) <- liveExecutors; rack <- getRackForHost(host)) {
-    hostsByRack.getOrElseUpdate(rack, new mutable.HashSet[String]()) += host
+  for ((execId, host) <- liveExecutors) {
+    hostsByRack.getOrElseUpdate(getRackForHost(host), new mutable.HashSet[String]()) += host
   }
 
   dagScheduler = new FakeDAGScheduler(sc, this)
@@ -116,7 +123,8 @@ class FakeTaskScheduler(sc: SparkContext, liveExecutors: (String, String)* /* ex
     val hostId = host.get
     val executorsOnHost = hostToExecutors(hostId)
     executorsOnHost -= execId
-    for (rack <- getRackForHost(hostId); hosts <- hostsByRack.get(rack)) {
+    val rack = getRackForHost(hostId)
+    for (hosts <- hostsByRack.get(rack)) {
       hosts -= hostId
       if (hosts.isEmpty) {
         hostsByRack -= rack
@@ -139,13 +147,14 @@ class FakeTaskScheduler(sc: SparkContext, liveExecutors: (String, String)* /* ex
     val executorsOnHost = hostToExecutors.getOrElseUpdate(host, new mutable.HashSet[String])
     executorsOnHost += execId
     executorIdToHost += execId -> host
-    for (rack <- getRackForHost(host)) {
-      hostsByRack.getOrElseUpdate(rack, new mutable.HashSet[String]()) += host
-    }
+    hostsByRack.getOrElseUpdate(getRackForHost(host), new mutable.HashSet[String]()) += host
   }
 
+  override val defaultRackValue: String = "default"
 
-  override def getRackForHost(value: String): Option[String] = FakeRackUtil.getRackForHost(value)
+  override def doGetRacksForHosts(hosts: Seq[String]): Seq[String] = {
+    FakeRackUtil.getRacksForHosts(hosts)
+  }
 }
 
 /**
@@ -1316,7 +1325,7 @@ class TaskSetManagerSuite extends SparkFunSuite with LocalSparkContext with Logg
     val taskDesc = taskSetManagerSpy.resourceOffer(exec, host, TaskLocality.ANY)
 
     // Assert the task has been black listed on the executor it was last executed on.
-    when(taskSetManagerSpy.addPendingTask(anyInt())).thenAnswer(
+    when(taskSetManagerSpy.addPendingTask(anyInt(), any())).thenAnswer(
       new Answer[Unit] {
         override def answer(invocationOnMock: InvocationOnMock): Unit = {
           val task: Int = invocationOnMock.getArgument(0)
@@ -1330,7 +1339,7 @@ class TaskSetManagerSuite extends SparkFunSuite with LocalSparkContext with Logg
     val e = new ExceptionFailure("a", "b", Array(), "c", None)
     taskSetManagerSpy.handleFailedTask(taskDesc.get.taskId, TaskState.FAILED, e)
 
-    verify(taskSetManagerSpy, times(1)).addPendingTask(anyInt())
+    verify(taskSetManagerSpy, times(1)).addPendingTask(anyInt(), any())
   }
 
   test("SPARK-21563 context's added jars shouldn't change mid-TaskSet") {
@@ -1602,4 +1611,60 @@ class TaskSetManagerSuite extends SparkFunSuite with LocalSparkContext with Logg
     verify(sched.dagScheduler).taskEnded(manager.tasks(3), Success, result.value(),
       result.accumUpdates, info3)
   }
+
+  test("SPARK-27038 Rack Resolution is done with a batch of de-duped hosts") {
+    val conf = new SparkConf()
+      .set(config.LOCALITY_WAIT.key, "0")
+      .set(config.LOCALITY_WAIT_RACK.key, "1s")
+    sc = new SparkContext("local", "test", conf)
+    // Create a cluster with 20 racks, with hosts spread out among them
+    val execAndHost = (0 to 199).map { i =>
+      FakeRackUtil.assignHostToRack("host" + i, "rack" + (i % 20))
+      ("exec" + i, "host" + i)
+    }
+    sched = new FakeTaskScheduler(sc, execAndHost: _*)
+    // make a taskset with preferred locations on the first 100 hosts in our cluster
+    val locations = new ArrayBuffer[Seq[TaskLocation]]()
+    for (i <- 0 to 99) {
+      locations += Seq(TaskLocation("host" + i))
+    }
+    val taskSet = FakeTask.createTaskSet(100, locations: _*)
+    val clock = new ManualClock
+    val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES, clock = clock)
+    // with rack locality, reject an offer on a host with an unknown rack
+    assert(manager.resourceOffer("otherExec", "otherHost", TaskLocality.RACK_LOCAL).isEmpty)
+    (0 until 20).foreach { rackIdx =>
+      (0 until 5).foreach { offerIdx =>
+        // if we offer hosts which are not in preferred locations,
+        // we'll reject them at NODE_LOCAL level,
+        // but accept them at RACK_LOCAL level if they're on OK racks
+        val hostIdx = 100 + rackIdx
+        assert(manager.resourceOffer("exec" + hostIdx, "host" + hostIdx, TaskLocality.NODE_LOCAL)
+          .isEmpty)
+        assert(manager.resourceOffer("exec" + hostIdx, "host" + hostIdx, TaskLocality.RACK_LOCAL)
+          .isDefined)
+      }
+    }
+    assert(FakeRackUtil.numBatchInvocation === 1)
+  }
+
+  test("SPARK-27038 Rack resolving is skipped when spark.locality.wait is zero") {
+    val conf = new SparkConf().set(config.LOCALITY_WAIT.key, "0")
+    sc = new SparkContext("local", "test", conf)
+    for (i <- 0 to 99) {
+      FakeRackUtil.assignHostToRack("host" + i, "rack" + (i % 20))
+    }
+    sched = new FakeTaskScheduler(sc,
+      ("execA", "host1"), ("execB", "host2"), ("execC", "host3"))
+    val locations = new ArrayBuffer[Seq[TaskLocation]]()
+    for (i <- 0 to 99) {
+      locations += Seq(TaskLocation("host" + i))
+    }
+    val taskSet = FakeTask.createTaskSet(100, locations: _*)
+    val clock = new ManualClock
+    val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES, clock = clock)
+    // verify the total number not changed with SPARK-27038
+    assert(manager.getPendingTasksForRack(sched.defaultRackValue).length === 100)
+    assert(FakeRackUtil.numBatchInvocation === 0)
+  }
 }

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

@@ -138,9 +138,8 @@ private[yarn] class LocalityPreferredContainerPlacementStrategy(
         // Only filter out the ratio which is larger than 0, which means the current host can
         // still be allocated with new container request.
         val hosts = preferredLocalityRatio.filter(_._2 > 0).keys.toArray
-        val racks = hosts.map { h =>
-          resolver.resolve(yarnConf, h)
-        }.toSet
+        val racks = resolver.resolve(yarnConf, hosts).map(_.getNetworkLocation)
+          .filter(_ != null).toSet
         containerLocalityPreferences += ContainerLocalityPreferences(hosts, racks.toArray)
 
         // Minus 1 each time when the host is used. When the current ratio is 0,