[query] HailContext is vanquished

Author: ehigham

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

@@ -194,6 +194,8 @@ object BatchQueryDriver extends HttpLikeRpc with Logging {
         jobConfig,
       )
 
+    Logging.preamble()
+
     // FIXME: when can the classloader be shared? (optimizer benefits!)
     try runRpc(
         Env(

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

@@ -275,6 +275,34 @@ final class Py4JQueryDriver(backend: Backend) extends Closeable {
       IRParser.parse_blockmatrix_ir(ctx, s)
     }._1
 
+  private[this] def fileAndLineCounts(
+    regex: String,
+    files: Seq[String],
+    maxLines: Int,
+  ): Map[String, Array[WithContext[String]]] =
+    synchronized {
+      val regexp = regex.r
+      backend.asSpark.sc
+        .textFilesLines(tmpFileManager.fs.globAll(files).map(_.getPath))
+        .filter(line => regexp.findFirstIn(line.value).isDefined)
+        .take(maxLines)
+        .groupBy(_.source.file)
+    }
+
+  def pyGrepPrint(regex: String, files: Seq[String], maxLines: Int): Unit =
+    fileAndLineCounts(regex, files, maxLines).foreach { case (file, lines) =>
+      info(s"$file: ${lines.length} ${plural(lines.length, "match", "matches")}:")
+      lines.map(_.value).foreach { line =>
+        val (screen, logged) = line.truncatable().strings
+        log.info("\t" + logged)
+        println(s"\t$screen")
+      }
+    }
+
+  def pyGrepReturn(regex: String, files: Seq[String], maxLines: Int)
+    : Array[(String, Array[String])] =
+    fileAndLineCounts(regex, files, maxLines).mapValues(_.map(_.value)).toArray
+
   private[this] def addReference(rg: ReferenceGenome): Unit =
     ReferenceGenome.addFatalOnCollision(references, FastSeq(rg))
 

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

@@ -49,6 +49,7 @@ object LocalBackend extends Backend {
   ): LocalBackend.type =
     synchronized {
       if (!skipLoggingConfiguration) Logging.configureLogging(logFile, quiet, append)
+      Logging.preamble()
       this
     }
 

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

@@ -108,6 +108,7 @@ object SparkBackend {
   def createSparkConf(appName: String, master: String, local: String, blockSize: Long)
     : SparkConf = {
     require(blockSize >= 0)
+
     checkSparkCompatibility(is.hail.HAIL_SPARK_VERSION, org.apache.spark.SPARK_VERSION)
 
     val conf = new SparkConf().setAppName(appName)
@@ -219,6 +220,7 @@ object SparkBackend {
 
       // there should be only one SparkContext
       assert(sc == null || (sc eq theSparkBackend.sc))
+      Logging.preamble()
 
       val initializedMinBlockSize =
         theSparkBackend.sc.getConf.getLong(
@@ -258,6 +260,7 @@ object SparkBackend {
       if (!skipLoggingConfiguration)
         Logging.configureLogging(logFile, quiet, append)
 
+      Logging.preamble()
       var sc1 = sc
       if (sc1 == null)
         sc1 = configureAndCreateSparkContext(appName, master, local, minBlockSize)

hail/hail/src/is/hail/backend/spark/SparkBackend.scala

@@ -1,5 +1,7 @@
 package is.hail.utils
 
+import is.hail.HAIL_PRETTY_VERSION
+
 import java.util.Properties
 
 import org.apache.log4j.{LogManager, Logger, PropertyConfigurator}
@@ -94,4 +96,10 @@ object Logging {
     PropertyConfigurator.configure(logProps)
   }
 
+  def preamble(): Unit = {
+    info(s"Running Hail version $HAIL_PRETTY_VERSION")
+    val jreVersion = Runtime.version().feature()
+    if (jreVersion != 11) warn(s"Hail is tested against Java 11, found $jreVersion")
+  }
+
 }

hail/hail/src/is/hail/utils/Logging.scala

@@ -78,9 +78,8 @@ def __init__(
             append,
             skip_logging_configuration,
         )
-        jhc = hail_package.HailContext.apply()
 
-        super().__init__(self._gateway.jvm, jbackend, jhc, tmpdir, tmpdir)
+        super().__init__(self._gateway.jvm, jbackend, tmpdir, tmpdir)
         self.gcs_requester_pays_configuration = gcs_requester_pays_configuration
         self._fs = self._exit_stack.enter_context(
             RouterFS(gcs_kwargs={'gcs_requester_pays_configuration': gcs_requester_pays_configuration})

hail/python/hail/backend/local_backend.py

@@ -178,7 +178,6 @@ def __init__(
         self,
         jvm: JVMView,
         jbackend: JavaObject,
-        jhc: JavaObject,
         tmpdir: str,
         remote_tmpdir: str,
     ):
@@ -195,7 +194,6 @@ def decode_bytearray(encoded):
         self._jvm = jvm
         self._hail_package = getattr(self._jvm, 'is').hail
         self._utils_package_object = scala_package_object(self._hail_package.utils)
-        self._jhc = jhc
 
         self._jbackend = self._hail_package.backend.driver.Py4JQueryDriver(jbackend)
         self.local_tmpdir = tmpdir
@@ -337,8 +335,6 @@ def _stop_jhttp_server(self):
     def stop(self):
         self._stop_jhttp_server()
         self._jbackend.close()
-        self._jhc.stop()
-        self._jhc = None
         uninstall_exception_handler()
         super().stop()
 

hail/python/hail/backend/py4j_backend.py

@@ -117,7 +117,6 @@ def __init__(
                 skip_logging_configuration,
                 min_block_size,
             )
-            jhc = hail_package.HailContext.getOrCreate()
         else:
             jbackend = hail_package.backend.spark.SparkBackend.apply(
                 jsc,
@@ -130,17 +129,13 @@ def __init__(
                 skip_logging_configuration,
                 min_block_size,
             )
-            jhc = hail_package.HailContext.apply()
 
         self._jsc = jbackend.sc()
-        if sc:
-            self.sc = sc
-        else:
-            self.sc = pyspark.SparkContext(gateway=self._gateway, jsc=jvm.JavaSparkContext(self._jsc))
+        self.sc = sc if sc else pyspark.SparkContext(gateway=self._gateway, jsc=jvm.JavaSparkContext(self._jsc))
         self._jspark_session = jbackend.sparkSession().apply()
         self._spark_session = pyspark.sql.SparkSession(self.sc, self._jspark_session)
 
-        super().__init__(jvm, jbackend, jhc, local_tmpdir, tmpdir)
+        super().__init__(jvm, jbackend, local_tmpdir, tmpdir)
         self.gcs_requester_pays_configuration = gcs_requester_pays_config
 
         self._logger = None

hail/python/hail/backend/spark_backend.py

@@ -8,6 +8,7 @@
 
 import hail as hl
 from hail import ir
+from hail.backend.py4j_backend import Py4JBackend
 from hail.expr import (
     LocusExpression,
     StructExpression,
@@ -1062,24 +1063,21 @@ def grep(regex, path, max_count=100, *, show: bool = True, force: bool = False,
     ---
     :obj:`dict` of :class:`str` to :obj:`list` of :obj:`str`
     """
-    from hail.backend.spark_backend import SparkBackend
-
-    if isinstance(hl.current_backend(), SparkBackend):
-        jfs = Env.spark_backend('grep').fs._jfs
-        if show:
-            Env.backend()._jhc.grepPrint(jfs, regex, jindexed_seq_args(path), max_count)
-            return
-        else:
-            jarr = Env.backend()._jhc.grepReturn(jfs, regex, jindexed_seq_args(path), max_count)
-            return {x._1(): list(x._2()) for x in jarr}
+    backend = hl.current_backend()
+    if isinstance(backend, Py4JBackend):
+        jb = backend._jbackend
+        return (
+            jb.pyGrepPrint(regex, jindexed_seq_args(path), max_count)
+            if show
+            else {x._1(): list(x._2()) for x in jb.pyGrepReturn(regex, jindexed_seq_args(path), max_count)}
+        )
 
     ht = hl.import_lines(path, force=force, force_bgz=force_bgz)
     ht = ht.filter(ht.text.matches(regex))
     ht = ht.head(max_count)
     lines = ht.collect()
     if show:
-        print('\n'.join(line.file + ': ' + line.text for line in lines))
-        return
+        return print('\n'.join(line.file + ': ' + line.text for line in lines))
 
     results = defaultdict(list)
     for line in lines: