[query] remove branchingFactor from HailContext`

Author: ehigham

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

@@ -92,19 +92,13 @@ object HailContext {
     }
   }
 
-  def getOrCreate(backend: Backend, branchingFactor: Int = 50): HailContext = {
-    if (theContext == null)
-      return HailContext(backend, branchingFactor)
-
-    if (theContext.branchingFactor != branchingFactor)
-      warn(
-        s"Requested branchingFactor $branchingFactor, but already initialized to ${theContext.branchingFactor}.  Ignoring requested setting."
-      )
-
-    theContext
-  }
+  def getOrCreate(backend: Backend): HailContext =
+    synchronized {
+      if (isInitialized) theContext
+      else HailContext(backend)
+    }
 
-  def apply(backend: Backend, branchingFactor: Int = 50): HailContext = synchronized {
+  def apply(backend: Backend): HailContext = synchronized {
     require(theContext == null)
     checkJavaVersion()
 
@@ -122,7 +116,7 @@ object HailContext {
       )
     }
 
-    theContext = new HailContext(backend, branchingFactor)
+    theContext = new HailContext(backend)
 
     info(s"Running Hail version ${theContext.version}")
 
@@ -164,8 +158,7 @@ object HailContext {
 }
 
 class HailContext private (
-  var backend: Backend,
-  val branchingFactor: Int,
+  var backend: Backend
 ) {
   def stop(): Unit = HailContext.stop()
 

hail/hail/src/is/hail/HailFeatureFlags.scala

@@ -2,13 +2,12 @@ package is.hail
 
 import is.hail.backend.ExecutionCache
 import is.hail.backend.spark.SparkBackend
-import is.hail.expr.ir.Optimize
+import is.hail.expr.ir.{Optimize, agg}
 import is.hail.io.fs.RequesterPaysConfig
 import is.hail.types.encoded.EType
 import is.hail.utils._
 
 import scala.collection.mutable
-
 import org.json4s.JsonAST.{JArray, JObject, JString}
 
 object HailFeatureFlags {
@@ -48,6 +47,7 @@ object HailFeatureFlags {
     ),
     (Optimize.Flags.Optimize, "HAIL_QUERY_OPTIMIZE" -> "1"),
     (Optimize.Flags.MaxOptimizerIterations, "HAIL_OPTIMIZER_ITERATIONS" -> null),
+    (agg.Flags.BranchFactor, "HAIL_BRANCH_FACTOR" -> null),
   )
 
   def fromEnv(m: Map[String, String] = sys.env): HailFeatureFlags =

hail/hail/src/is/hail/expr/ir/agg/package.scala

@@ -0,0 +1,37 @@
+package is.hail.expr.ir
+
+import is.hail.backend.ExecuteContext
+import is.hail.utils.HailException
+
+package agg {
+  object Flags {
+    val BranchFactor = "branch_factor"
+  }
+}
+
+package object agg {
+  val DefaultBranchFactor: Int = 50
+
+  def branchFactor(ctx: ExecuteContext): Int = {
+    val factor =
+      ctx.flags.lookup(Flags.BranchFactor).map { s =>
+        val factor =
+          try s.toInt
+          catch {
+            case _: NumberFormatException =>
+              throw new HailException(
+                f"'${Flags.BranchFactor}' must be a positive integer, got '$s'."
+              )
+          }
+
+        if (factor < 0)
+          throw new HailException(
+            f"'${Flags.BranchFactor}' must be greater than 0, got '$factor'."
+          )
+
+        factor
+      }
+
+    factor.getOrElse(DefaultBranchFactor)
+  }
+}

hail/hail/src/is/hail/expr/ir/lowering/LowerTableIR.scala

@@ -1,6 +1,5 @@
 package is.hail.expr.ir.lowering
 
-import is.hail.HailContext
 import is.hail.backend.ExecuteContext
 import is.hail.expr.ir.{agg, TableNativeWriter, _}
 import is.hail.expr.ir.defs._
@@ -850,7 +849,7 @@ object LowerTableIR {
           InitFromSerializedValue(i, GetTupleElement(initStateRef, i), agg.state)
         })
 
-        val branchFactor = HailContext.get.branchingFactor
+        val branchFactor = agg.branchFactor(ctx)
         val useTreeAggregate = aggs.shouldTreeAggregate && branchFactor < lc.numPartitions
         val isCommutative = aggs.isCommutative
         log.info(s"Aggregate: useTreeAggregate=$useTreeAggregate")
@@ -1681,7 +1680,7 @@ object LowerTableIR {
           val initFromSerializedStates = Begin(aggs.aggs.zipWithIndex.map { case (agg, i) =>
             InitFromSerializedValue(i, GetTupleElement(initStateRef, i), agg.state)
           })
-          val branchFactor = HailContext.get.branchingFactor
+          val branchFactor = agg.branchFactor(ctx)
           val big = aggs.shouldTreeAggregate && branchFactor < lc.numPartitions
           val (partitionPrefixSumValues, transformPrefixSum): (IR, IR => IR) = if (big) {
             val tmpDir = ctx.createTmpPath("aggregate_intermediates/")

hail/hail/src/is/hail/rvd/RVD.scala

@@ -5,7 +5,7 @@ import is.hail.annotations._
 import is.hail.asm4s.{theHailClassLoaderForSparkWorkers, HailClassLoader}
 import is.hail.backend.{ExecuteContext, HailStateManager, HailTaskContext}
 import is.hail.backend.spark.{SparkBackend, SparkTaskContext}
-import is.hail.expr.ir.InferPType
+import is.hail.expr.ir.{agg, InferPType}
 import is.hail.expr.ir.PruneDeadFields.isSupertype
 import is.hail.io._
 import is.hail.io.index.IndexWriter
@@ -722,7 +722,7 @@ class RVD(
     }
 
     if (tree) {
-      val depth = treeAggDepth(getNumPartitions, HailContext.get.branchingFactor)
+      val depth = treeAggDepth(getNumPartitions, agg.branchFactor(execCtx))
       val scale = math.max(
         math.ceil(math.pow(getNumPartitions, 1.0 / depth)).toInt,
         2,

hail/python/hail/backend/backend.py

@@ -197,6 +197,7 @@ class Backend(abc.ABC):
         "write_ir_files": ("HAIL_WRITE_IR_FILES", None),
         "optimize": ("HAIL_QUERY_OPTIMIZE", "1"),
         "max_optimizer_iterations": ("HAIL_OPTIMIZER_ITERATIONS", None),
+        "branch_factor": ("HAIL_BRANCH_FACTOR", None),
     }
 
     def _valid_flags(self) -> AbstractSet[str]:

hail/python/hail/backend/local_backend.py

@@ -3,7 +3,7 @@
 import os
 import sys
 from contextlib import ExitStack
-from typing import List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 from py4j.java_gateway import GatewayParameters, JavaGateway, launch_gateway
 
@@ -78,7 +78,7 @@ def __init__(
             append,
             skip_logging_configuration,
         )
-        jhc = hail_package.HailContext.apply(jbackend, branching_factor)
+        jhc = hail_package.HailContext.apply(jbackend)
 
         super().__init__(self._gateway.jvm, jbackend, jhc, tmpdir, tmpdir)
         self.gcs_requester_pays_configuration = gcs_requester_pays_configuration
@@ -87,7 +87,12 @@ def __init__(
         )
 
         self._logger = None
-        self._initialize_flags({})
+
+        flags: Dict[str, str] = {}
+        if branching_factor is not None:
+            flags['branch_factor'] = str(branching_factor)
+
+        self._initialize_flags(flags)
 
     def validate_file(self, uri: str) -> None:
         async_to_blocking(validate_file(uri, self._fs.afs))

hail/python/hail/backend/service_backend.py

@@ -130,6 +130,7 @@ async def create(
         regions: Optional[List[str]] = None,
         gcs_requester_pays_configuration: Optional[GCSRequesterPaysConfiguration] = None,
         gcs_bucket_allow_list: Optional[List[str]] = None,
+        branching_factor: Optional[int] = None,
     ):
         async_exit_stack = AsyncExitStack()
         billing_project = configuration_of(ConfigVariable.BATCH_BILLING_PROJECT, billing_project, None)
@@ -183,6 +184,9 @@ async def create(
                 disable_progress_bar = len(disable_progress_bar_str) > 0
 
         flags = flags or {}
+        if branching_factor is not None:
+            flags['branch_factor'] = str(branching_factor)
+
         if 'gcs_requester_pays_project' in flags or 'gcs_requester_pays_buckets' in flags:
             raise ValueError(
                 'Specify neither gcs_requester_pays_project nor gcs_requester_'

hail/python/hail/backend/spark_backend.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from typing import Any, Optional
+from typing import Any, Dict, Optional
 
 import orjson
 import pyspark
@@ -117,7 +117,7 @@ def __init__(
                 skip_logging_configuration,
                 min_block_size,
             )
-            jhc = hail_package.HailContext.getOrCreate(jbackend, branching_factor)
+            jhc = hail_package.HailContext.getOrCreate(jbackend)
         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, branching_factor)
+            jhc = hail_package.HailContext.apply(jbackend)
 
         self._jsc = jbackend.sc()
         if sc:
@@ -152,7 +152,11 @@ def __init__(
 
             jbackend.pyStartProgressBar()
 
-        self._initialize_flags({})
+        flags: Dict[str, str] = {}
+        if branching_factor is not None:
+            flags['branch_factor'] = str(branching_factor)
+
+        self._initialize_flags(flags)
 
         self._router_async_fs = RouterAsyncFS(
             gcs_kwargs={"gcs_requester_pays_configuration": gcs_requester_pays_config}

hail/python/hail/context.py

@@ -386,6 +386,7 @@ def init(
                 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':
@@ -522,6 +523,7 @@ def init_spark(
     gcs_requester_pays_configuration=nullable(oneof(str, sized_tupleof(str, sequenceof(str)))),
     regions=nullable(sequenceof(str)),
     gcs_bucket_allow_list=nullable(sequenceof(str)),
+    braching_factor=nullable(int),
 )
 async def init_batch(
     *,
@@ -545,6 +547,7 @@ async def init_batch(
     gcs_requester_pays_configuration: Optional[GCSRequesterPaysConfiguration] = None,
     regions: Optional[List[str]] = None,
     gcs_bucket_allow_list: Optional[List[str]] = None,
+    branching_factor: Optional[int] = None,
 ):
     from hail.backend.service_backend import ServiceBackend
 
@@ -563,6 +566,7 @@ async def init_batch(
         regions=regions,
         gcs_requester_pays_configuration=gcs_requester_pays_configuration,
         gcs_bucket_allow_list=gcs_bucket_allow_list,
+        branching_factor=branching_factor,
     )
 
     log = _get_log(log)