Optimize implementation of getAggregateRawMetrics in core-tools #1468

amahussein · 2024-12-17T15:33:20Z

Signed-off-by: Ahmed Hussein (amahussein) [email protected]

Contributes to #1461

This commit improves the implementation of aggregation accross raw metrics by replacing the builtin scala collections with accumulators.

In legacy implementation:

The cost of aggregation was too high:

O(m X n): where n is the number of the tasks and m is the number of metric fields in TaskModel
Using filterKeys turned to be very expensive because
- It would eventually visit all the keys in a hashMap. For aggrgation this means that we do O(m X n^2) where m is the number of keys in the hashMap and n is the number of stages per SQL/Stage
- Internally, it will create a complete hashMap which adds to the memory stress

In new implementation:

Using a single case class with var fields is way cheaper because:
- It costs a single allocation
- It reduces cache-pollution because the accumulator object is likely to stay in the cache Vs. visiting all the tasks for each time we accumulate a field.
- This reduced the cost to O(Kn) where n is the number of tasks, and K is a constant (number of fields)
Replacing filterKeys with a loop on the stageIDs and using filter and contains:
- This guarantees that we only visit the keys we are interested in.
- There are no implicit allocations.

Impact of performance

total CPU Time: 305,576 ms ->  down by 80%
total time: 2,228,968 ms       ->  down by 80% (from 23-28 minutes to 5-6 minutes)
total allocation: 360.43 GB   ->  down by 94%

getAggRawMetrics: 
	Memory 2.55 GB
	CPU 740
    Time 740

Code changes in details

This pull request introduces significant changes to the AppSparkMetricsAnalyzer class and related utility classes to improve the aggregation of Spark metrics. The changes include the addition of helper classes for accumulating metrics and the refactoring of existing methods to use these helpers. The most important changes are outlined below:

Refactoring and Code Simplification:

Refactored the AppSparkMetricsAnalyzer class to use the new AggAccumHelper and AggAccumPhotonHelper classes for aggregating metrics, simplifying the code and improving readability. [1] [2] [3] [4]

New Helper Classes:

Added AggAccumHelper class to facilitate the accumulation of aggregate metrics, allowing for future customization and parallel processing.
Added AggAccumPhotonHelper class to extend AggAccumHelper for Photon-specific metrics, handling shuffle write values and peak memory values.

New Accumulator Classes:

Introduced JobAggAccum class to optimize the aggregation of job-level metrics by avoiding the use of the Scala collections API on each field for the entire number of tasks/stages in a job.
Introduced SQLAggAccum class to optimize the aggregation of SQL-level metrics, including the calculation of executor CPU ratio and average input bytes read.

Import and Dependency Cleanup:

Cleaned up imports in AppSparkMetricsAnalyzer.scala, removing unused imports and consolidating others for better organization. [1] [2]

These changes collectively improve the maintainability and performance of the AppSparkMetricsAnalyzer class by leveraging the new helper and accumulator classes.

Signed-off-by: Ahmed Hussein (amahussein) <[email protected]> Contributes to NVIDIA#1461 This commit improves the implementation of aggregation accross raw metrics by replacing the builtin scala collections with accumulators.

amahussein · 2024-12-17T15:36:03Z

core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/AppSparkMetricsAnalyzer.scala

@@ -517,12 +485,4 @@ object AppSparkMetricsAnalyzer  {
      arr.max
    }
  }
-


Removed because it is not used anymore

amahussein · 2024-12-17T15:36:11Z

