[query] run ServiceBackend in Py4jQueryDriver

Author: ehigham

batch/batch/front_end/front_end.py

@@ -1842,11 +1842,15 @@ async def create_update(request, userdata):
     update_id, start_job_group_id, start_job_id = await _create_batch_update(
         batch_id, update_spec['token'], n_jobs, n_job_groups, user, db
     )
-    return json_response({
+
+    response = {
         'update_id': update_id,
         'start_job_group_id': start_job_group_id,
         'start_job_id': start_job_id,
-    })
+    }
+
+    log.info(response)
+    return json_response(response)
 
 
 async def _create_batch_update(
@@ -1866,7 +1870,7 @@ async def update(tx: Transaction):
         )
 
         if record:
-            return (record['update_id'], record['start_job_id'], record['start_job_group_id'])
+            return (record['update_id'], record['start_job_group_id'], record['start_job_id'])
 
         # We use FOR UPDATE so that we serialize batch update insertions
         # This is necessary to reserve job id and job group id ranges.

hail/hail/src/is/hail/backend/service/ServiceBackend.scala

@@ -1,5 +1,6 @@
 package is.hail.backend.service
 
+import is.hail.HAIL_REVISION
 import is.hail.backend._
 import is.hail.backend.Backend.PartitionFn
 import is.hail.backend.local.LocalTaskContext
@@ -12,6 +13,7 @@ import is.hail.expr.ir.analyses.SemanticHash
 import is.hail.expr.ir.lowering._
 import is.hail.services._
 import is.hail.services.JobGroupStates.{Cancelled, Failure, Success}
+import is.hail.services.oauth2.{CloudCredentials, HailCredentials}
 import is.hail.types._
 import is.hail.types.physical._
 import is.hail.utils._
@@ -26,8 +28,65 @@ import scala.util.control.NonFatal
 import java.io._
 import java.util.concurrent.Executors
 
+import com.fasterxml.jackson.core.StreamReadConstraints
+
 object ServiceBackend {
   val MaxAvailableGcsConnections = 1000
+
+  // See https://github.com/hail-is/hail/issues/14580
+  StreamReadConstraints.overrideDefaultStreamReadConstraints(
+    StreamReadConstraints.builder().maxStringLength(Integer.MAX_VALUE).build()
+  )
+
+  def pyServiceBackend(
+    name: String,
+    batchId_ : Integer,
+    billingProject: String,
+    deployConfigFile: String,
+    workerCores: String,
+    workerMemory: String,
+    storage: String,
+    cloudfuse: Array[CloudfuseConfig],
+    regions: Array[String],
+  ): ServiceBackend = {
+    val credentials: CloudCredentials =
+      HailCredentials().getOrElse(CloudCredentials(keyPath = None))
+
+    val client =
+      BatchClient(
+        DeployConfig.fromConfigFile(deployConfigFile),
+        credentials,
+      )
+
+    val batchId =
+      Option(batchId_).map(_.toInt).getOrElse {
+        client.newBatch(
+          BatchRequest(
+            billing_project = billingProject,
+            token = tokenUrlSafe,
+            n_jobs = 0,
+            attributes = Map("name" -> name),
+          )
+        )
+      }
+
+    val workerConfig =
+      BatchJobConfig(
+        workerCores,
+        workerMemory,
+        storage,
+        cloudfuse,
+        regions,
+      )
+
+    new ServiceBackend(
+      name,
+      client,
+      GitRevision(HAIL_REVISION),
+      BatchConfig(batchId, 0),
+      workerConfig,
+    )
+  }
 }
 
 case class BatchJobConfig(
@@ -122,10 +181,16 @@ class ServiceBackend(
           )
 
         stageCount += 1
-
-        Thread.sleep(600) // it is not possible for the batch to be finished in less than 600ms
-        val response = batchClient.waitForJobGroup(batchConfig.batchId, jobGroupId)
-        (response, startJobId)
+        try {
+          Thread.sleep(600) // it is not possible for the batch to be finished in less than 600ms
+          val response = batchClient.waitForJobGroup(batchConfig.batchId, jobGroupId)
+          (response, startJobId)
+        } catch {
+          case _: InterruptedException =>
+            batchClient.cancelJobGroup(batchConfig.batchId, jobGroupId)
+            Thread.currentThread().interrupt()
+            throw new CancellationException()
+        }
       }
 
       override def mapCollectPartitions(

hail/hail/src/is/hail/services/BatchClient.scala

@@ -7,7 +7,7 @@ import is.hail.services.BatchClient.{
 }
 import is.hail.services.JobGroupStates.isTerminal
 import is.hail.services.oauth2.CloudCredentials
-import is.hail.services.requests.Requester
+import is.hail.services.requests.{ClientResponseException, Requester}
 import is.hail.utils._
 
 import scala.collection.immutable.Stream.cons
@@ -270,22 +270,33 @@ case class BatchClient private (req: Requester) extends Logging with AutoCloseab
     batchId
   }
 
