[SPARK-24248][K8S] Use level triggering and state reconciliation in scheduling and lifecycle

dev/deps/spark-deps-hadoop-2.6

@@ -172,6 +172,8 @@ parquet-jackson-1.10.0.jar
 protobuf-java-2.5.0.jar
 py4j-0.10.7.jar
 pyrolite-4.13.jar
+reactive-streams-1.0.2.jar
+rxjava-2.1.13.jar
 scala-compiler-2.11.8.jar
 scala-library-2.11.8.jar
 scala-parser-combinators_2.11-1.0.4.jar

dev/deps/spark-deps-hadoop-2.7

@@ -173,6 +173,8 @@ parquet-jackson-1.10.0.jar
 protobuf-java-2.5.0.jar
 py4j-0.10.7.jar
 pyrolite-4.13.jar
+reactive-streams-1.0.2.jar
+rxjava-2.1.13.jar
 scala-compiler-2.11.8.jar
 scala-library-2.11.8.jar
 scala-parser-combinators_2.11-1.0.4.jar

dev/deps/spark-deps-hadoop-3.1

@@ -192,6 +192,8 @@ protobuf-java-2.5.0.jar
 py4j-0.10.7.jar
 pyrolite-4.13.jar
 re2j-1.1.jar
+reactive-streams-1.0.2.jar
+rxjava-2.1.13.jar
 scala-compiler-2.11.8.jar
 scala-library-2.11.8.jar
 scala-parser-combinators_2.11-1.0.4.jar

pom.xml

@@ -748,6 +748,11 @@
         <artifactId>jline</artifactId>
         <version>2.12.1</version>
       </dependency>
+      <dependency>
+        <groupId>io.reactivex.rxjava2</groupId>
+        <artifactId>rxjava</artifactId>
+        <version>2.1.13</version>
+      </dependency>
       <dependency>
         <groupId>org.scalatest</groupId>
         <artifactId>scalatest_${scala.binary.version}</artifactId>
@@ -760,6 +765,12 @@
         <version>1.10.19</version>
         <scope>test</scope>
       </dependency>
+      <dependency>
+        <groupId>org.jmock</groupId>
+        <artifactId>jmock-junit4</artifactId>
+        <scope>test</scope>
+        <version>2.8.4</version>
+      </dependency>
       <dependency>
         <groupId>org.scalacheck</groupId>
         <artifactId>scalacheck_${scala.binary.version}</artifactId>

resource-managers/kubernetes/core/pom.xml

@@ -77,16 +77,27 @@
     </dependency>
     <!-- End of shaded deps. -->
 
+    <dependency>
+      <groupId>com.squareup.okhttp3</groupId>
+      <artifactId>okhttp</artifactId>
+      <version>3.8.1</version>
+    </dependency>
+
+	<dependency>
+      <groupId>io.reactivex.rxjava2</groupId>
+      <artifactId>rxjava</artifactId>
+    </dependency>
+
     <dependency>
       <groupId>org.mockito</groupId>
       <artifactId>mockito-core</artifactId>
       <scope>test</scope>
     </dependency>
 
     <dependency>
-      <groupId>com.squareup.okhttp3</groupId>
-      <artifactId>okhttp</artifactId>
-      <version>3.8.1</version>
+      <groupId>org.jmock</groupId>
+      <artifactId>jmock-junit4</artifactId>
+      <scope>test</scope>
     </dependency>
 
   </dependencies>

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

