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

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

@@ -814,6 +814,8 @@ private[spark] class TaskSchedulerImpl(
   // By default, rack is unknown
   def getRackForHost(value: String): Option[String] = None
 
+  def getRacksForHosts(values: List[String]): List[Option[String]] = values.map(getRackForHost)
+
   private def waitBackendReady(): Unit = {
     if (backend.isReady) {
       return

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

@@ -184,11 +184,23 @@ private[spark] class TaskSetManager(
     t.epoch = epoch
   }
 
+  // An array to store preferred location and its task index
+  private val locationWithTaskIndex: ArrayBuffer[(String, Int)] = new ArrayBuffer[(String, Int)]()
+  private val addTaskStartTime = System.nanoTime()
   // 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)
+    addPendingTask(i, true)
   }
+  // Convert preferred location list to rack list in one invocation and zip with the origin index
+  private val rackWithTaskIndex = sched.getRacksForHosts(locationWithTaskIndex.map(_._1).toList)
+    .zip(locationWithTaskIndex.map(_._2))
+  rackWithTaskIndex.foreach {
+    case (Some(rack), index) =>
+      pendingTasksForRack.getOrElseUpdate(rack, new ArrayBuffer) += index
+    case _ =>
+  }
+  logInfo(s"Adding pending tasks take ${(System.nanoTime() - addTaskStartTime) / 1e9} seconds")
 
   /**
    * Track the set of locality levels which are valid given the tasks locality preferences and
@@ -214,7 +226,7 @@ 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, initializing: Boolean = false) {
     for (loc <- tasks(index).preferredLocations) {
       loc match {
         case e: ExecutorCacheTaskLocation =>
@@ -234,8 +246,13 @@ private[spark] class TaskSetManager(
         case _ =>
       }
       pendingTasksForHost.getOrElseUpdate(loc.host, new ArrayBuffer) += index
-      for (rack <- sched.getRackForHost(loc.host)) {
-        pendingTasksForRack.getOrElseUpdate(rack, new ArrayBuffer) += index
+
+      if (initializing) {
+        locationWithTaskIndex += ((loc.host, index))
+      } else {
+        for (rack <- sched.getRackForHost(loc.host)) {
+          pendingTasksForRack.getOrElseUpdate(rack, new ArrayBuffer) += index
+        }
       }
     }
 
@@ -249,24 +266,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())
   }
 

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

@@ -22,7 +22,7 @@ import java.util.{Properties, Random}
 import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
 
-import org.mockito.ArgumentMatchers.{any, anyInt, anyString}
+import org.mockito.ArgumentMatchers.{any, anyBoolean, anyInt, anyString}
 import org.mockito.Mockito.{mock, never, spy, times, verify, when}
 import org.mockito.invocation.InvocationOnMock
 import org.mockito.stubbing.Answer
@@ -69,18 +69,48 @@ class FakeDAGScheduler(sc: SparkContext, taskScheduler: FakeTaskScheduler)
 // Get the rack for a given host
 object FakeRackUtil {
   private val hostToRack = new mutable.HashMap[String, String]()
+  var loopCount = 0
 
   def cleanUp() {
     hostToRack.clear()
+    loopCount = 0
   }
 
   def assignHostToRack(host: String, rack: String) {
     hostToRack(host) = rack
   }
 
   def getRackForHost(host: String): Option[String] = {
+    loopCount = simulateRunResolveCommand(Seq(host))
     hostToRack.get(host)
   }
+
+  def getRacksForHosts(hosts: List[String]): List[Option[String]] = {
+    loopCount = simulateRunResolveCommand(hosts)
+    hosts.map(hostToRack.get)
+  }
+
+  /**
+   * This is a simulation of building and executing the resolution command.
+   * Simulate function `runResolveCommand()` in [[org.apache.hadoop.net.ScriptBasedMapping]].
+   * If Seq has 100 elements, it returns 4. If Seq has 1 elements, it returns 1.
+   * @param args a list of arguments
+   * @return script execution times
+   */
+  private def simulateRunResolveCommand(args: Seq[String]): Int = {
+    val maxArgs = 30 // Simulate NET_TOPOLOGY_SCRIPT_NUMBER_ARGS_DEFAULT
+    var numProcessed = 0
+    var loopCount = 0
+    while (numProcessed != args.size) {
+      var start = maxArgs * loopCount
+      numProcessed = start
+      while (numProcessed < (start + maxArgs) && numProcessed < args.size) {
+        numProcessed += 1
+      }
+      loopCount += 1
+    }
+    loopCount
+  }
 }
 
 /**
@@ -97,6 +127,7 @@ class FakeTaskScheduler(sc: SparkContext, liveExecutors: (String, String)* /* ex
   val finishedManagers = new ArrayBuffer[TaskSetManager]
   val taskSetsFailed = new ArrayBuffer[String]
   val speculativeTasks = new ArrayBuffer[Int]
+  var slowRackResolve = false
 
   val executors = new mutable.HashMap[String, String]
   for ((execId, host) <- liveExecutors) {
@@ -145,7 +176,11 @@ class FakeTaskScheduler(sc: SparkContext, liveExecutors: (String, String)* /* ex
   }
 
 
-  override def getRackForHost(value: String): Option[String] = FakeRackUtil.getRackForHost(value)
+  override def getRackForHost(value: String): Option[String] =
+    FakeRackUtil.getRackForHost(value)
+
+  override def getRacksForHosts(values: List[String]): List[Option[String]] =
+    FakeRackUtil.getRacksForHosts(values)
 }
 
 /**
@@ -1316,7 +1351,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(), anyBoolean())).thenAnswer(
       new Answer[Unit] {
         override def answer(invocationOnMock: InvocationOnMock): Unit = {
           val task: Int = invocationOnMock.getArgument(0)
@@ -1330,7 +1365,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(), anyBoolean())
   }
 
   test("SPARK-21563 context's added jars shouldn't change mid-TaskSet") {
@@ -1602,4 +1637,27 @@ class TaskSetManagerSuite extends SparkFunSuite with LocalSparkContext with Logg
     verify(sched.dagScheduler).taskEnded(manager.tasks(3), Success, result.value(),
       result.accumUpdates, info3)
   }
+
+  test("SPARK-27038: Verify the rack resolving time has been reduced") {
+    sc = new SparkContext("local", "test")
+    for (i <- 1 to 100) {
+      FakeRackUtil.assignHostToRack("host" + i, "rack" + i)
+    }
+    sched = new FakeTaskScheduler(sc,
+      ("execA", "host1"), ("execB", "host2"), ("execC", "host3"))
+    sched.slowRackResolve = true
+    val locations = new ArrayBuffer[Seq[TaskLocation]]()
+    for (i <- 1 to 100) {
+      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)
+    var total = 0
+    for (i <- 1 to 100) {
+      total += manager.getPendingTasksForRack("rack" + i).length
+    }
+    assert(total === 100) // verify the total number not changed with SPARK-27038
+    assert(FakeRackUtil.loopCount == 4) // verify script execution loop count decreased
+  }
 }

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

@@ -17,10 +17,18 @@
 
 package org.apache.spark.deploy.yarn
 
+import scala.collection.mutable.ArrayBuffer
+
+import com.google.common.base.Strings
 import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.fs.CommonConfigurationKeysPublic
+import org.apache.hadoop.net._
+import org.apache.hadoop.util.ReflectionUtils
 import org.apache.hadoop.yarn.util.RackResolver
 import org.apache.log4j.{Level, Logger}
 
+import org.apache.spark.internal.Logging
+
 /**
  * Wrapper around YARN's [[RackResolver]]. This allows Spark tests to easily override the
  * default behavior, since YARN's class self-initializes the first time it's called, and
@@ -38,3 +46,72 @@ private[yarn] class SparkRackResolver {
   }
 
 }
+
+/**
+ * Utility to resolve the rack for hosts in an efficient manner.
+ * It will cache the rack for individual hosts to avoid
+ * repeatedly performing the same expensive lookup.
+ *
+ * This will be unnecessary in hadoop releases with YARN-9332.
+ * With that, we could just directly use [[org.apache.hadoop.yarn.util.RackResolver]].
+ * In the meantime, this is a re-implementation for spark's use.
+ */
+object SparkRackResolver extends Logging {
+  private var dnsToSwitchMapping: DNSToSwitchMapping = _
+  private var initCalled = false
+  // advisory count of arguments for rack script, less than this will get a warning
+  private val ADVISORY_MINIMUM_NUMBER_SCRIPT_ARGS = 10000
+
+  def init(conf: Configuration): Unit = {
+    if (!initCalled) {
+      initCalled = true
+      val dnsToSwitchMappingClass =
+        conf.getClass(CommonConfigurationKeysPublic.NET_TOPOLOGY_NODE_SWITCH_MAPPING_IMPL_KEY,
+          classOf[ScriptBasedMapping], classOf[DNSToSwitchMapping])
+      if (classOf[ScriptBasedMapping].isAssignableFrom(dnsToSwitchMappingClass)) {
+        val numArgs = conf.getInt(CommonConfigurationKeysPublic.NET_TOPOLOGY_SCRIPT_NUMBER_ARGS_KEY,
+          CommonConfigurationKeysPublic.NET_TOPOLOGY_SCRIPT_NUMBER_ARGS_DEFAULT)
+        if (numArgs < ADVISORY_MINIMUM_NUMBER_SCRIPT_ARGS) {
+          logWarning(s"Increasing the value of" +
+            s" ${CommonConfigurationKeysPublic.NET_TOPOLOGY_SCRIPT_NUMBER_ARGS_KEY} could reduce" +
+            s" the time of rack resolving when submits a stage with a mass of tasks." +
+            s" Current number is $numArgs")
+        }
+      }
+      try {
+        val newInstance = ReflectionUtils.newInstance(dnsToSwitchMappingClass, conf)
+          .asInstanceOf[DNSToSwitchMapping]
+        dnsToSwitchMapping = newInstance match {
+          case _: CachedDNSToSwitchMapping => newInstance
+          case _ => new CachedDNSToSwitchMapping(newInstance)
+        }
+      } catch {
+        case e: Exception =>
+          throw new RuntimeException(e)
+      }
+    }
+  }
+
+  def resolveRacks(conf: Configuration, hostNames: List[String]): List[Node] = {
+    init(conf)
+    val nodes = new ArrayBuffer[Node]
+    val rNameList = dnsToSwitchMapping.resolve(hostNames.toList.asJava).asScala
+    if (rNameList == null || rNameList.isEmpty) {
+      hostNames.foreach(nodes += new NodeBase(_, NetworkTopology.DEFAULT_RACK))
+      logInfo(s"Got an error when resolve hostNames. " +
+        s"Falling back to ${NetworkTopology.DEFAULT_RACK} for all")
+    } else {
+      for ((hostName, rName) <- hostNames.zip(rNameList)) {
+        if (Strings.isNullOrEmpty(rName)) {
+          // fallback to use default rack
+          nodes += new NodeBase(hostName, NetworkTopology.DEFAULT_RACK)
+          logDebug(s"Could not resolve $hostName. " +
+            s"Falling back to ${NetworkTopology.DEFAULT_RACK}")
+        } else {
+          nodes += new NodeBase(hostName, rName)
+        }
+      }
+    }
+    nodes.toList
+  }
+}

resource-managers/yarn/src/main/scala/org/apache/spark/scheduler/cluster/YarnScheduler.scala

@@ -17,10 +17,13 @@
 
 package org.apache.spark.scheduler.cluster
 
+import org.apache.hadoop.net.NetworkTopology
 import org.apache.hadoop.yarn.util.RackResolver
 import org.apache.log4j.{Level, Logger}
 
 import org.apache.spark._
+import org.apache.spark.deploy.yarn.SparkRackResolver
+import org.apache.spark.internal.config.LOCALITY_WAIT
 import org.apache.spark.scheduler.TaskSchedulerImpl
 import org.apache.spark.util.Utils
 
@@ -31,9 +34,35 @@ private[spark] class YarnScheduler(sc: SparkContext) extends TaskSchedulerImpl(s
     Logger.getLogger(classOf[RackResolver]).setLevel(Level.WARN)
   }
 
+  // Add a on-off switch to save time for rack resolving
+  private val skipRackResolving = sc.conf.getTimeAsMs(
+    "spark.locality.wait.rack", sc.conf.get(LOCALITY_WAIT).toString) == 0
+
   // By default, rack is unknown
   override def getRackForHost(hostPort: String): Option[String] = {
-    val host = Utils.parseHostPort(hostPort)._1
-    Option(RackResolver.resolve(sc.hadoopConfiguration, host).getNetworkLocation)
+    if (skipRackResolving) {
+      Option(NetworkTopology.DEFAULT_RACK)
+    } else {
+      val host = Utils.parseHostPort(hostPort)._1
+      Option(RackResolver.resolve(sc.hadoopConfiguration, host).getNetworkLocation)
+    }
+  }
+
+  /**
+   * Get racks info for a list of host.
+   * Use [[SparkRackResolver]] to resolve racks instead of [[RackResolver]]
+   * before the dependency Hadoop applied YARN-9332
+   * @param hostPorts host list to resolve
+   * @return rack list
+   */
+  override def getRacksForHosts(hostPorts: List[String]): List[Option[String]] = {
+    val hosts = hostPorts.map(Utils.parseHostPort(_)._1)
+    if (skipRackResolving) {
+      hosts.map(_ => Option(NetworkTopology.DEFAULT_RACK))
+    } else {
+      SparkRackResolver.resolveRacks(sc.hadoopConfiguration, hosts).map { node =>
+        Option(node.getNetworkLocation)
+      }
+    }
   }
 }