Calculate task metric aggregates on-the-fly to reduce memory usage

[sayedbilalbari]

Signed off by: Sayed Bilal Bari ( sayedbilalbari )

Fixes #1545

This PR updates the AccumInfo logic to not store the updates associated with each task and maintain rolling Statistics at the stage level. This improves the overall memory usage by qualification run.

Why this change?

• This change aims to reduce overall heap usage of the tools qualification run.
• The qualification tool parses Spark event logs, collates data in memory, and outputs aggregate metrics (SQL metrics, Job level metrics, etc.) to generate a recommendation report.
• Major bottleneck identified:
  • The tool stores all update values coming in for each Task towards an Accumulable.
  • Because Tasks are numerous, the HashMap storing these update values can become very large.
    • For example, out of 14G of usage for an event log, the update map uses ~10G.
  • These update values are primarily used to calculate min, median, max, and total contribution of a stage towards an Accumulable (e.g., total shuffle).
• Since the final usage is just min, median, max, and total, this PR changes the previous behavior:
  • On-the-fly calculation of these metrics without storing all update values in memory.

Change

1. AccumInfo:
   • Updated to remove the taskUpdateMap.
   • The stageUpdateMap now:
     • Stores a map of values associated with each task.
     • Maintains a StatisticMetric object that is updated with the latest task metrics on a rolling basis.
2. AppSQLPlanAnalyzer:
   • Previously, there was significant aggregation at the stage level.
   • Updated to use a new method for calculating metrics at the stage level.
   • Removed other methods to favor the StageLevel metrics object.
3. Profiler:
   • Timeline generation is no longer used - will be tackled in next PR
   • Updated code for TimelineGeneration to remove taskUpdate usage

Testing

• The output reports generated by this change were compared against the original output reports:
  • No change in the recommendations.
  • The only difference in raw metrics is in median values, since a rolling average is now used instead of a strict median.
    • Observed no major outliers or disparities.
• Local test cases failing only due to the slight change in median values.
  • Very small difference observed.

Improvement

Metrics are for a 34G Databricks ZSTD log:

• Total runtime improvement of 37%
• Peak heap usage improvement of 68% from 15G to 5G.
• Min heap requirement improvement of 85% from 14G to 3G

BenchmarkSuite output before change:

Benchmark :                               Best Time(ms)   Avg Time(ms)   Stdev(ms)      Avg GC Time(ms)       Avg GC Count     Stdev GC Count    Max GC Time(ms)       Max GC Count   Relative
--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
Enable_Per_SQL_Arg_Qualification                 722587         724549        2774              20098.0              364.0                 63              20601                396      1.00X

BenchmarkSuite output after change:

Benchmark` :                               Best Time(ms)   Avg Time(ms)   Stdev(ms)      Avg GC Time(ms)       Avg GC Count     Stdev GC Count    Max GC Time(ms)       Max GC Count   Relative
--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
Enable_Per_SQL_Arg_Qualification                 449096         450413        1862               2951.0              198.0                  1               2955                199      1.00X

Metrics by running with dev checkpoint enabled
Ran with JVM params ->
-Xms=16G
-Xmx=16G

Before -

Memory Marker: Post processing events, 
jvm.heap.max -> 17.17G
jvm.heap.total -> 17.17G
jvm.heap.free -> 4.2G

After -

Memory Marker: Post processing events, 
jvm.heap.max -> 17.17G
jvm.heap.total -> 17.17G
jvm.heap.free -> 14.77G

This pull request includes multiple changes aimed at improving the performance and accuracy of the Spark RAPIDS tool. The most significant changes involve the removal of unused methods, updates to the handling of accumulator statistics, and improvements to the profiling and analysis components.

Codebase simplification and removal of unused methods:

• Removed the getStageTaskIds and filterAccumTaskUpdatesForStage methods from the AppSQLPlanAnalyzer class as they are no longer needed. [1] [2] [3] [4]

Enhancements to accumulator statistics:

• Updated the AccumInfo class to use stageValuesMap for storing stage-level metrics and added methods to calculate and retrieve statistics. This change improves the accuracy of the accumulator statistics by leveraging stage-level metrics. [1] [2] [3]
• Modified the StatisticsMetrics case class to use mutable fields for max and total to support incremental updates.

Profiling and analysis improvements:

• Updated the GenerateTimeline object to use stageValuesMap for calculating semaphore wait times and other metrics, improving the accuracy of the profiling data. [1] [2]
• Adjusted the AppSparkMetricsAnalyzer class to leverage stage-level metrics for peak memory and shuffle write time calculations, enhancing the profiling of Photon applications.

Miscellaneous updates:

• Updated the copyright year in several files to 2025. [1] [2]
• Modified the profiling expectations CSV file to reflect the changes in the accumulator statistics.

These changes collectively improve the performance, accuracy, and maintainability of the Spark RAPIDS tool.

[amahussein]

Profiler - timeline generation is no longer used. Removed code for TimelineGeneration

was it possible to do some workarounds to keep that functionality for now? if we can avoid removing the feature in this PR it will be better. Otherwise, we will have to fix the CI/CDs and the QA tests

Metrics mentioned are for a 34G databricks zstd log
Total runtime comes down 27% ( from 9 min 20 sec to 6 min 50 sec )
Memory utilisation comes down 68% ( from 15G to 5G )

Is the runtime evaluation done through the benchmarkSuite?

What is meant here by "Memory utilisation"? Does it mean minimum required heap to run the eventlog? or the total allocated bytes as measured by the Java's profiler?

[sayedbilalbari]

@amahussein Thanks for the review !

• Have updated the code to keep GenerateTimeline as of now
• Added benchmarkSuite runtime and updated the exact improvement metric
• Have added precise memoryUtilisation metrics as well ( peak heap usage and min heap usage )

[amahussein]

Thanks @sayedbilalbari
This is going to be a great milestone in the optimizations.

There is a few things we need to address:

• When evaluating the memory savings, I like to see using the dev mode to show the heap size after reading the eventlog. Then we compare that with the dev branch. This method is way more reliable and consistent. It helps us to set a reference for the subsequent changes in the tools metrics. In addition, not all devs/reviewers would be using the same Java profiler tool. While the JProfiler is handy, we want to be able to converge to standard means to evaluate performance.
  • It is a bonus if there is a script that builds in dev mode to gather these numbers, but I understand if we want to keep that in a different PR.
• We want to review how incomplete stages are handled. IIRC, taskEnd events are triggered before the stageCompletion. Also, we create stageModel only during the handling of stageSubmitted. We can do that in an offline discussion.
• since we are doing average instead of median. Probably we may consider renaming the fields to average instead. we can also discuss that offline.

[sayedbilalbari]

@amahussein Thanks for the review again !

• Have added the metrics from running the dev checkpoint. Added it to the description. Regarding the script, we can add that in a separate PR
• For handling missing StageCompleted events, for each TaskEnd, the stageUpdate map is updated with the cummulative update values.
• We create the StageModel when we get a StageSubmitted and update that on StageCompleted.
• If we miss out on the StageCompleted event, the metric that we might miss out on would be the completionTime, rest should be updated with the StageSubmitted event

[amahussein]

Overall. looks great!
Some comments to make it even better :)

[amahussein]

LGTME! Thanks @sayedbilalbari !
Lets move fwd with this to make it easy to investigate the correctness of metrics