[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[String] = Nil
+
   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))
+  for ((rack, index) <- rackWithTaskIndex) {
+    pendingTasksForRack.getOrElseUpdate(rack, new ArrayBuffer) += index
+  }
+  // visible for testing
+  private[scheduler] val addTaskElapsedTime = (System.nanoTime() - addTaskStartTime) / 1e9
+  logInfo(s"Adding pending task takes $addTaskElapsedTime 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, initialing: 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 (initialing) {
+        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
@@ -78,9 +78,19 @@ object FakeRackUtil {
     hostToRack(host) = rack
   }
 
-  def getRackForHost(host: String): Option[String] = {
+  def getRackForHost(host: String, slow: Boolean = false): Option[String] = {
+    if (slow) {
+      Thread.sleep(100) // assume resolving one host takes 100 ms
+    }
     hostToRack.get(host)
   }
+
+  def getRacksForHosts(hosts: List[String], slow: Boolean = false): List[String] = {
+    if (slow) {
+      Thread.sleep(500) // assume resolving multiple hosts takes 500 ms
+    }
+    hosts.flatMap(hostToRack.get)
+  }
 }
 
 /**
@@ -97,6 +107,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 +156,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, slowRackResolve)
+
+  override def getRacksForHosts(values: List[String]): List[String] =
+    FakeRackUtil.getRacksForHosts(values, slowRackResolve)
 }
 
 /**
@@ -1316,7 +1331,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 +1345,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 +1617,28 @@ 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 and result when initialing TaskSetManager") {
+    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 always equals 100 with/without SPARK-27038
+    // verify elapsed time should be less than 1s, without SPARK-27038, it should be larger 10s
+    assert(manager.addTaskElapsedTime < 1)
+  }
 }

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

@@ -0,0 +1,94 @@
+/*
+ * 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.scheduler.cluster
+
+import scala.collection.JavaConverters._
+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.spark.internal.Logging
+
+/**
+ * Added in SPARK-27038. Before using the higher Hadoop version which applied YARN-9332,
+ * we construct [[YarnRackResolver]] instead of [[org.apache.hadoop.yarn.util.RackResolver]]
+ * to revolve the rack info.
+ */
+object YarnRackResolver extends Logging {
+  private var dnsToSwitchMapping: DNSToSwitchMapping = _
+  private var initCalled = false
+  // advisory count of arguments for rack script
+  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,12 @@
 
 package org.apache.spark.scheduler.cluster
 
+import org.apache.hadoop.net.{NetworkTopology, NodeBase}
 import org.apache.hadoop.yarn.util.RackResolver
 import org.apache.log4j.{Level, Logger}
 
 import org.apache.spark._
+import org.apache.spark.internal.config.LOCALITY_WAIT
 import org.apache.spark.scheduler.TaskSchedulerImpl
 import org.apache.spark.util.Utils
 
@@ -31,9 +33,33 @@ 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 [[YarnRackResolver]] 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[String] = {
+    val hosts = hostPorts.map(Utils.parseHostPort(_)._1)
+    if (skipRackResolving) {
+      hosts.map(new NodeBase(_, NetworkTopology.DEFAULT_RACK)).map(_.getNetworkLocation)
+    } else {
+      YarnRackResolver.resolveRacks(sc.hadoopConfiguration, hosts).map(_.getNetworkLocation)
+    }
   }
 }