Qualification tool: Add penalty for row conversions

core/src/main/scala/org/apache/spark/sql/rapids/tool/qualification/QualificationAppInfo.scala

@@ -51,6 +51,7 @@ class QualificationAppInfo(
     HashMap.empty[Long, StageTaskQualificationSummary]
   val stageIdToTaskEndSum: HashMap[Long, StageTaskQualificationSummary] =
     HashMap.empty[Long, StageTaskQualificationSummary]
+  val stageIdToGpuCpuTransitions: HashMap[Int, Int] = HashMap.empty[Int, Int]
 
   val stageIdToSqlID: HashMap[Int, Long] = HashMap.empty[Int, Long]
   val sqlIDtoFailures: HashMap[Long, ArrayBuffer[String]] = HashMap.empty[Long, ArrayBuffer[String]]
@@ -241,13 +242,49 @@ class QualificationAppInfo(
     stages.map { stageId =>
       val stageTaskTime = stageIdToTaskEndSum.get(stageId)
         .map(_.totalTaskDuration).getOrElse(0L)
+      val transitionsTime = stageIdToGpuCpuTransitions.getOrElse(stageId, 0) match {
+        case gpuCpuTransitions if gpuCpuTransitions > 0 =>
+          // Duration to transfer data from GPU to CPU and vice versa.
+          // Assuming it's a PCI-E Gen3, but also assuming that some of the result could be
+          // spilled to disk.
+          // Duration in Spark metrics is in millisecond, so multiply this by 1000 to make
+          // it consistent
+          val totalBytesRead =
+            stageIdToTaskEndSum.get(stageId).map(_.totalbytesRead).getOrElse(0L).toDouble
+          if (totalBytesRead > 0) {
+            val fallback_duration = (totalBytesRead / QualificationAppInfo.CPU_GPU_TRANSFER_RATE) *
+              QualificationAppInfo.SECONDS_TO_MILLISECONDS * gpuCpuTransitions
+            fallback_duration.toLong
+          } else {
+            0L
+          }
+        case _ => 0L
+      }
+      // Update totaltaskduration of stageIdToTaskEndSum to include transitions time
+      val stageIdToTasksMetrics = stageIdToTaskEndSum.get(stageId).orElse(None)
+      if (stageIdToTasksMetrics.isDefined) {
+        stageIdToTasksMetrics.get.totalTaskDuration += transitionsTime
+      }
       StageQualSummaryInfo(stageId, allSpeedupFactorAvg, stageTaskTime,
-        eachStageUnsupported, estimated)
+        eachStageUnsupported + transitionsTime, estimated)
     }.toSet
   }
 
   def summarizeStageLevel(execInfos: Seq[ExecInfo], sqlID: Long): Set[StageQualSummaryInfo] = {
     val (allStagesToExecs, execsNoStage) = getStageToExec(execInfos)
+
+    // Get the total number of transitions between CPU and GPU for each stage and
+    // store it in a Map.
+    allStagesToExecs.foreach { case (stageId, execs) =>
+      val topLevelExecs = execs.filterNot(x => x.exec.startsWith("WholeStage"))
+      val childrenExecs = execs.flatMap(_.children).flatten
+      val allExecs = topLevelExecs ++ childrenExecs
+      val transitions = allExecs.zip(allExecs.drop(1)).count {
+        case (exec1, exec2) =>
+          (exec1.isSupported && !exec2.isSupported) || (!exec1.isSupported && exec2.isSupported)
+      }
+      stageIdToGpuCpuTransitions(stageId) = transitions
+    }
     if (allStagesToExecs.isEmpty) {
       // use job level
       // also get the job ids associated with the SQLId
@@ -670,7 +707,8 @@ class StageTaskQualificationSummary(
     val stageAttemptId: Int,
     var executorRunTime: Long,
     var executorCPUTime: Long,
-    var totalTaskDuration: Long)
+    var totalTaskDuration: Long,
+    var totalbytesRead: Long)
 
 case class QualApplicationInfo(
     appName: String,
@@ -746,6 +784,8 @@ object QualificationAppInfo extends Logging {
   val NOT_APPLICABLE = "Not Applicable"
   val LOWER_BOUND_RECOMMENDED = 1.3
   val LOWER_BOUND_STRONGLY_RECOMMENDED = 2.5
+  val CPU_GPU_TRANSFER_RATE = 10000000L
+  val SECONDS_TO_MILLISECONDS = 1000L
 
   private def handleException(e: Exception, path: EventLogInfo): String = {
     val message: String = e match {

core/src/main/scala/org/apache/spark/sql/rapids/tool/qualification/QualificationEventProcessor.scala

@@ -72,16 +72,26 @@ class QualificationEventProcessor(app: QualificationAppInfo, perSqlOnly: Boolean
     super.doSparkListenerTaskEnd(app, event)
     // keep all stage task times to see for nonsql duration
     val taskSum = app.stageIdToTaskEndSum.getOrElseUpdate(event.stageId, {
-      new StageTaskQualificationSummary(event.stageId, event.stageAttemptId, 0, 0, 0)
+      new StageTaskQualificationSummary(event.stageId, event.stageAttemptId, 0, 0, 0, 0)
     })
     taskSum.executorRunTime += event.taskMetrics.executorRunTime
     taskSum.executorCPUTime += NANOSECONDS.toMillis(event.taskMetrics.executorCpuTime)
     taskSum.totalTaskDuration += event.taskInfo.duration
+    // Add the total bytes read from the task if it's available. This is from inputMetrics if
+    // it is reading from datasource, or shuffleReadMetrics if it is reading from shuffle.
+    val inputMetrics = event.taskMetrics.inputMetrics
+    if (inputMetrics != null) {
+      taskSum.totalbytesRead += inputMetrics.bytesRead
+    }
+    val shuffleReadMetrics = event.taskMetrics.shuffleReadMetrics
+    if (shuffleReadMetrics != null) {
+      taskSum.totalbytesRead += shuffleReadMetrics.totalBytesRead
+    }
 
     // Adds in everything (including failures)
     app.stageIdToSqlID.get(event.stageId).foreach { sqlID =>
       val taskSum = app.sqlIDToTaskEndSum.getOrElseUpdate(sqlID, {
-        new StageTaskQualificationSummary(event.stageId, event.stageAttemptId, 0, 0, 0)
+        new StageTaskQualificationSummary(event.stageId, event.stageAttemptId, 0, 0, 0, 0)
       })
       taskSum.executorRunTime += event.taskMetrics.executorRunTime
       taskSum.executorCPUTime += NANOSECONDS.toMillis(event.taskMetrics.executorCpuTime)