Migrate CollectTimeIterator to use new NvtxIdWithMetrics class

docs/dev/nvtx_ranges.md

@@ -19,17 +19,26 @@ See [nvtx_profiling.md](https://nvidia.github.io/spark-rapids/docs/dev/nvtx_prof
 
 Name | Description
 -----|-------------
-getMapSizesByExecId|Call to internal Spark API for retrieving size and location of shuffle map output blocks
 gpuKudoSerialize|Perform kudo serialization on the gpu
+getMapSizesByExecId|Call to internal Spark API for retrieving size and location of shuffle map output blocks
+probe right|Probing the right side of a join input iterator to get the data size for preparing the join
+fetch join stream|IO time on the stream side data for the following join
+shuffled join stream|GpuShuffledHashJoinExec op is preparing build batches for join
 Acquire GPU|Time waiting for GPU semaphore to be acquired
 Release GPU|Releasing the GPU semaphore
+GpuCoalesceBatches: collect|GPU combining of small batches post-kernel processing
 gpuKudoSliceBuffers|slice kudo serialized buffers on host into partitions
+broadcast join stream|time it takes to materialize a broadcast batch on the host
 CommitShuffle|After all temporary shuffle writes are done, produce a single file (shuffle_[map_id]_0) in the commit phase
 ParallelDeserializerIterator.next|Calling next on the MT shuffle reader iterator
 queueFetched|MT shuffle manager is using the RapidsShuffleBlockFetcherIterator to queue the next set of fetched results
 WaitingForWrites|Rapids Shuffle Manager (multi threaded) is waiting for any queued writes to finish before finalizing the map output writer
+AbstractGpuCoalesceIterator|Default range for a code path in the AbstractGpuCoalesceIterator for an op which is not explicitly documented in its own range
 ThreadedWriter.write|Rapids Shuffle Manager (multi threaded) writing
 ThreadedReader.read|Rapids Shuffle Manager (multi threaded) reading
 gpuKudoCopyToHost|copy gpu kudo serialized outputs back to the host
+hash join build|IO time on the build side data for the following join
+probe left|Probing the left side of a join input iterator to get the data size for preparing the join
+build batch: collect|Perform a join where the build side fits in a single GPU batch
 BatchWait|Rapids Shuffle Manager (multi threaded) reader blocked waiting for batches to finish decoding
 RapidsCachingWriter.write|Rapids Shuffle Manager (ucx) writing

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuCoalesceBatches.scala

@@ -245,6 +245,15 @@ case class BatchedByKey(gpuOrder: Seq[SortOrder])(val cpuOrder: Seq[SortOrder])
   override def children: Seq[Expression] = gpuOrder
 }
 
+object OpNameNvtxMap {
+  private val map = Map(
+    "GpuCoalesceBatches: collect" -> NvtxRegistry.GPU_COALESCE_BATCHES_COLLECT,
+    "build batch: collect" -> NvtxRegistry.BUILD_BATCH_COLLECT
+  )
+
+  def get(opName: String): Option[NvtxId] = map.get(opName)
+}
+
 abstract class AbstractGpuCoalesceIterator(
     inputIter: Iterator[ColumnarBatch],
     goal: CoalesceSizeGoal,
@@ -261,7 +270,9 @@ abstract class AbstractGpuCoalesceIterator(
     case NoopMetric => new LocalGpuMetric
     case _ => streamTimeOrNoop
   }
-  private val iter = new CollectTimeIterator(s"$opName: collect", inputIter, streamTime)
+  private val iter = new CollectTimeIterator(
+    OpNameNvtxMap.get(s"$opName: collect").getOrElse(NvtxRegistry.GPU_COALESCE_ITERATOR),
+    inputIter, streamTime)
 
   private var batchInitialized: Boolean = false
 

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuMetrics.scala

@@ -18,8 +18,6 @@ package com.nvidia.spark.rapids
 
 import scala.collection.immutable.TreeMap
 
-import ai.rapids.cudf.NvtxColor
-import com.nvidia.spark.rapids.Arm.withResource
 import com.nvidia.spark.rapids.metrics.GpuBubbleTimerManager
 
 import org.apache.spark.{SparkContext, TaskContext}
@@ -423,17 +421,17 @@ final class LocalGpuMetric extends GpuMetric {
 }
 
 class CollectTimeIterator[T](
-    nvtxName: String,
+    nvtxId: NvtxId,
     it: Iterator[T],
     collectTime: GpuMetric) extends Iterator[T] {
   override def hasNext: Boolean = {
-    withResource(new NvtxWithMetrics(nvtxName, NvtxColor.BLUE, collectTime)) { _ =>
+    NvtxIdWithMetrics(nvtxId, collectTime) {
       it.hasNext
     }
   }
 
   override def next(): T = {
-    withResource(new NvtxWithMetrics(nvtxName, NvtxColor.BLUE, collectTime)) { _ =>
+    NvtxIdWithMetrics(nvtxId, collectTime) {
       it.next
     }
   }

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuShuffledHashJoinExec.scala

@@ -241,7 +241,7 @@ case class GpuShuffledHashJoinExec(
       (streamIter, buildIter) => {
         val (buildData, maybeBufferedStreamIter) =
           GpuShuffledHashJoinExec.prepareBuildBatchesForJoin(buildIter,
-            new CollectTimeIterator("shuffled join stream", streamIter, streamTime),
+            new CollectTimeIterator(NvtxRegistry.SHUFFLED_JOIN_STREAM, streamIter, streamTime),
             realTarget, localBuildOutput, buildGoal, subPartConf, coalesceMetrics, readOption)
 
         buildData match {
@@ -371,7 +371,7 @@ object GpuShuffledHashJoinExec extends Logging {
             } else {
               logDebug("Return multiple batches as the build side data for the following " +
                 "sub-partitioning join")
-              Right(new CollectTimeIterator("hash join build", gpuBuildIter, buildTime))
+              Right(new CollectTimeIterator(NvtxRegistry.HASH_JOIN_BUILD, gpuBuildIter, buildTime))
             }
           }
           buildTime += System.nanoTime() - startTime
@@ -421,7 +421,7 @@ object GpuShuffledHashJoinExec extends Logging {
         val safeIter = GpuSubPartitionHashJoin.safeIteratorFromSeq(spillBuf.toSeq).map { sp =>
           withRetryNoSplit(sp)(_.getColumnarBatch())
         } ++ filteredIter
-        Right(new CollectTimeIterator("hash join build", safeIter, buildTime))
+        Right(new CollectTimeIterator(NvtxRegistry.HASH_JOIN_BUILD, safeIter, buildTime))
       } else {
         // The size after filtering is within the target size or sub-partitioning is disabled.
         while(filteredIter.hasNext) {

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuShuffledSizedHashJoinExec.scala

@@ -672,8 +672,10 @@ object GpuShuffledSymmetricHashJoinExec {
       val rightTime = new LocalGpuMetric
       val buildTime = metrics(BUILD_TIME)
       val streamTime = metrics(STREAM_TIME)
-      val leftIter = new CollectTimeIterator("probe left", setupForProbe(rawLeftIter), leftTime)
-      val rightIter = new CollectTimeIterator("probe right", setupForProbe(rawRightIter), rightTime)
+      val leftIter = new CollectTimeIterator(NvtxRegistry.PROBE_LEFT,
+        setupForProbe(rawLeftIter), leftTime)
+      val rightIter = new CollectTimeIterator(NvtxRegistry.PROBE_RIGHT,
+        setupForProbe(rawRightIter), rightTime)
       closeOnExcept(mutable.Queue.empty[T]) { leftQueue =>
         closeOnExcept(mutable.Queue.empty[T]) { rightQueue =>
           var leftSize = 0L
@@ -723,7 +725,7 @@ object GpuShuffledSymmetricHashJoinExec {
           } else {
             baseBuildIter
           }
-          val streamIter = new CollectTimeIterator("fetch join stream",
+          val streamIter = new CollectTimeIterator(NvtxRegistry.FETCH_JOIN_STREAM,
             setupForJoin(streamQueue, rawStreamIter, exprs.streamTypes, gpuBatchSizeBytes, metrics),
             streamTime)
           JoinInfo(joinType, buildSide, buildIter, buildSize, None, streamIter, exprs)

sql-plugin/src/main/scala/com/nvidia/spark/rapids/NvtxRangeWithDoc.scala

@@ -134,6 +134,34 @@ object NvtxRegistry {
   val GET_MAP_SIZES_BY_EXEC_ID: NvtxId = NvtxId("getMapSizesByExecId", NvtxColor.CYAN,
     "Call to internal Spark API for retrieving size and location of shuffle map output blocks")
 
+  val GPU_COALESCE_BATCHES_COLLECT: NvtxId = NvtxId("GpuCoalesceBatches: collect", NvtxColor.BLUE,
+    "GPU combining of small batches post-kernel processing")
+
+  val BUILD_BATCH_COLLECT: NvtxId = NvtxId("build batch: collect", NvtxColor.BLUE,
+    "Perform a join where the build side fits in a single GPU batch")
+
+  val GPU_COALESCE_ITERATOR: NvtxId = NvtxId("AbstractGpuCoalesceIterator", NvtxColor.BLUE,
+    "Default range for a code path in the AbstractGpuCoalesceIterator for an op which " +
+      "is not explicitly documented in its own range")
+
+  val SHUFFLED_JOIN_STREAM: NvtxId = NvtxId("shuffled join stream", NvtxColor.BLUE,
+    "GpuShuffledHashJoinExec op is preparing build batches for join")
+
+  val HASH_JOIN_BUILD: NvtxId = NvtxId("hash join build", NvtxColor.BLUE,
+    "IO time on the build side data for the following join")
+
+  val PROBE_LEFT: NvtxId = NvtxId("probe left", NvtxColor.BLUE,
+    "Probing the left side of a join input iterator to get the data size for preparing the join")
+
+  val PROBE_RIGHT: NvtxId = NvtxId("probe right", NvtxColor.BLUE,
+    "Probing the right side of a join input iterator to get the data size for preparing the join")
+
+  val FETCH_JOIN_STREAM: NvtxId = NvtxId("fetch join stream", NvtxColor.BLUE,
+    "IO time on the stream side data for the following join")
+
+  val BROADCAST_JOIN_STREAM: NvtxId = NvtxId("broadcast join stream", NvtxColor.BLUE,
+    "time it takes to materialize a broadcast batch on the host")
+
   val GPU_KUDO_SERIALIZE: NvtxId = NvtxId("gpuKudoSerialize", NvtxColor.YELLOW,
     "Perform kudo serialization on the gpu")
 
@@ -158,6 +186,15 @@ object NvtxRegistry {
     register(QUEUE_FETCHED)
     register(RAPIDS_CACHING_WRITER_WRITE)
     register(GET_MAP_SIZES_BY_EXEC_ID)
+    register(GPU_COALESCE_BATCHES_COLLECT)
+    register(BUILD_BATCH_COLLECT)
+    register(GPU_COALESCE_ITERATOR)
+    register(SHUFFLED_JOIN_STREAM)
+    register(HASH_JOIN_BUILD)
+    register(PROBE_LEFT)
+    register(PROBE_RIGHT)
+    register(FETCH_JOIN_STREAM)
+    register(BROADCAST_JOIN_STREAM)
     register(GPU_KUDO_SERIALIZE)
     register(GPU_KUDO_COPY_TO_HOST)
     register(GPU_KUDO_SLICE_BUFFERS)

sql-plugin/src/main/scala/com/nvidia/spark/rapids/NvtxWithMetrics.scala

@@ -53,6 +53,29 @@ class NvtxWithMetrics(name: String, color: NvtxColor, val metrics: Seq[GpuMetric
   }
 }
 
+/**
+ *  NvtxId with option to pass one or more nano timing metric(s) that are updated upon close
+ *  by the amount of time spent in the range
+ */
+object NvtxIdWithMetrics {
+
+  def apply[V](nvtxId: NvtxId, metrics: GpuMetric*)(block: => V): V = {
+    val needTracks = metrics.map(_.tryActivateTimer(Seq.empty))
+    val start = System.nanoTime()
+
+    try {
+      nvtxId.apply(block)
+    } finally {
+      val time = System.nanoTime() - start
+      metrics.toSeq.zip(needTracks).foreach { pair =>
+        if (pair._2) {
+          pair._1.deactivateTimer(time, Seq.empty)
+        }
+      }
+    }
+  }
+}
+
 class MetricRange(val metrics: Seq[GpuMetric], val excludeMetric: Seq[GpuMetric] = Seq.empty)
   extends AutoCloseable {
 

sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastHashJoinExecBase.scala

@@ -151,7 +151,7 @@ abstract class GpuBroadcastHashJoinExecBase(
         GpuBroadcastHelper.getBroadcastBuiltBatchAndStreamIter(
           broadcastRelation,
           buildSchema,
-          new CollectTimeIterator("broadcast join stream", it, streamTime))
+          new CollectTimeIterator(NvtxRegistry.BROADCAST_JOIN_STREAM, it, streamTime))
       // builtBatch will be closed in doJoin
       doJoin(builtBatch, streamIter, targetSize, numOutputRows, numOutputBatches, opTime, joinTime)
     }