Adds aggregation across metrics for failed/succeeded and non completed stages

[sayedbilalbari]

Fixes #1552

Currently we store the stageInfo using the stageModelManager class where we map incoming stage information during the following events -

1. doSparkListenerStageCompleted
   

   spark-rapids-tools/core/src/main/scala/org/apache/spark/sql/rapids/tool/EventProcessorBase.scala

   Line 475 in 1f037fa

   app.getOrCreateStage(event.stageInfo)

2. doSparkListenerStageSubmitted. -
   

   spark-rapids-tools/core/src/main/scala/org/apache/spark/sql/rapids/tool/EventProcessorBase.scala

   Line 464 in 1f037fa

   app.getOrCreateStage(event.stageInfo)

So a stage information is updated once when a stage is submitted and once during completion.
A stageCompleted event comes for all attempts of a stage ( eg - there will be two stage Submitted and StageCompleted events for stage that fails on first attempt and succeeds on attempt 2)
This PR changes that behavior to aggregate all attempts for a stage ( failed + succeeded )

Changes -

This pull request includes several changes to improve the handling of stage attempts and task metrics in the Spark RAPIDS tool. The most important changes include adding logic to handle multiple stage attempts, modifying methods to aggregate metrics for these attempts, and updating the AccumManager to simplify task accumulation.

Handling multiple stage attempts:

• core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/AppSparkMetricsAnalyzer.scala: Added logic to handle multiple stage attempts by aggregating metrics for each attempt.
• core/src/main/scala/com/nvidia/spark/rapids/tool/profiling/ProfileClassWarehouse.scala: Introduced aggregateStageProfileMetric method to combine metrics for multiple attempts of the same stage.

Simplifying task accumulation:

• core/src/main/scala/org/apache/spark/sql/rapids/tool/EventProcessorBase.scala: Modified addAccToTask method to remove the taskId parameter and simplify task accumulation.
• core/src/main/scala/org/apache/spark/sql/rapids/tool/store/AccumInfo.scala: Updated addAccumToTask method to remove the taskId parameter.
• core/src/main/scala/org/apache/spark/sql/rapids/tool/store/AccumManager.scala: Simplified addAccToTask method by removing the taskId parameter.

Testing

This change has been tested against internal event logs and integration tests have been updated to ensure this behavior is tested for the future

[amahussein]

Did you see any change in the output as an impact of that change? If you a diff between the outputs of before/after is there any change in the metrics?

[eordentlich]

@leewyang can you take a look at this wrt qualx?

[amahussein]

Thanks @sayedbilalbari

It is interesting that this change did not trigger any change in the expected results of the unit tests. If that's the case, it will be great if we can add a test that shows this version is actually generating something different compared to the previous version.
Otherwise, it might be possible that the code had no impact at all.

[leewyang]

For qualx purposes, would there be a way to get the metrics associated with the failed stages/tasks, or will these be dropped entirely? For customers who care about the total time/resources used, it would be useful to get the metrics associated with successful stages/tasks and the failed stages/tasks, so we could tell that a job spent 90% of it's time/resources on failed stages/tasks. Otherwise, the job would appear to only take 10% of the actual time/resources used in reality.

[sayedbilalbari]

Thanks @amahussein for reviewing.
For the killed/failed tasks, as mentioned before, the way accumulable come in is different from how we are currently parsing it. The logic is as follows-

1. For killed tasks, the accumUpdates show up in the object event.reason.accumUpdates but not all of them are relevant. Ob further analysis, majority of them are zeros. They just denote the accumulables this task was supposed to deal with.
2. But the accumulables where it actually made updates show up in the taskInfo.accumulables ( as mentioned above like peakExecutoinMemory, jvmGCTime etc. )
3. So no extra logic is needed to parse the accumulableUpdates that come in with killedTasks, since we already parse the TaskInfo and use that to store relevant information.

As for no failure scenarios in the test, this change tackles two major scenarios -

1. Multiple successful attempts for a stage - will update the test case for this ( have already tested that this gives correct and different results for event logs )
2. Failed or non completed stages - test event logs did not have scenarios where a stage does not complete at all. There are scenarios where a primary attempt will fail but the secondary attempt will override the primary failure entry and hence no change in aggregated files is seen.

We can discuss this offline about what would be the best way to generate event logs like this for the test case.

[leewyang]

For qualx purposes, would there be a way to get the metrics associated with the failed stages/tasks, or will these be dropped entirely? For customers who care about the total time/resources used, it would be useful to get the metrics associated with successful stages/tasks and teh failed stages/tasks, so we could tell that a job spent 90% of it's time/resources on failed stages/stasks. Otherwise, the job would appear to only take 10% of the actual time/resources used in reality.

FYI, we have existing code that is trying to look at metrics for failed stages.

[winningsix]

For sm.duration, can we output both all stages and successFulStageAttempts? It's quite common to have failed stages for an in-product runs. For duration, we want to see the cost there.

[sayedbilalbari]

I think a better way of adding this would be another view which incorporates both failed and successful tasks/stages. Doing that segregation on a metric level does not seem like a nice idea. That would require a another view creation. We can take care of this in another issue/PR.

[wjxiz1992]

Found a case that no "sql_level_agg_task_metrics.csv" is producde with this change while 24.12.3 release version does produce. I'll file an internal issue.

[sayedbilalbari]

@wjxiz1992 Thanks for testing this out. This is happening because of the fact that the events logs had no successful attempt. Hence some of the files that dealt with the stage_aggregated results were not generated.
However post discussion with @amahussein , will change this behavior to incorporate a combined output as this can lead to confusion. Will update the PR

[leewyang]

After discussion with @sayedbilalbari and @amahussein, I think this PR is fine for qualx purposes. 👍

Note: there may be changes to the models after this PR is merged, since the stage metrics will be different than before (but theoretically more correct).

[amahussein]

Thanks @sayedbilalbari !
QQ: since this PR fixes the bug caused by iterating on all stages without considerations to the attempt.
is there any other calls to stageManager.getAllStages that might exhibit the same bug? If so, can you enumerate them so we target them in followups?

[sayedbilalbari]

Thanks @amahussein for reviewing. I see three more usages of getAllStages -

• One in the generateTimeline - this I am assuming will be taken care in the issue for dead code removal
• getStagesWithMlFunctions - this gets all stages and checks for Ml operations, Now this will basically have duplicate entries of the same stage coming in with the use of this function. No breaking change related to aggregation but might be a good idea to look into the function to check for duplicate handling
• getFailedStages - this needs no change as this filters out the attempts of a stage that might have failed. So should be good

[parthosa]

Thanks @sayedbilalbari. LGTME.