WX-1595 GCP Batch backend refactor to include the PAPI request manager

.sdkmanrc

@@ -1,3 +1,4 @@
 # Enable auto-env through the sdkman_auto_env config
 # Add key=value pairs of SDKs to use below
-java=11.0.11.hs-adpt
+#java=11.0.11.hs-adpt
+java=11.0.21-amzn

core/src/main/resources/reference.conf

@@ -625,6 +625,10 @@ load-control {
   # Google requests to the Pipelines API are also queued and batched
   papi-requests = 10000
 
+  ## Backend specific ##
+  # Google requests to the Batch API are also queued and batched
+  batch-requests = 10000
+
   ## Misc. ##
   # How often each actor should update its perceived load
   monitoring-frequency = 5 seconds

core/src/main/scala/cromwell/core/LoadConfig.scala

@@ -19,4 +19,5 @@ object LoadConfig {
   val MetadataWriteThreshold = conf.as[Int]("metadata-write")
   val MonitoringFrequency = conf.as[FiniteDuration]("monitoring-frequency")
   val PAPIThreshold = conf.as[Int]("papi-requests")
+  val BatchThreshold = conf.as[Int]("batch-requests")
 }

core/src/test/scala/cromwell/core/LoadConfigSpec.scala

@@ -22,5 +22,6 @@ class LoadConfigSpec extends AnyFlatSpec with CromwellTimeoutSpec with Matchers
     LoadConfig.IoNormalWindowMaximum shouldBe 60.seconds
     LoadConfig.MonitoringFrequency shouldBe 5.seconds
     LoadConfig.PAPIThreshold shouldBe 10000
+    LoadConfig.BatchThreshold shouldBe 10000
   }
 }

supportedBackends/google/batch/src/main/scala/cromwell/backend/google/batch/GcpBatchBackendLifecycleActorFactory.scala

@@ -2,26 +2,22 @@
 
 import akka.actor.{ActorRef, Props}
 import com.google.api.client.util.ExponentialBackOff
+import com.google.api.gax.rpc.FixedHeaderProvider
+import com.google.cloud.batch.v1.BatchServiceSettings
+import com.google.common.collect.ImmutableMap
 import com.typesafe.scalalogging.StrictLogging
+import cromwell.backend._
 import cromwell.backend.google.batch.GcpBatchBackendLifecycleActorFactory.{
   preemptionCountKey,
   robustBuildAttributes,
   unexpectedRetryCountKey
 }
 import cromwell.backend.google.batch.actors._
-import cromwell.backend.google.batch.api.{GcpBatchApiRequestHandler, GcpBatchRequestFactoryImpl}
-import cromwell.backend.google.batch.models.{GcpBatchConfiguration, GcpBatchConfigurationAttributes}
+import cromwell.backend.google.batch.api.request.{BatchRequestExecutor, RequestHandler}
 import cromwell.backend.google.batch.callcaching.{BatchBackendCacheHitCopyingActor, BatchBackendFileHashingActor}
+import cromwell.backend.google.batch.models.{GcpBatchConfiguration, GcpBatchConfigurationAttributes}
 import cromwell.backend.standard._
 import cromwell.backend.standard.callcaching.{StandardCacheHitCopyingActor, StandardFileHashingActor}
-import cromwell.backend.{
-  BackendConfigurationDescriptor,
-  BackendInitializationData,
-  BackendWorkflowDescriptor,
-  Gcp,
-  JobExecutionMap,
-  Platform
-}
 import cromwell.cloudsupport.gcp.GoogleConfiguration
 import cromwell.core.CallOutputs
 import wom.graph.CommandCallNode
@@ -70,7 +66,6 @@
 
   override def workflowFinalizationActorParams(workflowDescriptor: BackendWorkflowDescriptor,
                                                ioActor: ActorRef,
-                                               // batchConfiguration: GcpBatchConfiguration,
                                                calls: Set[CommandCallNode],
                                                jobExecutionMap: JobExecutionMap,
                                                workflowOutputs: CallOutputs,
@@ -93,10 +88,19 @@
   )
 
   override def backendSingletonActorProps(serviceRegistryActor: ActorRef): Option[Props] = {
-    val requestHandler = new GcpBatchApiRequestHandler
-    val requestFactory = new GcpBatchRequestFactoryImpl()(batchConfiguration.batchAttributes.gcsTransferConfiguration)
+    implicit val requestHandler: RequestHandler = new RequestHandler
+
+    val batchSettings = BatchServiceSettings.newBuilder
+      .setHeaderProvider(FixedHeaderProvider.create(ImmutableMap.of("user-agent", "cromwell")))
+      .build
+
     Option(
-      GcpBatchBackendSingletonActor.props(requestFactory, serviceRegistryActor = serviceRegistryActor)(requestHandler)
+      GcpBatchBackendSingletonActor.props(
+        qps = batchConfiguration.batchAttributes.qps,
+        requestWorkers = batchConfiguration.batchAttributes.requestWorkers,
+        serviceRegistryActor = serviceRegistryActor,
+        batchRequestExecutor = new BatchRequestExecutor.CloudImpl(batchSettings)
+      )(requestHandler)
     )
   }
 

