Migrate CollectTimeIterator to use new NvtxIdWithMetrics class (#12900)

We are in the process of migrating all of our nvtx ranges to use the new
documented Nvtx api. Many of the existing ranges are already wrapped in
other classes such as the `NvtxWithMetrics` class. This pr establishes a
replacement for `NvtxWithMetrics` called `NvtxIdWithMetrics` which
provides the same wrapping functionality, but using the documented Nvtx
class as input. Then, the PR begins migrating some of the many uses of
this wrapper, namely the ranges that go through the
`CollectTimeIterator`

---------

Signed-off-by: Zach Puller <zpuller@nvidia.com>

Author: zpuller

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)
     }