use add for watermark accumulator

Signed-off-by: Zach Puller <[email protected]>
NVIDIA · Sep 17, 2024 · 4697e40 · 4697e40
1 parent fecf8a7
commit 4697e40
Showing 1 changed file with 1 addition and 5 deletions.
diff --git a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/GpuTaskMetrics.scala b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/GpuTaskMetrics.scala
@@ -96,10 +96,6 @@ class WatermarkAccumulator extends AccumulatorV2[jl.Long, Long] {
     _value += v
   }
 
-  def setMaxIfGreater(v: Long): Unit = {
-    _value = _value.max(v)
-  }
-
   override def merge(other: AccumulatorV2[jl.Long, Long]): Unit = other match {
     case wa: WatermarkAccumulator =>
       _value = _value.max(wa._value)
@@ -216,7 +212,7 @@ class GpuTaskMetrics extends Serializable {
   def updateMaxMemory(taskAttemptId: Long): Unit = {
     val maxMem = RmmSpark.getAndResetMaxMemoryAllocated(taskAttemptId)
     if (maxMem > 0) {
-      maxDeviceMemoryBytes.setMaxIfGreater(maxMem)
+      maxDeviceMemoryBytes.add(maxMem)
     }
   }
 }