supportedBackends/google/batch/src/main/scala/cromwell/backend/google/batch/actors/BatchApiAbortClient.scala

@@ -2,21 +2,47 @@
 
 import akka.actor.{Actor, ActorLogging, ActorRef}
 import com.google.cloud.batch.v1.JobName
+import cromwell.backend.google.batch.api.BatchApiRequestManager.{BatchAbortRequest, BatchApiAbortQueryFailed}
+import cromwell.backend.google.batch.api.GcpBatchRequestFactory
 import cromwell.backend.google.batch.monitoring.BatchInstrumentation
+import cromwell.backend.standard.StandardAsyncJob
+import cromwell.core.WorkflowId
+import cromwell.core.logging.JobLogging
 
-trait BatchApiAbortClient { this: Actor with ActorLogging with BatchInstrumentation =>
+object BatchApiAbortClient {
+  sealed trait BatchAbortRequestSuccess
+  case class BatchAbortRequestSuccessful(jobId: String) extends BatchAbortRequestSuccess
 
-  def abortJob(jobName: JobName, backendSingletonActor: ActorRef): Unit =
-    backendSingletonActor ! GcpBatchBackendSingletonActor.Action.AbortJob(jobName)
+  // The operation is no longer running. Maybe it was already cancelled, maybe it finished on its own. We don't know
+  // the details and for abort they don't really matter.
+  case class BatchOperationIsAlreadyTerminal(jobId: String) extends BatchAbortRequestSuccess
+}
+
+trait BatchApiAbortClient { this: Actor with ActorLogging with JobLogging with BatchInstrumentation =>
+  import BatchApiAbortClient._
+
+  def abortJob(workflowId: WorkflowId,
+               jobName: JobName,
+               backendSingletonActor: ActorRef,
+               requestFactory: GcpBatchRequestFactory
+  ): Unit =
+    backendSingletonActor ! BatchAbortRequest(
+      workflowId = workflowId,
+      requester = self,
+      httpRequest = requestFactory.abortRequest(jobName),
+      jobId = StandardAsyncJob(jobName.toString)
+    )
 
   def abortActorClientReceive: Actor.Receive = {
-    case GcpBatchBackendSingletonActor.Event.JobAbortRequestSent(job) =>
-      log.info(s"Job aborted on GCP: ${job.getName}")
+    case BatchAbortRequestSuccessful(jobId) =>
       abortSuccess()
+      jobLogger.info(s"Successfully requested cancellation of $jobId")
+
+    // In this case we could immediately return an aborted handle and spare ourselves a round of polling
+    case BatchOperationIsAlreadyTerminal(jobId) =>
+      jobLogger.info(s"Job $jobId has already finished")
 
-    case GcpBatchBackendSingletonActor.Event.ActionFailed(jobName, cause) =>
-      val msg = s"Failed to abort job ($jobName) from GCP"
-      log.error(cause, msg)
-      abortFailed()
+    case BatchApiAbortQueryFailed(jobId, e) =>
+      jobLogger.error(s"Could not request cancellation of job $jobId", e)
   }
 }

supportedBackends/google/batch/src/main/scala/cromwell/backend/google/batch/actors/BatchApiRunCreationClient.scala

@@ -1,31 +1,34 @@
 package cromwell.backend.google.batch.actors
 
 import akka.actor.{Actor, ActorLogging, ActorRef}
+import cromwell.backend.google.batch.api.{BatchApiRequestManager, GcpBatchRequestFactory}
+import cromwell.backend.google.batch.api.BatchApiRequestManager.{
+  BatchApiRunCreationQueryFailed,
+  SystemBatchApiException
+}
 import cromwell.backend.google.batch.models.GcpBatchRequest
 import cromwell.backend.google.batch.monitoring.BatchInstrumentation
 import cromwell.backend.standard.StandardAsyncJob
+import cromwell.core.logging.JobLogger
 
 import scala.concurrent.{Future, Promise}
 import scala.util.{Failure, Success, Try}
 
 /**
-  * Handles the flow for submitting a single job to GCP, we can't do anything when that fails
+  * Handles the flow for submitting a single job to GCP
   */
 trait BatchApiRunCreationClient { this: Actor with ActorLogging with BatchInstrumentation =>
   private var runCreationClientPromise: Option[Promise[StandardAsyncJob]] = None
 
   // handles messages produced from GcpBatchBackendSingletonActor
   def runCreationClientReceive: Actor.Receive = {
-    case GcpBatchBackendSingletonActor.Event.JobSubmitted(job) =>
-      log.info(s"Job submitted to GCP: ${job.getName}")
+    case job: StandardAsyncJob =>
+      log.info(s"runCreationClientReceive -> StandardAsyncJob: ${job.jobId}")
       runSuccess()
-      completePromise(Success(StandardAsyncJob(job.getName)))
-
-    case GcpBatchBackendSingletonActor.Event.ActionFailed(jobName, cause) =>
-      val msg = s"Failed to submit job ($jobName) to GCP"
-      log.error(cause, msg)
-      runFailed()
-      completePromise(Failure(cause))
+      completePromise(Success(job))
+    case BatchApiRunCreationQueryFailed(_, e) =>
+      log.error(e, s"runCreationClientReceive -> BatchApiRunCreationQueryFailed: ${e.getMessage}")
+      completePromise(Failure(e))
   }
 
   private def completePromise(job: Try[StandardAsyncJob]): Unit = {
@@ -35,15 +38,37 @@
     runCreationClientPromise = None
   }
 
-  def runBatchJob(request: GcpBatchRequest, backendSingletonActor: ActorRef): Future[StandardAsyncJob] =
+  def runBatchJob(
+    request: GcpBatchRequest,
+    backendSingletonActor: ActorRef,
+    requestFactory: GcpBatchRequestFactory,
+    jobLogger: JobLogger
+  ): Future[StandardAsyncJob] =
     runCreationClientPromise match {
       case Some(p) =>
         p.future
       case None =>
-        log.info(s"Asking singleton actor to submit a job: ${request.jobName}")
-        backendSingletonActor ! GcpBatchBackendSingletonActor.Action.SubmitJob(request)
+        jobLogger.info(s"Asking singleton actor to submit a job: ${request.jobName}")
+
+        backendSingletonActor ! BatchApiRequestManager.BatchRunCreationRequest(
+          request.workflowId,
+          self,
+          requestFactory.submitRequest(request)
+        )
         val newPromise = Promise[StandardAsyncJob]()
         runCreationClientPromise = Option(newPromise)
         newPromise.future
     }
 }
+
+object BatchApiRunCreationClient {
+
+  /**
+   * Exception used to represent the fact that a job was aborted before a creation attempt was made.
+   * Meaning it was in the queue when the abort request was made, so it was just removed from the queue.
+   */
+  case object JobAbortedException
+      extends SystemBatchApiException(
+        new Exception("The job was removed from the queue before a Batch creation request was made")
+      )
+}

supportedBackends/google/batch/src/main/scala/cromwell/backend/google/batch/actors/BatchApiStatusRequestClient.scala

@@ -0,0 +1,57 @@
+package cromwell.backend.google.batch.actors
+
+import akka.actor.{Actor, ActorLogging, ActorRef}
+import com.google.cloud.batch.v1.JobName
+import cromwell.backend.google.batch.api.BatchApiRequestManager.{BatchApiStatusQueryFailed, BatchStatusPollRequest}
+import cromwell.backend.google.batch.api.GcpBatchRequestFactory
+import cromwell.backend.google.batch.models.RunStatus
+import cromwell.backend.google.batch.monitoring.BatchInstrumentation
+import cromwell.backend.standard.StandardAsyncJob
+import cromwell.core.WorkflowId
+
+import scala.concurrent.{Future, Promise}
+import scala.util.{Failure, Success, Try}
+
+/**
+  * Allows fetching a job status
+  */
+trait BatchApiStatusRequestClient { this: Actor with ActorLogging with BatchInstrumentation =>
+
+  private var pollingActorClientPromise: Option[Promise[RunStatus]] = None
+
+  def pollingActorClientReceive: Actor.Receive = {
+    case status: RunStatus =>
+      log.debug(s"Polled status received: ${status}")
+      pollSuccess()
+      completePromise(Success(status))
+    case BatchApiStatusQueryFailed(_, e) =>
+      log.debug("JES poll failed!")
+      completePromise(Failure(e))
+  }
+
+  private def completePromise(result: Try[RunStatus]): Unit = {
+    pollingActorClientPromise foreach { _.complete(result) }
+    pollingActorClientPromise = None
+  }
+
+  def pollStatus(
+    workflowId: WorkflowId,
+    jobName: JobName,
+    backendSingletonActor: ActorRef,
+    requestFactory: GcpBatchRequestFactory
+  ): Future[RunStatus] =
+    pollingActorClientPromise match {
+      case Some(p) => p.future
+      case None =>
+        backendSingletonActor ! BatchStatusPollRequest(
+          workflowId,
+          self,
+          requestFactory.queryRequest(jobName),
+          StandardAsyncJob(jobName.toString)
+        )
+
+        val newPromise = Promise[RunStatus]()
+        pollingActorClientPromise = Option(newPromise)
+        newPromise.future
+    }
+}