Use the new Armada Scala Client job watcher API method (#95)

Replace the local logic for watching a job with the new `jobWatcher()`
method from the Armada Scala Client.

Also, add the `-p` option to the `nc` command invocation for checking
for server reachability/readiness - this fixes an issue with running in
C3 where some older versions of `nc` did not correctly use the given
port number - with this, the `nc` listener process we run works with
both newer and older versions of nc.

---------

Signed-off-by: Rich Scott <richscott@sent.com>

Author: richscott

.github/workflows/build-snapshots.yaml

@@ -14,7 +14,7 @@ jobs:
       matrix:
         include:
           - scala_version: "2.13.16"
-            spark_version: "4.1.0-SNAPSHOT"
+            spark_version: "4.1.2-SNAPSHOT"
             java_version: "17"
 
     steps:
@@ -25,4 +25,4 @@ jobs:
         with:
           spark_version: ${{ matrix.spark_version }}
           scala_version: ${{ matrix.scala_version }}
-          java_version: ${{ matrix.java_version }}
+          java_version: ${{ matrix.java_version }}

pom.xml

@@ -153,7 +153,7 @@
         <dependency>
           <groupId>io.armadaproject</groupId>
           <artifactId>armada-scala-client_${scala.binary.version}</artifactId>
-          <version>0.2.0</version>
+          <version>0.3.0</version>
         </dependency>
 
         <!-- Test dependencies -->

src/main/scala/org/apache/spark/scheduler/cluster/armada/ArmadaClusterManagerBackend.scala

@@ -27,7 +27,7 @@ import api.submit.JobCancelRequest
 import io.armadaproject.armada.ArmadaClient
 import io.grpc.netty.NettyChannelBuilder
 import io.grpc.stub.MetadataUtils
-import io.grpc.{ManagedChannel, ManagedChannelBuilder, Metadata}
+import io.grpc.Metadata
 
 import org.apache.spark.SparkContext
 import org.apache.spark.deploy.armada.Config._
@@ -85,7 +85,6 @@ private[spark] class ArmadaClusterManagerBackend(
     }
 
   /** Armada connection details */
-  private var grpcChannel: Option[ManagedChannel]      = None
   private var armadaClient: Option[ArmadaClient]       = None
   private var eventWatcher: Option[ArmadaEventWatcher] = None
   private var queue: Option[String]                    = None
@@ -207,10 +206,8 @@ private[spark] class ArmadaClusterManagerBackend(
         channelBuilderWithTls.build()
     }
 
-    grpcChannel = Some(channel)
-
     // Start event watcher
-    val watcher = new ArmadaEventWatcher(channel, q, jsId, this, jobIdToExecutor, client)
+    val watcher = new ArmadaEventWatcher(q, jsId, this, jobIdToExecutor, client)
     eventWatcher = Some(watcher)
     watcher.start()
     logInfo(s"Armada Event Watcher started for queue=$q jobSetId=$jsId at $host:$port")
@@ -262,15 +259,6 @@ private[spark] class ArmadaClusterManagerBackend(
       ThreadUtils.shutdown(executorService)
     }
 
-    // Shutdown gRPC channel
-    Utils.tryLogNonFatalError {
-      grpcChannel.foreach { ch =>
-        ch.shutdown()
-        ch.awaitTermination(5, TimeUnit.SECONDS)
-      }
-    }
-
-    grpcChannel = None
     eventWatcher = None
     armadaClient = None
     executorAllocator = None

src/main/scala/org/apache/spark/scheduler/cluster/armada/ArmadaEventWatcher.scala

@@ -23,8 +23,6 @@ import scala.concurrent.Await
 import scala.concurrent.duration._
 
 import api.event.{
-  EventGrpc,
-  WatchRequest,
   EventMessage,
   JobSubmittedEvent,
   JobQueuedEvent,
@@ -35,16 +33,14 @@ import api.event.{
   JobCancelledEvent,
   JobPreemptingEvent,
   JobPreemptedEvent,
-  JobTerminatedEvent,
-  JobUnableToScheduleEvent
+  JobTerminatedEvent
 }
 import io.grpc.ManagedChannel
 import org.apache.spark.internal.Logging
 
 /** Keeps track of the armada event stream for the jobset
   */
 private[spark] class ArmadaEventWatcher(
-    grpcChannel: ManagedChannel,
     queue: String,
     jobSetId: String,
     backend: ArmadaClusterManagerBackend,
@@ -89,17 +85,11 @@ private[spark] class ArmadaEventWatcher(
 
     while (running) {
       try {
-        logInfo(s"Connecting to Armada event stream: queue=$queue, jobSetId=$jobSetId")
+        val watcher = submitClient.jobWatcher(queue, jobSetId, lastMessageId)
 
-        val eventStub    = EventGrpc.blockingStub(grpcChannel)
-        val watchRequest = WatchRequest(queue = queue, jobSetId = jobSetId, fromId = lastMessageId)
-
-        logInfo("Connected to Armada event stream (blocking)")
-
-        val iterator = eventStub.watch(watchRequest)
-        while (running && iterator.hasNext) {
+        while (running && watcher.hasNext) {
           try {
-            val streamMessage = iterator.next()
+            val streamMessage = watcher.next()
             if (streamMessage != null) {
               lastMessageId = streamMessage.id
               streamMessage.message.foreach(processEventMessage)
@@ -183,9 +173,6 @@ private[spark] class ArmadaEventWatcher(
       case EventMessage.Events.Preempted(event) =>
         handlePreempted(event)
 
-      case EventMessage.Events.UnableToSchedule(event) =>
-        handleUnableToSchedule(event)
-
       case EventMessage.Events.Leased(_) =>
         logInfo("Job leased event received")
 
@@ -311,20 +298,4 @@ private[spark] class ArmadaEventWatcher(
     }
   }
 
-  private def handleUnableToSchedule(event: JobUnableToScheduleEvent): Unit = {
-    val jobId      = event.jobId
-    val executorId = jobIdToExecutor.get(jobId)
-    if (executorId != null) {
-      val reason = Option(event.reason).filter(_.nonEmpty).getOrElse("Unable to schedule")
-      val nodeName =
-        Option(event.nodeName).filter(_.nonEmpty).map(n => s" on node $n").getOrElse("")
-      val podName = Option(event.podName).filter(_.nonEmpty).map(p => s" (pod: $p)").getOrElse("")
-      val fullReason = s"$reason$nodeName$podName"
-      logWarning(s"Executor $executorId (job $jobId) unable to schedule: $fullReason")
-      backend.onExecutorUnableToSchedule(jobId, executorId, fullReason)
-    } else {
-      logInfo(s"Received unable to schedule event for unknown job $jobId")
-    }
-  }
-
 }

src/test/scala/org/apache/spark/deploy/armada/submit/ArmadaUtilsSuite.scala

@@ -43,7 +43,7 @@ private class ArmadaUtilsSuite
   test("Confirm initContainer sh command succeeds with server") {
     // start server
     val serverPort    = "54525"
-    val serverCommand = Seq("nc", "-l", serverPort)
+    val serverCommand = Seq("nc", "-l", "-p", serverPort)
     val server        = Process.apply(serverCommand).run
 
     // start client