[query] remove global Backend field

Author: ehigham

hail/hail/src/is/hail/backend/Backend.scala

@@ -1,6 +1,7 @@
 package is.hail.backend
 
 import is.hail.asm4s.HailClassLoader
+import is.hail.backend.Backend.PartitionFn
 import is.hail.backend.spark.SparkBackend
 import is.hail.expr.ir.{IR, LoweringAnalyses, SortField, TableIR, TableReader}
 import is.hail.expr.ir.lowering.{TableStage, TableStageDependency}
@@ -44,43 +45,31 @@ object Backend {
     codec.encode(ctx, elementType, t.loadField(off, 0), os)
   }
 
-  // Currently required by `BackendUtils.collectDArray`
-  private[this] var instance: Backend = _
-
-  def set(b: Backend): Unit =
-    synchronized { instance = b }
-
-  def get: Backend =
-    synchronized {
-      assert(instance != null)
-      instance
-    }
+  type PartitionFn = (Array[Byte], Array[Byte], HailTaskContext, HailClassLoader, FS) => Array[Byte]
 }
 
 abstract class BroadcastValue[T] { def value: T }
 
-trait BackendContext {
-  def executionCache: ExecutionCache
-}
-
-abstract class Backend extends Closeable {
-
-  def defaultParallelism: Int
-
-  def canExecuteParallelTasksOnDriver: Boolean = true
+abstract class DriverContext {
 
-  def broadcast[T: ClassTag](value: T): BroadcastValue[T]
+  def executionCache: ExecutionCache
 
-  def parallelizeAndComputeWithIndex(
-    backendContext: BackendContext,
-    fs: FS,
+  def collectDistributedArray(
+    globals: Array[Byte],
     contexts: IndexedSeq[Array[Byte]],
     stageIdentifier: String,
     dependency: Option[TableStageDependency] = None,
     partitions: Option[IndexedSeq[Int]] = None,
   )(
-    f: (Array[Byte], HailTaskContext, HailClassLoader, FS) => Array[Byte]
+    f: PartitionFn
   ): (Option[Throwable], IndexedSeq[(Array[Byte], Int)])
+}
+
+abstract class Backend extends Closeable {
+
+  def defaultParallelism: Int
+
+  def broadcast[T: ClassTag](value: T): BroadcastValue[T]
 
   def asSpark(implicit E: Enclosing): SparkBackend =
     fatal(s"${getClass.getSimpleName}: ${E.value} requires SparkBackend")
@@ -118,5 +107,5 @@ abstract class Backend extends Closeable {
 
   def execute(ctx: ExecuteContext, ir: IR): Either[Unit, (PTuple, Long)]
 
-  def backendContext(ctx: ExecuteContext): BackendContext
+  def driverContext(ctx: ExecuteContext): DriverContext
 }

hail/hail/src/is/hail/backend/BackendUtils.scala

@@ -27,9 +27,7 @@ class BackendUtils(
   def getModule(id: String): (HailClassLoader, FS, HailTaskContext, Region) => F = loadedModules(id)
 
   def collectDArray(
-    backendContext: BackendContext,
-    theDriverHailClassLoader: HailClassLoader,
-    fs: FS,
+    ctx: DriverContext,
     modID: String,
     contexts: Array[Array[Byte]],
     globals: Array[Byte],
@@ -42,7 +40,7 @@ class BackendUtils(
       semhash
         .map { s =>
           log.info(s"[collectDArray|$stageName]: querying cache for $s")
-          val cachedResults = backendContext.executionCache.lookup(s)
+          val cachedResults = ctx.executionCache.lookup(s)
           log.info(s"[collectDArray|$stageName]: found ${cachedResults.length} entries for $s.")
           cachedResults
         }
@@ -51,40 +49,21 @@ class BackendUtils(
     val remainingPartitions =
       contexts.indices.filterNot(k => cachedResults.containsOrdered[Int](k, _ < _, _._2))
 
-    val backend = Backend.get
     val mod = getModule(modID)
     val t = System.nanoTime()
     val (failureOpt, successes) =
       remainingPartitions match {
         case Seq() =>
           (None, IndexedSeq.empty)
-        case Seq(k) if backend.canExecuteParallelTasksOnDriver =>
-          try
-            using(new LocalTaskContext(k, 0)) { htc =>
-              using(htc.getRegionPool().getRegion()) { r =>
-                val f = mod(theDriverHailClassLoader, fs, htc, r)
-                val res = retryTransientErrors(f(r, contexts(k), globals))
-                (None, FastSeq(res -> k))
-              }
-            }
-          catch {
-            case NonFatal(ex) =>
-              (Some(ex), IndexedSeq.empty)
-          }
+
         case partitions =>
-          val globalsBC = backend.broadcast(globals)
-          val fsConfigBC = backend.broadcast(fs.getConfiguration())
-          backend.parallelizeAndComputeWithIndex(
-            backendContext,
-            fs,
+          ctx.collectDistributedArray(
+            globals,
             contexts,
             stageName,
             tsd,
             Some(partitions),
-          ) { (ctx, htc, theHailClassLoader, fs) =>
-            val fsConfig = fsConfigBC.value
-            val gs = globalsBC.value
-            fs.setConfiguration(fsConfig)
+          ) { (gs, ctx, htc, theHailClassLoader, fs) =>
             htc.getRegionPool().scopedRegion { region =>
               mod(theHailClassLoader, fs, htc, region)(region, ctx, gs)
             }
@@ -102,7 +81,7 @@ class BackendUtils(
         _._2 < _._2,
       )
 
-    semhash.foreach(s => backendContext.executionCache.put(s, results))
+    semhash.foreach(s => ctx.executionCache.put(s, results))
     failureOpt.foreach(throw _)
 
     results.map(_._1).toArray

hail/hail/src/is/hail/backend/ExecuteContext.scala

@@ -148,11 +148,7 @@ class ExecuteContext(
     f: (HailClassLoader, FS, HailTaskContext, Region) => T
   )(implicit E: Enclosing
   ): T =
-    using(new LocalTaskContext(0, 0)) { tc =>
-      Backend.set(backend)
-      try time(f(theHailClassLoader, fs, tc, r))
-      finally Backend.set(null)
-    }
+    using(new LocalTaskContext(0, 0))(tc => time(f(theHailClassLoader, fs, tc, r)))
 
   def createTmpPath(prefix: String, extension: String = null, local: Boolean = false): String =
     tempFileManager.newTmpPath(if (local) localTmpdir else tmpdir, prefix, extension)

hail/hail/src/is/hail/backend/local/LocalBackend.scala

@@ -1,23 +1,22 @@
 package is.hail.backend.local
 
-import is.hail.CancellingExecutorService
-import is.hail.asm4s._
 import is.hail.backend._
+import is.hail.backend.Backend.PartitionFn
 import is.hail.expr.Validate
 import is.hail.expr.ir._
 import is.hail.expr.ir.analyses.SemanticHash
 import is.hail.expr.ir.lowering._
-import is.hail.io.fs._
 import is.hail.types._
 import is.hail.types.physical.PTuple
 import is.hail.utils._
 
+import scala.collection.mutable
 import scala.reflect.ClassTag
+import scala.util.control.NonFatal
 
 import java.io.PrintWriter
 
 import com.fasterxml.jackson.core.StreamReadConstraints
-import com.google.common.util.concurrent.MoreExecutors
 
 class LocalBroadcastValue[T](val value: T) extends BroadcastValue[T] with Serializable
 
@@ -53,9 +52,6 @@ object LocalBackend extends Backend {
       this
     }
 
-  private case class Context(hcl: HailClassLoader, override val executionCache: ExecutionCache)
-      extends BackendContext
-
   def broadcast[T: ClassTag](value: T): BroadcastValue[T] = new LocalBroadcastValue[T](value)
 
   private[this] var stageIdx: Int = 0
@@ -67,32 +63,54 @@ object LocalBackend extends Backend {
       current
     }
 
-  override def parallelizeAndComputeWithIndex(
-    ctx: BackendContext,
-    fs: FS,
-    contexts: IndexedSeq[Array[Byte]],
-    stageIdentifier: String,
-    dependency: Option[TableStageDependency],
-    partitions: Option[IndexedSeq[Int]],
-  )(
-    f: (Array[Byte], HailTaskContext, HailClassLoader, FS) => Array[Byte]
-  ): (Option[Throwable], IndexedSeq[(Array[Byte], Int)]) = {
-
-    val stageId = nextStageId()
-    val hcl = ctx.asInstanceOf[Context].hcl
-    runAllKeepFirstError(new CancellingExecutorService(MoreExecutors.newDirectExecutorService())) {
-      partitions.getOrElse(contexts.indices).map { i =>
-        (
-          () => using(new LocalTaskContext(i, stageId))(f(contexts(i), _, hcl, fs)),
-          i,
-        )
+  override def driverContext(ctx: ExecuteContext): DriverContext = {
+    new DriverContext {
+
+      override val executionCache: ExecutionCache =
+        ExecutionCache.fromFlags(ctx.flags, ctx.fs, ctx.localTmpdir)
+
+      override def collectDistributedArray(
+        globals: Array[Byte],
+        contexts: IndexedSeq[Array[Byte]],
+        stageIdentifier: String,
+        dependency: Option[TableStageDependency],
+        partitions: Option[IndexedSeq[Int]],
+      )(
+        f: PartitionFn
+      ): (Option[Throwable], IndexedSeq[(Array[Byte], Int)]) = {
+
+        val todo: IndexedSeq[Int] =
+          partitions.getOrElse(contexts.indices)
+
+        val results: mutable.ArrayBuffer[(Array[Byte], Int)] =
+          new mutable.ArrayBuffer(todo.length)
+
+        var failure: Option[Throwable] =
+          None
+
+        val stageId = nextStageId()
+
+        try
+          for (idx <- todo) {
+            using(new LocalTaskContext(idx, stageId)) { tx =>
+              results += {
+                (
+                  f(globals, contexts(idx), tx, ctx.theHailClassLoader, ctx.fs),
+                  idx,
+                )
+              }
+            }
+          }
+        catch {
+          case NonFatal(t) =>
+            failure = Some(t)
+        }
+
+        (failure, results)
       }
     }
   }
 
-  override def backendContext(ctx: ExecuteContext): BackendContext =
-    Context(ctx.theHailClassLoader, ExecutionCache.fromFlags(ctx.flags, ctx.fs, ctx.tmpdir))
-
   def defaultParallelism: Int = 1
 
   def close(): Unit =