diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuExec.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuExec.scala index a3fc1bff0c9..1fb2ff6f415 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuExec.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuExec.scala @@ -158,12 +158,34 @@ sealed abstract class GpuMetric extends Serializable { def +=(v: Long): Unit def add(v: Long): Unit + private var isTimerActive = false + + final def tryActivateTimer(): Boolean = { + if (!isTimerActive) { + isTimerActive = true + true + } else { + false + } + } + + final def deactivateTimer(duration: Long): Unit = { + if (isTimerActive) { + isTimerActive = false + add(duration) + } + } + final def ns[T](f: => T): T = { - val start = System.nanoTime() - try { + if (tryActivateTimer()) { + val start = System.nanoTime() + try { + f + } finally { + deactivateTimer(System.nanoTime() - start) + } + } else { f - } finally { - add(System.nanoTime() - start) } } } diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/NvtxWithMetrics.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/NvtxWithMetrics.scala index 92a11f56123..538f117e50f 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/NvtxWithMetrics.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/NvtxWithMetrics.scala @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2021, NVIDIA CORPORATION. + * Copyright (c) 2019-2024, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -32,26 +32,32 @@ object NvtxWithMetrics { * by the amount of time spent in the range */ class NvtxWithMetrics(name: String, color: NvtxColor, val metrics: GpuMetric*) - extends NvtxRange(name, color) { + extends NvtxRange(name, color) { + val needTracks = metrics.map(_.tryActivateTimer()) private val start = System.nanoTime() override def close(): Unit = { val time = System.nanoTime() - start - metrics.foreach { metric => - metric += time + metrics.toSeq.zip(needTracks).foreach { pair => + if (pair._2) { + pair._1.deactivateTimer(time) + } } super.close() } } class MetricRange(val metrics: GpuMetric*) extends AutoCloseable { + val needTracks = metrics.map(_.tryActivateTimer()) private val start = System.nanoTime() override def close(): Unit = { val time = System.nanoTime() - start - metrics.foreach { metric => - metric += time + metrics.toSeq.zip(needTracks).foreach { pair => + if (pair._2) { + pair._1.deactivateTimer(time) + } } } } diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/MetricsSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/MetricsSuite.scala new file mode 100644 index 00000000000..580c5a2ed55 --- /dev/null +++ b/tests/src/test/scala/com/nvidia/spark/rapids/MetricsSuite.scala @@ -0,0 +1,68 @@ +/* + * Copyright (c) 2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.nvidia.spark.rapids + +import ai.rapids.cudf.NvtxColor +import com.nvidia.spark.rapids.Arm.withResource +import org.scalatest.funsuite.AnyFunSuite + +class MetricsSuite extends AnyFunSuite { + + test("GpuMetric.ns: duplicate timing on the same metrics") { + val m1 = new LocalGpuMetric() + m1.ns( + m1.ns( + Thread.sleep(100) + ) + ) + // if the timing is duplicated, the value should be around 200,000,000 + assert(m1.value < 100000000 * 1.5) + assert(m1.value > 100000000 * 0.5) + } + + test("MetricRange: duplicate timing on the same metrics") { + val m1 = new LocalGpuMetric() + val m2 = new LocalGpuMetric() + withResource(new MetricRange(m1, m2)) { _ => + withResource(new MetricRange(m2, m1)) { _ => + Thread.sleep(100) + } + } + + // if the timing is duplicated, the value should be around 200,000,000 + assert(m1.value < 100000000 * 1.5) + assert(m1.value > 100000000 * 0.5) + assert(m2.value < 100000000 * 1.5) + assert(m2.value > 100000000 * 0.5) + } + + test("NvtxWithMetrics: duplicate timing on the same metrics") { + val m1 = new LocalGpuMetric() + val m2 = new LocalGpuMetric() + withResource(new NvtxWithMetrics("a", NvtxColor.BLUE, m1, m2)) { _ => + withResource(new NvtxWithMetrics("b", NvtxColor.BLUE, m2, m1)) { _ => + Thread.sleep(100) + } + } + + // if the timing is duplicated, the value should be around 200,000,000 + assert(m1.value < 100000000 * 1.5) + assert(m1.value > 100000000 * 0.5) + assert(m2.value < 100000000 * 1.5) + assert(m2.value > 100000000 * 0.5) + } +}