-  def newJobGroup(req: JobGroupRequest): (Int, Int) = {
-    val nJobs = req.jobs.length
-    val (updateId, startJobGroupId, startJobId) = beginUpdate(req.batch_id, req.token, nJobs)
-    log.info(s"Began update '$updateId' for batch '${req.batch_id}'.")
-
-    createJobGroup(updateId, req)
-    log.info(s"Created job group $startJobGroupId for batch ${req.batch_id}")
-
-    createJobsIncremental(req.batch_id, updateId, req.jobs)
-    log.info(s"Submitted $nJobs in job group $startJobGroupId for batch ${req.batch_id}")
-
-    commitUpdate(req.batch_id, updateId)
-    log.info(s"Committed update $updateId for batch ${req.batch_id}.")
-
-    (startJobGroupId, startJobId)
-  }
+  def newJobGroup(req: JobGroupRequest): (Int, Int) =
+    retryable { attempts =>
+      try {
+        val nJobs = req.jobs.length
+        val (updateId, startJobGroupId, startJobId) = beginUpdate(req.batch_id, req.token, nJobs)
+        log.info(s"Began update '$updateId' for batch '${req.batch_id}'.")
+
+        createJobGroup(updateId, req)
+        log.info(s"Created job group $startJobGroupId for batch ${req.batch_id}")
+
+        createJobsIncremental(req.batch_id, updateId, req.jobs)
+        log.info(s"Submitted $nJobs in job group $startJobGroupId for batch ${req.batch_id}")
+
+        commitUpdate(req.batch_id, updateId)
+        log.info(s"Committed update $updateId for batch ${req.batch_id}.")
+
+        (startJobGroupId, startJobId)
+      } catch {
+        case e: ClientResponseException
+            if attempts < 5 &&
+              e.status == 400 &&
+              e.getMessage.contains("job group specs were not submitted in order") =>
+          Thread.sleep(delayMsForTry(attempts))
+          log.info("retry", e)
+          retry
+      }
+    }
 
   def getJobGroup(batchId: Int, jobGroupId: Int): JobGroupResponse =
     req

hail/hail/src/is/hail/services/package.scala

@@ -4,6 +4,7 @@ import is.hail.services.requests.ClientResponseException
 import is.hail.shadedazure.com.azure.storage.common.implementation.Constants
 import is.hail.utils._
 
+import scala.annotation.tailrec
 import scala.util.Random
 
 import java.io._
