Weight stage duration average by task count

[WilliamK112]

Fixes #2145

Summary

• derive the merged stage durationAvg from the pooled durationSum and numTasks
• avoid giving small retry attempts the same weight as large attempts
• keep merged-stage rounding aligned with single-attempt, job, and SQL aggregation

Impact and scope

• job-level and SQL-level duration_avg are unaffected because their aggregation already recomputes from pooled sums and task counts
• qualx consumes the corrected stage-level column as its duration_mean model feature under the default stage-level filter
• QualX is unaffected: under the default stage filter, QualX renames stage duration_avg to duration_mean but then overwrites it with duration_sum / numTasks_sum before modeling.
• the zero-task durationMin behavior remains out of scope and is tracked separately by [BUG] duration_min reports 0 when a stage produced no tasks #2151

Review updates

• rebased onto current dev (36c1c56)
• moved the regression into AnalysisSuite
• use three pooled tasks totaling 1000 ms and assert 333.3, which pins one-decimal precision
• removed assertions unrelated to the defect

Validation

Using JDK 17:

• mvn -Dbuildver=357 -DwildcardSuites=com.nvidia.spark.rapids.tool.profiling.AnalysisSuite test: 27/27 passed
• mvn -q -Dbuildver=357 verify: 993/993 passed (Spark 3.5.7, Scala 2.12)
• mvn -q -Dbuildver=321 verify: 974/974 passed (Spark 3.2.1, Scala 2.12)
• mvn -q -Pscala213 -Dbuildver=357 verify: 993/993 passed (Spark 3.5.7, Scala 2.13)
• ScalaStyle: 326 files, 0 errors, 0 warnings

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge; the corrected weighted average is consistent with existing per-stage calculation semantics and is covered by a targeted regression test.

Summary

• Introduces local merged task-count and duration-sum values for consistent aggregation.
• Calculates durationAvg using the same averaging and rounding helper as individual stage metrics.
• Adds regression coverage demonstrating that attempts with different task counts receive proportional weight.

[amahussein]

The change looks right, and the pooled form matches TaskMetricsAccumRec.finalizeAggregation exactly, so a merged stage now agrees with a single-attempt one. Two asks before merge.

1. First, could the new test move into AnalysisSuite rather than a new suite file? Same package, and it already covers this class of bug on the GPU side, where SQL and app avg pool the reporting tasks rather than re-averaging stage means is the same averaging-of-averages defect this fixes for stage attempts. Several tests there build result rows directly without a SparkSession, so the pattern fits.
2. Second, the validation: the run quoted in the description executes 1 test out of 993, never exercises Scala 2.13 or a non-default Spark version, and the 325-file scalastyle count is stale now that dev has 327. Worth redoing against current dev.

Two things the description could add, neither blocking: job-level and SQL-level duration_avg are unaffected, since TaskMetricsAccumRec recomputes from the pooled sum and count rather than reading the merged value; and the qualx featurizer consumes this column as its duration_mean model feature under the default stage-level filter, so a model input moves with it.

[amahussein]

P0: Could you move this into AnalysisSuite rather than adding a new file? Same package, and it already covers this class of bug on the GPU side: SQL and app avg pool the reporting tasks rather than re-averaging stage means is the same averaging-of-averages defect that this fixes for stage attempts. Several tests there already build result rows directly without a SparkSession, so the pattern fits.

[WilliamK112]

Done — I moved the regression into AnalysisSuite, adjacent to the existing pooled-average coverage, and removed the standalone suite.

[amahussein]

The 1 passed to calculateAverage is the rounding precision, and this assertion cannot see a change to it: 2000/101 rounds to 19.8 at one decimal place and 19.80 at two, which are the same Double. That matters because TaskMetricsAccumRec.finalizeAggregation hardcodes the same 1 for a single attempt, and the value reaches the CSV through a bare toString, so a drift would print merged rows at a different precision from unmerged ones in the same column. An input whose roundings differ, say numTasks = 3, durationSum = 1000 asserting 333.3, would pin it.

[WilliamK112]

Done — the fixture now pools two tasks totaling 800 ms with one retry task totaling 200 ms, and asserts 1000 / 3 = 333.3. This fails as 300.0 under the old averaging-of-averages implementation and pins the one-decimal precision.

[amahussein]

These two assertions were already true before the change; only the durationAvg one exercises the fix.

[WilliamK112]

Done — I removed the redundant numTasks and durationSum assertions, leaving only the durationAvg assertion that exercises this fix.

[amahussein]

Flagging for awareness; no change requested here. An attempt that ran zero tasks still produces a row, and finalizeAggregation calls resetFields() when numTasks < 1, which sets durationMin to 0. Math.min(0, x) then publishes 0 as the stage minimum. Measured: a stage whose only task took 907 ms reports duration_min of 0 once a second attempt is submitted and aborted before any task end, in an otherwise complete log. Your fix makes durationAvg immune to this, which leaves durationMin the last field in this merge that a zero-task attempt corrupts.

Filed separately as #2151.

[WilliamK112]

Thanks for isolating and filing this. I have kept the zero-task durationMin behavior out of this PR and referenced #2151 in the updated description.

[amahussein]

This is character-for-character what TaskMetricsAccumRec.finalizeAggregation uses for a single attempt, so a merged stage now publishes what one attempt carrying the same tasks would. Job and SQL rows already pooled this way, so the stage row was the only member of the family that did not.

[WilliamK112]

Addressed the review in 2394a4d and rebased onto current dev:

• moved the regression into AnalysisSuite
• changed it to the rounding-sensitive 3-task / 1000 ms case (333.3)
• removed the unrelated assertions
• kept the separate zero-task durationMin issue out of scope

I also ran the full JDK 17 matrix locally: Spark 3.5.7 / Scala 2.12 (993 tests), Spark 3.2.1 / Scala 2.12 (974 tests), and Spark 3.5.7 / Scala 2.13 (993 tests), all passing; ScalaStyle checked 326 files with no errors or warnings.

When convenient, could you please take another look?

[amahussein]

LGTM The regression is now adjacent to the existing pooled-average coverage in AnalysisSuite

• One process note for future updates: please avoid rebasing an active PR, because the force-push discards review context. Merge dev into the PR branch instead when it needs an update.
• I updated the PR description because it has icorrectly stated that the QualX inherits the fix. The corrected fact is that QualX renames stage duration_avg to duration_mean but then overwrites it with duration_sum / numTasks_sum before modeling.