[CARMEL-7374][CARMEL-5347] Fair task scheduler consider user resource usage (apache#118)

Author: wangyum

core/src/main/scala/org/apache/spark/internal/config/package.scala

@@ -2102,6 +2102,14 @@ package object config {
       .timeConf(TimeUnit.MILLISECONDS)
       .createOptional
 
+  private[spark] val SCHEDULER_FAIR_USER_LEVEL_ENABLED =
+    ConfigBuilder("spark.scheduler.fair.userLevel")
+      .internal()
+      .doc("If set true, and scheduler mode set to fair, " +
+        "task scheduler will consider user running tasks instead of only task set level")
+      .booleanConf
+      .createWithDefault(true)
+
   private[spark] val SPECULATION_ENABLED =
     ConfigBuilder("spark.speculation")
       .version("0.6.0")

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

@@ -38,6 +38,7 @@ import org.apache.spark.internal.config
 import org.apache.spark.internal.config._
 import org.apache.spark.resource.ResourceUtils
 import org.apache.spark.scheduler.TaskLocality.TaskLocality
+import org.apache.spark.scheduler.TaskSchedulerImpl.SCHEDULER_MODE_PROPERTY
 import org.apache.spark.storage.BlockManagerId
 import org.apache.spark.util.{AccumulatorV2, SystemClock, ThreadUtils, Utils}
 
@@ -100,6 +101,26 @@ private[spark] class AnalyticsTaskSchedulerImpl(
   private val FURTHER_SCHEDULE_LOCK_WAIT_TIME = conf.
     getTimeAsMs("spark.task.furtherScheduledLockWaitTime", "5s")
 
+  private val userLevelFairScheduledEnabled = conf.get(SCHEDULER_FAIR_USER_LEVEL_ENABLED)
+
+  override def initialize(backend: SchedulerBackend): Unit = {
+    this.backend = backend
+    schedulableBuilder = {
+      schedulingMode match {
+        case SchedulingMode.FIFO =>
+          new FIFOSchedulableBuilder(rootPool)
+        case SchedulingMode.FAIR =>
+          val userResourceManager =
+            if (userLevelFairScheduledEnabled) sc.userResourceManager else None
+          new FairSchedulableBuilder(rootPool, sc, userResourceManager)
+        case _ =>
+          throw new IllegalArgumentException(s"Unsupported $SCHEDULER_MODE_PROPERTY: " +
+            s"$schedulingMode")
+      }
+    }
+    schedulableBuilder.buildPools()
+  }
+
   override def start(): Unit = {
     backend.start()
 

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

@@ -72,6 +72,8 @@ private[spark] class AnalyticsTaskSetManager(
 
   override def userInfo(): Option[UserInfo] = _userInfo
 
+  override def user: Option[String] = _userInfo.map(_.user)
+
   // 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.
   addPendingTasks()

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

@@ -32,7 +32,8 @@ private[spark] class Pool(
     val poolName: String,
     val schedulingMode: SchedulingMode,
     initMinShare: Int,
-    initWeight: Int)
+    initWeight: Int,
+    userResourceManager: Option[UserResourceManager] = None)
   extends Schedulable with Logging {
 
   val schedulableQueue = new ConcurrentLinkedQueue[Schedulable]
@@ -50,7 +51,7 @@ private[spark] class Pool(
   private val taskSetSchedulingAlgorithm: SchedulingAlgorithm = {
     schedulingMode match {
       case SchedulingMode.FAIR =>
-        new FairSchedulingAlgorithm()
+        new FairSchedulingAlgorithm(userResourceManager)
       case SchedulingMode.FIFO =>
         new FIFOSchedulingAlgorithm()
       case _ =>

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

@@ -56,7 +56,8 @@ private[spark] class FIFOSchedulableBuilder(val rootPool: Pool)
   }
 }
 
-private[spark] class FairSchedulableBuilder(val rootPool: Pool, sc: SparkContext)
+private[spark] class FairSchedulableBuilder(
+    val rootPool: Pool, sc: SparkContext, userResourceManager: Option[UserResourceManager] = None)
   extends SchedulableBuilder with Logging {
 
   val schedulerAllocFile = sc.conf.get(SCHEDULER_ALLOCATION_FILE)
@@ -135,7 +136,8 @@ private[spark] class FairSchedulableBuilder(val rootPool: Pool, sc: SparkContext
       val weight = getIntValue(poolNode, poolName, WEIGHT_PROPERTY,
         DEFAULT_WEIGHT, fileName)
 
-      rootPool.addSchedulable(new Pool(poolName, schedulingMode, minShare, weight))
+      rootPool.addSchedulable(
+        new Pool(poolName, schedulingMode, minShare, weight, userResourceManager))
 
       logInfo("Created pool: %s, schedulingMode: %s, minShare: %d, weight: %d".format(
         poolName, schedulingMode, minShare, weight))

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

@@ -40,7 +40,8 @@ private[spark] class FIFOSchedulingAlgorithm extends SchedulingAlgorithm {
   }
 }
 
-private[spark] class FairSchedulingAlgorithm extends SchedulingAlgorithm {
+private[spark] class FairSchedulingAlgorithm(
+    userResourceManager: Option[UserResourceManager]) extends SchedulingAlgorithm {
   override def comparator(s1: Schedulable, s2: Schedulable): Boolean = {
     val minShare1 = s1.minShare
     val minShare2 = s2.minShare
@@ -61,7 +62,13 @@ private[spark] class FairSchedulingAlgorithm extends SchedulingAlgorithm {
     } else if (s1Needy && s2Needy) {
       compare = minShareRatio1.compareTo(minShareRatio2)
     } else {
-      compare = taskToWeightRatio1.compareTo(taskToWeightRatio2)
+      if (userResourceManager.isDefined && s1.user.isDefined && s2.user.isDefined) {
+        val r = userResourceManager.get
+        compare = r.occupiedCpus(s1.user.get).compareTo(r.occupiedCpus(s2.user.get))
+      }
+      if (compare == 0) {
+        compare = taskToWeightRatio1.compareTo(taskToWeightRatio2)
+      }
     }
     if (compare < 0) {
       true