@@ -0,0 +1,120 @@
+/*
+ * 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.k8s
+
+import java.util.concurrent.atomic.{AtomicInteger, AtomicLong}
+
+import io.fabric8.kubernetes.api.model.{Pod, PodBuilder}
+import io.fabric8.kubernetes.client.KubernetesClient
+import scala.collection.mutable
+
+import org.apache.spark.{SparkConf, SparkException}
+import org.apache.spark.deploy.k8s.Config._
+import org.apache.spark.deploy.k8s.Constants._
+import org.apache.spark.deploy.k8s.KubernetesConf
+import org.apache.spark.internal.Logging
+
+private[spark] class ExecutorPodsAllocator(
+    conf: SparkConf,
+    executorBuilder: KubernetesExecutorBuilder,
+    kubernetesClient: KubernetesClient,
+    eventQueue: ExecutorPodsEventQueue) extends Logging {
+
+  private val EXECUTOR_ID_COUNTER = new AtomicLong(0L)
+
+  private val totalExpectedExecutors = new AtomicInteger(0)
+
+  private val podAllocationSize = conf.get(KUBERNETES_ALLOCATION_BATCH_SIZE)
+
+  private val podAllocationDelay = conf.get(KUBERNETES_ALLOCATION_BATCH_DELAY)
+
+  private val kubernetesDriverPodName = conf
+    .get(KUBERNETES_DRIVER_POD_NAME)
+    .getOrElse(throw new SparkException("Must specify the driver pod name"))
+
+  private val driverPod = kubernetesClient.pods()
+    .withName(kubernetesDriverPodName)
+    .get()
+
+  // Use sets of ids instead of counters to be able to handle duplicate events.
+
+  // Executor IDs that have been requested from Kubernetes but are not running yet.
+  private val pendingExecutors = mutable.Set.empty[Long]
+
+  // We could use CoarseGrainedSchedulerBackend#totalRegisteredExecutors here for tallying the
+  // executors that are running. But, here we choose instead to maintain all state within this
+  // class from the persecptive of the k8s API. Therefore whether or not this scheduler loop
+  // believes an executor is running is dictated by the K8s API rather than Spark's RPC events.
+  // We may need to consider where these perspectives may differ and which perspective should
+  // take precedence.
+  private val runningExecutors = mutable.Set.empty[Long]
+
+  def start(applicationId: String): Unit = {
+    eventQueue.addSubscriber(podAllocationDelay) { updatedPods =>
+      processUpdatedPodEvents(applicationId, updatedPods)
+    }
+  }
+
+  def setTotalExpectedExecutors(total: Int): Unit = totalExpectedExecutors.set(total)
+
+  private def processUpdatedPodEvents(applicationId: String, updatedPods: Seq[Pod]): Unit = {
+    updatedPods.foreach { updatedPod =>
+      val execId = updatedPod.getMetadata.getLabels.get(SPARK_EXECUTOR_ID_LABEL).toLong
+      val phase = updatedPod.getStatus.getPhase.toLowerCase
+      phase match {
+        case "running" =>
+          pendingExecutors -= execId
+          runningExecutors += execId
+        case "failed" | "succeeded" | "error" =>
+          pendingExecutors -= execId
+          runningExecutors -= execId
+      }
+    }
+
+    val currentRunningExecutors = runningExecutors.size
+    val currentTotalExpectedExecutors = totalExpectedExecutors.get
+    if (pendingExecutors.isEmpty && currentRunningExecutors < currentTotalExpectedExecutors) {
+      val numExecutorsToAllocate = math.min(
+        currentTotalExpectedExecutors - currentRunningExecutors, podAllocationSize)
+      logInfo(s"Going to request $numExecutorsToAllocate executors from Kubernetes.")
+      val newExecutorIds = mutable.Buffer.empty[Long]
+      val podsToAllocate = mutable.Buffer.empty[Pod]
+      for ( _ <- 0 until numExecutorsToAllocate) {
+        val newExecutorId = EXECUTOR_ID_COUNTER.incrementAndGet()
+        val executorConf = KubernetesConf.createExecutorConf(
+          conf,
+          newExecutorId.toString,
+          applicationId,
+          driverPod)
+        val executorPod = executorBuilder.buildFromFeatures(executorConf)
+        val podWithAttachedContainer = new PodBuilder(executorPod.pod)
+          .editOrNewSpec()
+          .addToContainers(executorPod.container)
+          .endSpec()
+          .build()
+        kubernetesClient.pods().create(podWithAttachedContainer)
+        pendingExecutors += newExecutorId
+      }
+    } else if (currentRunningExecutors == currentTotalExpectedExecutors) {
+      logDebug("Current number of running executors is equal to the number of requested" +
+        " executors. Not scaling up further.")
+    } else if (pendingExecutors.nonEmpty) {
+      logInfo(s"Still waiting for ${pendingExecutors.size} executors to begin running before" +
+        s" requesting for more executors.")
+    }
+  }
+}

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

@@ -0,0 +1,29 @@
+/*
+ * 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.k8s
+
+import io.fabric8.kubernetes.api.model.Pod
+
+private[spark] trait ExecutorPodsEventQueue {
+
+  def addSubscriber(processBatchIntervalMillis: Long)(onNextBatch: Seq[Pod] => Unit): Unit
+
+  def stopProcessingEvents(): Unit
+
+  def pushPodUpdate(updatedPod: Pod): Unit
+
+}

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

@@ -0,0 +1,77 @@
+/*
+ * 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.k8s
+
+import java.util.concurrent.{Executor, ScheduledExecutorService, TimeUnit}
+
+import com.google.common.collect.Lists
+import io.fabric8.kubernetes.api.model.Pod
+import io.reactivex.disposables.Disposable
+import io.reactivex.functions.Consumer
+import io.reactivex.schedulers.Schedulers
+import io.reactivex.subjects.PublishSubject
+import scala.collection.JavaConverters._
+import scala.collection.mutable
+
+import org.apache.spark.util.Utils
+
+private[spark] class ExecutorPodsEventQueueImpl(
+    bufferEventsExecutor: ScheduledExecutorService,
+    executeSubscriptionsExecutor: Executor)
+  extends ExecutorPodsEventQueue {
+
+  private val eventsObservable = PublishSubject.create[Pod]()
+  private val observedDisposables = mutable.Buffer.empty[Disposable]
+
+  def addSubscriber(processBatchIntervalMillis: Long)(onNextBatch: Seq[Pod] => Unit): Unit = {
+    observedDisposables += eventsObservable
+      // Group events in the time window given by the caller. These buffers are then sent
+      // to the caller's lambda at the given interval, with the pod updates that occurred
+      // in that given interval.
+      .buffer(
+        processBatchIntervalMillis,
+        TimeUnit.MILLISECONDS,
+        // For testing - specifically use the given scheduled executor service to trigger
+        // buffer boundaries. Allows us to inject a deterministic scheduler here.
+        Schedulers.from(bufferEventsExecutor))
+      // Trigger an event cycle immediately. Not strictly required to be fully correct, but
+      // in particular the pod allocator should try to request executors immediately instead
+      // of waiting for one pod allocation delay.
+      .startWith(Lists.newArrayList[Pod]())
+      // Force all triggered events - both the initial event above and the buffered ones in
+      // the following time windows - to execute asynchronously to this call's thread.
+      .observeOn(Schedulers.from(executeSubscriptionsExecutor))
+      .subscribe(toReactivexConsumer { (pods: java.util.List[Pod]) =>
+        Utils.tryLogNonFatalError {
+          onNextBatch(pods.asScala)
+        }
+      })
+  }
+
+  def stopProcessingEvents(): Unit = {
+    observedDisposables.foreach(_.dispose())
+    eventsObservable.onComplete()
+  }
+
+  def pushPodUpdate(updatedPod: Pod): Unit = eventsObservable.onNext(updatedPod)
+
+  private def toReactivexConsumer[T](consumer: T => Unit): Consumer[T] = {
+    new Consumer[T] {
+      override def accept(item: T): Unit = consumer(item)
+    }
+  }
+}