Spark rapids tools 1120 fea args

core/src/main/scala/org/apache/spark/rapids/tool/benchmarks/Benchmark.scala

@@ -20,7 +20,7 @@ import java.io.{OutputStream, PrintStream}
 
 import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
-import scala.concurrent.duration.{DurationInt, FiniteDuration, NANOSECONDS}
+import scala.concurrent.duration.NANOSECONDS
 
 import org.apache.commons.io.output.TeeOutputStream
 
@@ -39,9 +39,8 @@ import org.apache.spark.sql.rapids.tool.util.{RuntimeUtil, ToolsTimer}
 class Benchmark(
     name: String,
     valuesPerIteration: Long,
-    minNumIters: Int = 2,
-    warmupTime: FiniteDuration = 2.seconds,
-    minTime: FiniteDuration = 2.seconds,
+    minNumIters: Int,
+    warmUpIterations: Int,
     outputPerIteration: Boolean = false,
     output: Option[OutputStream] = None) {
   import Benchmark._
@@ -121,16 +120,16 @@ class Benchmark(
    */
   def measure(num: Long, overrideNumIters: Int)(f: ToolsTimer => Unit): Result = {
     System.gc()  // ensures garbage from previous cases don't impact this one
-    val warmupDeadline = warmupTime.fromNow
-    while (!warmupDeadline.isOverdue) {
+    var wi = 0
+    while (wi < warmUpIterations) {
       f(new ToolsTimer(-1))
+      wi += 1
     }
     val minIters = if (overrideNumIters != 0) overrideNumIters else minNumIters
-    val minDuration = if (overrideNumIters != 0) 0 else minTime.toNanos
     val runTimes = ArrayBuffer[Long]()
     var totalTime = 0L
     var i = 0
-    while (i < minIters || totalTime < minDuration) {
+    while (i < minIters) {
       val timer = new ToolsTimer(i)
       f(timer)
       val runTime = timer.totalTime()

core/src/main/scala/org/apache/spark/rapids/tool/benchmarks/BenchmarkArgs.scala

@@ -0,0 +1,29 @@
+package org.apache.spark.rapids.tool.benchmarks
+
+import org.rogach.scallop.{ScallopConf, ScallopOption}
+
+class BenchmarkArgs(arguments: Seq[String]) extends ScallopConf(arguments) {
+
+    banner("""
+Benchmarker class for running various benchmarks.
+    """)
+
+  val iterations: ScallopOption[Int] = opt[Int](name = "iterations", short = 'i', default = Some(5),
+    descr = "Total iterations to run")
+
+  val warmupIterations: ScallopOption[Int] = opt[Int](name = "warmupIterations", short = 'w' ,
+    default = Some(3), descr = "Number of warmup iterations to run")
+
+  val outputDirectory: ScallopOption[String] = opt[String](name = "outputDirectory", short = 'd',
+    default = Some("."), descr = "Directory to write output to")
+
+  val outputFormat: ScallopOption[String] = opt[String](name = "outputFormat", short = 'o',
+    default = Some("tbl"), descr = "Format of output ( tbl, json)")
+
+  val extraArgs: ScallopOption[String] = opt[String](name = "extraArgs" , short = 'a',
+    required = false,
+    descr = "Extra arguments to pass to the benchmark")
+
+  verify()
+
+}

core/src/main/scala/org/apache/spark/rapids/tool/benchmarks/BenchmarkBase.scala

@@ -31,7 +31,10 @@ abstract class BenchmarkBase {
    * Implementations of this method are supposed to use the wrapper method `runBenchmark`
    * for each benchmark scenario.
    */
-  def runBenchmarkSuite(mainArgs: Array[String]): Unit
+  def runBenchmarkSuite(iterations: Int,
+      warmUpIterations: Int,
+      outputFormat: String,
+      mainArgs: Array[String]): Unit
 
   final def runBenchmark(benchmarkName: String)(func: => Any): Unit = {
     val separator = "=" * 96
@@ -49,8 +52,10 @@ abstract class BenchmarkBase {
   def afterAll(): Unit = {}
 
   def main(args: Array[String]): Unit = {
+
+    val conf = new BenchmarkArgs(args)
     // TODO: get the dirRoot from the arguments instead
-    val dirRoot = ""
+    val dirRoot = conf.outputDirectory().stripSuffix("/")
     val resultFileName = "results.txt"
     val dir = new File(s"$dirRoot/$prefix/")
     if (!dir.exists()) {
@@ -61,7 +66,10 @@ abstract class BenchmarkBase {
       file.createNewFile()
     }
     output = Some(new FileOutputStream(file))
-    runBenchmarkSuite(args)
+    runBenchmarkSuite(conf.iterations(),
+      conf.warmupIterations(),
+      conf.outputFormat(),
+      conf.extraArgs().split("\\s+").filter(_.nonEmpty))
     output.foreach { o =>
       if (o != null) {
         o.close()

core/src/main/scala/org/apache/spark/rapids/tool/benchmarks/QualificationBenchmark.scala

@@ -0,0 +1,27 @@
+package org.apache.spark.rapids.tool.benchmarks
+
+import com.nvidia.spark.rapids.tool.qualification.QualificationArgs
+import com.nvidia.spark.rapids.tool.qualification.QualificationMain.mainInternal
+
+object QualificationBenchmark extends BenchmarkBase {
+  override def runBenchmarkSuite(iterations: Int,
+    warmUpIterations: Int,
+    outputFormat: String,
+    mainArgs: Array[String]): Unit = {
+    runBenchmark("QualificationBenchmark") {
+      val benchmarker =
+        new Benchmark(
+          "QualificationBenchmark",
+          2,
+          output = output,
+          outputPerIteration = true,
+          warmUpIterations = warmUpIterations,
+          minNumIters = iterations)
+      benchmarker.addCase("QualificationBenchmark") { _ =>
+        mainInternal(new QualificationArgs(mainArgs),
+          printStdout = true, enablePB = true)
+      }
+      benchmarker.run()
+    }
+  }
+}