More tests for the scheduler backend.

Author: mccheah

resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/kubernetes/KubernetesClusterManager.scala

@@ -120,13 +120,16 @@ private[spark] class KubernetesClusterManager extends ExternalClusterManager wit
         kubernetesShuffleManager)
     val allocatorExecutor = ThreadUtils
         .newDaemonSingleThreadScheduledExecutor("kubernetes-pod-allocator")
+    val requestExecutorsService = ThreadUtils.newDaemonCachedThreadPool(
+        "kubernetes-request-executors")
     new KubernetesClusterSchedulerBackend(
         scheduler.asInstanceOf[TaskSchedulerImpl],
         sc.env.rpcEnv,
         executorPodFactory,
         kubernetesShuffleManager,
         kubernetesClient,
-        allocatorExecutor)
+        allocatorExecutor,
+        requestExecutorsService)
   }
 
   override def initialize(scheduler: TaskScheduler, backend: SchedulerBackend): Unit = {

resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/kubernetes/KubernetesClusterSchedulerBackend.scala

@@ -19,7 +19,7 @@ package org.apache.spark.scheduler.cluster.kubernetes
 import java.io.Closeable
 import java.net.InetAddress
 import java.util.Collections
-import java.util.concurrent.{ConcurrentHashMap, ScheduledExecutorService, TimeUnit}
+import java.util.concurrent.{ConcurrentHashMap, ExecutorService, ScheduledExecutorService, ThreadPoolExecutor, TimeUnit}
 import java.util.concurrent.atomic.{AtomicInteger, AtomicLong, AtomicReference}
 
 import io.fabric8.kubernetes.api.model._
@@ -36,19 +36,21 @@ import org.apache.spark.rpc.{RpcAddress, RpcCallContext, RpcEndpointAddress, Rpc
 import org.apache.spark.scheduler.{ExecutorExited, SlaveLost, TaskSchedulerImpl}
 import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages.{RetrieveSparkAppConfig, SparkAppConfig}
 import org.apache.spark.scheduler.cluster.CoarseGrainedSchedulerBackend
-import org.apache.spark.util.{ThreadUtils, Utils}
+import org.apache.spark.util.Utils
 
 private[spark] class KubernetesClusterSchedulerBackend(
     scheduler: TaskSchedulerImpl,
     rpcEnv: RpcEnv,
     executorPodFactory: ExecutorPodFactory,
     shuffleManager: Option[KubernetesExternalShuffleManager],
     kubernetesClient: KubernetesClient,
-    allocatorExecutor: ScheduledExecutorService)
+    allocatorExecutor: ScheduledExecutorService,
+    requestExecutorsService: ExecutorService)
   extends CoarseGrainedSchedulerBackend(scheduler, rpcEnv) {
 
   import KubernetesClusterSchedulerBackend._
 
+  private val EXECUTOR_ID_COUNTER = new AtomicLong(0L)
   private val RUNNING_EXECUTOR_PODS_LOCK = new Object
   // Indexed by executor IDs and guarded by RUNNING_EXECUTOR_PODS_LOCK.
   private val runningExecutorsToPods = new mutable.HashMap[String, Pod]
@@ -71,7 +73,7 @@ private[spark] class KubernetesClusterSchedulerBackend(
       throw new SparkException("Must specify the driver pod name"))
 
   private implicit val requestExecutorContext = ExecutionContext.fromExecutorService(
-      ThreadUtils.newDaemonCachedThreadPool("kubernetes-executor-requests"))
+      requestExecutorsService)
 
   private val driverPod = try {
     kubernetesClient.pods().inNamespace(kubernetesNamespace).
@@ -212,7 +214,7 @@ private[spark] class KubernetesClusterSchedulerBackend(
             .watch(new ExecutorPodsWatcher()))
 
     allocatorExecutor.scheduleWithFixedDelay(
-        allocatorRunnable, 0, podAllocationInterval, TimeUnit.SECONDS)
+        allocatorRunnable, 0L, podAllocationInterval, TimeUnit.SECONDS)
     shuffleManager.foreach(_.start(applicationId()))
 
     if (!Utils.isDynamicAllocationEnabled(conf)) {
@@ -315,11 +317,13 @@ private[spark] class KubernetesClusterSchedulerBackend(
   override def doKillExecutors(executorIds: Seq[String]): Future[Boolean] = Future[Boolean] {
     RUNNING_EXECUTOR_PODS_LOCK.synchronized {
       for (executor <- executorIds) {
-        runningExecutorsToPods.remove(executor) match {
-          case Some(pod) =>
-            kubernetesClient.pods().delete(pod)
-            runningPodsToExecutors.remove(pod.getMetadata.getName)
-          case None => logWarning(s"Unable to remove pod for unknown executor $executor")
+        val maybeRemovedExecutor = runningExecutorsToPods.remove(executor)
+        maybeRemovedExecutor.foreach { executorPod =>
+          kubernetesClient.pods().delete(executorPod)
+          runningPodsToExecutors.remove(executorPod.getMetadata.getName)
+        }
+        if (maybeRemovedExecutor.isEmpty) {
+          logWarning(s"Unable to remove pod for unknown executor $executor")
         }
       }
     }
@@ -445,7 +449,7 @@ private[spark] class KubernetesClusterSchedulerBackend(
       new PartialFunction[Any, Unit]() {
         override def isDefinedAt(msg: Any): Boolean = {
           msg match {
-            case RetrieveSparkAppConfig(executorId) =>
+            case RetrieveSparkAppConfig(_) =>
               shuffleManager.isDefined
             case _ => false
           }
@@ -475,7 +479,6 @@ private[spark] class KubernetesClusterSchedulerBackend(
 
 private object KubernetesClusterSchedulerBackend {
   private val DEFAULT_STATIC_PORT = 10000
-  private val EXECUTOR_ID_COUNTER = new AtomicLong(0L)
   private val VMEM_EXCEEDED_EXIT_CODE = -103
   private val PMEM_EXCEEDED_EXIT_CODE = -104
   private val UNKNOWN_EXIT_CODE = -111

resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/kubernetes/KubernetesClusterSchedulerBackendSuite.scala

@@ -16,20 +16,25 @@
  */
 package org.apache.spark.scheduler.cluster.kubernetes
 
-import java.util.concurrent.ScheduledExecutorService
+import java.util.concurrent.{ExecutorService, ScheduledExecutorService, TimeUnit}
 
 import io.fabric8.kubernetes.api.model.{DoneablePod, Pod, PodBuilder, PodList}
 import io.fabric8.kubernetes.client.{KubernetesClient, Watch, Watcher}
-import io.fabric8.kubernetes.client.dsl.{FilterWatchListDeletable, MixedOperation, NonNamespaceOperation, PodResource, Resource}
-import org.mockito.{ArgumentCaptor, Mock, MockitoAnnotations}
-import org.mockito.Mockito.{verify, when}
+import io.fabric8.kubernetes.client.dsl.{FilterWatchListDeletable, MixedOperation, NonNamespaceOperation, PodResource}
+import org.mockito.{AdditionalAnswers, ArgumentCaptor, Mock, MockitoAnnotations}
+import org.mockito.Matchers.{any, eq => mockitoEq}
+import org.mockito.Mockito.{doNothing, never, verify, when}
 import org.scalatest.BeforeAndAfter
+import org.scalatest.mock.MockitoSugar._
+import scala.collection.JavaConverters._
 
 import org.apache.spark.{SparkConf, SparkContext, SparkFunSuite}
 import org.apache.spark.deploy.kubernetes.config._
 import org.apache.spark.deploy.kubernetes.constants._
-import org.apache.spark.rpc.RpcEnv
-import org.apache.spark.scheduler.TaskSchedulerImpl
+import org.apache.spark.rpc.{RpcCallContext, RpcEndpoint, RpcEndpointAddress, RpcEndpointRef, RpcEnv}
+import org.apache.spark.scheduler.{LiveListenerBus, TaskSchedulerImpl}
+import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages.RegisterExecutor
+import org.apache.spark.scheduler.cluster.CoarseGrainedSchedulerBackend
 
 private[spark] class KubernetesClusterSchedulerBackendSuite
     extends SparkFunSuite with BeforeAndAfter {
@@ -39,22 +44,31 @@ private[spark] class KubernetesClusterSchedulerBackendSuite
   private val NAMESPACE = "test-namespace"
   private val SPARK_DRIVER_HOST = "localhost"
   private val SPARK_DRIVER_PORT = 7077
+  private val POD_ALLOCATION_INTERVAL = 60L
+  private val DRIVER_URL = RpcEndpointAddress(
+      SPARK_DRIVER_HOST, SPARK_DRIVER_PORT, CoarseGrainedSchedulerBackend.ENDPOINT_NAME).toString
 
   private type PODS = MixedOperation[Pod, PodList, DoneablePod, PodResource[Pod, DoneablePod]]
   private type LABELLED_PODS = FilterWatchListDeletable[
-        Pod, PodList, java.lang.Boolean, Watch, Watcher[Pod]]
+      Pod, PodList, java.lang.Boolean, Watch, Watcher[Pod]]
   private type IN_NAMESPACE_PODS = NonNamespaceOperation[
-        Pod, PodList, DoneablePod, PodResource[Pod, DoneablePod]]
+      Pod, PodList, DoneablePod, PodResource[Pod, DoneablePod]]
 
   @Mock
   private var sparkContext: SparkContext = _
 
+  @Mock
+  private var listenerBus: LiveListenerBus = _
+
   @Mock
   private var taskSchedulerImpl: TaskSchedulerImpl = _
 
   @Mock
   private var allocatorExecutor: ScheduledExecutorService = _
 
+  @Mock
+  private var requestExecutorsService: ExecutorService = _
+
   @Mock
   private var executorPodFactory: ExecutorPodFactory = _
 
@@ -79,11 +93,17 @@ private[spark] class KubernetesClusterSchedulerBackendSuite
   @Mock
   private var rpcEnv: RpcEnv = _
 
+  @Mock
+  private var driverEndpointRef: RpcEndpointRef = _
+
   @Mock
   private var executorPodsWatch: Watch = _
 
   private var sparkConf: SparkConf = _
   private var executorPodsWatcherArgument: ArgumentCaptor[Watcher[Pod]] = _
+  private var allocatorRunnable: ArgumentCaptor[Runnable] = _
+  private var requestExecutorRunnable: ArgumentCaptor[Runnable] = _
+  private var driverEndpoint: ArgumentCaptor[RpcEndpoint] = _
 
   private val driverPod = new PodBuilder()
     .withNewMetadata()
@@ -101,8 +121,13 @@ private[spark] class KubernetesClusterSchedulerBackendSuite
         .set(KUBERNETES_NAMESPACE, NAMESPACE)
         .set("spark.driver.host", SPARK_DRIVER_HOST)
         .set("spark.driver.port", SPARK_DRIVER_PORT.toString)
+        .set(KUBERNETES_ALLOCATION_BATCH_DELAY, POD_ALLOCATION_INTERVAL)
     executorPodsWatcherArgument = ArgumentCaptor.forClass(classOf[Watcher[Pod]])
+    allocatorRunnable = ArgumentCaptor.forClass(classOf[Runnable])
+    requestExecutorRunnable = ArgumentCaptor.forClass(classOf[Runnable])
+    driverEndpoint = ArgumentCaptor.forClass(classOf[RpcEndpoint])
     when(sparkContext.conf).thenReturn(sparkConf)
+    when(sparkContext.listenerBus).thenReturn(listenerBus)
     when(taskSchedulerImpl.sc).thenReturn(sparkContext)
     when(kubernetesClient.pods()).thenReturn(podOperations)
     when(podOperations.withLabel(SPARK_APP_ID_LABEL, APP_ID)).thenReturn(podsWithLabelOperations)
@@ -111,26 +136,164 @@ private[spark] class KubernetesClusterSchedulerBackendSuite
     when(podOperations.inNamespace(NAMESPACE)).thenReturn(podsInNamespace)
     when(podsInNamespace.withName(DRIVER_POD_NAME)).thenReturn(podsWithDriverName)
     when(podsWithDriverName.get()).thenReturn(driverPod)
+    when(allocatorExecutor.scheduleWithFixedDelay(
+        allocatorRunnable.capture(),
+        mockitoEq(0L),
+        mockitoEq(POD_ALLOCATION_INTERVAL),
+        mockitoEq(TimeUnit.SECONDS))).thenReturn(null)
+    // Creating Futures in Scala backed by a Java executor service resolves to running
+    // ExecutorService#execute (as opposed to submit)
+    doNothing().when(requestExecutorsService).execute(requestExecutorRunnable.capture())
+    when(rpcEnv.setupEndpoint(
+        mockitoEq(CoarseGrainedSchedulerBackend.ENDPOINT_NAME), driverEndpoint.capture()))
+        .thenReturn(driverEndpointRef)
   }
 
   test("Basic lifecycle expectations when starting and stopping the scheduler.") {
     val scheduler = newSchedulerBackend(true)
     scheduler.start()
     verify(shuffleManager).start(APP_ID)
     assert(executorPodsWatcherArgument.getValue != null)
+    assert(allocatorRunnable.getValue != null)
     scheduler.stop()
     verify(shuffleManager).stop()
     verify(executorPodsWatch).close()
   }
 
+  test("Static allocation should request executors upon first allocator run.") {
+    sparkConf
+        .set(KUBERNETES_ALLOCATION_BATCH_SIZE, 2)
+        .set(org.apache.spark.internal.config.EXECUTOR_INSTANCES, 2)
+    val scheduler = newSchedulerBackend(true)
+    scheduler.start()
+    requestExecutorRunnable.getValue.run()
+    val firstExecutorPod = new PodBuilder()
+      .withNewMetadata()
+        .withName("pod1")
+        .endMetadata()
+      .build()
+    val secondExecutorPod = new PodBuilder()
+      .withNewMetadata()
+        .withName("pod2")
+        .endMetadata()
+      .build()
+    when(podOperations.create(firstExecutorPod))
+        .thenReturn(firstExecutorPod)
+    when(podOperations.create(secondExecutorPod))
+      .thenReturn(secondExecutorPod)
+    when(executorPodFactory.createExecutorPod(
+        "1",
+        APP_ID,
+        DRIVER_URL,
+        sparkConf.getExecutorEnv,
+        driverPod,
+        Map.empty)).thenReturn(firstExecutorPod)
+    when(executorPodFactory.createExecutorPod(
+        "2",
+        APP_ID,
+        DRIVER_URL,
+        sparkConf.getExecutorEnv,
+        driverPod,
+        Map.empty)).thenReturn(secondExecutorPod)
+    allocatorRunnable.getValue.run()
+    verify(podOperations).create(firstExecutorPod)
+    verify(podOperations).create(secondExecutorPod)
+  }
+
+  test("Killing executors deletes the executor pods") {
+    sparkConf
+        .set(KUBERNETES_ALLOCATION_BATCH_SIZE, 2)
+        .set(org.apache.spark.internal.config.EXECUTOR_INSTANCES, 2)
+    val scheduler = newSchedulerBackend(true)
+    scheduler.start()
+    requestExecutorRunnable.getValue.run()
+    val firstExecutorPod = new PodBuilder()
+      .withNewMetadata()
+      .withName("pod1")
+      .endMetadata()
+      .build()
+    val secondExecutorPod = new PodBuilder()
+      .withNewMetadata()
+      .withName("pod2")
+      .endMetadata()
+      .build()
+    when(podOperations.create(any(classOf[Pod])))
+        .thenAnswer(AdditionalAnswers.returnsFirstArg())
+    when(executorPodFactory.createExecutorPod(
+        "1",
+        APP_ID,
+        DRIVER_URL,
+        sparkConf.getExecutorEnv,
+        driverPod,
+        Map.empty)).thenReturn(firstExecutorPod)
+    when(executorPodFactory.createExecutorPod(
+        "2",
+        APP_ID,
+        DRIVER_URL,
+        sparkConf.getExecutorEnv,
+        driverPod,
+        Map.empty)).thenReturn(secondExecutorPod)
+    allocatorRunnable.getValue.run()
+    scheduler.doKillExecutors(Seq("1", "2"))
+    requestExecutorRunnable.getAllValues.asScala.last.run()
+    verify(podOperations).delete(firstExecutorPod)
+    verify(podOperations).delete(secondExecutorPod)
+  }
+
+  test("Executors should be requested in batches.") {
+    sparkConf
+      .set(KUBERNETES_ALLOCATION_BATCH_SIZE, 1)
+      .set(org.apache.spark.internal.config.EXECUTOR_INSTANCES, 2)
+    val scheduler = newSchedulerBackend(true)
+    scheduler.start()
+    requestExecutorRunnable.getValue.run()
+    val firstExecutorPod = new PodBuilder()
+      .withNewMetadata()
+      .withName("pod1")
+      .endMetadata()
+      .build()
+    val secondExecutorPod = new PodBuilder()
+      .withNewMetadata()
+      .withName("pod2")
+      .endMetadata()
+      .build()
+    when(podOperations.create(any(classOf[Pod])))
+      .thenAnswer(AdditionalAnswers.returnsFirstArg())
+    when(executorPodFactory.createExecutorPod(
+        "1",
+        APP_ID,
+        DRIVER_URL,
+        sparkConf.getExecutorEnv,
+        driverPod,
+        Map.empty)).thenReturn(firstExecutorPod)
+    when(executorPodFactory.createExecutorPod(
+        "2",
+        APP_ID,
+        DRIVER_URL,
+        sparkConf.getExecutorEnv,
+        driverPod,
+        Map.empty)).thenReturn(secondExecutorPod)
+    allocatorRunnable.getValue.run()
+    verify(podOperations).create(firstExecutorPod)
+    verify(podOperations, never()).create(secondExecutorPod)
+    val registerFirstExecutorMessage = RegisterExecutor(
+        "1", mock[RpcEndpointRef], "localhost", 1, Map.empty[String, String])
+    when(taskSchedulerImpl.resourceOffers(any())).thenReturn(Seq.empty)
+    driverEndpoint.getValue.receiveAndReply(mock[RpcCallContext])
+        .apply(registerFirstExecutorMessage)
+    allocatorRunnable.getValue.run()
+    verify(podOperations).create(secondExecutorPod)
+  }
+
   private def newSchedulerBackend(externalShuffle: Boolean): KubernetesClusterSchedulerBackend = {
     new KubernetesClusterSchedulerBackend(
         taskSchedulerImpl,
         rpcEnv,
         executorPodFactory,
         if (externalShuffle) Some(shuffleManager) else None,
         kubernetesClient,
-        allocatorExecutor)
+        allocatorExecutor,
+        requestExecutorsService)
   }
 
 }