[query] remove Backend from HailContext

Author: ehigham

hail/hail/src/is/hail/HailContext.scala

@@ -1,23 +1,14 @@
 package is.hail
 
-import is.hail.backend.Backend
 import is.hail.backend.spark.SparkBackend
 import is.hail.expr.ir.functions.IRFunctionRegistry
 import is.hail.io.fs.FS
 import is.hail.utils._
 
-import org.apache.spark._
-
 object HailContext {
 
   private var theContext: HailContext = _
 
-  def get: HailContext = synchronized {
-    assert(TaskContext.get() == null, "HailContext not available on worker")
-    assert(theContext != null, "HailContext not initialized")
-    theContext
-  }
-
   def checkJavaVersion(): Unit = {
     val javaVersion = raw"(\d+)\.(\d+)\.(\d+).*".r
     val versionString = System.getProperty("java.version")
@@ -36,13 +27,13 @@ object HailContext {
     }
   }
 
-  def getOrCreate(backend: Backend): HailContext =
+  def getOrCreate: HailContext =
     synchronized {
       if (theContext != null) theContext
-      else HailContext(backend)
+      else apply
     }
 
-  def apply(backend: Backend): HailContext = synchronized {
+  def apply: HailContext = synchronized {
     require(theContext == null)
     checkJavaVersion()
 
@@ -60,7 +51,7 @@ object HailContext {
       )
     }
 
-    theContext = new HailContext(backend)
+    theContext = new HailContext
 
     info(s"Running Hail version $HAIL_PRETTY_VERSION")
 
@@ -70,14 +61,11 @@ object HailContext {
   def stop(): Unit =
     synchronized {
       IRFunctionRegistry.clearUserFunctions()
-      theContext.backend.close()
       theContext = null
     }
 }
 
-class HailContext private (
-  var backend: Backend
-) {
+class HailContext {
   def stop(): Unit = HailContext.stop()
 
   private[this] def fileAndLineCounts(

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

@@ -43,6 +43,18 @@ object Backend {
     assert(t.isFieldDefined(off, 0))
     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
+    }
 }
 
 abstract class BroadcastValue[T] { def value: T }

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

@@ -1,6 +1,5 @@
 package is.hail.backend
 
-import is.hail.HailContext
 import is.hail.annotations.Region
 import is.hail.asm4s._
 import is.hail.backend.local.LocalTaskContext
@@ -52,7 +51,7 @@ class BackendUtils(
     val remainingPartitions =
       contexts.indices.filterNot(k => cachedResults.containsOrdered[Int](k, _ < _, _._2))
 
-    val backend = HailContext.get.backend
+    val backend = Backend.get
     val mod = getModule(modID)
     val t = System.nanoTime()
     val (failureOpt, successes) =

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

@@ -149,9 +149,9 @@ class ExecuteContext(
   )(implicit E: Enclosing
   ): T =
     using(new LocalTaskContext(0, 0)) { tc =>
-      time {
-        f(theHailClassLoader, fs, tc, r)
-      }
+      Backend.set(backend)
+      try time(f(theHailClassLoader, fs, tc, r))
+      finally Backend.set(null)
     }
 
   def createTmpPath(prefix: String, extension: String = null, local: Boolean = false): String =

hail/hail/src/is/hail/backend/driver/BatchQueryDriver.scala

@@ -194,7 +194,7 @@ object BatchQueryDriver extends HttpLikeRpc with Logging {
         jobConfig,
       )
 
-    HailContext(backend)
+    HailContext.getOrCreate
     log.info("HailContext initialized.")
 
     // FIXME: when can the classloader be shared? (optimizer benefits!)
@@ -211,7 +211,10 @@ object BatchQueryDriver extends HttpLikeRpc with Logging {
           payload,
         )
       )
-    finally HailContext.stop()
+    finally {
+      backend.close()
+      HailContext.stop()
+    }
   }
 }
 

hail/hail/src/is/hail/backend/driver/Py4JQueryDriver.scala

@@ -284,6 +284,7 @@ final class Py4JQueryDriver(backend: Backend) extends Closeable {
       compiledCodeCache.clear()
       irCache.clear()
       coercerCache.clear()
+      backend.close()
     }
 
   private[this] def removeReference(name: String): Unit =

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

@@ -177,8 +177,7 @@ object Worker {
     timer.end("readInputs")
     timer.start("executeFunction")
 
-    // FIXME: workers should not have backends, but some things do need hail contexts
-    HailContext.getOrCreate(new ServiceBackend(null, null, null, null, null))
+    HailContext.getOrCreate
     val result =
       try
         using(new ServiceTaskContext(i)) { htc =>

hail/hail/test/src/is/hail/HailSuite.scala

@@ -54,24 +54,26 @@ class HailSuite extends TestNGSuite with TestUtils {
   def setupHailContext(): Unit = {
     Logging.configureLogging("/tmp/hail.log", quiet = false, append = false)
     RVD.CheckRvdKeyOrderingForTesting = true
-    val backend = SparkBackend(
-      sc = new SparkContext(
-        SparkBackend.createSparkConf(
-          appName = "Hail.TestNG",
-          master = System.getProperty("hail.master"),
-          local = "local[2]",
-          blockSize = 0,
-        )
-          .set("spark.unsafe.exceptionOnMemoryLeak", "true")
-      ),
-      skipLoggingConfiguration = true,
+    Backend.set(
+      SparkBackend(
+        sc = new SparkContext(
+          SparkBackend.createSparkConf(
+            appName = "Hail.TestNG",
+            master = System.getProperty("hail.master"),
+            local = "local[2]",
+            blockSize = 0,
+          )
+            .set("spark.unsafe.exceptionOnMemoryLeak", "true")
+        ),
+        skipLoggingConfiguration = true,
+      )
     )
-    HailSuite.hc_ = HailContext(backend)
+    HailSuite.hc_ = HailContext.getOrCreate
   }
 
   @BeforeClass
   def setupExecuteContext(): Unit = {
-    val backend = HailSuite.hc_.backend.asSpark
+    val backend = Backend.get.asSpark
     val conf = new Configuration(backend.sc.hadoopConfiguration)
     val fs = new HadoopFS(new SerializableHadoopConfiguration(conf))
     val pool = RegionPool()
@@ -103,7 +105,7 @@ class HailSuite extends TestNGSuite with TestUtils {
 
     hadoop.fs.FileSystem.closeAll()
 
-    if (HailSuite.hc_.backend.asSpark.sc.isStopped)
+    if (SparkBackend.sparkContext.isStopped)
       throw new RuntimeException(s"'${context.getName}' stopped spark context!")
   }
 

hail/python/hail/backend/local_backend.py

@@ -78,7 +78,7 @@ def __init__(
             append,
             skip_logging_configuration,
         )
-        jhc = hail_package.HailContext.apply(jbackend)
+        jhc = hail_package.HailContext.apply()
 
         super().__init__(self._gateway.jvm, jbackend, jhc, tmpdir, tmpdir)
         self.gcs_requester_pays_configuration = gcs_requester_pays_configuration

hail/python/hail/backend/spark_backend.py

@@ -117,7 +117,7 @@ def __init__(
                 skip_logging_configuration,
                 min_block_size,
             )
-            jhc = hail_package.HailContext.getOrCreate(jbackend)
+            jhc = hail_package.HailContext.getOrCreate()
         else:
             jbackend = hail_package.backend.spark.SparkBackend.apply(
                 jsc,
@@ -130,7 +130,7 @@ def __init__(
                 skip_logging_configuration,
                 min_block_size,
             )
-            jhc = hail_package.HailContext.apply(jbackend)
+            jhc = hail_package.HailContext.apply()
 
         self._jsc = jbackend.sc()
         if sc: