[query] run ServiceBackend in Py4jQueryDriver

Author: ehigham

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.compat.immutable.LazyList
@@ -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 e.status == 400
+              && e.getMessage.contains("job group specs were not submitted in order") =>
+          val delay = delayMsForTry(Random.nextInt(attempts + 1))
+          log.warn(f"Another process is updating batch ${req.batch_id} - retrying in $delay ms.", e)
+          Thread.sleep(delay)
+          retry
+      }
+    }
 
   def getJobGroup(batchId: Int, jobGroupId: Int): JobGroupResponse =
     req
@@ -338,7 +349,7 @@ case class BatchClient private (req: Requester) extends Logging with AutoCloseab
       Thread.sleep(d.toLong)
     }
 
-    throw new AssertionError("unreachable")
+    unreachable
   }
 
   override def close(): Unit =

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/package.scala

@@ -9,7 +9,7 @@ import scala.collection.compat._
 import scala.collection.mutable
 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
@@ -379,6 +379,8 @@ package object utils
 
   def uninitialized[T]: T = null.asInstanceOf[T]
 
+  def unreachable[A]: A = throw new AssertionError("unreachable")
+
   private object mapAccumulateInstance extends MapAccumulate[Nothing, Nothing]
 
   def mapAccumulate[C[_], U] =
@@ -958,6 +960,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
+      }
+
+    unreachable
+  }
 }
 
 class CancellingExecutorService(delegate: ExecutorService) extends AbstractExecutorService {

hail/python/hail/context.py

@@ -366,26 +366,47 @@ def init(
         backend = 'batch'
 
     if backend == 'batch':
-        return hail_event_loop().run_until_complete(
-            init_batch(
+        if os.getenv('HAIL_QUERY_USE_LOCAL_DRIVER') is not None:
+            return hail.experimental.init(
+                backend=backend,
+                app_name=app_name,
                 log=log,
                 quiet=quiet,
                 append=append,
-                tmpdir=tmp_dir,
+                tmp_dir=tmp_dir,
                 default_reference=default_reference,
                 global_seed=global_seed,
                 driver_cores=driver_cores,
                 driver_memory=driver_memory,
                 worker_cores=worker_cores,
                 worker_memory=worker_memory,
                 batch_id=batch_id,
-                name_prefix=app_name,
                 gcs_requester_pays_configuration=gcs_requester_pays_configuration,
                 regions=regions,
                 gcs_bucket_allow_list=gcs_bucket_allow_list,
                 branching_factor=branching_factor,
             )
-        )
+        else:
+            return hail_event_loop().run_until_complete(
+                init_batch(
+                    log=log,
+                    quiet=quiet,
+                    append=append,
+                    tmpdir=tmp_dir,
+                    default_reference=default_reference,
+                    global_seed=global_seed,
+                    driver_cores=driver_cores,
+                    driver_memory=driver_memory,
+                    worker_cores=worker_cores,
+                    worker_memory=worker_memory,
+                    batch_id=batch_id,
+                    name_prefix=app_name,
+                    gcs_requester_pays_configuration=gcs_requester_pays_configuration,
+                    regions=regions,
+                    gcs_bucket_allow_list=gcs_bucket_allow_list,
+                    branching_factor=branching_factor,
+                )
+            )
     if backend == 'spark':
         return init_spark(
             sc=sc,

hail/python/hail/experimental/__init__.py

@@ -1,3 +1,4 @@
+from .context import init
 from .datasets import load_dataset
 from .db import DB
 from .export_entries_by_col import export_entries_by_col
@@ -35,6 +36,7 @@
     'hail_metadata',
     'haplotype_freq_em',
     'import_gtf',
+    'init',
     'ld_score',
     'ld_score_regression',
     'load_dataset',