core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/AppSparkMetricsAnalyzer.scala

      stageLevelSparkMetrics(index).put(sm.stageInfo.stageId, stageRow)
    }
  }
 }


 object AppSparkMetricsAnalyzer  {
-  def getDurations(tcs: Iterable[TaskModel]): (Long, Long, Long, Double) = {


Removed because it is not used anymore

parthosa

Thanks @amahussein for this design refactor. Minor comments.

parthosa · 2024-12-17T21:39:08Z

core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/AppSparkMetricsAnalyzer.scala

@@ -182,66 +176,55 @@ class AppSparkMetricsAnalyzer(app: AppBase) extends AppAnalysisBase(app) {
      if (app.sqlIdToStages.contains(sqlId)) {
        val stagesInSQL = app.sqlIdToStages(sqlId)
        // TODO: Should we only consider successful tasks?
-        val cachedResBySQL = stageLevelSparkMetrics(index).filterKeys(stagesInSQL.contains).values


nit: Can we combine filter and map and use collect to process in single pass?

parthosa · 2024-12-17T21:40:01Z

core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/AppSparkMetricsAnalyzer.scala

-        if (profResultsInJob.isEmpty) {
+        val jobAggAccumulator = new AggAccumHelper()
+        val perJobRec = jobAggAccumulator.accumPerJob(
+          jc.stageIds.filter(stageLevelSparkMetrics(index).contains)


Similarly, can we replace filter and map by collect?

parthosa · 2024-12-17T22:15:24Z

core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/util/TaskMetricsAccumRec.scala

+   */
+  def isEmptyAggregates: Boolean = numTasks == 0
+
+  def resetFields(): Unit = {


nit: Can we add a comment on why do we need to reset fields here?

Done! Also refactored the code to do that within the class which is better OOP

core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/util/TaskMetricsAccumRec.scala

Signed-off-by: Ahmed Hussein (amahussein) <[email protected]>

amahussein

Thanks @parthosa and @nartal1 for your comments.
I addressed your comments in the last commit and did a small change in the aggDiagnostics.

In order to test that this PR did not change the behavior, I did a diff between the output folders of the Profiler tool. the output matched perfectly.

core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/util/TaskMetricsAccumRec.scala

amahussein · 2024-12-18T15:41:31Z

core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/util/TaskMetricsAccumRec.scala

+   */
+  def isEmptyAggregates: Boolean = numTasks == 0
+
+  def resetFields(): Unit = {


Done! Also refactored the code to do that within the class which is better OOP

amahussein · 2024-12-18T15:42:09Z

core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/AppSparkMetricsAnalyzer.scala

@@ -182,66 +176,55 @@ class AppSparkMetricsAnalyzer(app: AppBase) extends AppAnalysisBase(app) {
      if (app.sqlIdToStages.contains(sqlId)) {
        val stagesInSQL = app.sqlIdToStages(sqlId)
        // TODO: Should we only consider successful tasks?
-        val cachedResBySQL = stageLevelSparkMetrics(index).filterKeys(stagesInSQL.contains).values


amahussein · 2024-12-18T15:42:26Z

core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/AppSparkMetricsAnalyzer.scala

-        if (profResultsInJob.isEmpty) {
+        val jobAggAccumulator = new AggAccumHelper()
+        val perJobRec = jobAggAccumulator.accumPerJob(
+          jc.stageIds.filter(stageLevelSparkMetrics(index).contains)


amahussein · 2024-12-18T15:43:18Z

core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/AppSparkMetricsAnalyzer.scala

-    }
-  }
-
-  def minWithEmptyHandling(arr: Iterable[Long]): Long = {


Removed because it is unused

amahussein · 2024-12-18T15:49:24Z

core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/AppSparkMetricsAnalyzer.scala

+      val nodeNames = sqlAnalyzer.stageToNodeNames.getOrElse(sm.stageInfo.stageId, emptyNodeNames)
+      val diagnosticMetricsMap =
+        sqlAnalyzer.stageToDiagnosticMetrics
+          .getOrElse(sm.stageInfo.stageId, emptyDiagnosticMetrics)


Reformated the code because it was not easy to read that withDefaultValue is applied on getOrElse

amahussein · 2024-12-18T15:49:54Z

core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/AppSparkMetricsAnalyzer.scala

+      AccumProfileResults(0, 0, AccumMetaRef.EMPTY_ACCUM_META_REF, 0L, 0L, 0L, 0L)
+    val emptyNodeNames = Seq.empty[String]
+    val emptyDiagnosticMetrics = HashMap.empty[String, AccumProfileResults]


@cindyyuanjiang

It is better to avoid creating metrics/nodeNames with empty Strings. Because it is harder to notice them and then it could lead to other problems in the CSV files or on joining based on metric names when the string is empty. That's why I replaced empty string with "N/A"

Moved the creation of default values outside the map block.

thanks @amahussein!

nartal1

Thanks @amahussein ! This is a great refactor. Runtime down by 80% and memory usage optimization is nice indeed.

parthosa

Thanks @amahussein. LGTME.

parthosa · 2024-12-18T17:27:37Z

core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/util/AggAccumHelper.scala

@@ -45,22 +40,19 @@ class AggAccumHelper {

  def accumPerStage(taskRecords: Iterable[TaskModel]): TaskMetricsAccumRec = {
    val resRec = createStageAccumRecord()
-    initializeRecord(resRec, taskRecords)


I was also wondering the need to initializeRecord() before.

amahussein added the core_tools Scope the core module (scala) label Dec 17, 2024

amahussein self-assigned this Dec 17, 2024

amahussein commented Dec 17, 2024

View reviewed changes

amahussein mentioned this pull request Dec 17, 2024

[BUG] Investigate long execution time of eventlogs #1461

Closed

amahussein requested review from parthosa, cindyyuanjiang, tgravescs and nartal1 December 17, 2024 16:07

parthosa reviewed Dec 17, 2024

View reviewed changes

nartal1 reviewed Dec 17, 2024

View reviewed changes

core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/util/TaskMetricsAccumRec.scala Show resolved Hide resolved

address reviews and fix issues in aggregateDiagnostic

6fe341c

Signed-off-by: Ahmed Hussein (amahussein) <[email protected]>

amahussein commented Dec 18, 2024

View reviewed changes

amahussein requested review from nartal1 and parthosa December 18, 2024 15:53

amahussein mentioned this pull request Dec 18, 2024

[FEA] Add IO diagnostic output for GPU slowness in Profiler tool #1451

Open

nartal1 approved these changes Dec 18, 2024

View reviewed changes

parthosa approved these changes Dec 18, 2024

View reviewed changes

amahussein merged commit 18b0472 into NVIDIA:dev Dec 18, 2024
15 checks passed

amahussein deleted the rapids-tools-1461-part02 branch December 18, 2024 17:31

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Optimize implementation of getAggregateRawMetrics in core-tools #1468

Optimize implementation of getAggregateRawMetrics in core-tools #1468

amahussein commented Dec 17, 2024 •

edited

Loading

amahussein Dec 17, 2024

amahussein Dec 17, 2024

parthosa left a comment

parthosa Dec 17, 2024

amahussein Dec 18, 2024

parthosa Dec 17, 2024

amahussein Dec 18, 2024

parthosa Dec 17, 2024

amahussein Dec 18, 2024

amahussein left a comment

amahussein Dec 18, 2024

amahussein Dec 18, 2024

amahussein Dec 18, 2024

amahussein Dec 18, 2024

amahussein Dec 18, 2024

amahussein Dec 18, 2024

cindyyuanjiang Dec 18, 2024

nartal1 left a comment

parthosa left a comment

parthosa Dec 18, 2024

@@ @@ -517,12 +485,4 @@ object AppSparkMetricsAnalyzer { @@
                     arr.max
                   }
                 }

Optimize implementation of getAggregateRawMetrics in core-tools #1468

Optimize implementation of getAggregateRawMetrics in core-tools #1468

Conversation

amahussein commented Dec 17, 2024 • edited Loading

In legacy implementation:

In new implementation:

Impact of performance

Refactoring and Code Simplification:

New Helper Classes:

New Accumulator Classes:

Import and Dependency Cleanup:

Choose a reason for hiding this comment

Choose a reason for hiding this comment

parthosa left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

amahussein left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nartal1 left a comment

Choose a reason for hiding this comment

parthosa left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

amahussein commented Dec 17, 2024 •

edited

Loading