@@ -30,31 +31,24 @@ package object services {
 
   private[this] val LOG_2_MAX_MULTIPLIER =
     30 // do not set larger than 30 due to integer overflow calculating multiplier
-  private[this] val DEFAULT_MAX_DELAY_MS = 60000
-  private[this] val DEFAULT_BASE_DELAY_MS = 1000
+  private[this] val DEFAULT_MAX_DELAY_MS = 60000L
+  private[this] val DEFAULT_BASE_DELAY_MS = 1000L
 
   def delayMsForTry(
     tries: Int,
-    baseDelayMs: Int = DEFAULT_BASE_DELAY_MS,
-    maxDelayMs: Int = DEFAULT_MAX_DELAY_MS,
-  ): Int = {
+    baseDelayMs: Long = DEFAULT_BASE_DELAY_MS,
+    maxDelayMs: Long = DEFAULT_MAX_DELAY_MS,
+  ): Long = {
     // Based on AWS' recommendations:
     // - https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/
     /* -
      * https://github.com/aws/aws-sdk-java/blob/master/aws-java-sdk-core/src/main/java/com/amazonaws/retry/PredefinedBackoffStrategies.java */
     val multiplier = 1L << math.min(tries, LOG_2_MAX_MULTIPLIER)
-    val ceilingForDelayMs = math.min(baseDelayMs * multiplier, maxDelayMs.toLong).toInt
-    val proposedDelayMs = ceilingForDelayMs / 2 + Random.nextInt(ceilingForDelayMs / 2 + 1)
-    return proposedDelayMs
+    val ceilingForDelayMs = math.min(baseDelayMs * multiplier, maxDelayMs).toInt
+    ceilingForDelayMs / 2L + Random.nextInt(ceilingForDelayMs / 2 + 1)
   }
 
-  def sleepBeforTry(
-    tries: Int,
-    baseDelayMs: Int = DEFAULT_BASE_DELAY_MS,
-    maxDelayMs: Int = DEFAULT_MAX_DELAY_MS,
-  ) =
-    Thread.sleep(delayMsForTry(tries, baseDelayMs, maxDelayMs).toLong)
-
+  @tailrec
   def isLimitedRetriesError(_e: Throwable): Boolean = {
     // An exception is a "retry once error" if a rare, known bug in a dependency or in a cloud
     // provider can manifest as this exception *and* that manifestation is indistinguishable from a
@@ -94,6 +88,7 @@ package object services {
     }
   }
 
+  @tailrec
   def isTransientError(_e: Throwable): Boolean = {
     // ReactiveException is package private inside reactore.core.Exception so we cannot access
     // it directly for an isInstance check. AFAICT, this is the only way to check if we received
@@ -185,14 +180,11 @@ package object services {
     }
   }
 
-  def retryTransientErrors[T](f: => T, reset: Option[() => Unit] = None): T = {
-    var tries = 0
-    while (true) {
-      try
-        return f
+  def retryTransientErrors[T](f: => T, reset: Option[() => Unit] = None): T =
+    retryable { tries =>
+      try f
       catch {
         case e: Exception =>
-          tries += 1
           val delay = delayMsForTry(tries)
           if (tries <= 5 && isLimitedRetriesError(e)) {
             log.warn(
@@ -205,14 +197,12 @@ package object services {
           } else if (tries % 10 == 0) {
             log.warn(s"Encountered $tries transient errors, most recent one was $e.")
           }
-          Thread.sleep(delay.toLong)
+          Thread.sleep(delay)
+          reset.foreach(_())
+          retry
       }
-      reset.foreach(_())
     }
 
-    throw new AssertionError("unreachable")
-  }
-
   def formatException(e: Throwable): String = {
     using(new StringWriter()) { sw =>
       using(new PrintWriter(sw)) { pw =>

hail/hail/src/is/hail/utils/StringSocketAppender.scala

@@ -13,12 +13,16 @@ object StringSocketAppender {
   // low reconnection delay because everything is local
   val DEFAULT_RECONNECTION_DELAY = 100
 
-  var theAppender: StringSocketAppender = _
+  private var theAppender: StringSocketAppender = _
 
-  def get(): StringSocketAppender = theAppender
+  def get(): StringSocketAppender =
+    synchronized {
+      if (theAppender == null) theAppender = new StringSocketAppender
+      theAppender
+    }
 }
 
-class StringSocketAppender() extends AppenderSkeleton {
+class StringSocketAppender extends AppenderSkeleton {
   private var address: InetAddress = _
   private var port: Int = _
   private var os: OutputStream = _
@@ -39,7 +43,7 @@ class StringSocketAppender() extends AppenderSkeleton {
 
   override def close(): Unit = {
     if (closed) return
-    this.closed = true
+    closed = true
     cleanUp()
   }
 
@@ -58,6 +62,10 @@ class StringSocketAppender() extends AppenderSkeleton {
       connector.interrupted = true
       connector = null // allow gc
     }
+
+    StringSocketAppender.synchronized {
+      StringSocketAppender.theAppender = null
+    }
   }
 
   private def connect(address: InetAddress, port: Int): Unit = {

hail/hail/src/is/hail/utils/package.scala

@@ -9,7 +9,7 @@ import scala.collection.{mutable, GenTraversableOnce, TraversableOnce}
 import scala.collection.compat._
 import scala.collection.mutable.ArrayBuffer
 import scala.reflect.ClassTag
-import scala.util.control.NonFatal
+import scala.util.control.{ControlThrowable, NonFatal}
 
 import java.io._
 import java.lang.reflect.Method
@@ -1038,6 +1038,23 @@ package object utils
 
   def jsonToBytes(v: JValue): Array[Byte] =
     JsonMethods.compact(v).getBytes(StandardCharsets.UTF_8)
+
+  private[this] object Retry extends ControlThrowable
+
+  def retry[A]: A = throw Retry
+
+  def retryable[A](f: Int => A): A = {
+    var attempts: Int = 0
+
+    while (true)
+      try return f(attempts)
+      catch {
+        case Retry =>
+          attempts += 1
+      }
+
+    uninitialized
+  }
 }
 
 class CancellingExecutorService(delegate: ExecutorService) extends AbstractExecutorService {