[query] use a single regionpool per thread

Author: ehigham

hail/hail/src/is/hail/annotations/RegionPool.scala

@@ -118,9 +118,7 @@ final class RegionPool private (strictMemoryCheck: Boolean, threadName: String,
     }
   }
 
-  def getRegion(): Region = getRegion(Region.REGULAR)
-
-  def getRegion(size: Int): Region = {
+  def getRegion(size: Int = Region.REGULAR): Region = {
     val r = new Region(size, this)
     r.memory = getMemory(size)
     r
@@ -154,11 +152,29 @@ final class RegionPool private (strictMemoryCheck: Boolean, threadName: String,
 
   def report(context: String): Unit = {
     val inBlocks = bytesInBlocks()
+    val (chunksAllocated, cacheHits) = chunkCache.getUsage()
 
     logger.info(
-      s"RegionPool: $context: ${readableBytes(totalAllocatedBytes)} allocated (${readableBytes(inBlocks)} blocks / " +
-        s"${readableBytes(totalAllocatedBytes - inBlocks)} chunks), regions.size = ${regions.size}, " +
-        s"$numJavaObjects current java objects, thread $threadID: $threadName"
+      s"""RegionPool: $context
+         |  thread:
+         |    id:   $threadID
+         |    name: $threadName
+         |  objects:  $numJavaObjects
+         |  allocations:
+         |    peak:     $getHighestTotalUsage
+         |    total:    ${readableBytes(totalAllocatedBytes)}
+         |    blocks:   ${readableBytes(inBlocks)}
+         |    chunks:   ${readableBytes(totalAllocatedBytes - inBlocks)}
+         |  regions:
+         |    total:  ${regions.size}
+         |    free:   ${freeRegions.size}
+         |  blocks:
+         |    total:  ${blocks.sum}
+         |    free:   ${freeBlocks.view.map(_.size).sum}
+         |  chunks:
+         |    total:  $chunksAllocated
+         |    reused: $cacheHits
+         | """.stripMargin
     )
 //    logger.info("-----------STACK_TRACES---------")
 //    val stacks: String = regions.result().toIndexedSeq.flatMap(r => r.stackTrace.map((r.getTotalChunkMemory(), _))).foldLeft("")((a: String, b) => a + "\n" + b.toString())
@@ -170,8 +186,6 @@ final class RegionPool private (strictMemoryCheck: Boolean, threadName: String,
   def scopedSmallRegion[T](f: Region => T): T = using(Region(Region.SMALL, pool = this))(f)
   def scopedTinyRegion[T](f: Region => T): T = using(Region(Region.TINY, pool = this))(f)
 
-  override def finalize(): Unit = close()
-
   private[this] var closed: Boolean = false
 
   override def close(): Unit = {

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

@@ -1,12 +1,10 @@
 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.{Compiled, IR, LoweringAnalyses, SortField, TableIR, TableReader}
 import is.hail.expr.ir.lowering.{TableStage, TableStageDependency}
 import is.hail.io.{BufferSpec, TypedCodecSpec}
-import is.hail.io.fs.FS
 import is.hail.types.RTable
 import is.hail.types.encoded.EType
 import is.hail.types.physical.PTuple
@@ -45,7 +43,7 @@ object Backend {
     codec.encode(ctx, elementType, t.loadField(off, 0), os)
   }
 
-  type PartitionFn = (Array[Byte], Array[Byte], HailTaskContext, HailClassLoader, FS) => Array[Byte]
+  type PartitionFn = Compiled[(Array[Byte], Array[Byte]) => Array[Byte]]
 }
 
 abstract class BroadcastValue[T] { def value: T }

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

@@ -3,25 +3,18 @@ package is.hail.backend
 import is.hail.annotations.Region
 import is.hail.asm4s._
 import is.hail.collection.implicits.toRichIndexedSeq
+import is.hail.expr.ir.Compiled
 import is.hail.expr.ir.analyses.SemanticHash
 import is.hail.expr.ir.lowering.TableStageDependency
-import is.hail.io.fs._
 import is.hail.utils._
 
 object BackendUtils {
   type F = AsmFunction3[Region, Array[Byte], Array[Byte], Array[Byte]]
 }
 
-class BackendUtils(
-  mods: Array[(String, (HailClassLoader, FS, HailTaskContext, Region) => BackendUtils.F)]
-) extends Logging {
+class BackendUtils(mods: Array[(String, Compiled[BackendUtils.F])]) extends Logging {
 
-  import BackendUtils.F
-
-  private[this] val loadedModules
-    : Map[String, (HailClassLoader, FS, HailTaskContext, Region) => F] = mods.toMap
-
-  def getModule(id: String): (HailClassLoader, FS, HailTaskContext, Region) => F = loadedModules(id)
+  private[this] val getModule = mods.toMap
 
   def collectDArray(
     ctx: DriverRuntimeContext,
@@ -82,7 +75,7 @@ class BackendUtils(
     tsd: Option[TableStageDependency],
   ): (Option[Throwable], IndexedSeq[(Array[Byte], Int)]) = {
 
-    val mod = getModule(modID)
+    val loadFn = getModule(modID)
     val start = System.nanoTime()
 
     val r = rtx.mapCollectPartitions(
@@ -91,10 +84,10 @@ class BackendUtils(
       stageName,
       tsd,
       partitions,
-    ) { (gs, ctx, htc, theHailClassLoader, fs) =>
-      htc.getRegionPool().scopedRegion { region =>
-        mod(theHailClassLoader, fs, htc, region)(region, ctx, gs)
-      }
+    ) {
+      (hcl, fs, htc, r) =>
+        val fn = loadFn(hcl, fs, htc, r)
+        fn(r, _, _)
     }
 
     val elapsed = System.nanoTime() - start

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

@@ -3,8 +3,7 @@ package is.hail.backend
 import is.hail.HailFeatureFlags
 import is.hail.annotations.{Region, RegionPool}
 import is.hail.asm4s.HailClassLoader
-import is.hail.backend.local.LocalTaskContext
-import is.hail.expr.ir.{BaseIR, CodeCacheKey, CompiledFunction}
+import is.hail.expr.ir.{BaseIR, CompileCache, Compiled}
 import is.hail.expr.ir.LoweredTableReader.LoweredTableReaderCoercer
 import is.hail.expr.ir.lowering.IrMetadata
 import is.hail.io.fs.FS
@@ -67,12 +66,12 @@ object ExecuteContext {
     flags: HailFeatureFlags,
     irMetadata: IrMetadata,
     blockMatrixCache: mutable.Map[String, BlockMatrix],
-    codeCache: mutable.Map[CodeCacheKey, CompiledFunction[_]],
+    compileCache: CompileCache,
     irCache: mutable.Map[Int, BaseIR],
     coercerCache: mutable.Map[Any, LoweredTableReaderCoercer],
   )(
     f: ExecuteContext => T
-  ): T = {
+  ): T =
     RegionPool.scoped { pool =>
       pool.scopedRegion { region =>
         using(new ExecuteContext(
@@ -88,13 +87,12 @@ object ExecuteContext {
           flags,
           irMetadata,
           blockMatrixCache,
-          codeCache,
+          compileCache,
           irCache,
           coercerCache,
         ))(f(_))
       }
     }
-  }
 
   def createTmpPathNoCleanup(tmpdir: String, prefix: String, extension: String = null): String = {
     val random = new SecureRandom()
@@ -113,17 +111,17 @@ class ExecuteContext(
   val backend: Backend,
   val references: Map[String, ReferenceGenome],
   val fs: FS,
-  val r: Region,
+  override val r: Region,
   val timer: ExecutionTimer,
   val tempFileManager: TempFileManager,
   val theHailClassLoader: HailClassLoader,
   val flags: HailFeatureFlags,
   val irMetadata: IrMetadata,
   val BlockMatrixCache: mutable.Map[String, BlockMatrix],
-  val CodeCache: mutable.Map[CodeCacheKey, CompiledFunction[_]],
+  val CompileCache: CompileCache,
   val PersistedIrCache: mutable.Map[Int, BaseIR],
   val PersistedCoercerCache: mutable.Map[Any, LoweredTableReaderCoercer],
-) extends Closeable {
+) extends HailTaskContext with Closeable {
 
   val rngNonce: Long =
     try
@@ -142,13 +140,14 @@ class ExecuteContext(
 
   val memo: mutable.Map[Any, Any] = new mutable.HashMap[Any, Any]()
 
-  val taskContext: HailTaskContext = new LocalTaskContext(0, 0)
+  private[this] val onCloseTasks = mutable.ArrayBuffer.empty[() => Unit]
+  override def onClose(f: () => Unit): Unit = onCloseTasks += f
 
-  def scopedExecution[T](
-    f: (HailClassLoader, FS, HailTaskContext, Region) => T
-  )(implicit E: Enclosing
-  ): T =
-    using(new LocalTaskContext(0, 0))(tc => time(f(theHailClassLoader, fs, tc, r)))
+  def run[A](f: Compiled[A])(implicit E: Enclosing): A =
+    time(f(theHailClassLoader, fs, this, r))
+
+  def scopedExecution[T](f: Compiled[T])(implicit E: Enclosing): T =
+    r.pool.scopedRegion(r => local(r = r)(_.run(f)))
 
   def createTmpPath(prefix: String, extension: String = null, local: Boolean = false): String =
     tempFileManager.newTmpPath(if (local) localTmpdir else tmpdir, prefix, extension)
@@ -162,8 +161,8 @@ class ExecuteContext(
   def shouldLogIR(): Boolean = !shouldNotLogIR()
 
   override def close(): Unit = {
+    onCloseTasks.foreach(_())
     tempFileManager.close()
-    taskContext.close()
   }
 
   def time[A](block: => A)(implicit E: Enclosing): A =
@@ -182,7 +181,7 @@ class ExecuteContext(
     flags: HailFeatureFlags = this.flags,
     irMetadata: IrMetadata = this.irMetadata,
     blockMatrixCache: mutable.Map[String, BlockMatrix] = this.BlockMatrixCache,
-    codeCache: mutable.Map[CodeCacheKey, CompiledFunction[_]] = this.CodeCache,
+    compileCache: CompileCache = this.CompileCache,
     persistedIrCache: mutable.Map[Int, BaseIR] = this.PersistedIrCache,
     persistedCoercerCache: mutable.Map[Any, LoweredTableReaderCoercer] = this.PersistedCoercerCache,
   )(
@@ -201,7 +200,7 @@ class ExecuteContext(
       flags,
       irMetadata,
       blockMatrixCache,
-      codeCache,
+      compileCache,
       persistedIrCache,
       persistedCoercerCache,
     ))(f)

hail/hail/src/is/hail/backend/HailTaskContext.scala

@@ -1,50 +1,35 @@
 package is.hail.backend
 
-import is.hail.annotations.RegionPool
-import is.hail.utils._
+import is.hail.annotations.{Region, RegionPool}
+import is.hail.utils.using
 
 import scala.collection.mutable
 
-import java.io.Closeable
+trait HailTaskContext {
 
-class TaskFinalizer {
-  val closeables = mutable.ArrayBuffer.empty[Closeable]
+  /** region whose lifetime is at least as long as this task */
+  def r: Region
 
-  def clear(): Unit =
-    closeables.clear()
-
-  def addCloseable(c: Closeable): Unit =
-    closeables += c
-
-  def closeAll(): Unit = closeables.foreach(_.close())
+  /** register an action that will be called when this task completes */
+  def onClose(f: () => Unit): Unit
 }
 
-abstract class HailTaskContext extends AutoCloseable with Logging {
-  def stageId(): Int
-
-  def partitionId(): Int
-
-  def attemptNumber(): Int
-
-  private lazy val thePool = RegionPool()
-
-  def getRegionPool(): RegionPool = thePool
+object HailTaskContext {
+  def runPartition[A](partId: Int)(f: HailTaskContext => A): A =
+    using(new PartitionContext(partId))(f)
+}
 
-  val finalizers = mutable.ArrayBuffer.empty[TaskFinalizer]
+class PartitionContext(partId: Int) extends HailTaskContext with AutoCloseable {
+  private[this] val onCloseTasks = mutable.ArrayBuffer.empty[() => Unit]
 
-  def newFinalizer(): TaskFinalizer = {
-    val f = new TaskFinalizer
-    finalizers += f
-    f
-  }
+  private[this] val pool = RegionPool()
+  override val r: Region = Region(pool = pool)
+  override def onClose(f: () => Unit): Unit = onCloseTasks += f
 
   override def close(): Unit = {
-    logger.info(
-      s"TaskReport: stage=${stageId()}, partition=${partitionId()}, attempt=${attemptNumber()}, " +
-        s"peakBytes=${thePool.getHighestTotalUsage}, peakBytesReadable=${formatSpace(thePool.getHighestTotalUsage)}, " +
-        s"chunks requested=${thePool.getUsage._1}, cache hits=${thePool.getUsage._2}"
-    )
-    finalizers.foreach(_.closeAll())
-    thePool.close()
+    onCloseTasks.foreach(_())
+    r.close()
+    pool.logStats(s"Partition $partId")
+    pool.close()
   }
 }

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

@@ -113,7 +113,7 @@ object BatchQueryDriver extends HttpLikeRpc with Logging {
           flags = env.flags,
           irMetadata = new IrMetadata(),
           blockMatrixCache = ImmutableMap.empty,
-          codeCache = ImmutableMap.empty,
+          compileCache = ImmutableMap.empty,
           irCache = ImmutableMap.empty,
           coercerCache = ImmutableMap.empty,
         )(f)

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

@@ -17,6 +17,7 @@ import is.hail.io.fs._
 import is.hail.io.reference.{IndexedFastaSequenceFile, LiftOver}
 import is.hail.sparkextras.implicits._
 import is.hail.types.physical.PStruct
+import is.hail.types.physical.stypes.SingleCodeType
 import is.hail.types.virtual.{TArray, TInterval}
 import is.hail.types.virtual.Kinds.{BlockMatrix, Matrix, Table, Value}
 import is.hail.utils._
@@ -46,7 +47,10 @@ final class Py4JQueryDriver(backend: Backend) extends Closeable with Logging {
   private[this] val hcl = new HailClassLoader(getClass.getClassLoader)
   private[this] val references = mutable.Map(ReferenceGenome.builtinReferences().toSeq: _*)
   private[this] val blockMatrixCache = mutable.Map[String, linalg.BlockMatrix]()
-  private[this] val compiledCodeCache = new Cache[CodeCacheKey, CompiledFunction[_]](50)
+
+  private[this] val compiledCodeCache =
+    new Cache[CompileCacheKey, (Option[SingleCodeType], Compiled[_])](50)
+
   private[this] val irCache = mutable.Map[Int, BaseIR]()
   private[this] val coercerCache = new Cache[Any, LoweredTableReaderCoercer](32)
   private[this] var irID: Int = 0
@@ -347,7 +351,7 @@ final class Py4JQueryDriver(backend: Backend) extends Closeable with Logging {
           flags = flags,
           irMetadata = new IrMetadata(),
           blockMatrixCache = blockMatrixCache,
-          codeCache = compiledCodeCache,
+          compileCache = compiledCodeCache,
           irCache = irCache,
           coercerCache = coercerCache,
         )(f)