diff --git a/core/src/main/resources/bootstrap/tuningConfigs.yaml b/core/src/main/resources/bootstrap/tuningConfigs.yaml index 1f0e76049..1e22b2ceb 100644 --- a/core/src/main/resources/bootstrap/tuningConfigs.yaml +++ b/core/src/main/resources/bootstrap/tuningConfigs.yaml @@ -218,6 +218,35 @@ default: default: 2 usedBy: spark.sql.shuffle.partitions + # Downward shuffle-partition tuning configs. + # These drive the final downward-only pass that lowers an oversized shuffle partition + # recommendation based on the total uncompressed shuffle input of the worst consumer stage. + - name: DOWNWARD_SHUFFLE_ENABLED + description: >- + Enables the final downward-only shuffle partition pass. When disabled, the AutoTuner + never lowers the shuffle partition recommendation produced by the normal tuning passes. + Disabled by default, so enabling it is a per-run opt-in. What has to be settled before it + defaults to true is the target partition size, not a count of validated applications: on a + large workload the hand-tuned value, the ColumnarExchange bound and this pass at a 1 GiB + target disagree, and this pass asks for the most partitions of the three. + default: false + usedBy: spark.sql.shuffle.partitions, spark.sql.adaptive.coalescePartitions.initialPartitionNum + + - name: DOWNWARD_SHUFFLE_TARGET_PARTITION_SIZE + description: >- + Target amount of estimated GPU shuffle input processed by a single partition. The + per-stage partition requirement is the estimated stage input divided by this value. + default: 1g + usedBy: spark.sql.shuffle.partitions, spark.sql.adaptive.coalescePartitions.initialPartitionNum + + - name: DOWNWARD_SHUFFLE_INPUT_SIZE_FACTOR + description: >- + Factor applied to the measured uncompressed shuffle input to estimate the GPU input size. + Profiling reads measured GpuColumnarExchange data, so it uses 1.0. Qualification reads CPU + Exchange data and overrides this in the qualification section. + default: 1.0 + usedBy: spark.sql.shuffle.partitions, spark.sql.adaptive.coalescePartitions.initialPartitionNum + - name: WORKER_GPU_COUNT description: >- Default number of GPUs per worker node @@ -450,6 +479,13 @@ qualification: default: 1g usedBy: spark.rapids.sql.batchSizeBytes + - name: DOWNWARD_SHUFFLE_INPUT_SIZE_FACTOR + description: >- + Qualification reads uncompressed CPU Exchange data, so the GPU input size is estimated + rather than measured. This factor is an initial conservative estimate. + default: 0.8 + usedBy: spark.sql.shuffle.partitions, spark.sql.adaptive.coalescePartitions.initialPartitionNum + # Profiling tool specific tuning configs profiling: - name: BATCH_SIZE_BYTES diff --git a/core/src/main/scala/com/nvidia/spark/rapids/tool/AppSummaryInfoBaseProvider.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/AppSummaryInfoBaseProvider.scala index 99d5b717e..d54036fd7 100644 --- a/core/src/main/scala/com/nvidia/spark/rapids/tool/AppSummaryInfoBaseProvider.scala +++ b/core/src/main/scala/com/nvidia/spark/rapids/tool/AppSummaryInfoBaseProvider.scala @@ -17,10 +17,12 @@ package com.nvidia.spark.rapids.tool import com.nvidia.spark.rapids.tool.analysis.AggRawMetricsResult +import com.nvidia.spark.rapids.tool.analysis.AppSQLPlanAnalyzer import com.nvidia.spark.rapids.tool.profiling.{AppInfoColumnarExchangeMetrics, AppInfoJobStageAggMetricsVisitor, AppInfoPropertyGetter, AppInfoReadMetrics, - AppInfoSqlTaskAggMetricsVisitor, AppInfoSQLTaskInputSizes, BaseProfilingAppSummaryInfoProvider, - DataSourceProfileResult, ProfilerResult, PySparkMemoryEvidence, SingleAppSummaryInfoProvider} + AppInfoShuffleStageInputMetrics, AppInfoSqlTaskAggMetricsVisitor, AppInfoSQLTaskInputSizes, + BaseProfilingAppSummaryInfoProvider, DataSourceProfileResult, ProfilerResult, + PySparkMemoryEvidence, SingleAppSummaryInfoProvider} import com.nvidia.spark.rapids.tool.tuning.QualAppSummaryInfoProvider import org.apache.spark.sql.rapids.tool.ToolUtils @@ -35,7 +37,8 @@ class AppSummaryInfoBaseProvider extends AppInfoPropertyGetter with AppInfoSqlTaskAggMetricsVisitor with AppInfoSQLTaskInputSizes with AppInfoReadMetrics - with AppInfoColumnarExchangeMetrics { + with AppInfoColumnarExchangeMetrics + with AppInfoShuffleStageInputMetrics { def isAppInfoAvailable = false override def getAllProperties: Map[String, String] = Map[String, String]() override def getSparkProperty(propKey: String): Option[String] = None @@ -89,12 +92,15 @@ object AppSummaryInfoBaseProvider { * tool. * @param appInfo * @param appAggStats optional aggregate of application stats + * @param sqlAnalyzer the SQL plan analyzer already built for this application, reused so the + * plans are not traversed a second time * @return object that can be used by the AutoTuner to calculate the recommendations */ def fromQualAppInfo(appInfo: QualificationAppInfo, appAggStats: Option[QualificationSummaryInfo] = None, rawAggMetrics: AggRawMetricsResult, - dsInfo: Seq[DataSourceProfileResult]): AppSummaryInfoBaseProvider = { - new QualAppSummaryInfoProvider(appInfo, appAggStats, rawAggMetrics, dsInfo) + dsInfo: Seq[DataSourceProfileResult], + sqlAnalyzer: Option[AppSQLPlanAnalyzer] = None): AppSummaryInfoBaseProvider = { + new QualAppSummaryInfoProvider(appInfo, appAggStats, rawAggMetrics, dsInfo, sqlAnalyzer) } } diff --git a/core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/AppSQLPlanAnalyzer.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/AppSQLPlanAnalyzer.scala index 0d57edece..3db738213 100644 --- a/core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/AppSQLPlanAnalyzer.scala +++ b/core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/AppSQLPlanAnalyzer.scala @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024-2025, NVIDIA CORPORATION. + * Copyright (c) 2024-2026, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -20,7 +20,7 @@ import scala.collection.mutable.{AbstractSet, ArrayBuffer, HashMap, LinkedHashSe import com.nvidia.spark.rapids.tool.analysis.util.IOAccumDiagnosticMetrics._ import com.nvidia.spark.rapids.tool.analysis.util.StageAccumDiagnosticMetrics._ -import com.nvidia.spark.rapids.tool.profiling.{AccumProfileResults, IODiagnosticResult, SQLAccumProfileResults, SQLMetricInfoCase, SQLStageInfoProfileResult, UnsupportedSQLPlan, WholeStageCodeGenResults} +import com.nvidia.spark.rapids.tool.profiling.{AccumProfileResults, IODiagnosticResult, ShuffleStageInputAnalysis, SQLAccumProfileResults, SQLMetricInfoCase, SQLStageInfoProfileResult, UnsupportedSQLPlan, WholeStageCodeGenResults} import org.apache.spark.sql.rapids.tool.{AppBase, RDDCheckHelper} import org.apache.spark.sql.rapids.tool.plangraph.{SparkPlanGraphCluster, SparkPlanGraphNode, ToolsPlanGraph} @@ -68,6 +68,13 @@ class AppSQLPlanAnalyzer(app: AppBase) val IODiagnosticMetricsMap: HashMap[(Long, Long), ArrayBuffer[SQLAccumProfileResults]] = HashMap.empty[(Long, Long), ArrayBuffer[SQLAccumProfileResults]] + /** + * Raw consumer-stage shuffle input inventory used by the AutoTuner's downward shuffle-partition + * pass. It is computed lazily and cached so that both tools reuse this analyzer's plan graphs + * instead of traversing the SQL plans a second time. + */ + lazy val shuffleStageInputAnalysis: ShuffleStageInputAnalysis = ShuffleStageInputAnalyzer(app) + /** * Updates the stageToDiagnosticMetrics mapping with the provided AccumProfileResults. * @param accum AccumProfileResults instance containing diagnostic metrics to be added diff --git a/core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/ShuffleStageInputAnalyzer.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/ShuffleStageInputAnalyzer.scala new file mode 100644 index 000000000..7416e01ab --- /dev/null +++ b/core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/ShuffleStageInputAnalyzer.scala @@ -0,0 +1,415 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.nvidia.spark.rapids.tool.analysis + +import scala.collection.mutable +import scala.util.Try + +import com.nvidia.spark.rapids.tool.profiling.{ShuffleInputProvenance, ShuffleStageInputAnalysis, ShuffleStageInputIncompleteReason, ShuffleStageInputRecord} + +import org.apache.spark.internal.Logging +import org.apache.spark.sql.rapids.tool.AppBase +import org.apache.spark.sql.rapids.tool.plangraph.{SparkPlanGraphNode, ToolsPlanGraph} +import org.apache.spark.sql.rapids.tool.store.{StageModel, TaskModel} + +/** + * Builds the raw consumer-stage shuffle input inventory used by the downward shuffle-partition + * pass. + * + * Sizing shuffle exchanges in isolation is not enough, because a join stage can consume several + * shuffled inputs at once. This analyzer therefore attributes each exchange's uncompressed + * `data size` to every downstream branch that consumes it and totals those branches per consumer + * stage. An exchange reused through two branches into the same stage contributes twice, because + * each branch processes the data. + * + * Lowering a global partition count is more dangerous than overestimating it, so the analysis is + * all-or-nothing: an unsupported shuffle node, a missing size metric, an ambiguous consumer + * mapping, a SQL execution that did not terminate successfully, or a consumer stage without a + * completed successful attempt marks the whole application incomplete rather than silently + * dropping the affected stage. + */ +class ShuffleStageInputAnalyzer(app: AppBase) extends Logging { + + import ShuffleStageInputAnalyzer._ + + /** Accumulated bytes and branch count for one (SQL execution, consumer stage). */ + private case class BranchTotals(bytes: Long, branches: Int) { + // Saturating rather than wrapping: a total that overflowed to a small positive value would + // understate the requirement, which is the one direction this pass must never fail in. + def add(moreBytes: Long): BranchTotals = { + val sum = bytes + moreBytes + val saturated = if (sum < bytes) Long.MaxValue else sum + BranchTotals(saturated, branches + 1) + } + } + + private val provenance: ShuffleInputProvenance = { + if (app.gpuMode) ShuffleInputProvenance.Measured else ShuffleInputProvenance.Estimated + } + + /** + * Shuffle exchange node names this analyzer can size for the current application. + * + * A GPU application must expose GPU columnar exchanges; a plain CPU `Exchange` that executed in + * a GPU plan means the GPU metrics do not cover all the shuffled data, which is exactly the + * situation that makes a downward decision unsafe. + * + * TODO(#2133): a CPU exchange inside a GPU run could in principle be sized from its CPU data + * size, the way qualification already sizes an all-CPU run. It is not handled today because the + * input-size factor that converts CPU bytes to estimated GPU bytes is application-level, so a + * single run cannot apply one factor to its CPU exchanges and another to its GPU ones. + */ + private def isSupportedShuffleExchange(nodeName: String): Boolean = { + if (app.gpuMode) { + nodeName.contains(COLUMNAR_EXCHANGE_NAME) + } else { + nodeName == EXCHANGE_NAME + } + } + + def build(): ShuffleStageInputAnalysis = { + val records = mutable.ArrayBuffer.empty[ShuffleStageInputRecord] + val reasons = mutable.ArrayBuffer.empty[ShuffleStageInputIncompleteReason] + + app.sqlManager.sqlPlans.values.foreach { planModel => + Try(planModel.getToolsPlanGraph).toOption match { + case Some(graph) => analyzeSql(planModel.id, graph, records, reasons) + case None => + // No final plan graph means there is nothing trustworthy to size. + reasons += ShuffleStageInputIncompleteReason.IncompleteSqlExecution(planModel.id) + } + } + ShuffleStageInputAnalysis(records.toSeq, reasons.toSeq, provenance, + appHasFailedStage = app.stageManager.getFailedStages.nonEmpty) + } + + /** + * Edge index for one plan graph. + * + * `ToolsPlanGraph.getSinkNodes` rescans every edge per call, so the walk builds both directions + * once. The incoming direction is what identifies an exchange's producing stage without relying + * on stage-number ordering. + */ + private class EdgeIndex(graph: ToolsPlanGraph) { + private val outgoing: Map[Long, Seq[Long]] = + graph.edges.groupBy(_.fromId).map { case (from, es) => from -> es.map(_.toId).toSeq }.toMap + private val incoming: Map[Long, Seq[Long]] = + graph.edges.groupBy(_.toId).map { case (to, es) => to -> es.map(_.fromId).toSeq }.toMap + + /** Downstream consumer branches of a node. A reused exchange has one entry per reuse site. */ + def sinksOf(nodeId: Long): Seq[Long] = outgoing.getOrElse(nodeId, Seq.empty) + def sourcesOf(nodeId: Long): Seq[Long] = incoming.getOrElse(nodeId, Seq.empty) + } + + /** + * Stages that write into an exchange, taken from the stage assignment of its child nodes. + * + * The exchange's own assignment spans both the writing and the reading side, so it cannot + * distinguish them. Walking one edge upstream identifies the producing side from the graph + * itself rather than from stage-number ordering. + */ + private def producerStagesOf( + graph: ToolsPlanGraph, edgeIndex: EdgeIndex, nodeId: Long): Set[Int] = { + edgeIndex.sourcesOf(nodeId).flatMap(graph.getNodeStageLogicalAssignment).toSet + } + + private def analyzeSql( + sqlId: Long, + graph: ToolsPlanGraph, + records: mutable.ArrayBuffer[ShuffleStageInputRecord], + reasons: mutable.ArrayBuffer[ShuffleStageInputIncompleteReason]): Unit = { + // A node that was never assigned to a stage did not execute (for example a branch that AQE + // replaced), so it carries no shuffle data and is not evidence of missing coverage. + val executedShuffleNodes = graph.allNodes.filter { node => + isShuffleInputCandidate(node.name) && graph.getNodeStageRawAssignment(node.id).nonEmpty + } + if (executedShuffleNodes.isEmpty) { + return + } + if (!app.sqlIdToInfo.get(sqlId).exists(_.isTerminalSuccess)) { + reasons += ShuffleStageInputIncompleteReason.IncompleteSqlExecution(sqlId) + return + } + + val sqlStageIds = stageIdsOfSql(sqlId) + val edgeIndex = new EdgeIndex(graph) + val branchTotals = mutable.LinkedHashMap.empty[Int, BranchTotals] + + executedShuffleNodes.foreach { node => + val producerStages = producerStagesOf(graph, edgeIndex, node.id) + if (!isSupportedShuffleExchange(node.name)) { + reasons += ShuffleStageInputIncompleteReason.UnsupportedShuffleNode( + sqlId, node.id, node.name) + } else if (producerStages.isEmpty) { + // Without a producing side there is no way to tell which end of the exchange a + // downstream stage sits on, so the branch cannot be attributed safely. + reasons += ShuffleStageInputIncompleteReason.UnresolvedConsumerStage( + sqlId, node.id, node.name) + } else { + // Resolve the exchange metric once, before attributing it to any branch, so that a + // reused exchange cannot be counted from a different stage's accumulator rows. + resolveExchangeDataSize(node, producerStages) match { + case None => + reasons += ShuffleStageInputIncompleteReason.MissingExchangeMetric( + sqlId, node.id, node.name) + case Some(dataSize) => + attributeBranches(sqlId, graph, edgeIndex, node, producerStages, sqlStageIds, + dataSize, branchTotals, reasons) + } + } + } + + branchTotals.foreach { case (stageId, totals) => + selectStageAttempt(stageId) match { + case None => + reasons += ShuffleStageInputIncompleteReason.NoSuccessfulStageAttempt(sqlId, stageId) + case Some(stageModel) => + records += buildRecord(sqlId, stageId, stageModel, totals) + } + } + } + + /** + * Attributes an exchange's resolved bytes to every distinct downstream branch. + * + * A reused exchange is represented as extra outgoing edges from the single original node, so the + * branch count here is what preserves reuse multiplicity, including two branches that land in + * the same consumer stage. + */ + private def attributeBranches( + sqlId: Long, + graph: ToolsPlanGraph, + edgeIndex: EdgeIndex, + node: SparkPlanGraphNode, + producerStages: Set[Int], + sqlStageIds: Set[Int], + dataSize: Long, + branchTotals: mutable.LinkedHashMap[Int, BranchTotals], + reasons: mutable.ArrayBuffer[ShuffleStageInputIncompleteReason]): Unit = { + val branches = edgeIndex.sinksOf(node.id) + val walkedStages = if (branches.isEmpty) { + // The topmost exchange of a plan has no downstream node at all. + Seq(Set.empty[Int]) + } else { + branches.map(resolveConsumerStages(graph, edgeIndex, _, producerStages, sqlStageIds)) + } + // A branch that reaches no assigned stage still has a reading stage: near the root of a plan + // the downstream nodes often carry no metrics, so the walk dead-ends. That stage is + // recoverable from the exchange's own assignment, which spans both sides of the shuffle, + // once the writing side is removed. + val fallbackStages = consumerStagesFromOwnAssignment(graph, node.id, producerStages, + sqlStageIds) + walkedStages.foreach { walked => + val consumerStages = if (walked.nonEmpty) walked else fallbackStages + if (consumerStages.isEmpty) { + reasons += ShuffleStageInputIncompleteReason.UnresolvedConsumerStage( + sqlId, node.id, node.name) + } else { + consumerStages.foreach { consumerStageId => + val current = branchTotals.getOrElse(consumerStageId, BranchTotals(0L, 0)) + branchTotals(consumerStageId) = current.add(dataSize) + } + } + } + } + + /** + * Consumer stages derived from the exchange's own stage assignment. + * + * A shuffle exchange is assigned to both the stage that writes it (from its write metrics) and + * the stage that reads it (from its read metrics), so removing the producing side leaves the + * reading side. This is the fallback for branches whose graph walk reaches no assigned node. + */ + private def consumerStagesFromOwnAssignment( + graph: ToolsPlanGraph, + nodeId: Long, + producerStages: Set[Int], + sqlStageIds: Set[Int]): Set[Int] = { + graph.getNodeStageRawAssignment(nodeId).filter { stageId => + !producerStages.contains(stageId) && sqlStageIds.contains(stageId) + } + } + + /** + * Walks one outgoing branch to the first downstream node assigned to a stage other than the + * exchange's producing stage. + * + * The walk follows graph edges rather than comparing stage numbers, because stage ids are not + * ordered by data flow. A candidate stage must also be confirmed by the node's own raw + * assignment and must belong to this SQL execution's jobs. + * + * A branch can legitimately reach several stages at once: when AQE splits a skewed join, the + * one consuming operator is assigned to every split stage. The full exchange size is then + * attributed to each of them. That deliberately overstates each split stage, because AQE divided + * the data between them, and overstating can only raise the partition requirement and make a + * reduction less likely. Understating it is the outcome that would be unsafe. + * + * TODO(#2133): the split could be estimated instead of duplicated, for example by apportioning + * the exchange size across the split stages by their per-stage shuffle read metrics. That would + * tighten the requirement on skewed joins, where duplication is at its most conservative. + * + * @return the consumer stages of this branch, or an empty set when none could be resolved + */ + private def resolveConsumerStages( + graph: ToolsPlanGraph, + edgeIndex: EdgeIndex, + sinkId: Long, + producerStages: Set[Int], + sqlStageIds: Set[Int]): Set[Int] = { + val visited = mutable.HashSet.empty[Long] + var frontier = List(sinkId) + var depth = 0 + while (frontier.nonEmpty && depth < MAX_CONSUMER_WALK_DEPTH) { + val candidates = frontier.flatMap { nodeId => + val raw = graph.getNodeStageRawAssignment(nodeId) + graph.getNodeStageLogicalAssignment(nodeId).filter { stageId => + !producerStages.contains(stageId) && sqlStageIds.contains(stageId) && + raw.contains(stageId) + } + }.toSet + if (candidates.nonEmpty) { + return candidates + } + visited ++= frontier + frontier = frontier.flatMap(edgeIndex.sinksOf).distinct.filterNot(visited.contains) + depth += 1 + } + Set.empty + } + + /** + * Reads the uncompressed `data size` of a single exchange node. + * + * Stage-scoped accumulator evidence is preferred so that a reused exchange is not inflated by + * rows belonging to another stage. A present zero is a real measurement and is returned as + * such; only the complete absence of accumulator evidence returns None. + */ + private def resolveExchangeDataSize( + node: SparkPlanGraphNode, producerStages: Set[Int]): Option[Long] = { + node.metrics.find(_.name == DATA_SIZE_METRIC).flatMap { metric => + val accumInfoOpt = app.accumManager.accumInfoMap.get(metric.accumulatorId) + val fromStages = accumInfoOpt.flatMap { accumInfo => + val stageValues = producerStages.toSeq.sorted.flatMap(accumInfo.getTotalForStage) + // Fall back to the accumulator's own maximum when the producing stage cannot be pinned + // down, which keeps the existing driver/task maximum semantics. + stageValues.reduceOption(_ max _).orElse(accumInfo.getMaxTotalAcrossStages) + } + // Local-mode plans report exchange sizes through driver accumulator updates instead. + val fromDriver = app.driverAccumMap.get(metric.accumulatorId) + .flatMap(_.map(_.value).reduceOption(_ max _)) + Seq(fromStages, fromDriver).flatten.reduceOption(_ max _) + } + } + + /** + * Picks the attempt whose evidence represents the consumer stage: the highest completed attempt + * that did not fail. Earlier failed attempts are ignored, so a stage that failed once and then + * succeeded cleanly is still eligible. + */ + private def selectStageAttempt(stageId: Int): Option[StageModel] = { + app.stageManager.getStagesByIds(Seq(stageId)) + .filter(sm => !sm.hasFailed && sm.stageInfo.completionTime.isDefined) + .reduceOption((left, right) => if (right.getAttemptId > left.getAttemptId) right else left) + } + + private def buildRecord( + sqlId: Long, + stageId: Int, + stageModel: StageModel, + totals: BranchTotals): ShuffleStageInputRecord = { + val attemptId = stageModel.getAttemptId + val tasks = app.taskManager.getTasks(stageId, attemptId, Some(countsTowardTotals)).toSeq + val spillCandidates = + app.taskManager.getTasks(stageId, attemptId, Some(countsTowardSpillGate)).toSeq + val hasTaskSpill = spillCandidates.exists(spilled) + val hasSpill = hasTaskSpill || app.accumManager.hasGpuSpillEvidence(stageId, attemptId) + ShuffleStageInputRecord( + sqlId = sqlId, + stageId = stageId, + stageAttemptId = attemptId, + totalShuffleInputBytes = totals.bytes, + numShuffleBranches = totals.branches, + numTasks = stageModel.stageInfo.numTasks, + hasPositiveSpill = hasSpill, + hasSkew = hasShuffleReadSkew(tasks.map(_.sr_totalBytesRead))) + } + + /** + * Matches the existing shuffle-skew heuristic: a task reading more than three times the attempt + * average and more than 100 MB. + */ + private def hasShuffleReadSkew(shuffleReadBytes: Seq[Long]): Boolean = { + if (shuffleReadBytes.isEmpty) { + false + } else { + val average = shuffleReadBytes.sum.toDouble / shuffleReadBytes.size + shuffleReadBytes.exists(bytes => bytes > 3 * average && bytes > SKEW_MIN_READ_BYTES) + } + } + + /** Stage ids reachable from the jobs that belong to this SQL execution. */ + private def stageIdsOfSql(sqlId: Long): Set[Int] = { + app.jobIdToInfo.values.collect { + case job if job.sqlID.contains(sqlId) => job.stageIds + }.flatten.toSet + } +} + +object ShuffleStageInputAnalyzer { + /** Uncompressed shuffle bytes written by an exchange, as reported by Spark and RAPIDS. */ + val DATA_SIZE_METRIC = "data size" + private val EXCHANGE_NAME = "Exchange" + private val COLUMNAR_EXCHANGE_NAME = "ColumnarExchange" + private val BROADCAST_MARKER = "Broadcast" + /** Bound on how far a branch walk looks for a consumer stage before failing closed. */ + private val MAX_CONSUMER_WALK_DEPTH = 16 + /** Matches AppSparkMetricsAnalyzer.shuffleSkewCheck. */ + private val SKEW_MIN_READ_BYTES = 100L * 1024L * 1024L + + /** + * Tasks whose metrics describe the work the recommendation governs. Speculative duplicates and + * failed attempts did not produce the stage's output, so their bytes must not inflate the + * requirement the candidate is sized from. + */ + private[analysis] def countsTowardTotals(task: TaskModel): Boolean = { + task.successful && !task.speculative + } + + /** + * Tasks whose spill evidence blocks a reduction. Deliberately broader than + * [[countsTowardTotals]]: a task that spilled and then failed is exactly the memory pressure + * this pass must not reduce into, so excluding it would hide the signal the gate exists to see. + * Speculative duplicates are still excluded, since they re-run work already counted elsewhere. + */ + private[analysis] def countsTowardSpillGate(task: TaskModel): Boolean = !task.speculative + + private[analysis] def spilled(task: TaskModel): Boolean = { + task.memoryBytesSpilled > 0L || task.diskBytesSpilled > 0L + } + + /** + * Any exchange that shuffles data. Broadcast exchanges are excluded because they replicate a + * small side rather than partitioning it, so they do not drive the partition count. + */ + def isShuffleInputCandidate(nodeName: String): Boolean = { + nodeName.contains(EXCHANGE_NAME) && !nodeName.contains(BROADCAST_MARKER) + } + + def apply(app: AppBase): ShuffleStageInputAnalysis = { + new ShuffleStageInputAnalyzer(app).build() + } +} diff --git a/core/src/main/scala/com/nvidia/spark/rapids/tool/profiling/ApplicationSummaryInfo.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/profiling/ApplicationSummaryInfo.scala index 35b569ccc..0f4b7acdf 100644 --- a/core/src/main/scala/com/nvidia/spark/rapids/tool/profiling/ApplicationSummaryInfo.scala +++ b/core/src/main/scala/com/nvidia/spark/rapids/tool/profiling/ApplicationSummaryInfo.scala @@ -102,6 +102,18 @@ trait AppInfoColumnarExchangeMetrics { def getMaxColumnarExchangeDataSizeBytes: Option[Long] = None } +/** + * Exposes the raw consumer-stage shuffle input inventory to the AutoTuner. + * + * Both tools report the same record shape; only the provenance differs. The default fails closed + * so a provider that cannot produce an analysis can never enable a downward recommendation. + */ +trait AppInfoShuffleStageInputMetrics { + def getShuffleStageInputAnalysis: ShuffleStageInputAnalysis = { + ShuffleStageInputAnalysis.empty(ShuffleInputProvenance.Estimated) + } +} + /** * Base class for Profiling App Summary Info Provider. */ @@ -257,6 +269,11 @@ class SingleAppSummaryInfoProvider( SingleAppSummaryInfoProvider.computeMaxColumnarExchangeDataSizeBytes(app.sqlMetrics) } + // Reuses the plan analysis the profiler already built rather than traversing the plans again. + override def getShuffleStageInputAnalysis: ShuffleStageInputAnalysis = { + appInfo.planMetricProcessor.shuffleStageInputAnalysis + } + override def getClassPathEntries: Map[String, String] = { appInfo.classpathEntries } diff --git a/core/src/main/scala/com/nvidia/spark/rapids/tool/profiling/ProfileClassWarehouse.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/profiling/ProfileClassWarehouse.scala index a73d43d87..f68b73937 100644 --- a/core/src/main/scala/com/nvidia/spark/rapids/tool/profiling/ProfileClassWarehouse.scala +++ b/core/src/main/scala/com/nvidia/spark/rapids/tool/profiling/ProfileClassWarehouse.scala @@ -312,9 +312,20 @@ class SQLExecutionInfoClass( // Empty for Spark 3.2.x event logs or when no runtime config changes were made. // This stores only the overrides, not the full config set. val modifiedConfigs: Map[String, String] = Map.empty) { + // Failure text reported by SparkListenerSQLExecutionEnd, when the Spark version records one. + // A recorded end time alone does not prove the execution succeeded, so consumers that need + // terminal success must consult this as well. + var failureReason: Option[String] = None + def setDsOrRdd(value: Boolean): Unit = { hasDatasetOrRDD = value } + + /** + * True only when the execution reached a terminal end event without a reported failure. + * Used by analyses that must fail closed on incomplete or failed SQL executions. + */ + def isTerminalSuccess: Boolean = endTime.isDefined && failureReason.isEmpty } case class SQLAccumProfileResults( diff --git a/core/src/main/scala/com/nvidia/spark/rapids/tool/profiling/ShuffleStageInputMetrics.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/profiling/ShuffleStageInputMetrics.scala new file mode 100644 index 000000000..09bd925e3 --- /dev/null +++ b/core/src/main/scala/com/nvidia/spark/rapids/tool/profiling/ShuffleStageInputMetrics.scala @@ -0,0 +1,152 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.nvidia.spark.rapids.tool.profiling + +/** + * Describes how the shuffle input bytes of a consumer stage were obtained. + * + * Profiling reads measured `GpuColumnarExchange` data sizes, while qualification reads regular + * CPU `Exchange` data sizes and must estimate the GPU input from them. + */ +sealed abstract class ShuffleInputProvenance(val label: String) + +object ShuffleInputProvenance { + /** GPU event logs expose the exchange data size that the GPU actually processed. */ + case object Measured extends ShuffleInputProvenance("measured") + /** CPU event logs expose CPU exchange data size, which only estimates the GPU input. */ + case object Estimated extends ShuffleInputProvenance("estimated") +} + +/** + * Total uncompressed shuffle input entering a single consumer-stage execution, together with the + * safety evidence of the stage attempt that was selected to represent it. + * + * One record is produced per (SQL execution, consumer stage). `totalShuffleInputBytes` sums every + * distinct shuffle branch that feeds the stage, so an exchange reused through two branches into the + * same stage contributes twice. + * + * The record is intentionally raw: it carries measured bytes only. The tool-specific input-size + * factor belongs to AutoTuner configuration and is applied by the downward policy, not here. + * + * @param sqlId SQL execution the consumer stage belongs to + * @param stageId consumer stage id + * @param stageAttemptId attempt selected as the representative successful attempt + * @param totalShuffleInputBytes sum of the uncompressed `data size` of every incoming branch + * @param numShuffleBranches number of distinct incoming shuffle branches that were summed + * @param numTasks declared task count of the selected attempt + * @param hasPositiveSpill true when any successful task of the selected attempt spilled to + * host/memory or disk, or when GPU SQL spill was recorded for it + * @param hasSkew true when the selected attempt shows shuffle read skew + */ +case class ShuffleStageInputRecord( + sqlId: Long, + stageId: Int, + stageAttemptId: Int, + totalShuffleInputBytes: Long, + numShuffleBranches: Int, + numTasks: Int, + hasPositiveSpill: Boolean, + hasSkew: Boolean) + +/** + * Reason the shuffle-stage input inventory could not be proven complete. + * + * A global reduction is more dangerous than an overestimate, so any single reason makes the whole + * analysis unusable rather than only dropping the affected stage. + */ +sealed abstract class ShuffleStageInputIncompleteReason(val description: String) + +object ShuffleStageInputIncompleteReason { + /** An executed non-broadcast shuffle input had no supported `data size` metric. */ + case class MissingExchangeMetric(sqlId: Long, nodeId: Long, nodeName: String) + extends ShuffleStageInputIncompleteReason( + s"SQL $sqlId node $nodeId ($nodeName) has no supported shuffle 'data size' metric") + + /** A shuffle node type that this analysis cannot size was executed. */ + case class UnsupportedShuffleNode(sqlId: Long, nodeId: Long, nodeName: String) + extends ShuffleStageInputIncompleteReason( + s"SQL $sqlId node $nodeId ($nodeName) is an unsupported shuffle input") + + /** The downstream consumer stage of an exchange branch could not be resolved unambiguously. */ + case class UnresolvedConsumerStage(sqlId: Long, nodeId: Long, nodeName: String) + extends ShuffleStageInputIncompleteReason( + s"SQL $sqlId node $nodeId ($nodeName) has no reliable consumer-stage mapping") + + /** The SQL execution never reached a terminal successful end event. */ + case class IncompleteSqlExecution(sqlId: Long) + extends ShuffleStageInputIncompleteReason( + s"SQL $sqlId did not complete successfully") + + /** No completed, non-failed attempt exists for a consumer stage. */ + case class NoSuccessfulStageAttempt(sqlId: Long, stageId: Int) + extends ShuffleStageInputIncompleteReason( + s"SQL $sqlId consumer stage $stageId has no completed successful attempt") +} + +/** + * Raw result of the consumer-stage shuffle input analysis for one application. + * + * An application that genuinely executed no shuffle exchange is still `analyzed`: it produces no + * records and no incomplete reasons, and the downward pass simply has nothing to size. That is a + * different state from a provider that never produced an analysis at all, which must fail closed. + * + * @param records one record per (SQL execution, consumer stage) + * @param incompleteReasons non-empty when the inventory is not trustworthy; the downward pass + * must then keep the normal recommendation + * @param provenance whether `records` carry measured or estimated shuffle bytes + * @param analyzed false when no analysis was produced for the application at all + * @param appHasFailedStage true when any stage in the application failed, anywhere. This is an + * application-wide signal, deliberately broader than the per-record + * attempt evidence: a run that failed a stage is not a run to size a + * global reduction from. + */ +case class ShuffleStageInputAnalysis( + records: Seq[ShuffleStageInputRecord], + incompleteReasons: Seq[ShuffleStageInputIncompleteReason], + provenance: ShuffleInputProvenance, + analyzed: Boolean = true, + appHasFailedStage: Boolean = false) { + + def isComplete: Boolean = analyzed && incompleteReasons.isEmpty + + /** First few reasons, used to keep the user-facing diagnostic concise. */ + def incompleteSummary(maxReasons: Int = 2): String = { + if (!analyzed) { + ShuffleStageInputAnalysis.notAnalyzedSummary + } else { + val shown = incompleteReasons.take(maxReasons).map(_.description) + val suffix = if (incompleteReasons.size > shown.size) { + s" (and ${incompleteReasons.size - shown.size} more)" + } else { + "" + } + shown.mkString("; ") + suffix + } + } +} + +object ShuffleStageInputAnalysis { + val notAnalyzedSummary = "no shuffle input analysis was produced for the application" + + /** + * An analysis that carries no evidence at all, used by providers that cannot produce one. + * It always fails the completeness gate. + */ + def empty(provenance: ShuffleInputProvenance): ShuffleStageInputAnalysis = { + ShuffleStageInputAnalysis(Seq.empty, Seq.empty, provenance, analyzed = false) + } +} diff --git a/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala index 37589c7f5..ad72551a1 100644 --- a/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala +++ b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala @@ -255,9 +255,28 @@ abstract class AutoTuner( // Recommendations for these properties will not be computed, ensuring that dependent properties // are also affected correctly. private val skippedRecommendations: mutable.HashSet[String] = mutable.HashSet[String]() + // Properties that have already contributed a "was not set" comment, so a later pass recommending + // the same key does not repeat it. + private val keysWithMissingComment: mutable.HashSet[String] = mutable.HashSet[String]() // Set of properties for which only source application values are used and // no calculations are performed. protected val limitedLogicRecommendations: mutable.HashSet[String] = mutable.HashSet[String]() + // Reasons the normal tuning passes raised the shuffle partition recommendation. The final + // downward pass must preserve every one of them, so they are recorded where they are applied + // instead of being inferred afterwards. + private val shufflePartitionUpwardReasons: mutable.LinkedHashSet[String] = + mutable.LinkedHashSet[String]() + + /** + * True when the application hit an out-of-memory failure. Only GPU profiling carries this + * evidence, so the base implementation reports false and the profiling AutoTuner overrides it. + */ + protected def applicationHadOom: Boolean = false + + /** Records an upward shuffle-partition decision so the downward pass cannot undo it. */ + protected def recordShufflePartitionUpwardReason(reason: String): Unit = { + shufflePartitionUpwardReasons += reason + } // When enabled, the profiler recommendations should only include updated settings. private var filterByUpdatedPropertiesEnabled: Boolean = true // Non-executor memory (reserved for OS, resource manager, etc), configurable via tuning configs @@ -459,13 +478,19 @@ abstract class AutoTuner( /** * Append a comment to the list by looking up the missing comment if any in the tuningEntry * table. + * + * A property is either set in the source application or it is not, so the comment is emitted at + * most once even when several passes recommend a value for the same key. + * * @param key the property set by the autotuner. */ private def appendMissingComment(key: String): Unit = { - val missingComment = finalTuningTable.get(key) - .flatMap(_.getMissingComment()) - .getOrElse(s"was not set.") - appendComment(key, missingComment) + if (keysWithMissingComment.add(key)) { + val missingComment = finalTuningTable.get(key) + .flatMap(_.getMissingComment()) + .getOrElse(s"was not set.") + appendComment(key, missingComment) + } } /** @@ -1811,6 +1836,9 @@ abstract class AutoTuner( // and doesn't capture CPU shuffle data. val columnarExchangeRatio = (maxDataSize.toDouble / gpuBatchSize).ceil.toInt if (columnarExchangeRatio > finalPartitionValue) { + recordShufflePartitionUpwardReason( + s"the GPU ColumnarExchange batch-size bound raised partitions to " + + s"$columnarExchangeRatio") appendComment(s"'$initialPartitionNumKey' adjusted from " + s"$finalPartitionValue to $columnarExchangeRatio based on " + s"ColumnarExchange data size (${maxDataSize} bytes) and " + @@ -1998,6 +2026,7 @@ abstract class AutoTuner( inputShufflePartitions *= configProvider.getEntry("SHUFFLE_PARTITION_MULTIPLIER").getDefault.toInt // Could be memory instead of partitions + recordShufflePartitionUpwardReason("spilling was detected in shuffle stages") appendComment(shufflePartitionsCommentForSpilling) } } @@ -2033,6 +2062,197 @@ abstract class AutoTuner( } } + /** + * Final downward-only shuffle-partition pass. + * + * Runs after the normal job-level and cluster-level recommendations so that it observes the + * effective partition value they produced and the recommended cluster shape. It sizes the worst + * consumer stage from the total uncompressed shuffle input entering it, quantizes the result to + * whole execution waves of that cluster, and lowers the recommendation only when the reduction + * is material and every safety gate passes. + * + * Lowering a global partition count is riskier than leaving it too high, so anything unproven + * keeps the normal recommendation. It never raises a value and never touches the AQE advisory + * partition size. + */ + private def recommendDownwardShufflePartitions(): Unit = { + val configResult = DownwardShufflePolicyConfig.fromProvider(configProvider) + val enabled = configResult.exists(_.enabled) + // Building the analysis walks every SQL plan, so it is only worth doing once the pass is + // known to be enabled. The pass ships off, which makes this the common path. + val slotCount = if (enabled) downwardShuffleSlotCount else None + val analysis = + if (enabled) appInfoProvider.getShuffleStageInputAnalysis else emptyShuffleStageInputAnalysis + val decision = DownwardShufflePartitionsPolicy.decide( + configResult, + shufflePartitionValue, + slotCount, + analysis) + decision match { + case DownwardShuffleDecision.InvalidConfig(errors) => + // Fail closed as one decision: no property is touched when any policy input is invalid. + logWarning("Skipping the downward shuffle partition pass because its configuration is " + + s"invalid: ${errors.mkString("; ")}") + appendComment(downwardShufflePartitionsInvalidConfigComment) + case DownwardShuffleDecision.Skipped(reason) => + reportDownwardShuffleSkip(reason) + case applied: DownwardShuffleDecision.Applied => + applyDownwardShufflePartitions(applied) + } + } + + /** + * Task slots of the cluster this run is recommending, which is the size of one execution wave. + * + * The executor count is read from the recommendation map rather than through `getPropertyValue`, + * because that helper falls back to the source properties and would silently yield the source CPU + * executor count on a platform that excludes 'spark.executor.instances'. It falls back to + * `recommendedClusterInfo.numExecutors` only when no recommendation exists, since + * `recommendDynamicAllocationConfigs` rescales the executor counts it recommends without ever + * updating the cluster record. + * + * @return the slot count, or None when the recommended cluster shape or the per-executor + * multiplier cannot be resolved + */ + private def downwardShuffleSlotCount: Option[Int] = { + platform.recommendedClusterInfo.flatMap { clusterInfo => + val executors = + recommendedIntValue("spark.executor.instances").getOrElse(clusterInfo.numExecutors) + // Cores, not GPU task concurrency: concurrency is auto-tuned by recent plugin versions, and + // sizing against it would badly under-use a cluster running mixed CPU and GPU stages. + val coresPerExecutor = clusterInfo.coresPerExecutor + if (executors > 0 && coresPerExecutor > 0) { + val slots = executors.toLong * coresPerExecutor.toLong + if (slots > Int.MaxValue.toLong) None else Some(slots.toInt) + } else { + None + } + } + } + + /** Tuned value of a recommended property parsed as an Int, ignoring source-property fallbacks. */ + private def recommendedIntValue(property: String): Option[Int] = { + recommendations.get(property).flatMap(_.tunedValue) + .flatMap(value => Try(value.trim.toInt).toOption) + } + + /** + * Stand-in used when the pass is off or misconfigured, so no analysis has to be built. The + * policy short-circuits on the config before it reads any of this. + */ + private def emptyShuffleStageInputAnalysis: ShuffleStageInputAnalysis = { + ShuffleStageInputAnalysis.empty(ShuffleInputProvenance.Measured) + } + + /** + * Runs the AutoTuner-owned gates that the pure policy cannot see, then updates every required + * partition property or none of them. + */ + private def applyDownwardShufflePartitions( + applied: DownwardShuffleDecision.Applied): Unit = { + downwardShuffleBlockingReason(applied) match { + case Some(reason) => reportDownwardShuffleSkip(reason) + case None => + // Clamp each property against its own current value rather than against the single + // effective value the decision was made from, which is the maximum across them. Today the + // AQE pass has already levelled the two, so this is an invariant guard rather than a live + // difference: it keeps the pass downward-only per property if that ordering ever changes. + val updates = requiredPartitionProperties.map { property => + val current = getPropertyValue(property).flatMap(v => Try(v.trim.toInt).toOption) + property -> current.map(_ min applied.selectedValue).getOrElse(applied.selectedValue) + } + updates.foreach { case (property, value) => appendRecommendation(property, value) } + appendComment( + DownwardShufflePartitionsPolicy.appliedComment(updates, applied)) + } + } + + /** + * Properties that must all be updated together for the recommendation to be coherent. + * When AQE coalescing is disabled this is only 'spark.sql.shuffle.partitions'; otherwise the + * active AQE partition property must move with it. + */ + private def requiredPartitionProperties: Seq[String] = { + // 'spark.sql.shuffle.partitions' always applies. The AQE partition property is updated only + // when a value for it already exists, from the source application or an earlier + // recommendation: this pass lowers a partition count, it does not introduce a property the + // application never carried. + val aqeProperty = aqePartitionProperty.filter(getPropertyValue(_).isDefined) + (aqeProperty.toSeq :+ "spark.sql.shuffle.partitions").distinct + } + + /** + * First reason the reduction cannot be applied, or None when every gate passes. + * + * The gates are evaluated before any property is written, because `appendRecommendation` + * silently ignores protected properties and would otherwise leave the two partition properties + * disagreeing with each other. + */ + private def downwardShuffleBlockingReason( + applied: DownwardShuffleDecision.Applied): Option[DownwardShuffleSkipReason] = { + // 1. Any increase the normal passes made for spill, OOM, or the GPU batch-size bound wins. + shufflePartitionUpwardReasons.headOption + .map(DownwardShuffleSkipReason.UpwardSafetyReason) + // 2. A platform that keeps optimizing shuffle at runtime may ignore the manual value. + .orElse { + if (isDatabricksAutoOptimizeShuffleActive) { + Some(DownwardShuffleSkipReason.PlatformControlledShuffle) + } else { + None + } + } + // 3. A run that failed a stage or ran out of memory is not evidence to reduce from. + .orElse { + val analysis = appInfoProvider.getShuffleStageInputAnalysis + if (analysis.appHasFailedStage) { + Some(DownwardShuffleSkipReason.ApplicationHasFailedStage) + } else if (applicationHadOom) { + Some(DownwardShuffleSkipReason.ApplicationHadOom) + } else { + None + } + } + // 4. Every affected consumer stage must be free of skew and spill. + .orElse { + appInfoProvider.getShuffleStageInputAnalysis.records.collectFirst { + case record if record.hasPositiveSpill => + DownwardShuffleSkipReason.StagePressure(record.stageId, "positive spill") + case record if record.hasSkew => + DownwardShuffleSkipReason.StagePressure(record.stageId, "shuffle read skew") + } + } + // 5. Every property this decision must write has to be writable. + .orElse { + requiredPartitionProperties.collectFirst { + case property if ignoreRecommendation(property) || !isCalculationEnabled(property) => + DownwardShuffleSkipReason.PropertyNotMutable(property) + } + } + } + + /** + * True while Databricks automatic shuffle optimization is still effectively enabled after + * normal tuning, which means the runtime, not this recommendation, governs partitioning. + */ + private def isDatabricksAutoOptimizeShuffleActive: Boolean = { + platform.isInstanceOf[DatabricksPlatform] && + getPropertyValue("spark.databricks.adaptive.autoOptimizeShuffle.enabled") + .exists(_.trim.equalsIgnoreCase("true")) + } + + /** + * Reports a no-op. Ordinary policy and safety decisions stay log-only so that enabling this + * pass does not add comments to every application; only actionable states earn a comment. + */ + private def reportDownwardShuffleSkip(reason: DownwardShuffleSkipReason): Unit = { + if (reason.isWarning) { + logWarning(s"Skipping the downward shuffle partition pass: ${reason.description}") + appendComment(downwardShufflePartitionsIncompleteEvidenceComment) + } else { + logInfo(s"Skipping the downward shuffle partition pass: ${reason.description}") + } + } + /** * Analyzes unsupported driver logs and generates recommendations for configuration properties. */ @@ -2284,6 +2504,9 @@ abstract class AutoTuner( recommendPluginProps() calculateJobLevelRecommendations() calculateClusterLevelRecommendations() + // Final downward-only pass. It runs last so that it sees the effective normal shuffle + // partition recommendation, and it can only lower that value, never raise it. + recommendDownwardShufflePartitions() // Add all platform specific recommendations platform.platformSpecificRecommendations.collect { @@ -2565,6 +2788,11 @@ class ProfilingAutoTuner( } } + override protected def applicationHadOom: Boolean = { + appInfoProvider.scanStagesWithGpuOom.nonEmpty || + appInfoProvider.gpuShuffleStagesWithContainerOom.nonEmpty + } + /** * Overrides the calculation for 'spark.sql.shuffle.partitions'. * This method checks for task OOM errors in shuffle stages and recommends to increase @@ -2576,6 +2804,7 @@ class ProfilingAutoTuner( // Shuffle Stages with Task OOM detected. We may want to increase shuffle partitions. val recShufflePartitions = shufflePartitionValue * configProvider.getEntry("SHUFFLE_PARTITION_MULTIPLIER").getDefault.toInt + recordShufflePartitionUpwardReason("task OOM was detected in shuffle stages") appendComment(shufflePartitionsCommentForGpuOOM) math.max(calculatedValue, recShufflePartitions) } else { @@ -2816,6 +3045,16 @@ trait AutoTunerStaticComments { "Shuffle partitions should be increased since task OOM occurred in shuffle stages." } + def downwardShufflePartitionsInvalidConfigComment: String = { + "Shuffle partitions were not lowered because the downward shuffle partition tuning " + + "configuration is invalid. See the tool logs for the offending entries." + } + + def downwardShufflePartitionsIncompleteEvidenceComment: String = { + "Shuffle partitions were not lowered because the shuffle input of every consumer stage " + + "could not be measured. See the tool logs for details." + } + /** * Comment for missing GPU discovery script. * Since this comment is conditional, it is not included in the diff --git a/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/DownwardShufflePartitionsPolicy.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/DownwardShufflePartitionsPolicy.scala new file mode 100644 index 000000000..207f0afa8 --- /dev/null +++ b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/DownwardShufflePartitionsPolicy.scala @@ -0,0 +1,442 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.nvidia.spark.rapids.tool.tuning + +import scala.util.Try + +import com.nvidia.spark.rapids.tool.profiling.{ShuffleInputProvenance, ShuffleStageInputAnalysis, ShuffleStageInputRecord} +import com.nvidia.spark.rapids.tool.tuning.config.TuningConfigProvider + +import org.apache.spark.network.util.ByteUnit +import org.apache.spark.sql.rapids.tool.util.StringUtils + +/** + * Validated policy inputs of the downward-only shuffle partition pass. + * + * Every field is user-overridable through the tuning-config mechanism. The values are validated as + * one unit before any recommendation can be mutated, so a partially valid configuration can never + * produce a partially applied reduction. + * + * @param enabled master switch for the downward pass + * @param targetPartitionSizeBytes estimated GPU input a single partition should process + * @param inputSizeFactor factor converting measured shuffle bytes to estimated GPU bytes + */ +case class DownwardShufflePolicyConfig( + enabled: Boolean, + targetPartitionSizeBytes: Long, + inputSizeFactor: Double) + +object DownwardShufflePolicyConfig { + val ENABLED_KEY = "DOWNWARD_SHUFFLE_ENABLED" + val TARGET_PARTITION_SIZE_KEY = "DOWNWARD_SHUFFLE_TARGET_PARTITION_SIZE" + val INPUT_SIZE_FACTOR_KEY = "DOWNWARD_SHUFFLE_INPUT_SIZE_FACTOR" + + /** Config used when the feature is switched off. The remaining fields are never read. */ + val disabled: DownwardShufflePolicyConfig = DownwardShufflePolicyConfig( + enabled = false, + targetPartitionSizeBytes = 0L, + inputSizeFactor = 0.0) + + /** + * Reads and validates every policy entry from the tuning-config provider. + * + * All problems are collected so that an invalid configuration surfaces as one fail-closed + * decision. When the feature is explicitly disabled the remaining entries are not validated, + * which keeps a disabled run quiet even if its unused policy entries are malformed. + * + * @return the validated config, or the list of validation errors + */ + def fromProvider( + configProvider: TuningConfigProvider): Either[Seq[String], DownwardShufflePolicyConfig] = { + rawValue(configProvider, ENABLED_KEY).flatMap(parseStrictBoolean(ENABLED_KEY, _)) match { + case Left(err) => Left(Seq(err)) + case Right(false) => Right(disabled) + case Right(true) => parseEnabledConfig(configProvider) + } + } + + private def parseEnabledConfig( + configProvider: TuningConfigProvider): Either[Seq[String], DownwardShufflePolicyConfig] = { + val targetSize = parseMemoryBytes(configProvider, TARGET_PARTITION_SIZE_KEY) + val factor = parseDouble(configProvider, INPUT_SIZE_FACTOR_KEY, min = 0.0, minInclusive = false) + + val errors = Seq(targetSize, factor).collect { + case Left(err) => err + } + (targetSize, factor) match { + case (Right(size), Right(f)) if errors.isEmpty => + Right(DownwardShufflePolicyConfig( + enabled = true, + targetPartitionSizeBytes = size, + inputSizeFactor = f)) + case _ => Left(errors) + } + } + + private def rawValue( + configProvider: TuningConfigProvider, key: String): Either[String, String] = { + Try(configProvider.getEntry(key).getDefault).toOption + .filter(v => v != null && v.trim.nonEmpty) + .toRight(s"'$key' is not defined") + } + + /** Accepts only an exact 'true' or 'false'; anything else fails closed. */ + private def parseStrictBoolean(key: String, value: String): Either[String, Boolean] = { + value.trim.toLowerCase match { + case "true" => Right(true) + case "false" => Right(false) + case other => Left(s"'$key' must be exactly 'true' or 'false' but was '$other'") + } + } + + private def parseMemoryBytes( + configProvider: TuningConfigProvider, key: String): Either[String, Long] = { + rawValue(configProvider, key).flatMap { raw => + Try(StringUtils.convertMemorySizeToBytes(raw, Some(ByteUnit.BYTE))).toOption + .filter(_ > 0L) + .toRight(s"'$key' must be a positive memory size but was '$raw'") + } + } + + private def parseDouble( + configProvider: TuningConfigProvider, + key: String, + min: Double, + minInclusive: Boolean): Either[String, Double] = { + val bound = if (minInclusive) s">= $min" else s"> $min" + rawValue(configProvider, key).flatMap { raw => + Try(raw.trim.toDouble).toOption + .filter(v => !v.isNaN && !v.isInfinite && (if (minInclusive) v >= min else v > min)) + .toRight(s"'$key' must be a finite number $bound but was '$raw'") + } + } +} + +/** Reason the downward pass left the normal recommendation unchanged. */ +sealed abstract class DownwardShuffleSkipReason(val description: String) { + /** + * Warning-level reasons point at something the user may want to fix, so they earn one concise + * comment. Ordinary policy and safety decisions stay log-only to avoid broad output churn. + */ + def isWarning: Boolean = false +} + +object DownwardShuffleSkipReason { + case object Disabled + extends DownwardShuffleSkipReason("the downward shuffle partition pass is disabled") + + /** + * The analysis ran but found a gap it cannot reason about. This is actionable, so it earns one + * concise comment in addition to the log. + */ + case class IncompleteEvidence(summary: String) + extends DownwardShuffleSkipReason( + s"shuffle input evidence is incomplete: $summary") { + override def isWarning: Boolean = true + } + + /** + * No analysis was produced for the application at all, for example because the caller does not + * wire one up. There is nothing for the user to act on, so this stays log-only. + */ + case object NoAnalysisAvailable + extends DownwardShuffleSkipReason( + "no shuffle input analysis is available for this application") + + case object NoStageEvidence + extends DownwardShuffleSkipReason("no consumer stage shuffle input was found") + + /** + * The worst stage needs more partitions than a Spark partition count can express, or covering it + * with whole waves would. Failing closed here is safer than recommending a partition count below + * the requirement. + */ + case class RequirementOutOfRange(requirement: Long) + extends DownwardShuffleSkipReason( + s"the worst consumer stage requires $requirement partitions, which cannot be rounded up to" + + s" whole cluster waves within the largest representable partition count" + + s" (${Int.MaxValue})") { + override def isWarning: Boolean = true + } + + /** + * The recommended cluster shape is not known, so there is no wave to quantize to. This is the + * common qualification case rather than something the user can act on, so it stays log-only. + */ + case object SlotCountUnavailable + extends DownwardShuffleSkipReason( + "the recommended cluster's task slot count could not be determined") + + case class NotDownward(candidate: Int, normalValue: Int) + extends DownwardShuffleSkipReason( + s"candidate $candidate does not lower the current recommendation $normalValue") + + /** + * The application failed a stage or hit an OOM anywhere. Sizing a global reduction from a run + * that did not complete cleanly is unsafe even when the consumer stages themselves look healthy. + */ + case object ApplicationHasFailedStage + extends DownwardShuffleSkipReason("the application had a failed stage") + + case object ApplicationHadOom + extends DownwardShuffleSkipReason("the application had an out-of-memory failure") + + case class UpwardSafetyReason(reason: String) + extends DownwardShuffleSkipReason( + s"an existing upward recommendation must be preserved: $reason") + + case class StagePressure(stageId: Int, reason: String) + extends DownwardShuffleSkipReason( + s"consumer stage $stageId shows $reason") + + case class PropertyNotMutable(property: String) + extends DownwardShuffleSkipReason( + s"partition property '$property' cannot be updated by the AutoTuner") + + case object PlatformControlledShuffle + extends DownwardShuffleSkipReason( + "Databricks automatic shuffle optimization is still effectively enabled") +} + +/** Outcome of the pure downward policy. */ +sealed trait DownwardShuffleDecision + +object DownwardShuffleDecision { + /** + * A reduction that passed every policy gate. + * + * @param normalValue the effective recommendation produced by normal tuning + * @param selectedValue the wave-quantized partition count to recommend instead + * @param determiningRecord the consumer stage that produced the worst requirement + * @param estimatedInputBytes determining stage bytes after the input-size factor + * @param rawRequirement partition count before rounding up to whole waves + * @param slotCount task slots of the recommended cluster; one execution wave + * @param waveCount whole waves the selected value spans + * @param provenance whether the bytes were measured or estimated + * @param inputSizeFactor factor that was applied + */ + case class Applied( + normalValue: Int, + selectedValue: Int, + determiningRecord: ShuffleStageInputRecord, + estimatedInputBytes: Long, + rawRequirement: Long, + slotCount: Int, + waveCount: Int, + provenance: ShuffleInputProvenance, + inputSizeFactor: Double) extends DownwardShuffleDecision + + case class Skipped(reason: DownwardShuffleSkipReason) extends DownwardShuffleDecision + + case class InvalidConfig(errors: Seq[String]) extends DownwardShuffleDecision +} + +/** + * Pure, deterministic calculator for the downward-only shuffle partition pass. + * + * This object never reads or mutates AutoTuner state. It converts raw consumer-stage shuffle input + * records into either an applied reduction or a typed no-op reason, using overflow-safe arithmetic + * and quantizing the result to whole execution waves of the recommended cluster. + */ +object DownwardShufflePartitionsPolicy { + + /** A consumer stage together with its derived estimated bytes and partition requirement. */ + private case class StageRequirement( + record: ShuffleStageInputRecord, + estimatedBytes: Long, + requirement: Long) + + /** + * Evaluates the policy for one application. + * + * @param configResult validated policy config, or its validation errors + * @param normalValue effective shuffle partition recommendation after normal tuning + * @param slotCount task slots of the recommended cluster, or None when it is unknown + * @param analysis raw consumer-stage shuffle input analysis + */ + def decide( + configResult: Either[Seq[String], DownwardShufflePolicyConfig], + normalValue: Int, + slotCount: Option[Int], + analysis: ShuffleStageInputAnalysis): DownwardShuffleDecision = { + configResult match { + case Left(errors) => DownwardShuffleDecision.InvalidConfig(errors) + case Right(config) if !config.enabled => + DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.Disabled) + case Right(config) => decideWithConfig(config, normalValue, slotCount, analysis) + } + } + + private def decideWithConfig( + config: DownwardShufflePolicyConfig, + normalValue: Int, + slotCount: Option[Int], + analysis: ShuffleStageInputAnalysis): DownwardShuffleDecision = { + if (!analysis.analyzed) { + return DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.NoAnalysisAvailable) + } + if (!analysis.isComplete) { + return DownwardShuffleDecision.Skipped( + DownwardShuffleSkipReason.IncompleteEvidence(analysis.incompleteSummary())) + } + if (analysis.records.isEmpty) { + return DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.NoStageEvidence) + } + + // A slot count is what a wave is measured in, so without one there is nothing to quantize to. + val slots = slotCount.filter(_ > 0).getOrElse { + return DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.SlotCountUnavailable) + } + + val worst = selectWorstStage(config, analysis.records) + waveQuantized(worst.requirement, slots) match { + case None => + // No representable whole-wave count covers the requirement: fail closed rather than + // recommend a partition count that is known to be too small. + DownwardShuffleDecision.Skipped( + DownwardShuffleSkipReason.RequirementOutOfRange(worst.requirement)) + case Some(candidate) if candidate >= normalValue => + // Also covers the case where the normal value is already at or below one wave, because + // every candidate is at least the slot count. + DownwardShuffleDecision.Skipped( + DownwardShuffleSkipReason.NotDownward(candidate, normalValue)) + case Some(candidate) => + DownwardShuffleDecision.Applied( + normalValue = normalValue, + selectedValue = candidate, + determiningRecord = worst.record, + estimatedInputBytes = worst.estimatedBytes, + rawRequirement = worst.requirement, + slotCount = slots, + waveCount = candidate / slots, + provenance = analysis.provenance, + inputSizeFactor = config.inputSizeFactor) + } + } + + /** + * Picks the consumer stage with the largest partition requirement. + * + * Ties are broken by the larger estimated input, then by ascending SQL id, stage id, and attempt + * id, so the determining stage is reproducible across runs regardless of record ordering. + */ + private def selectWorstStage( + config: DownwardShufflePolicyConfig, + records: Seq[ShuffleStageInputRecord]): StageRequirement = { + records.map { record => + val estimated = estimateGpuInputBytes(record.totalShuffleInputBytes, config.inputSizeFactor) + StageRequirement(record, estimated, partitionRequirement(estimated, + config.targetPartitionSizeBytes)) + }.reduceLeft { (best, candidate) => + if (isWorseThan(candidate, best)) candidate else best + } + } + + /** Strict "is a worse (larger) requirement than" comparison implementing the tie-break order. */ + private def isWorseThan(left: StageRequirement, right: StageRequirement): Boolean = { + if (left.requirement != right.requirement) { + left.requirement > right.requirement + } else if (left.estimatedBytes != right.estimatedBytes) { + left.estimatedBytes > right.estimatedBytes + } else if (left.record.sqlId != right.record.sqlId) { + left.record.sqlId < right.record.sqlId + } else if (left.record.stageId != right.record.stageId) { + left.record.stageId < right.record.stageId + } else { + left.record.stageAttemptId < right.record.stageAttemptId + } + } + + /** + * Applies the tool-specific input-size factor. The product is computed in `Double` and clamped so + * that a huge stage total can never wrap around to a small positive `Long`. + */ + private[tuning] def estimateGpuInputBytes(totalBytes: Long, factor: Double): Long = { + if (totalBytes <= 0L) { + 0L + } else { + val scaled = math.ceil(totalBytes.toDouble * factor) + if (scaled >= Long.MaxValue.toDouble) Long.MaxValue else scaled.toLong + } + } + + /** + * Ceiling of estimated bytes divided by the target partition size, with a lower bound of 1 so a + * stage that carries any shuffle input always requires at least one partition. + */ + private[tuning] def partitionRequirement(estimatedBytes: Long, targetBytes: Long): Long = { + if (estimatedBytes <= 0L) { + 1L + } else { + // targetBytes is validated to be positive, so this cannot divide by zero and neither + // branch can overflow because estimatedBytes is at most Long.MaxValue. + val quotient = estimatedBytes / targetBytes + if (estimatedBytes % targetBytes == 0L) math.max(quotient, 1L) else quotient + 1L + } + } + + /** + * Rounds a raw requirement up to a whole number of execution waves of the recommended cluster. + * + * The slot count is a hard floor as well as the quantum, so a stage needing fewer partitions than + * the cluster has slots still gets exactly one full wave. Every arithmetic step stays in `Long` + * and the result is range-checked, because `ceil(raw / slots) * slots` can exceed the largest + * partition count Spark can express even when the requirement itself does not. + * + * @return the wave-quantized candidate, or None when no whole-wave count within `Int.MaxValue` + * covers the requirement + */ + private[tuning] def waveQuantized(requirement: Long, slots: Int): Option[Int] = { + if (slots <= 0 || requirement > Int.MaxValue.toLong) { + return None + } + // Both operands are at most Int.MaxValue here, so neither the sum nor the product overflows. + val raw = math.max(slots.toLong, requirement) + val candidate = ((raw + slots - 1L) / slots) * slots + if (candidate > Int.MaxValue.toLong) None else Some(candidate.toInt) + } + + /** + * Builds the single user-facing comment for an applied reduction. It names every input of the + * decision so the recommendation can be audited without re-running the tool. + */ + def appliedComment( + partitionUpdates: Seq[(String, Int)], + decision: DownwardShuffleDecision.Applied): String = { + val record = decision.determiningRecord + // Normally every property lands on the same value, and the comment reads as one sentence about + // both. They can only differ if a property was already below the candidate and got clamped to + // its own value, in which case each is named with what it actually became. + val appliedValues = partitionUpdates.map(_._2).distinct + val lowered = if (appliedValues.size == 1) { + s"${partitionUpdates.map(p => s"'${p._1}'").mkString(" and ")} lowered from " + + s"${decision.normalValue} to ${appliedValues.head}" + } else { + partitionUpdates.map { case (property, value) => s"'$property' lowered to $value" } + .mkString(" and ") + s", from an effective ${decision.normalValue}" + } + lowered + s" based on the " + + s"${decision.provenance.label} shuffle input of the worst consumer stage " + + s"(SQL ${record.sqlId}, stage ${record.stageId}, attempt ${record.stageAttemptId}): " + + s"${record.totalShuffleInputBytes} bytes across ${record.numShuffleBranches} shuffle " + + s"branch(es), input size factor ${decision.inputSizeFactor}, " + + s"${decision.estimatedInputBytes} estimated bytes, " + + s"raw requirement ${decision.rawRequirement} partitions rounded up to " + + s"${decision.waveCount} execution wave(s) of ${decision.slotCount} cluster task slots " + + s"(${decision.selectedValue} partitions)." + } +} diff --git a/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/QualAppSummaryInfoProvider.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/QualAppSummaryInfoProvider.scala index acf08c417..9d48d0226 100644 --- a/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/QualAppSummaryInfoProvider.scala +++ b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/QualAppSummaryInfoProvider.scala @@ -17,8 +17,8 @@ package com.nvidia.spark.rapids.tool.tuning import com.nvidia.spark.rapids.tool.AppSummaryInfoBaseProvider -import com.nvidia.spark.rapids.tool.analysis.AggRawMetricsResult -import com.nvidia.spark.rapids.tool.profiling.{DataSourceProfileResult, PySparkMemoryEvidence} +import com.nvidia.spark.rapids.tool.analysis.{AggRawMetricsResult, AppSQLPlanAnalyzer} +import com.nvidia.spark.rapids.tool.profiling.{DataSourceProfileResult, PySparkMemoryEvidence, ShuffleInputProvenance, ShuffleStageInputAnalysis} import org.apache.spark.internal.Logging import org.apache.spark.sql.rapids.tool.plangraph.ToolsPlanGraph @@ -31,12 +31,16 @@ import org.apache.spark.sql.rapids.tool.qualification.{QualificationAppInfo, Qua * need to feed the autotuner with values from the aggregates. * @param rawAggMetrics the raw profiler aggregation metrics * @param dsInfo Data source information + * @param sqlAnalyzer the SQL plan analyzer already created for this application. It is reused + * rather than rebuilt so the SQL plans are traversed only once. */ class QualAppSummaryInfoProvider( val appInfo: QualificationAppInfo, val appAggStats: Option[QualificationSummaryInfo], val rawAggMetrics: AggRawMetricsResult, - val dsInfo: Seq[DataSourceProfileResult]) extends AppSummaryInfoBaseProvider with Logging { + val dsInfo: Seq[DataSourceProfileResult], + val sqlAnalyzer: Option[AppSQLPlanAnalyzer] = None) + extends AppSummaryInfoBaseProvider with Logging { // Group data source info by location for distinct location calculations. // TODO: Should We only consider data sources from the final plan? // And also drop the location is unknown. @@ -151,6 +155,16 @@ class QualAppSummaryInfoProvider( appInfo.classpathEntries } + /** + * CPU event logs expose regular Exchange data sizes, so the GPU input is estimated from them. + * Without an analyzer there is no evidence at all, which fails closed. + */ + override def getShuffleStageInputAnalysis: ShuffleStageInputAnalysis = { + sqlAnalyzer + .map(_.shuffleStageInputAnalysis) + .getOrElse(ShuffleStageInputAnalysis.empty(ShuffleInputProvenance.Estimated)) + } + override def getPySparkMemoryEvidence: Seq[PySparkMemoryEvidence] = { PySparkMemoryEvidence.fromApp(appInfo) } diff --git a/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/QualificationAutoTunerRunner.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/QualificationAutoTunerRunner.scala index c79f91c27..86644cf2c 100644 --- a/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/QualificationAutoTunerRunner.scala +++ b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/QualificationAutoTunerRunner.scala @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024-2025, NVIDIA CORPORATION. + * Copyright (c) 2024-2026, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -19,7 +19,7 @@ package com.nvidia.spark.rapids.tool.tuning import scala.util.{Failure, Success, Try} import com.nvidia.spark.rapids.tool.{AppSummaryInfoBaseProvider, Platform} -import com.nvidia.spark.rapids.tool.analysis.AggRawMetricsResult +import com.nvidia.spark.rapids.tool.analysis.{AggRawMetricsResult, AppSQLPlanAnalyzer} import com.nvidia.spark.rapids.tool.profiling.DataSourceProfileResult import com.nvidia.spark.rapids.tool.tuning.config.TuningConfiguration import com.nvidia.spark.rapids.tool.views.qualification.QualReportGenConfProvider @@ -92,10 +92,12 @@ object QualificationAutoTunerRunner extends Logging { appAggStats: Option[QualificationSummaryInfo], tunerContext: TunerContext, rawAggMetrics: AggRawMetricsResult, - dsInfo: Seq[DataSourceProfileResult]): Option[QualificationAutoTunerRunner] = { + dsInfo: Seq[DataSourceProfileResult], + sqlAnalyzer: Option[AppSQLPlanAnalyzer] = None): Option[QualificationAutoTunerRunner] = { Try { val qualInfoProvider: QualAppSummaryInfoProvider = - AppSummaryInfoBaseProvider.fromQualAppInfo(appInfo, appAggStats, rawAggMetrics, dsInfo) + AppSummaryInfoBaseProvider + .fromQualAppInfo(appInfo, appAggStats, rawAggMetrics, dsInfo, sqlAnalyzer) .asInstanceOf[QualAppSummaryInfoProvider] new QualificationAutoTunerRunner(qualInfoProvider, tunerContext) } match { diff --git a/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/TunerContext.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/TunerContext.scala index 4d8ee38d8..d8e21ed15 100644 --- a/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/TunerContext.scala +++ b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/TunerContext.scala @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024-2025, NVIDIA CORPORATION. + * Copyright (c) 2024-2026, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -98,7 +98,9 @@ case class TunerContext ( val sqlAnalyzer = AppSQLPlanAnalyzer(appInfo) val rawAggMetrics = QualSparkMetricsAggregator.getAggRawMetrics(appInfo, appIndex, Some(sqlAnalyzer)) - QualificationAutoTunerRunner(appInfo, appAggStats, this, rawAggMetrics, dsInfo).collect { + // Pass the analyzer instance along so the tuning path reuses its SQL plan analysis. + QualificationAutoTunerRunner(appInfo, appAggStats, this, rawAggMetrics, dsInfo, + Some(sqlAnalyzer)).collect { case qualTuner => Try { qualTuner.runAutoTuner(platform, getUserProvidedTuningConfigs) diff --git a/core/src/main/scala/org/apache/spark/sql/rapids/tool/EventProcessorBase.scala b/core/src/main/scala/org/apache/spark/sql/rapids/tool/EventProcessorBase.scala index 1eb7e68bd..ae80a3695 100644 --- a/core/src/main/scala/org/apache/spark/sql/rapids/tool/EventProcessorBase.scala +++ b/core/src/main/scala/org/apache/spark/sql/rapids/tool/EventProcessorBase.scala @@ -192,6 +192,9 @@ abstract class EventProcessorBase[T <: AppBase](app: T) extends SparkListener wi app.sqlIdToInfo.get(event.executionId).foreach { sql => sql.endTime = Some(event.time) sql.duration = ProfileUtils.OptionLongMinusLong(sql.endTime, sql.startTime) + // An end time alone does not prove the execution succeeded. Record the failure text so + // analyses that need terminal success can tell a completed execution from a failed one. + sql.failureReason = EventUtils.readErrorMessageFromSQLEndEvent(event) } } @@ -399,11 +402,10 @@ abstract class EventProcessorBase[T <: AppBase](app: T) extends SparkListener wi def doSparkListenerTaskEnd( app: T, event: SparkListenerTaskEnd): Unit = { - // TODO: this implementation needs to be updated to use attemptID // Parse task accumulables for (res <- event.taskInfo.accumulables) { try { - app.accumManager.addAccToTask(event.stageId, res) + app.accumManager.addAccToTask(event.stageId, event.stageAttemptId, res) } catch { case NonFatal(e) => logWarning("Exception when parsing accumulables on task-completed " diff --git a/core/src/main/scala/org/apache/spark/sql/rapids/tool/store/AccumManager.scala b/core/src/main/scala/org/apache/spark/sql/rapids/tool/store/AccumManager.scala index bd3b287fe..d46a88432 100644 --- a/core/src/main/scala/org/apache/spark/sql/rapids/tool/store/AccumManager.scala +++ b/core/src/main/scala/org/apache/spark/sql/rapids/tool/store/AccumManager.scala @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024-2025, NVIDIA CORPORATION. + * Copyright (c) 2024-2026, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -21,6 +21,7 @@ import scala.collection.{mutable, Map} import com.nvidia.spark.rapids.tool.analysis.StatisticsMetrics import org.apache.spark.scheduler.AccumulableInfo +import org.apache.spark.sql.rapids.tool.util.EventUtils /** * A class that manages task/stage accumulables - @@ -31,6 +32,9 @@ class AccumManager { new mutable.HashMap[Long, AccumInfo]() } + // Stage attempts that reported positive GPU spill. See [[hasGpuSpillEvidence]]. + private val gpuSpillStageAttempts: mutable.HashSet[(Int, Int)] = mutable.HashSet.empty + private def getOrCreateAccumInfo(id: Long, name: Option[String]): AccumInfo = { accumInfoMap.getOrElseUpdate(id, AccumInfo(AccumMetaRef(id, name))) } @@ -40,9 +44,48 @@ class AccumManager { accumInfoRef.addAccumToStage(stageId, accumulableInfo) } - def addAccToTask(stageId: Int, accumulableInfo: AccumulableInfo): Unit = { + /** + * Records a task-level accumulable. + * + * The per-accumulator statistics are intentionally keyed by stage only (see [[AccumInfo]]). + * The stage attempt is additionally retained for the GPU spill metrics, because a downward + * tuning decision must be able to tell spill in a failed attempt apart from spill in the + * later successful attempt it selected. + */ + def addAccToTask(stageId: Int, stageAttemptId: Int, accumulableInfo: AccumulableInfo): Unit = { val accumInfoRef = getOrCreateAccumInfo(accumulableInfo.id, accumulableInfo.name) accumInfoRef.addAccumToTask(stageId, accumulableInfo) + recordGpuSpillAttempt(stageId, stageAttemptId, accumulableInfo) + } + + /** + * Notes the stage attempt when a GPU spill accumulable reports a positive update. + * + * Only the spill metric names are tracked, so this adds a bounded amount of state rather than + * duplicating every accumulator per attempt. + */ + private def recordGpuSpillAttempt( + stageId: Int, stageAttemptId: Int, accumulableInfo: AccumulableInfo): Unit = { + val isSpillMetric = + accumulableInfo.name.exists(AccumManager.GPU_SPILL_METRIC_NAMES.contains) + if (isSpillMetric) { + val positiveUpdate = accumulableInfo.update + .flatMap(EventUtils.parseAccumFieldToLong) + .exists(_ > 0L) + if (positiveUpdate) { + gpuSpillStageAttempts += ((stageId, stageAttemptId)) + } + } + } + + /** + * True when the given stage attempt reported any positive GPU spill activity. + * + * GPU host and disk spill are not visible in Spark's task metrics, so this is the only + * attempt-scoped spill evidence available for GPU event logs. + */ + def hasGpuSpillEvidence(stageId: Int, stageAttemptId: Int): Boolean = { + gpuSpillStageAttempts.contains((stageId, stageAttemptId)) } def getAccStageIds(id: Long): Set[Int] = { @@ -74,3 +117,15 @@ class AccumManager { accumInfoMap.values.foreach(f) } } + +object AccumManager { + /** + * RAPIDS accumulator names that indicate the GPU spilled. These are reported as durations, so + * only their positive/zero state is meaningful here, not their magnitude. + */ + val GPU_SPILL_METRIC_NAMES: Set[String] = Set( + "gpuSpillToHostTime", + "gpuSpillToDiskTime", + "gpuReadSpillFromHostTime", + "gpuReadSpillFromDiskTime") +} diff --git a/core/src/main/scala/org/apache/spark/sql/rapids/tool/util/EventUtils.scala b/core/src/main/scala/org/apache/spark/sql/rapids/tool/util/EventUtils.scala index 9959d5728..8fb62b769 100644 --- a/core/src/main/scala/org/apache/spark/sql/rapids/tool/util/EventUtils.scala +++ b/core/src/main/scala/org/apache/spark/sql/rapids/tool/util/EventUtils.scala @@ -26,7 +26,7 @@ import scala.util.matching.Regex import org.json4s.jackson.JsonMethods.parse import org.apache.spark.internal.Logging -import org.apache.spark.sql.execution.ui.SparkListenerSQLExecutionStart +import org.apache.spark.sql.execution.ui.{SparkListenerSQLExecutionEnd, SparkListenerSQLExecutionStart} /** * Utility containing the implementation of helpers used for parsing data from event. @@ -210,6 +210,18 @@ object EventUtils extends Logging { }.toOption.flatten.getOrElse(Map.empty) } + // Reads errorMessage via reflection (Spark 3.4+). This is the only failure state that + // SparkListenerSQLExecutionEnd serializes into the event log; its sibling `executionFailure` + // is a live-listener-only field and is always empty when replaying a log. + // Returns None on older versions or when the execution succeeded. + def readErrorMessageFromSQLEndEvent( + event: SparkListenerSQLExecutionEnd): Option[String] = { + Try { + Option(invokeMethodOnEvent(event, "errorMessage")) + .map(_.asInstanceOf[Option[String]]) + }.toOption.flatten.flatten.filter(_.nonEmpty) + } + // Reads jobTags via reflection (Spark 3.5+, introduced for Connect support). // Returns empty set on older versions. def readJobTagsFromSQLStartEvent( diff --git a/core/src/test/scala/com/nvidia/spark/rapids/tool/analysis/ShuffleStageInputMetricsSuite.scala b/core/src/test/scala/com/nvidia/spark/rapids/tool/analysis/ShuffleStageInputMetricsSuite.scala new file mode 100644 index 000000000..130b6f2b6 --- /dev/null +++ b/core/src/test/scala/com/nvidia/spark/rapids/tool/analysis/ShuffleStageInputMetricsSuite.scala @@ -0,0 +1,451 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.nvidia.spark.rapids.tool.analysis + +import java.io.File +import java.nio.charset.StandardCharsets +import java.nio.file.{Files, Paths} + +import com.nvidia.spark.rapids.tool.{EventLogPathProcessor, ToolTestUtils} +import com.nvidia.spark.rapids.tool.profiling.{ShuffleInputProvenance, ShuffleStageInputAnalysis, ShuffleStageInputIncompleteReason, ShuffleStageInputRecord} +import org.apache.hadoop.conf.Configuration +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.spark.internal.Logging +import org.apache.spark.sql.{DataFrame, SparkSession, TrampolineUtil} +import org.apache.spark.sql.rapids.tool.profiling.ApplicationInfo +import org.apache.spark.sql.rapids.tool.store.TaskModel + +/** + * Tests the consumer-stage shuffle input analysis on real Spark plan graphs. + * + * The positive cases run actual queries so the plan graphs, AQE rewrites, exchange reuse, and + * accumulator wiring are the real ones. The fail-closed cases start from the same real + * application and then remove one piece of evidence, which is the most direct way to prove that a + * single gap disables the whole analysis. + */ +class ShuffleStageInputMetricsSuite extends AnyFunSuite with Logging { + + private lazy val sparkSession: SparkSession = { + SparkSession.builder() + .master("local[*]") + .appName("Rapids Shuffle Stage Input Unit Tests") + .getOrCreate() + } + + private lazy val hadoopConf: Configuration = sparkSession.sparkContext.hadoopConfiguration + + private val profilingLogDir = ToolTestUtils.getTestResourcePath("spark-events-profiling") + private val qualificationLogDir = + ToolTestUtils.getTestResourcePath("spark-events-qualification") + + private def loadApp(eventLogPath: String): ApplicationInfo = { + val eventLogInfo = EventLogPathProcessor.getEventLogInfo(eventLogPath, hadoopConf).head._1 + new ApplicationInfo(hadoopConf, eventLogInfo) + } + + private def analyze(app: ApplicationInfo): ShuffleStageInputAnalysis = { + ShuffleStageInputAnalyzer(app) + } + + /** + * Runs a query through a local Spark session, writes its event log, and hands the resulting + * profiled application to the test body. + */ + private def withProfiledQuery(name: String)(query: SparkSession => DataFrame) + (body: ApplicationInfo => Unit): Unit = { + TrampolineUtil.withTempDir { eventLogDir => + val (eventLog, _) = ToolTestUtils.generateEventLog(eventLogDir, name)(query) + body(loadApp(eventLog)) + } + } + + private def joinQuery(spark: SparkSession): DataFrame = { + import spark.implicits._ + val left = spark.sparkContext.makeRDD(1 to 2000, 4).map(i => (i, s"l$i")).toDF("k", "lv") + val right = spark.sparkContext.makeRDD(1 to 2000, 4).map(i => (i, s"r$i")).toDF("k", "rv") + // Repartitioning both sides guarantees two shuffle branches into the join's consumer stage. + left.repartition(4, $"k").join(right.repartition(4, $"k"), "k").groupBy("k").count() + } + + private def singleShuffleQuery(spark: SparkSession): DataFrame = { + import spark.implicits._ + spark.sparkContext.makeRDD(1 to 2000, 4).map(i => (i % 17, i)).toDF("k", "v") + .groupBy("k").sum("v") + } + + private def selfJoinQuery(spark: SparkSession): DataFrame = { + import spark.implicits._ + val base = spark.sparkContext.makeRDD(1 to 2000, 4).map(i => (i % 31, i)).toDF("k", "v") + .groupBy("k").sum("v") + // Joining an aggregation to itself lets Spark reuse the same exchange on both sides. + base.as("a").join(base.as("b"), $"a.k" === $"b.k").select($"a.k") + } + + test("a single shuffle produces one complete consumer-stage record") { + withProfiledQuery("singleShuffle")(singleShuffleQuery) { app => + val analysis = analyze(app) + assert(analysis.isComplete, s"unexpected gaps: ${analysis.incompleteSummary(10)}") + assert(analysis.provenance == ShuffleInputProvenance.Estimated) + assert(analysis.records.nonEmpty) + analysis.records.foreach { record => + assert(record.numShuffleBranches >= 1) + assert(record.totalShuffleInputBytes >= 0L) + assert(record.numTasks > 0) + assert(record.stageAttemptId >= 0) + } + } + } + + test("a two-sided join sums both shuffle branches into one consumer stage") { + withProfiledQuery("twoSidedJoin")(joinQuery) { app => + val analysis = analyze(app) + assert(analysis.isComplete, s"unexpected gaps: ${analysis.incompleteSummary(10)}") + val multiBranch = analysis.records.filter(_.numShuffleBranches >= 2) + assert(multiBranch.nonEmpty, + s"expected a stage fed by several branches but got: ${analysis.records}") + // The stage total must exceed any single branch, which is the whole point of summing. + multiBranch.foreach { record => + assert(record.totalShuffleInputBytes > 0L) + } + } + } + + test("a reused exchange contributes once per consumer branch") { + withProfiledQuery("selfJoinReuse")(selfJoinQuery) { app => + val analysis = analyze(app) + assert(analysis.isComplete, s"unexpected gaps: ${analysis.incompleteSummary(10)}") + // Reuse shows up as extra branches, never as extra records for the same stage. + val branchCounts = analysis.records.map(_.numShuffleBranches) + assert(branchCounts.forall(_ >= 1)) + assert(analysis.records.map(r => (r.sqlId, r.stageId)).distinct.size == + analysis.records.size, "each (sql, consumer stage) must appear exactly once") + } + } + + test("records are deterministic across repeated analysis of the same application") { + withProfiledQuery("determinism")(joinQuery) { app => + assert(analyze(app) == analyze(app)) + } + } + + test("a GPU join event log yields measured per-consumer-stage totals") { + val app = loadApp(s"$profilingLogDir/rapids_join_eventlog.zstd") + val analysis = analyze(app) + assert(app.gpuMode) + assert(analysis.provenance == ShuffleInputProvenance.Measured) + assert(analysis.isComplete, s"unexpected gaps: ${analysis.incompleteSummary(10)}") + // This fixture joins two shuffled inputs into stage 2 and feeds the final stage 3 from one. + assert(analysis.records.toSet == Set( + ShuffleStageInputRecord(sqlId = 0L, stageId = 2, stageAttemptId = 0, + totalShuffleInputBytes = 80152384L, numShuffleBranches = 2, numTasks = 200, + hasPositiveSpill = false, hasSkew = false), + ShuffleStageInputRecord(sqlId = 0L, stageId = 3, stageAttemptId = 0, + totalShuffleInputBytes = 19600L, numShuffleBranches = 1, numTasks = 1, + hasPositiveSpill = false, hasSkew = false))) + } + + test("a failed SQL execution makes the whole analysis incomplete") { + withProfiledQuery("failedSql")(joinQuery) { app => + assert(analyze(app).isComplete) + // A recorded end time alone does not prove success; a failure must disable the analysis. + app.sqlIdToInfo.values.foreach(_.failureReason = Some("injected failure")) + val analysis = analyze(app) + assert(!analysis.isComplete) + assert(analysis.incompleteReasons.exists { + case _: ShuffleStageInputIncompleteReason.IncompleteSqlExecution => true + case _ => false + }) + } + } + + test("an unfinished SQL execution makes the whole analysis incomplete") { + withProfiledQuery("unfinishedSql")(joinQuery) { app => + app.sqlIdToInfo.values.foreach(_.endTime = None) + val analysis = analyze(app) + assert(!analysis.isComplete) + assert(analysis.incompleteReasons.exists { + case _: ShuffleStageInputIncompleteReason.IncompleteSqlExecution => true + case _ => false + }) + } + } + + test("a CPU exchange in a GPU application is an unsupported shuffle input") { + withProfiledQuery("mixedPlan")(joinQuery) { app => + assert(analyze(app).isComplete) + // Treat the CPU plan as a GPU one: its plain Exchange nodes now lack GPU size metrics, + // which is the mixed-plan situation that must never produce a reduction. + app.gpuMode = true + val analysis = analyze(app) + assert(!analysis.isComplete) + assert(analysis.incompleteReasons.exists { + case _: ShuffleStageInputIncompleteReason.UnsupportedShuffleNode => true + case _ => false + }) + assert(analysis.provenance == ShuffleInputProvenance.Measured) + } + } + + test("a missing exchange size metric makes the whole analysis incomplete") { + withProfiledQuery("missingMetric")(joinQuery) { app => + assert(analyze(app).isComplete) + // Drop the accumulator evidence behind every exchange 'data size' metric. + val dataSizeAccumIds = app.sqlManager.sqlPlans.values.flatMap { planModel => + planModel.getToolsPlanGraph.allNodes + .filter(n => ShuffleStageInputAnalyzer.isShuffleInputCandidate(n.name)) + .flatMap(_.metrics.filter(_.name == ShuffleStageInputAnalyzer.DATA_SIZE_METRIC)) + .map(_.accumulatorId) + }.toSet + assert(dataSizeAccumIds.nonEmpty) + dataSizeAccumIds.foreach(app.accumManager.removeAccumInfo) + val analysis = analyze(app) + assert(!analysis.isComplete) + assert(analysis.incompleteReasons.exists { + case _: ShuffleStageInputIncompleteReason.MissingExchangeMetric => true + case _ => false + }) + } + } + + test("a consumer stage without a completed successful attempt makes the analysis incomplete") { + withProfiledQuery("noSuccessfulAttempt")(joinQuery) { app => + val consumerStageIds = analyze(app).records.map(_.stageId) + assert(consumerStageIds.nonEmpty) + app.stageManager.removeStages(consumerStageIds) + val analysis = analyze(app) + assert(!analysis.isComplete) + assert(analysis.incompleteReasons.exists { + case _: ShuffleStageInputIncompleteReason.NoSuccessfulStageAttempt => true + case _ => false + }) + } + } + + test("a broadcast exchange is never treated as a shuffle input") { + assert(!ShuffleStageInputAnalyzer.isShuffleInputCandidate("BroadcastExchange")) + assert(!ShuffleStageInputAnalyzer.isShuffleInputCandidate("GpuBroadcastExchange")) + assert(ShuffleStageInputAnalyzer.isShuffleInputCandidate("Exchange")) + assert(ShuffleStageInputAnalyzer.isShuffleInputCandidate("GpuColumnarExchange")) + } + + /** + * Minimal task carrying only the fields the totals and spill predicates read. Every other field + * is zeroed, so a test that starts depending on one will fail loudly rather than silently. + */ + private def task( + successful: Boolean = true, + speculative: Boolean = false, + memoryBytesSpilled: Long = 0L, + diskBytesSpilled: Long = 0L): TaskModel = { + TaskModel( + stageId = 1, stageAttemptId = 0, taskType = "ShuffleMapTask", + endReason = if (successful) "Success" else "ExecutorLostFailure", + taskId = 1L, attempt = 0, launchTime = 0L, finishTime = 0L, duration = 0L, + successful = successful, taskStatus = if (successful) "SUCCESS" else "FAILED", + executorId = "1", host = "host1", taskLocality = "PROCESS_LOCAL", speculative = speculative, + gettingResultTime = 0L, executorDeserializeTime = 0L, executorDeserializeCPUTime = 0L, + executorRunTime = 0L, executorCPUTime = 0L, peakExecutionMemory = 0L, resultSize = 0L, + jvmGCTime = 0L, resultSerializationTime = 0L, + memoryBytesSpilled = memoryBytesSpilled, diskBytesSpilled = diskBytesSpilled, + sr_remoteBlocksFetched = 0L, sr_localBlocksFetched = 0L, sr_fetchWaitTime = 0L, + sr_remoteBytesRead = 0L, sr_remoteBytesReadToDisk = 0L, sr_localBytesRead = 0L, + sr_totalBytesRead = 0L, sw_bytesWritten = 0L, sw_writeTime = 0L, sw_recordsWritten = 0L, + input_bytesRead = 0L, input_recordsRead = 0L, output_bytesWritten = 0L, + output_recordsWritten = 0L) + } + + test("a failed task's spill still blocks a reduction even though its bytes are not counted") { + val failedSpiller = task(successful = false, memoryBytesSpilled = 1L) + // The totals must not absorb work that never produced output. + assert(!ShuffleStageInputAnalyzer.countsTowardTotals(failedSpiller)) + // The gate must still see it: a task that spilled and then died is the pressure this pass + // exists to avoid reducing into. + assert(ShuffleStageInputAnalyzer.countsTowardSpillGate(failedSpiller)) + assert(ShuffleStageInputAnalyzer.spilled(failedSpiller)) + } + + test("speculative duplicates are excluded from both the totals and the spill gate") { + val speculativeSpiller = task(speculative = true, diskBytesSpilled = 1L) + assert(!ShuffleStageInputAnalyzer.countsTowardTotals(speculativeSpiller)) + assert(!ShuffleStageInputAnalyzer.countsTowardSpillGate(speculativeSpiller)) + // A speculative task that also failed stays excluded. + assert(!ShuffleStageInputAnalyzer.countsTowardSpillGate( + task(successful = false, speculative = true))) + } + + test("a successful task counts toward both, and spill is either memory or disk") { + val clean = task() + assert(ShuffleStageInputAnalyzer.countsTowardTotals(clean)) + assert(ShuffleStageInputAnalyzer.countsTowardSpillGate(clean)) + assert(!ShuffleStageInputAnalyzer.spilled(clean)) + assert(ShuffleStageInputAnalyzer.spilled(task(memoryBytesSpilled = 1L))) + assert(ShuffleStageInputAnalyzer.spilled(task(diskBytesSpilled = 1L))) + // Zero spill on both counters is not spill evidence. + assert(!ShuffleStageInputAnalyzer.spilled( + task(memoryBytesSpilled = 0L, diskBytesSpilled = 0L))) + } + + test("GPU spill evidence is retained per stage attempt") { + TrampolineUtil.withTempDir { tmpDir => + // Attempt 0 of stage 10 spilled; attempt 1 of the same stage ran clean. + val logPath = Paths.get(tmpDir.getAbsolutePath, "gpu_spill_attempts_eventlog") + // scalastyle:off line.size.limit + val content = + """{"Event":"SparkListenerLogStart","Spark Version":"3.5.0"} + |{"Event":"SparkListenerApplicationStart","App Name":"GpuSpillAttempts","App ID":"local-1600000000000","Timestamp":123456,"User":"tester"} + |{"Event":"SparkListenerTaskEnd","Stage ID":10,"Stage Attempt ID":0,"Task Type":"ShuffleMapTask","Task End Reason":{"Reason":"Success"},"Task Info":{"Task ID":1,"Index":1,"Attempt":0,"Partition ID":1,"Launch Time":1712248533994,"Executor ID":"1","Host":"host1","Locality":"PROCESS_LOCAL","Speculative":false,"Getting Result Time":0,"Finish Time":1712248534994,"Failed":false,"Killed":false,"Accumulables":[{"ID":1018,"Name":"gpuSpillToHostTime","Update":"00:00:00.845","Value":"00:00:00.845","Internal":false,"Count Failed Values":true}]}} + |{"Event":"SparkListenerTaskEnd","Stage ID":10,"Stage Attempt ID":1,"Task Type":"ShuffleMapTask","Task End Reason":{"Reason":"Success"},"Task Info":{"Task ID":2,"Index":1,"Attempt":0,"Partition ID":1,"Launch Time":1712248535994,"Executor ID":"1","Host":"host1","Locality":"PROCESS_LOCAL","Speculative":false,"Getting Result Time":0,"Finish Time":1712248536994,"Failed":false,"Killed":false,"Accumulables":[{"ID":1010,"Name":"gpuSemaphoreWait","Update":"00:00:00.492","Value":"00:00:00.492","Internal":false,"Count Failed Values":true}]}}""".stripMargin + // scalastyle:on line.size.limit + Files.write(logPath, content.getBytes(StandardCharsets.UTF_8)) + val app = loadApp(logPath.toString) + assert(app.accumManager.hasGpuSpillEvidence(10, 0), + "the failed-style attempt 0 must retain its spill evidence") + assert(!app.accumManager.hasGpuSpillEvidence(10, 1), + "the clean attempt 1 must not inherit attempt 0's spill") + } + } + + test("a zero GPU spill update is not treated as spill evidence") { + TrampolineUtil.withTempDir { tmpDir => + val logPath = Paths.get(tmpDir.getAbsolutePath, "gpu_zero_spill_eventlog") + // scalastyle:off line.size.limit + val content = + """{"Event":"SparkListenerLogStart","Spark Version":"3.5.0"} + |{"Event":"SparkListenerApplicationStart","App Name":"GpuZeroSpill","App ID":"local-1600000000001","Timestamp":123456,"User":"tester"} + |{"Event":"SparkListenerTaskEnd","Stage ID":7,"Stage Attempt ID":0,"Task Type":"ShuffleMapTask","Task End Reason":{"Reason":"Success"},"Task Info":{"Task ID":1,"Index":1,"Attempt":0,"Partition ID":1,"Launch Time":1712248533994,"Executor ID":"1","Host":"host1","Locality":"PROCESS_LOCAL","Speculative":false,"Getting Result Time":0,"Finish Time":1712248534994,"Failed":false,"Killed":false,"Accumulables":[{"ID":1018,"Name":"gpuSpillToHostTime","Update":"00:00:00.000","Value":"00:00:00.000","Internal":false,"Count Failed Values":true}]}}""".stripMargin + // scalastyle:on line.size.limit + Files.write(logPath, content.getBytes(StandardCharsets.UTF_8)) + val app = loadApp(logPath.toString) + assert(!app.accumManager.hasGpuSpillEvidence(7, 0)) + } + } + + test("the profiling analyzer exposes the same analysis without a second traversal") { + withProfiledQuery("providerReuse")(joinQuery) { app => + val fromAnalyzer = app.planMetricProcessor.shuffleStageInputAnalysis + assert(fromAnalyzer == analyze(app)) + // The lazy val must be cached rather than recomputed on every access. + assert(app.planMetricProcessor.shuffleStageInputAnalysis eq fromAnalyzer) + } + } + + test("a CPU event log without a terminal SQL end event is incomplete") { + val app = loadApp(s"$qualificationLogDir/join_missing_sql_end") + val analysis = analyze(app) + assert(!analysis.isComplete) + assert(analysis.records.isEmpty) + assert(analysis.incompleteReasons == + Seq(ShuffleStageInputIncompleteReason.IncompleteSqlExecution(0L))) + } + + test("a CPU AQE shuffle event log records multi-branch totals and real spill pressure") { + val app = loadApp(s"$qualificationLogDir/aqeshuffle_eventlog.zstd") + val analysis = analyze(app) + assert(analysis.isComplete, s"unexpected gaps: ${analysis.incompleteSummary(10)}") + assert(analysis.provenance == ShuffleInputProvenance.Estimated) + val joinStage = analysis.records.find(_.stageId == 4) + assert(joinStage.isDefined, s"stage 4 missing from ${analysis.records}") + // Two shuffle branches feed stage 4, and its tasks really did spill in this fixture. + assert(joinStage.get.numShuffleBranches == 2) + assert(joinStage.get.totalShuffleInputBytes == 320000000L) + assert(joinStage.get.hasPositiveSpill, + "the fixture's spill evidence must reach the record so the pass stays blocked") + } + + test("a CPU query event log totals each consumer stage of its final plan") { + val app = loadApp(s"$qualificationLogDir/nds_q86_test") + val analysis = analyze(app) + assert(analysis.isComplete, s"unexpected gaps: ${analysis.incompleteSummary(10)}") + assert(analysis.records.toSet == Set( + ShuffleStageInputRecord(sqlId = 24L, stageId = 34, stageAttemptId = 0, + totalShuffleInputBytes = 5491688L, numShuffleBranches = 1, numTasks = 1024, + hasPositiveSpill = false, hasSkew = false), + ShuffleStageInputRecord(sqlId = 24L, stageId = 35, stageAttemptId = 0, + totalShuffleInputBytes = 17600L, numShuffleBranches = 1, numTasks = 1024, + hasPositiveSpill = false, hasSkew = false))) + } + + test("an event log with several stage attempts selects one completed successful attempt") { + val app = loadApp(s"$qualificationLogDir/multiple_attempts") + val analysis = analyze(app) + assert(analysis.isComplete, s"unexpected gaps: ${analysis.incompleteSummary(10)}") + assert(analysis.records.nonEmpty) + analysis.records.foreach { record => + val selected = app.stageManager.getStagesByIds(Seq(record.stageId)) + .find(_.getAttemptId == record.stageAttemptId) + assert(selected.isDefined, s"attempt ${record.stageAttemptId} missing for $record") + assert(!selected.get.hasFailed) + assert(selected.get.stageInfo.completionTime.isDefined) + // No later completed successful attempt may exist for the same stage. + assert(!app.stageManager.getStagesByIds(Seq(record.stageId)).exists { candidate => + candidate.getAttemptId > record.stageAttemptId && !candidate.hasFailed && + candidate.stageInfo.completionTime.isDefined + }) + } + } + + test("an exchange whose downstream nodes carry no stage is still attributed") { + // Ends in a repartition, so the topmost exchange's only sink is the plan root, which carries + // no metrics and therefore no stage assignment. Walking the graph dead-ends there, and the + // reading stage has to be recovered from the exchange's own assignment instead. + val confs = Map( + "spark.sql.adaptive.enabled" -> "true", + "spark.sql.autoBroadcastJoinThreshold" -> "-1", + "spark.sql.shuffle.partitions" -> "8") + TrampolineUtil.withTempDir { dir => + val (log, _) = ToolTestUtils.generateEventLog(dir, "rootExchange", Some(confs)) { spark => + import spark.implicits._ + val left = spark.sparkContext.makeRDD(1 to 2000, 4).map(i => (i, s"l$i")).toDF("k", "lv") + val right = spark.sparkContext.makeRDD(1 to 2000, 4).map(i => (i, s"r$i")).toDF("k", "rv") + left.join(right, "k").repartition(4) + } + val app = loadApp(log) + val analysis = analyze(app) + assert(analysis.isComplete, s"unexpected gaps: ${analysis.incompleteSummary(10)}") + + // Confirm the shape this test exists for: an exchange whose sink has no stage assignment. + val deadEnding = app.sqlManager.sqlPlans.values.flatMap { planModel => + val graph = planModel.getToolsPlanGraph + graph.allNodes + .filter(n => ShuffleStageInputAnalyzer.isShuffleInputCandidate(n.name)) + .filter { n => + val sinks = graph.getSinkNodes(n.id) + sinks.nonEmpty && sinks.forall(graph.getNodeStageRawAssignment(_).isEmpty) + } + } + assert(deadEnding.nonEmpty, + "fixture no longer produces an exchange whose sinks are unassigned") + + // Every executed exchange must contribute exactly one branch: the join's two inputs land + // in one consumer stage, and the trailing repartition feeds another. + val totalExchanges = app.sqlManager.sqlPlans.values.flatMap { planModel => + planModel.getToolsPlanGraph.allNodes + .filter(n => ShuffleStageInputAnalyzer.isShuffleInputCandidate(n.name)) + .filter(n => planModel.getToolsPlanGraph.getNodeStageRawAssignment(n.id).nonEmpty) + }.size + assert(analysis.records.map(_.numShuffleBranches).sum == totalExchanges, + s"not every exchange was attributed: ${analysis.records}") + assert(analysis.records.exists(_.numShuffleBranches == 2), + s"expected the join stage to sum both inputs: ${analysis.records}") + } + } + + test("test resources are available") { + assert(new File(s"$profilingLogDir/rapids_join_eventlog.zstd").exists()) + } +} diff --git a/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/BaseAutoTunerSuite.scala b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/BaseAutoTunerSuite.scala index 5abf99594..a3604bb17 100644 --- a/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/BaseAutoTunerSuite.scala +++ b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/BaseAutoTunerSuite.scala @@ -54,7 +54,9 @@ class AppInfoProviderMockTest(val maxInput: Double, val maxColumnarExchangeDataSizeBytes: Option[Long] = None, val maxFileScanInputOverride: Option[Option[Double]] = None, val pySparkMemoryEvidence: Seq[PySparkMemoryEvidence] = Seq.empty, - val hasSqlCache: Boolean = false) + val hasSqlCache: Boolean = false, + val shuffleStageInputAnalysis: ShuffleStageInputAnalysis = + ShuffleStageInputAnalysis.empty(ShuffleInputProvenance.Measured)) extends BaseProfilingAppSummaryInfoProvider { override def isAppInfoAvailable = true override def getMaxFileScanInput: Option[Double] = @@ -78,6 +80,7 @@ class AppInfoProviderMockTest(val maxInput: Double, override def getMaxColumnarExchangeDataSizeBytes: Option[Long] = maxColumnarExchangeDataSizeBytes override def getPySparkMemoryEvidence: Seq[PySparkMemoryEvidence] = pySparkMemoryEvidence override def hasSqlCacheEvidence: Boolean = hasSqlCache + override def getShuffleStageInputAnalysis: ShuffleStageInputAnalysis = shuffleStageInputAnalysis /** * Sets the spark master property in the properties map. @@ -148,12 +151,45 @@ abstract class BaseAutoTunerSuite extends AnyFunSuite with BeforeAndAfterEach maxColumnarExchangeDataSizeBytes: Option[Long] = None, maxFileScanInputOverride: Option[Option[Double]] = None, pySparkMemoryEvidence: Seq[PySparkMemoryEvidence] = Seq.empty, - hasSqlCache: Boolean = false): AppInfoProviderMockTest = { + hasSqlCache: Boolean = false, + shuffleStageInputAnalysis: ShuffleStageInputAnalysis = + ShuffleStageInputAnalysis.empty(ShuffleInputProvenance.Measured) + ): AppInfoProviderMockTest = { new AppInfoProviderMockTest(maxInput, spilledMetrics, jvmGCFractions, propsFromLog, sparkVersion, rapidsJars, distinctLocationPct, redundantReadSize, meanInput, meanShuffleRead, shuffleStagesWithPosSpilling, shuffleSkewStages, scanStagesWithGpuOom, gpuShuffleStagesWithContainerOom, maxColumnarExchangeDataSizeBytes, - maxFileScanInputOverride, pySparkMemoryEvidence, hasSqlCache) + maxFileScanInputOverride, pySparkMemoryEvidence, hasSqlCache, shuffleStageInputAnalysis) + } + + /** + * Builds a complete shuffle-stage input analysis with one record per (stage, bytes) pair. + * Used to drive the downward shuffle-partition pass from AutoTuner tests. + */ + protected def completeShuffleStageInputs( + stageInputs: Seq[(Int, Long)], + provenance: ShuffleInputProvenance = ShuffleInputProvenance.Measured, + numBranches: Int = 1, + hasPositiveSpill: Boolean = false, + hasSkew: Boolean = false, + appHasFailedStage: Boolean = false): ShuffleStageInputAnalysis = { + val records = stageInputs.map { case (stageId, bytes) => + ShuffleStageInputRecord(sqlId = 0L, stageId = stageId, stageAttemptId = 0, + totalShuffleInputBytes = bytes, numShuffleBranches = numBranches, numTasks = 200, + hasPositiveSpill = hasPositiveSpill, hasSkew = hasSkew) + } + ShuffleStageInputAnalysis(records, Seq.empty, provenance, + appHasFailedStage = appHasFailedStage) + } + + /** Builds an analysis that ran but found a gap, which must keep the normal recommendation. */ + protected def incompleteShuffleStageInputs( + provenance: ShuffleInputProvenance = ShuffleInputProvenance.Measured + ): ShuffleStageInputAnalysis = { + ShuffleStageInputAnalysis( + Seq.empty, + Seq(ShuffleStageInputIncompleteReason.MissingExchangeMetric(0L, 3L, "Exchange")), + provenance) } /** diff --git a/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/DownwardShufflePartitionsSuite.scala b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/DownwardShufflePartitionsSuite.scala new file mode 100644 index 000000000..94d42f6dd --- /dev/null +++ b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/DownwardShufflePartitionsSuite.scala @@ -0,0 +1,532 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.nvidia.spark.rapids.tool.tuning + +import scala.util.Try + +import com.nvidia.spark.rapids.tool.ToolTestUtils +import com.nvidia.spark.rapids.tool.profiling.{ShuffleInputProvenance, ShuffleStageInputAnalysis, ShuffleStageInputIncompleteReason, ShuffleStageInputRecord} +import com.nvidia.spark.rapids.tool.tuning.config.{ProfTuningConfigProvider, QualTuningConfigProvider, TuningConfigEntry, TuningConfigProvider} +import org.scalatest.funsuite.AnyFunSuite + +/** + * Focused tests for the configuration and the pure downward shuffle-partition policy. + * + * These tests never construct an AutoTuner: the policy layer must be provable on its own. + */ +class DownwardShufflePartitionsSuite extends AnyFunSuite { + + private val GiB = 1024L * 1024L * 1024L + + /** + * The pass ships disabled, so these providers opt in unless the test overrides the switch + * itself. Everything below the switch is then exercised against the shipped defaults. + */ + private def withOptIn(default: List[TuningConfigEntry]): List[TuningConfigEntry] = { + if (default.exists(_.name == DownwardShufflePolicyConfig.ENABLED_KEY)) { + default + } else { + TuningConfigEntry(name = DownwardShufflePolicyConfig.ENABLED_KEY, default = "true") :: default + } + } + + private def profProvider( + default: List[TuningConfigEntry] = List.empty): ProfTuningConfigProvider = { + TuningConfigProvider.builder + .withUserProvidedConfig(Some(ToolTestUtils.buildTuningConfigs(default = withOptIn(default)))) + .build[ProfTuningConfigProvider] + } + + private def qualProvider( + default: List[TuningConfigEntry] = List.empty, + qualification: List[TuningConfigEntry] = List.empty): QualTuningConfigProvider = { + TuningConfigProvider.builder + .withUserProvidedConfig(Some(ToolTestUtils.buildTuningConfigs( + default = withOptIn(default), qualification = qualification))) + .build[QualTuningConfigProvider] + } + + /** Provider with no user overrides at all, so it reflects exactly what the tool ships. */ + private def shippedProfProvider(): ProfTuningConfigProvider = { + TuningConfigProvider.builder.build[ProfTuningConfigProvider] + } + + private def record( + sqlId: Long = 0L, + stageId: Int = 1, + stageAttemptId: Int = 0, + totalBytes: Long, + numBranches: Int = 1, + numTasks: Int = 200, + hasSpill: Boolean = false, + hasSkew: Boolean = false): ShuffleStageInputRecord = { + ShuffleStageInputRecord(sqlId, stageId, stageAttemptId, totalBytes, numBranches, numTasks, + hasSpill, hasSkew) + } + + private def analysis( + records: Seq[ShuffleStageInputRecord], + provenance: ShuffleInputProvenance = ShuffleInputProvenance.Measured + ): ShuffleStageInputAnalysis = { + ShuffleStageInputAnalysis(records, Seq.empty, provenance) + } + + /** Config used by the arithmetic tests: 1 GiB target, cores basis, no reduction threshold. */ + private val baseConfig = DownwardShufflePolicyConfig( + enabled = true, + targetPartitionSizeBytes = GiB, + inputSizeFactor = 1.0) + + /** Slot count of the cluster the arithmetic tests recommend for: 125 executors x 16 cores. */ + private val slots = 2000 + + /** Unwraps a config result that the test expects to be valid. */ + private def expectValid( + result: Either[Seq[String], DownwardShufflePolicyConfig]): DownwardShufflePolicyConfig = { + result.fold(errors => fail(s"expected a valid config but got: ${errors.mkString(", ")}"), + identity) + } + + /** Unwraps a config result that the test expects to be invalid. */ + private def expectErrors( + result: Either[Seq[String], DownwardShufflePolicyConfig]): Seq[String] = { + result.fold(identity, config => fail(s"expected validation errors but got: $config")) + } + + private def decide( + normalValue: Int, + records: Seq[ShuffleStageInputRecord], + config: DownwardShufflePolicyConfig = baseConfig, + provenance: ShuffleInputProvenance = ShuffleInputProvenance.Measured, + slotCount: Option[Int] = Some(slots) + ): DownwardShuffleDecision = { + DownwardShufflePartitionsPolicy.decide(Right(config), normalValue, slotCount, + analysis(records, provenance)) + } + + // + // Configuration + // + + test("the pass ships disabled so enabling it is an explicit opt-in") { + assert(DownwardShufflePolicyConfig.fromProvider(shippedProfProvider()) == + Right(DownwardShufflePolicyConfig.disabled)) + } + + test("profiling defaults match the product contract") { + val config = expectValid(DownwardShufflePolicyConfig.fromProvider(profProvider())) + assert(config.enabled) + assert(config.targetPartitionSizeBytes == GiB) + assert(config.inputSizeFactor == 1.0) + } + + test("the retired sizing entries are gone and their removal does not break config loading") { + Seq("DOWNWARD_SHUFFLE_PARTITION_FLOOR", "DOWNWARD_SHUFFLE_RUNG_MULTIPLIER", + "DOWNWARD_SHUFFLE_SLOT_BASIS", "DOWNWARD_SHUFFLE_MIN_REDUCTION_FACTOR").foreach { key => + Seq(profProvider(), qualProvider()).foreach { provider => + assert(Try(provider.getEntry(key)).isFailure, s"'$key' should no longer be defined") + } + } + assert(DownwardShufflePolicyConfig.fromProvider(profProvider()).isRight) + assert(DownwardShufflePolicyConfig.fromProvider(qualProvider()).isRight) + } + + test("qualification overrides the input size factor in its own tool section") { + val config = expectValid(DownwardShufflePolicyConfig.fromProvider(qualProvider())) + assert(config.enabled) + assert(config.inputSizeFactor == 0.8) + // Everything else still comes from the shared defaults. + assert(config.targetPartitionSizeBytes == GiB) + } + + test("user overrides are honored in the default and tool sections") { + val fromDefault = DownwardShufflePolicyConfig.fromProvider(profProvider( + default = List( + TuningConfigEntry(name = "DOWNWARD_SHUFFLE_TARGET_PARTITION_SIZE", default = "512m")))) + assert(fromDefault == Right(baseConfig.copy( + targetPartitionSizeBytes = 512L * 1024L * 1024L))) + + val fromToolSection = DownwardShufflePolicyConfig.fromProvider(qualProvider( + qualification = List( + TuningConfigEntry(name = "DOWNWARD_SHUFFLE_INPUT_SIZE_FACTOR", default = "0.5")))) + assert(fromToolSection == Right(baseConfig.copy(inputSizeFactor = 0.5))) + } + + test("a disabled feature short-circuits without validating the other entries") { + val provider = profProvider(default = List( + TuningConfigEntry(name = "DOWNWARD_SHUFFLE_ENABLED", default = "false"), + // Deliberately invalid; it must not be read while the feature is off. + TuningConfigEntry(name = "DOWNWARD_SHUFFLE_TARGET_PARTITION_SIZE", default = "not-a-size"))) + val configResult = DownwardShufflePolicyConfig.fromProvider(provider) + assert(configResult == Right(DownwardShufflePolicyConfig.disabled)) + assert(DownwardShufflePartitionsPolicy.decide(configResult, 4000, Some(slots), + analysis(Seq(record(totalBytes = GiB)))) == + DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.Disabled)) + } + + test("the enabled switch requires an exact boolean") { + Seq("yes", "1", "TRUE!", "0").foreach { value => + val result = DownwardShufflePolicyConfig.fromProvider(profProvider( + default = List(TuningConfigEntry(name = "DOWNWARD_SHUFFLE_ENABLED", default = value)))) + assert(result.isLeft, s"'$value' should not parse as a boolean") + } + // Case-insensitive exact values are still accepted. + assert(DownwardShufflePolicyConfig.fromProvider(profProvider( + default = List(TuningConfigEntry(name = "DOWNWARD_SHUFFLE_ENABLED", default = "TRUE")))) + .exists(_.enabled)) + } + + test("every invalid numeric boundary is rejected") { + val invalidByKey = Seq( + "DOWNWARD_SHUFFLE_TARGET_PARTITION_SIZE" -> Seq("0", "-1g", "abc"), + "DOWNWARD_SHUFFLE_INPUT_SIZE_FACTOR" -> Seq("0", "-0.5", "abc")) + + invalidByKey.foreach { case (key, values) => + values.foreach { value => + val result = DownwardShufflePolicyConfig.fromProvider(profProvider( + default = List(TuningConfigEntry(name = key, default = value)))) + assert(result.isLeft, s"'$key' = '$value' should be rejected") + assert(result.swap.exists(_.exists(_.contains(key))), + s"the error for '$key' = '$value' should name the key") + } + } + } + + test("NaN and infinite values are rejected") { + Seq("DOWNWARD_SHUFFLE_INPUT_SIZE_FACTOR").foreach { key => + Seq("NaN", "Infinity", "-Infinity").foreach { value => + val result = DownwardShufflePolicyConfig.fromProvider(profProvider( + default = List(TuningConfigEntry(name = key, default = value)))) + assert(result.isLeft, s"'$key' = '$value' should be rejected") + } + } + } + + test("all configuration errors are reported as one fail-closed decision") { + val provider = profProvider(default = List( + TuningConfigEntry(name = "DOWNWARD_SHUFFLE_TARGET_PARTITION_SIZE", default = "-1g"), + TuningConfigEntry(name = "DOWNWARD_SHUFFLE_INPUT_SIZE_FACTOR", default = "-0.5"))) + val configResult = DownwardShufflePolicyConfig.fromProvider(provider) + val errors = expectErrors(configResult) + assert(errors.size == 2) + DownwardShufflePartitionsPolicy.decide(configResult, 4000, Some(slots), + analysis(Seq(record(totalBytes = GiB)))) match { + case DownwardShuffleDecision.InvalidConfig(reported) => assert(reported == errors) + case other => fail(s"expected an invalid-config decision but got $other") + } + } + + // + // Arithmetic + // + + test("estimateGpuInputBytes rounds up and clamps instead of overflowing") { + assert(DownwardShufflePartitionsPolicy.estimateGpuInputBytes(0L, 1.0) == 0L) + assert(DownwardShufflePartitionsPolicy.estimateGpuInputBytes(10L, 0.8) == 8L) + // 0.75 of 10 is 7.5 and must round up so the requirement is never understated. + assert(DownwardShufflePartitionsPolicy.estimateGpuInputBytes(10L, 0.75) == 8L) + assert(DownwardShufflePartitionsPolicy.estimateGpuInputBytes(Long.MaxValue, 2.0) == + Long.MaxValue) + } + + test("partitionRequirement is an exact ceiling with a lower bound of one") { + assert(DownwardShufflePartitionsPolicy.partitionRequirement(0L, GiB) == 1L) + assert(DownwardShufflePartitionsPolicy.partitionRequirement(1L, GiB) == 1L) + // Exact multiples of the target must not round up to an extra partition. + assert(DownwardShufflePartitionsPolicy.partitionRequirement(1000L * GiB, GiB) == 1000L) + assert(DownwardShufflePartitionsPolicy.partitionRequirement(1000L * GiB + 1L, GiB) == 1001L) + } + + test("waveQuantized rounds up to whole waves and treats the slot count as a floor") { + val cases = Seq( + // Below one wave still gets exactly one wave. + 1L -> 2000, 1999L -> 2000, + // Exactly one wave stays at one wave rather than rounding to two. + 2000L -> 2000, + // One partition above a wave boundary opens the next whole wave. + 2001L -> 4000, 4001L -> 6000, + // Several waves round up to the next whole wave, not to the requirement. + 5000L -> 6000, 6000L -> 6000) + cases.foreach { case (requirement, expected) => + assert(DownwardShufflePartitionsPolicy.waveQuantized(requirement, slots) == Some(expected), + s"requirement $requirement") + } + } + + test("waveQuantized handles a slot count of one and an absent-sized cluster") { + // With one slot per wave the candidate is just the requirement, bounded below by one. + assert(DownwardShufflePartitionsPolicy.waveQuantized(1L, 1) == Some(1)) + assert(DownwardShufflePartitionsPolicy.waveQuantized(37L, 1) == Some(37)) + // A non-positive slot count has no wave to quantize to. + assert(DownwardShufflePartitionsPolicy.waveQuantized(1000L, 0).isEmpty) + assert(DownwardShufflePartitionsPolicy.waveQuantized(1000L, -16).isEmpty) + } + + test("a requirement no representable wave count can cover fails closed") { + // Beyond the largest partition count Spark can express. + assert(DownwardShufflePartitionsPolicy.waveQuantized(Int.MaxValue.toLong + 1L, slots).isEmpty) + assert(DownwardShufflePartitionsPolicy.waveQuantized(Long.MaxValue, slots).isEmpty) + // Representable on its own, but rounding it up to a whole wave overshoots Int.MaxValue. + assert(DownwardShufflePartitionsPolicy.waveQuantized(Int.MaxValue.toLong, slots).isEmpty) + // The exact largest representable whole-wave count is still covered. + val largestWholeWave = (Int.MaxValue / slots) * slots + assert(DownwardShufflePartitionsPolicy.waveQuantized(largestWholeWave.toLong, slots) == + Some(largestWholeWave)) + } + + test("an out-of-range requirement skips instead of recommending an uncovering value") { + decide(Int.MaxValue, Seq(record(totalBytes = Long.MaxValue))) match { + case DownwardShuffleDecision.Skipped( + reason @ DownwardShuffleSkipReason.RequirementOutOfRange(_)) => + assert(reason.isWarning) + case other => fail(s"expected an out-of-range skip but got $other") + } + // Representable requirement whose wave rounding would wrap past Int.MaxValue. + decide(Int.MaxValue, Seq(record(totalBytes = Int.MaxValue.toLong * GiB))) match { + case DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.RequirementOutOfRange(_)) => + case other => fail(s"expected an out-of-range skip but got $other") + } + } + + test("a missing or non-positive slot count is a quiet no-op rather than a division error") { + Seq(None, Some(0), Some(-16)).foreach { slotCount => + assert(decide(8000, Seq(record(totalBytes = 900L * GiB)), slotCount = slotCount) == + DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.SlotCountUnavailable), + s"slot count $slotCount") + } + assert(!DownwardShuffleSkipReason.SlotCountUnavailable.isWarning) + } + + // + // Policy decisions + // + + test("AE1: a multi-input join stage is sized from its combined branches") { + // Two branches whose combined input needs 730 partitions at a 1 GiB target. + val joinStage = record(sqlId = 0L, stageId = 7, totalBytes = 729L * GiB + 1L, numBranches = 2) + decide(normalValue = 8000, records = Seq(joinStage)) match { + case applied: DownwardShuffleDecision.Applied => + assert(applied.rawRequirement == 730L) + // Below one wave, so the cluster's own slot count is the floor. + assert(applied.selectedValue == 2000) + assert(applied.slotCount == 2000) + assert(applied.waveCount == 1) + assert(applied.normalValue == 8000) + assert(applied.determiningRecord == joinStage) + assert(applied.provenance == ShuffleInputProvenance.Measured) + case other => fail(s"expected an applied reduction but got $other") + } + } + + test("the qualification factor lowers the requirement of the same stage") { + val stage = record(totalBytes = 4000L * GiB) + val qualConfig = baseConfig.copy(inputSizeFactor = 0.8) + decide(8000, Seq(stage), qualConfig, ShuffleInputProvenance.Estimated) match { + case applied: DownwardShuffleDecision.Applied => + assert(applied.estimatedInputBytes == 3200L * GiB) + assert(applied.rawRequirement == 3200L) + // 3200 spans two whole waves of 2000. + assert(applied.selectedValue == 4000) + assert(applied.waveCount == 2) + assert(applied.inputSizeFactor == 0.8) + assert(applied.provenance == ShuffleInputProvenance.Estimated) + case other => fail(s"expected an applied reduction but got $other") + } + } + + test("the worst consumer stage determines the candidate") { + val small = record(sqlId = 0L, stageId = 1, totalBytes = 100L * GiB) + val worst = record(sqlId = 1L, stageId = 9, totalBytes = 2500L * GiB) + Seq(Seq(small, worst), Seq(worst, small)).foreach { records => + decide(8000, records) match { + case applied: DownwardShuffleDecision.Applied => + assert(applied.determiningRecord == worst) + assert(applied.selectedValue == 4000) + case other => fail(s"expected an applied reduction but got $other") + } + } + } + + test("ties are broken by estimated bytes, then by ascending SQL and stage ids") { + // All three round to the same requirement of 2 partitions. + val moreBytes = record(sqlId = 5L, stageId = 3, totalBytes = GiB + 900L) + val lowIds = record(sqlId = 1L, stageId = 2, totalBytes = GiB + 1L) + val highIds = record(sqlId = 1L, stageId = 4, totalBytes = GiB + 1L) + + // Larger estimated bytes wins the tie even though its ids sort later. + decide(4000, Seq(lowIds, moreBytes, highIds)) match { + case applied: DownwardShuffleDecision.Applied => + assert(applied.determiningRecord == moreBytes) + case other => fail(s"expected an applied reduction but got $other") + } + // With equal bytes, the lowest (sqlId, stageId) wins regardless of input order. + Seq(Seq(lowIds, highIds), Seq(highIds, lowIds)).foreach { records => + decide(4000, records) match { + case applied: DownwardShuffleDecision.Applied => + assert(applied.determiningRecord == lowIds) + case other => fail(s"expected an applied reduction but got $other") + } + } + } + + test("an exact wave multiple does not round up to an extra wave") { + decide(8000, Seq(record(totalBytes = 4000L * GiB))) match { + case applied: DownwardShuffleDecision.Applied => + assert(applied.rawRequirement == 4000L) + assert(applied.selectedValue == 4000) + assert(applied.waveCount == 2) + case other => fail(s"expected an applied reduction but got $other") + } + } + + test("AE7: a candidate at or above the normal value never raises it") { + // The candidate wave equals the current value. + assert(decide(2000, Seq(record(totalBytes = 900L * GiB))) == + DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.NotDownward(2000, 2000))) + // The candidate wave is larger than the current value. + assert(decide(300, Seq(record(totalBytes = 900L * GiB))) == + DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.NotDownward(2000, 300))) + } + + test("AE7: a normal value below one wave is left alone") { + // Even a one-byte stage cannot pull the recommendation under a single full wave. + assert(decide(200, Seq(record(totalBytes = 1L))) == + DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.NotDownward(2000, 200))) + } + + test("any reduction to a whole wave applies, with no minimum-size threshold") { + // A 3000 -> 2000 reduction is only 1.5x. The wave quantum is the only size gate there is. + decide(3000, Seq(record(totalBytes = 900L * GiB))) match { + case applied: DownwardShuffleDecision.Applied => assert(applied.selectedValue == 2000) + case other => fail(s"expected an applied reduction but got $other") + } + // One partition above the wave is still a reduction and is still applied. + decide(2001, Seq(record(totalBytes = 900L * GiB))) match { + case applied: DownwardShuffleDecision.Applied => assert(applied.selectedValue == 2000) + case other => fail(s"expected an applied reduction but got $other") + } + } + + test("AE7: the candidate is always a whole wave, never above the normal value") { + val configs = Seq(baseConfig, baseConfig.copy(targetPartitionSizeBytes = 512L * 1024L * 1024L), + baseConfig.copy(inputSizeFactor = 0.8)) + val byteSizes = Seq(0L, 1L, GiB, 37L * GiB, 999L * GiB, 100000L * GiB, Long.MaxValue) + val slotCounts = Seq(1, 3, 16, 375, 2000, 100000) + val normalValues = Seq(1, 199, 200, 500, 501, 2000, 200000, Int.MaxValue) + for (config <- configs; bytes <- byteSizes; slotCount <- slotCounts; + normalValue <- normalValues) { + decide(normalValue, Seq(record(totalBytes = bytes)), config, slotCount = Some(slotCount)) + match { + case applied: DownwardShuffleDecision.Applied => + val context = s"config=$config bytes=$bytes slots=$slotCount normal=$normalValue" + assert(applied.selectedValue < normalValue, context) + assert(applied.selectedValue >= slotCount, context) + assert(applied.selectedValue % slotCount == 0, context) + assert(applied.selectedValue == applied.waveCount * slotCount, context) + case _ => // a skip is always safe + } + } + } + + // + // Evidence completeness + // + + test("incomplete evidence keeps the normal recommendation and warns") { + val incomplete = ShuffleStageInputAnalysis( + Seq(record(totalBytes = 900L * GiB)), + Seq(ShuffleStageInputIncompleteReason.MissingExchangeMetric(0L, 3L, "Exchange")), + ShuffleInputProvenance.Measured) + DownwardShufflePartitionsPolicy.decide(Right(baseConfig), 4000, Some(slots), + incomplete) match { + case DownwardShuffleDecision.Skipped(reason: DownwardShuffleSkipReason.IncompleteEvidence) => + assert(reason.isWarning) + assert(reason.description.contains("data size")) + case other => fail(s"expected an incomplete-evidence skip but got $other") + } + } + + test("an analysis that was never produced fails closed without a user comment") { + val notAnalyzed = ShuffleStageInputAnalysis.empty(ShuffleInputProvenance.Measured) + assert(!notAnalyzed.isComplete) + val decision = + DownwardShufflePartitionsPolicy.decide(Right(baseConfig), 4000, Some(slots), notAnalyzed) + assert(decision == + DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.NoAnalysisAvailable)) + // A provider that never produced an analysis is not something the user can act on. + assert(!DownwardShuffleSkipReason.NoAnalysisAvailable.isWarning) + } + + test("an application with no shuffle at all is a quiet no-op") { + val noShuffle = ShuffleStageInputAnalysis(Seq.empty, Seq.empty, ShuffleInputProvenance.Measured) + assert(noShuffle.isComplete) + val decision = + DownwardShufflePartitionsPolicy.decide(Right(baseConfig), 4000, Some(slots), noShuffle) + assert(decision == DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.NoStageEvidence)) + assert(!DownwardShuffleSkipReason.NoStageEvidence.isWarning) + } + + test("the incomplete summary stays concise") { + val reasons = (1 to 5).map { i => + ShuffleStageInputIncompleteReason.UnresolvedConsumerStage(0L, i.toLong, "Exchange") + } + val summary = ShuffleStageInputAnalysis(Seq.empty, reasons, ShuffleInputProvenance.Measured) + .incompleteSummary() + assert(summary.contains("and 3 more")) + assert(summary.contains("node 1")) + assert(summary.contains("node 2")) + assert(!summary.contains("node 3")) + } + + // + // Diagnostics + // + + test("the applied comment names the slot count and wave arithmetic, not rung rounding") { + val stage = record(sqlId = 2L, stageId = 7, stageAttemptId = 1, + totalBytes = 4500L * GiB, numBranches = 2) + val applied = decide(8000, Seq(stage)).asInstanceOf[DownwardShuffleDecision.Applied] + assert(applied.selectedValue == 6000 && applied.waveCount == 3) + val comment = DownwardShufflePartitionsPolicy.appliedComment( + Seq("spark.sql.shuffle.partitions" -> 6000, + "spark.sql.adaptive.coalescePartitions.initialPartitionNum" -> 6000), applied) + Seq("spark.sql.shuffle.partitions", + "spark.sql.adaptive.coalescePartitions.initialPartitionNum", + "8000", "6000", "measured", "SQL 2", "stage 7", "attempt 1", + (4500L * GiB).toString, "2 shuffle branch(es)", "1.0", "4500", + "3 execution wave(s)", "2000 cluster task slots").foreach { fragment => + assert(comment.contains(fragment), s"'$fragment' missing from: $comment") + } + assert(!comment.contains("rung"), s"the comment should no longer mention rungs: $comment") + } + + test("the applied comment names each property when the clamp gives them different values") { + val applied = decide(8000, Seq(record(totalBytes = 4500L * GiB))) + .asInstanceOf[DownwardShuffleDecision.Applied] + // A property already below the candidate keeps its own lower value, so one sentence about + // both would misreport it. + val comment = DownwardShufflePartitionsPolicy.appliedComment( + Seq("spark.sql.shuffle.partitions" -> 6000, + "spark.sql.adaptive.coalescePartitions.initialPartitionNum" -> 1000), applied) + assert(comment.contains("'spark.sql.shuffle.partitions' lowered to 6000")) + assert(comment.contains( + "'spark.sql.adaptive.coalescePartitions.initialPartitionNum' lowered to 1000")) + assert(comment.contains("from an effective 8000")) + // The wave arithmetic still describes the candidate that was computed. + assert(comment.contains("3 execution wave(s) of 2000 cluster task slots")) + } +} diff --git a/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/FileScanInputMetricsSuite.scala b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/FileScanInputMetricsSuite.scala index c4edef526..5cbf585ba 100644 --- a/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/FileScanInputMetricsSuite.scala +++ b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/FileScanInputMetricsSuite.scala @@ -22,8 +22,8 @@ import com.nvidia.spark.rapids.tool.{AppSummaryInfoBaseProvider, EventLogPathPro PlatformFactory, PlatformNames, ToolTestUtils} import com.nvidia.spark.rapids.tool.analysis.AggRawMetricsResult import com.nvidia.spark.rapids.tool.profiling.{ApplicationSummaryInfo, CollectInformation, - DataSourceProfileResult, ProfileArgs, PySparkMemoryEvidence, SingleAppSummaryInfoProvider, - StageAggTaskMetricsProfileResult} + DataSourceProfileResult, ProfileArgs, PySparkMemoryEvidence, ShuffleInputProvenance, + ShuffleStageInputAnalysis, SingleAppSummaryInfoProvider, StageAggTaskMetricsProfileResult} import com.nvidia.spark.rapids.tool.views.RawMetricProfilerView import org.apache.spark.sql.rapids.tool.AccumToStageRetriever @@ -188,6 +188,9 @@ class FileScanInputMetricsSuite extends ProfilingAutoTunerSuiteBase { override def getPySparkMemoryEvidence: Seq[PySparkMemoryEvidence] = Seq.empty override def scanStagesWithGpuOom: Set[Long] = Set.empty override def getClassPathEntries: Map[String, String] = Map.empty + // This stub carries no ApplicationInfo, so the plan analysis it would come from is absent. + override def getShuffleStageInputAnalysis: ShuffleStageInputAnalysis = + ShuffleStageInputAnalysis.empty(ShuffleInputProvenance.Measured) } private def maxPartitionRecommendation( diff --git a/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuiteV2.scala b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuiteV2.scala index 06ac8d6a8..c20b313f9 100644 --- a/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuiteV2.scala +++ b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuiteV2.scala @@ -20,7 +20,7 @@ import scala.collection.mutable import com.nvidia.spark.rapids.tool.{DynamicAllocationInfo, GpuTypes, NodeInstanceMapKey, PlatformFactory, PlatformInstanceTypes, PlatformNames, ToolTestUtils} import com.nvidia.spark.rapids.tool.profiling.{Profiler, PySparkMemoryEvidence, - RecommendedCommentResult} + RecommendedCommentResult, ShuffleStageInputAnalysis} import com.nvidia.spark.rapids.tool.tuning.config.{ConfTypeEnum, TuningConfigEntry, TuningConfiguration, TuningEntryDefinition} import org.scalatest.matchers.should.Matchers._ @@ -2661,6 +2661,359 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { compareOutput(expectedResults, autoTunerOutput) } + // + // Downward shuffle partition pass + // + + private val GiB = 1024L * 1024L * 1024L + private val SHUFFLE_PARTITIONS_KEY = "spark.sql.shuffle.partitions" + private val AQE_INITIAL_PARTITION_NUM_KEY = + "spark.sql.adaptive.coalescePartitions.initialPartitionNum" + private val ADVISORY_PARTITION_SIZE_KEY = "spark.sql.adaptive.advisoryPartitionSizeInBytes" + + /** + * Recommended cluster of the downward-pass fixture: 25 executors of 16 cores, so one execution + * wave is 400 task slots on the default cores basis. + */ + private val DOWNWARD_PASS_WORKERS = 25 + private val DOWNWARD_PASS_SLOTS = 400 + + /** + * Source properties of a GPU application whose normal recommendation lands on 8000 shuffle + * partitions, which is comfortably above the downward candidate the tests drive. + */ + private def downwardPassSourceProps( + extra: Map[String, String] = Map.empty): mutable.Map[String, String] = { + val base = mutable.LinkedHashMap[String, String]( + "spark.executor.cores" -> "16", + "spark.executor.instances" -> DOWNWARD_PASS_WORKERS.toString, + "spark.executor.memory" -> "80g", + "spark.executor.resource.gpu.amount" -> "1", + "spark.sql.adaptive.enabled" -> "true", + SHUFFLE_PARTITIONS_KEY -> "8000", + "spark.task.resource.gpu.amount" -> "0.25", + "spark.rapids.sql.enabled" -> "true", + "spark.plugins" -> "com.nvidia.spark.SQLPlugin", + "spark.rapids.sql.concurrentGpuTasks" -> "1") + extra.foreach { case (k, v) => base.put(k, v) } + base + } + + private def recommendedValue( + properties: Seq[TuningEntryTrait], key: String): Option[String] = { + properties.find(_.name == key).map(_.getTuneValue()) + } + + /** + * The downward pass ships disabled, so every test in this section opts in explicitly. Extra + * entries the caller supplies are merged on top of that opt-in. + */ + private def downwardPassConfigs( + enabled: Boolean = true, + extra: List[TuningConfigEntry] = List.empty): TuningConfiguration = { + ToolTestUtils.buildTuningConfigs(default = + TuningConfigEntry(name = "DOWNWARD_SHUFFLE_ENABLED", default = enabled.toString) :: extra) + } + + private def downwardPassEnabledConfigs: TuningConfiguration = downwardPassConfigs() + + /** + * Runs the AutoTuner over a GPU application with the given shuffle-stage evidence and returns + * its recommendations and comments. + */ + private def runDownwardPass( + shuffleStageInputAnalysis: ShuffleStageInputAnalysis, + sourceProps: mutable.Map[String, String] = downwardPassSourceProps(), + shuffleStagesWithPosSpilling: Set[Long] = Set(), + gpuShuffleStagesWithContainerOom: Set[Long] = Set(), + scanStagesWithGpuOom: Set[Long] = Set(), + meanInputOverride: Option[Double] = None, + maxColumnarExchangeDataSizeBytes: Option[Long] = None, + extraTuningConfigs: List[TuningConfigEntry] = List.empty, + platformName: String = PlatformNames.DATAPROC, + configureClusterInfo: Boolean = true, + enableDownwardPass: Boolean = true + ): (Seq[TuningEntryTrait], Seq[String]) = { + val userProvidedTuningConfigs = + Some(downwardPassConfigs(enableDownwardPass, extraTuningConfigs)) + val infoProvider = getMockInfoProvider( + maxInput = 1.0E9, + spilledMetrics = Seq(0, 0), + jvmGCFractions = Seq(0.1, 0.1), + propsFromLog = sourceProps, + sparkVersion = Some(testSparkVersion), + meanInput = meanInputOverride.getOrElse(1.0E9), + meanShuffleRead = 1.0E9, + shuffleStagesWithPosSpilling = shuffleStagesWithPosSpilling, + gpuShuffleStagesWithContainerOom = gpuShuffleStagesWithContainerOom, + scanStagesWithGpuOom = scanStagesWithGpuOom, + maxColumnarExchangeDataSizeBytes = maxColumnarExchangeDataSizeBytes, + shuffleStageInputAnalysis = shuffleStageInputAnalysis) + val platform = PlatformFactory.createInstance(platformName) + if (configureClusterInfo) { + configureEventLogClusterInfoForTest(platform, numCores = 16, + numWorkers = DOWNWARD_PASS_WORKERS, gpuCount = 1, sparkProperties = sourceProps.toMap) + } + val autoTuner = + buildAutoTunerForTests(infoProvider, platform, None, userProvidedTuningConfigs) + val (properties, comments) = autoTuner.getRecommendedProperties() + (properties, comments.map(_.comment)) + } + + /** A single consumer stage that needs 900 partitions at the default 1 GiB target. */ + private def worstStageNeeding900Partitions: ShuffleStageInputAnalysis = { + completeShuffleStageInputs(Seq(4 -> 900L * GiB)) + } + + test("Downward pass lowers both partition properties atomically when AQE coalescing is on") { + val (properties, comments) = runDownwardPass(worstStageNeeding900Partitions) + // 900 partitions rounds up to 3 whole waves of 400 slots. + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).contains("1200")) + assert(recommendedValue(properties, AQE_INITIAL_PARTITION_NUM_KEY).contains("1200")) + val appliedComment = comments.filter(_.contains("lowered from 8000 to 1200")) + assert(appliedComment.size == 1, s"expected exactly one applied comment in: $comments") + // The comment must identify every input of the decision so it can be audited. + Seq(SHUFFLE_PARTITIONS_KEY, AQE_INITIAL_PARTITION_NUM_KEY, "measured", "stage 4", + "900", "1.0", "3 execution wave(s)", s"$DOWNWARD_PASS_SLOTS cluster task slots") + .foreach { fragment => + assert(appliedComment.head.contains(fragment), + s"'$fragment' missing from: ${appliedComment.head}") + } + } + + test("Downward pass keeps the normal recommendation and stays quiet " + + "with no recommended cluster") { + val (properties, comments) = runDownwardPass(worstStageNeeding900Partitions, + configureClusterInfo = false) + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).forall(_ == "8000")) + // An unknown cluster shape is the common qualification case, so it must add no comment. + assert(!comments.exists(c => c.contains("lowered from") || c.contains("slot"))) + } + + test("Downward pass leaves the AQE advisory partition size untouched") { + val advisorySize = "64m" + val (properties, _) = runDownwardPass(worstStageNeeding900Partitions, + sourceProps = downwardPassSourceProps(Map(ADVISORY_PARTITION_SIZE_KEY -> advisorySize))) + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).contains("1200")) + assert(recommendedValue(properties, ADVISORY_PARTITION_SIZE_KEY).forall(_ == advisorySize)) + } + + test("Downward pass updates only the shuffle property when AQE coalescing is disabled") { + val (properties, comments) = runDownwardPass(worstStageNeeding900Partitions, + sourceProps = downwardPassSourceProps( + Map("spark.sql.adaptive.coalescePartitions.enabled" -> "false"))) + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).contains("1200")) + assert(recommendedValue(properties, AQE_INITIAL_PARTITION_NUM_KEY).isEmpty) + assert(comments.exists(_.contains("lowered from 8000 to 1200"))) + } + + test("Downward pass preserves an increase made for shuffle spilling") { + val (properties, comments) = runDownwardPass(worstStageNeeding900Partitions, + shuffleStagesWithPosSpilling = Set(1L)) + // The spill-driven doubling must survive; nothing may be lowered. + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).contains("16000")) + assert(comments.exists(_.contains("Shuffle partitions should be increased since spilling"))) + assert(!comments.exists(_.contains("lowered from"))) + } + + test("Downward pass preserves an increase made for shuffle task OOM") { + val (properties, comments) = runDownwardPass(worstStageNeeding900Partitions, + gpuShuffleStagesWithContainerOom = Set(1L)) + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).exists(_.toInt >= 8000)) + assert(!comments.exists(_.contains("lowered from"))) + } + + test("Downward pass preserves the existing GPU ColumnarExchange lower bound") { + // A 60000 GiB exchange against the ~2g batch size forces partitions above 8000. + val (properties, comments) = runDownwardPass(worstStageNeeding900Partitions, + maxColumnarExchangeDataSizeBytes = Some(60000L * GiB)) + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).exists(_.toInt > 8000)) + assert(!comments.exists(_.contains("lowered from"))) + } + + test("Downward pass is blocked by spill in an affected consumer stage") { + val withSpill = completeShuffleStageInputs(Seq(4 -> 900L * GiB), hasPositiveSpill = true) + val (properties, comments) = runDownwardPass(withSpill) + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).forall(_ == "8000")) + assert(!comments.exists(_.contains("lowered from"))) + } + + test("Downward pass is blocked by skew in an affected consumer stage") { + val withSkew = completeShuffleStageInputs(Seq(4 -> 900L * GiB), hasSkew = true) + val (properties, comments) = runDownwardPass(withSkew) + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).forall(_ == "8000")) + assert(!comments.exists(_.contains("lowered from"))) + } + + test("Downward pass warns once and changes nothing when evidence is incomplete") { + val (properties, comments) = runDownwardPass(incompleteShuffleStageInputs()) + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).forall(_ == "8000")) + val warnings = comments.filter(_.contains("could not be measured")) + assert(warnings.size == 1, s"expected exactly one incomplete-evidence comment in: $comments") + } + + test("Downward pass warns once and changes nothing when its configuration is invalid") { + val (properties, comments) = runDownwardPass(worstStageNeeding900Partitions, + extraTuningConfigs = List( + TuningConfigEntry(name = "DOWNWARD_SHUFFLE_TARGET_PARTITION_SIZE", default = "0"))) + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).forall(_ == "8000")) + val warnings = comments.filter(_.contains("tuning configuration is invalid")) + assert(warnings.size == 1, s"expected exactly one invalid-config comment in: $comments") + } + + test("Downward pass does nothing when disabled") { + val (properties, comments) = runDownwardPass(worstStageNeeding900Partitions, + enableDownwardPass = false) + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).forall(_ == "8000")) + // A disabled feature must stay completely silent in the user-facing output. + assert(!comments.exists(c => c.contains("lowered from") || c.contains("downward"))) + } + + test("Downward pass applies a reduction smaller than the retired 2x threshold") { + // 1500 partitions rounds up to 4 waves (1600), a 1.9x reduction from 3000. The wave quantum + // is the only size gate, so this applies where the retired threshold would have blocked it. + val (properties, comments) = runDownwardPass( + completeShuffleStageInputs(Seq(4 -> 1500L * GiB)), + sourceProps = downwardPassSourceProps(Map(SHUFFLE_PARTITIONS_KEY -> "3000"))) + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).contains("1600")) + assert(comments.exists(_.contains("lowered from 3000 to 1600"))) + } + + test("Downward pass never lowers below a single execution wave") { + val (properties, comments) = runDownwardPass(completeShuffleStageInputs(Seq(4 -> 1L))) + // Even a one-byte stage still gets one full wave of the recommended cluster. + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY) + .contains(DOWNWARD_PASS_SLOTS.toString)) + assert(comments.exists(_.contains(s"lowered from 8000 to $DOWNWARD_PASS_SLOTS"))) + assert(comments.exists(_.contains("1 execution wave(s)"))) + } + + test("Downward pass sizes the wave from cores, not GPU task concurrency") { + val (properties, comments) = runDownwardPass(worstStageNeeding900Partitions) + val concurrentGpuTasks = + recommendedValue(properties, "spark.rapids.sql.concurrentGpuTasks").map(_.toInt).getOrElse( + fail("expected a concurrentGpuTasks recommendation")) + val gpuSlots = DOWNWARD_PASS_WORKERS * concurrentGpuTasks + assert(gpuSlots != DOWNWARD_PASS_SLOTS, + "the two bases must differ for this test to prove anything") + assert(comments.exists(_.contains(s"$DOWNWARD_PASS_SLOTS cluster task slots"))) + assert(!comments.exists(_.contains(s"$gpuSlots cluster task slots"))) + } + + test("Downward pass changes neither property when one of them is enforced") { + val targetClusterInfo = ToolTestUtils.buildTargetClusterInfo( + enforcedSparkProperties = Map(AQE_INITIAL_PARTITION_NUM_KEY -> "4096")) + val sourceProps = downwardPassSourceProps() + val infoProvider = getMockInfoProvider( + maxInput = 1.0E9, + spilledMetrics = Seq(0, 0), + jvmGCFractions = Seq(0.1, 0.1), + propsFromLog = sourceProps, + sparkVersion = Some(testSparkVersion), + meanInput = 1.0E9, + meanShuffleRead = 1.0E9, + shuffleStageInputAnalysis = worstStageNeeding900Partitions) + val platform = PlatformFactory.createInstance(PlatformNames.DATAPROC, Some(targetClusterInfo)) + configureEventLogClusterInfoForTest(platform, numCores = 16, + numWorkers = DOWNWARD_PASS_WORKERS, gpuCount = 1, sparkProperties = sourceProps.toMap) + val autoTuner = buildAutoTunerForTests(infoProvider, platform, None, + Some(downwardPassEnabledConfigs)) + val (properties, comments) = autoTuner.getRecommendedProperties() + // The enforced AQE property blocks the whole update, so the shuffle property is untouched. + assert(recommendedValue(properties, AQE_INITIAL_PARTITION_NUM_KEY).contains("4096")) + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).forall(_ == "8000")) + assert(!comments.map(_.comment).exists(_.contains("lowered from"))) + } + + test("Downward pass is blocked while Databricks automatic shuffle optimization is active") { + val autoOptimizeKey = "spark.databricks.adaptive.autoOptimizeShuffle.enabled" + // The user enforces the Databricks setting, so normal tuning cannot turn it off and the + // runtime, not this recommendation, still governs partitioning. + val targetClusterInfo = ToolTestUtils.buildTargetClusterInfo( + enforcedSparkProperties = Map(autoOptimizeKey -> "true")) + val sourceProps = downwardPassSourceProps(Map(autoOptimizeKey -> "true")) + val infoProvider = getMockInfoProvider( + maxInput = 1.0E9, + spilledMetrics = Seq(0, 0), + jvmGCFractions = Seq(0.1, 0.1), + propsFromLog = sourceProps, + sparkVersion = Some(testDatabricksVersion), + meanInput = 1.0E9, + meanShuffleRead = 1.0E9, + shuffleStageInputAnalysis = worstStageNeeding900Partitions) + val platform = + PlatformFactory.createInstance(PlatformNames.DATABRICKS_AWS, Some(targetClusterInfo)) + configureEventLogClusterInfoForTest(platform, numCores = 16, + numWorkers = DOWNWARD_PASS_WORKERS, gpuCount = 1, sparkProperties = sourceProps.toMap) + val autoTuner = buildAutoTunerForTests(infoProvider, platform, None, + Some(downwardPassEnabledConfigs)) + val (properties, comments) = autoTuner.getRecommendedProperties() + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).forall(_ == "8000")) + assert(!comments.map(_.comment).exists(_.contains("lowered from"))) + } + + test("Downward pass falls back to the cluster record when the platform excludes instances") { + // Databricks excludes 'spark.executor.instances', so no recommendation exists for it and the + // slot count must come from the recommended cluster rather than the source CPU value. + val sourceProps = downwardPassSourceProps() + val infoProvider = getMockInfoProvider( + maxInput = 1.0E9, + spilledMetrics = Seq(0, 0), + jvmGCFractions = Seq(0.1, 0.1), + propsFromLog = sourceProps, + sparkVersion = Some(testDatabricksVersion), + meanInput = 1.0E9, + meanShuffleRead = 1.0E9, + shuffleStageInputAnalysis = worstStageNeeding900Partitions) + val platform = PlatformFactory.createInstance(PlatformNames.DATABRICKS_AWS) + configureEventLogClusterInfoForTest(platform, numCores = 16, + numWorkers = DOWNWARD_PASS_WORKERS, gpuCount = 1, sparkProperties = sourceProps.toMap) + val autoTuner = buildAutoTunerForTests(infoProvider, platform, None, + Some(downwardPassEnabledConfigs)) + val (properties, comments) = autoTuner.getRecommendedProperties() + assert(recommendedValue(properties, "spark.executor.instances").isEmpty, + "Databricks should not recommend an executor count") + val clusterInfo = platform.recommendedClusterInfo.getOrElse( + fail("expected a recommended cluster")) + val expectedSlots = clusterInfo.numExecutors * clusterInfo.coresPerExecutor + // Reading the source CPU executor cores instead would have produced a different wave. + assert(expectedSlots != DOWNWARD_PASS_SLOTS, + "the recommended cluster must differ from the source one for this test to prove anything") + val expected = math.ceil(900.0 / expectedSlots).toInt * expectedSlots + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).contains(expected.toString)) + assert(comments.map(_.comment).exists(_.contains(s"$expectedSlots cluster task slots"))) + } + + test("Downward pass is blocked when the application had a failed stage") { + val withFailure = completeShuffleStageInputs(Seq(4 -> 900L * GiB), appHasFailedStage = true) + val (properties, comments) = runDownwardPass(withFailure) + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).forall(_ == "8000")) + assert(!comments.exists(_.contains("lowered from"))) + } + + test("Downward pass is blocked when the application had an OOM") { + // A GPU OOM anywhere in the application disqualifies the run as sizing evidence, even + // though the consumer stages themselves look clean. + val (properties, comments) = runDownwardPass(worstStageNeeding900Partitions, + scanStagesWithGpuOom = Set(9L)) + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).forall(_ == "8000")) + assert(!comments.exists(_.contains("lowered from"))) + } + + test("Downward pass leaves the two partition properties in agreement") { + // The pass updates the AQE partition property only when a value for it already exists. In a + // normal AQE run recommendAQEProperties has already appended one, so the guard is inert and + // both properties move together; what users depend on is that they never disagree. + val sourceProps = downwardPassSourceProps() + sourceProps.remove(AQE_INITIAL_PARTITION_NUM_KEY) + val (properties, _) = runDownwardPass(worstStageNeeding900Partitions, + sourceProps = sourceProps) + val shuffle = recommendedValue(properties, SHUFFLE_PARTITIONS_KEY) + val aqe = recommendedValue(properties, AQE_INITIAL_PARTITION_NUM_KEY) + assert(shuffle.contains("1200")) + assert(aqe.forall(_ == shuffle.get), + s"partition properties disagree: shuffle=$shuffle aqe=$aqe") + } test("PySpark memory evidence without a positive source limit does not enable telemetry " + "by default") { val sourceProps = mutable.LinkedHashMap[String, String]( diff --git a/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/QualificationAutoTunerSuite.scala b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/QualificationAutoTunerSuite.scala index dc1da63d7..67c606e3e 100644 --- a/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/QualificationAutoTunerSuite.scala +++ b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/QualificationAutoTunerSuite.scala @@ -20,9 +20,10 @@ import java.nio.file.Paths import scala.collection.mutable -import com.nvidia.spark.rapids.tool.{DynamicAllocationInfo, GpuTypes, PlatformFactory, PlatformNames, ToolTestUtils} -import com.nvidia.spark.rapids.tool.profiling.{Profiler, PySparkMemoryEvidence} -import com.nvidia.spark.rapids.tool.qualification.{QualificationArgs, QualificationMain} +import com.nvidia.spark.rapids.tool.{DynamicAllocationInfo, EventLogPathProcessor, GpuTypes, PlatformFactory, PlatformNames, ToolTestUtils} +import com.nvidia.spark.rapids.tool.analysis.{AppSQLPlanAnalyzer, QualSparkMetricsAggregator} +import com.nvidia.spark.rapids.tool.profiling.{Profiler, PySparkMemoryEvidence, ShuffleInputProvenance, ShuffleStageInputAnalysis} +import com.nvidia.spark.rapids.tool.qualification.{PluginTypeChecker, QualificationArgs, QualificationMain} import com.nvidia.spark.rapids.tool.tuning.config.{CategoryEnum, ConfTypeEnum, LevelEnum, TuningConfigEntry, TuningEntryDefinition} import com.nvidia.spark.rapids.tool.views.CLUSTER_INFORMATION_LABEL import com.nvidia.spark.rapids.tool.views.qualification.QualReportGenConfProvider @@ -33,7 +34,8 @@ import org.scalatest.prop.TableFor3 import org.apache.spark.sql.TrampolineUtil import org.apache.spark.sql.rapids.tool.{MatchingInstanceTypeNotFoundException, RecommendedClusterInfo} -import org.apache.spark.sql.rapids.tool.util.FSUtils +import org.apache.spark.sql.rapids.tool.qualification.QualificationAppInfo +import org.apache.spark.sql.rapids.tool.util.{FSUtils, RapidsToolsConfUtil} /** * Suite to test the Qualification Tool's AutoTuner @@ -2374,6 +2376,123 @@ class QualificationAutoTunerSuite extends BaseAutoTunerSuite { } } + // + // Downward shuffle partition pass on CPU event logs + // + + private val QUAL_GiB = 1024L * 1024L * 1024L + + /** + * Builds a Qualification AutoTuner over a CPU application whose normal recommendation is 8000 + * shuffle partitions and whose worst consumer stage carries the given uncompressed input. + * + * The downward pass ships disabled, so every test here opts in explicitly. Extra entries the + * caller supplies are merged on top of that opt-in. + */ + private def buildDownwardPassAutoTuner( + shuffleStageInputAnalysis: ShuffleStageInputAnalysis, + extraDefaultConfigs: List[TuningConfigEntry] = List.empty, + extraQualificationConfigs: List[TuningConfigEntry] = List.empty): AutoTuner = { + val userProvidedTuningConfigs = Some(ToolTestUtils.buildTuningConfigs( + default = TuningConfigEntry(name = "DOWNWARD_SHUFFLE_ENABLED", default = "true") :: + extraDefaultConfigs, + qualification = extraQualificationConfigs)) + val sparkProps = defaultSparkProps ++ mutable.Map( + "spark.executor.memory" -> "212992MiB", + "spark.sql.adaptive.enabled" -> "true", + "spark.sql.shuffle.partitions" -> "8000") + val infoProvider = getMockInfoProvider(0, Seq(0), Seq(0.0), sparkProps, + Some(testSparkVersion), shuffleStageInputAnalysis = shuffleStageInputAnalysis) + val platform = PlatformFactory.createInstance(PlatformNames.EMR) + platform.configureClusterInfoFromEventLog( + coresPerExecutor = 32, execsPerNode = 4, numExecs = 20, numExecutorNodes = 5, + sparkProperties = sparkProps.toMap, systemProperties = Map.empty) + buildAutoTunerForTests(infoProvider, platform, None, userProvidedTuningConfigs) + } + + private def cpuStageInput(bytes: Long): ShuffleStageInputAnalysis = { + completeShuffleStageInputs(Seq(4 -> bytes), provenance = ShuffleInputProvenance.Estimated) + } + + /** Task slots of the cluster this fixture recommends: one executor of 32 cores. */ + private val QUAL_SLOTS = 32 + + test("test AutoTuner for Qualification lowers shuffle partitions using the CPU input factor") { + // 1000 GiB of CPU exchange data at the qualification factor of 0.8 estimates 800 GiB of GPU + // input, which needs 800 partitions: exactly 25 whole waves of 32 slots. + val autoTuner = buildDownwardPassAutoTuner(cpuStageInput(1000L * QUAL_GiB)) + val (properties, comments) = autoTuner.getRecommendedProperties(showOnlyUpdatedProps = + QualificationAutoTunerRunner.filterByUpdatedPropsEnabled) + val autoTunerOutput = Profiler.getAutoTunerResultsAsString(properties, comments) + assertExpectedLinesExist( + Seq("--conf spark.sql.shuffle.partitions=800", + "--conf spark.sql.adaptive.coalescePartitions.initialPartitionNum=800"), + autoTunerOutput) + // The comment must say the input was estimated, not measured, for a CPU event log, and it must + // name the wave arithmetic so the recommendation can be audited without re-running the tool. + val applied = comments.map(_.comment).filter(_.contains("lowered from 8000 to 800")) + assert(applied.size == 1, + s"expected exactly one applied comment in: ${comments.map(_.comment)}") + assert(applied.head.contains("estimated")) + assert(applied.head.contains("input size factor 0.8")) + assert(applied.head.contains(s"25 execution wave(s) of $QUAL_SLOTS cluster task slots")) + } + + test("test AutoTuner for Qualification honours a custom input size factor") { + // A factor of 0.4 estimates 400 GiB, which needs 400 partitions and rounds up to 13 waves. + val autoTuner = buildDownwardPassAutoTuner(cpuStageInput(1000L * QUAL_GiB), + extraQualificationConfigs = List( + TuningConfigEntry(name = "DOWNWARD_SHUFFLE_INPUT_SIZE_FACTOR", default = "0.4"))) + val (properties, comments) = autoTuner.getRecommendedProperties(showOnlyUpdatedProps = + QualificationAutoTunerRunner.filterByUpdatedPropsEnabled) + val autoTunerOutput = Profiler.getAutoTunerResultsAsString(properties, comments) + assertExpectedLinesExist(Seq("--conf spark.sql.shuffle.partitions=416"), autoTunerOutput) + val applied = comments.map(_.comment).filter(_.contains("lowered from 8000 to 416")) + assert(applied.size == 1, + s"expected exactly one applied comment in: ${comments.map(_.comment)}") + assert(applied.head.contains("input size factor 0.4")) + assert(applied.head.contains("raw requirement 400 partitions")) + assert(applied.head.contains(s"13 execution wave(s) of $QUAL_SLOTS cluster task slots")) + } + + test("test AutoTuner for Qualification keeps the recommendation when evidence is incomplete") { + val autoTuner = buildDownwardPassAutoTuner( + incompleteShuffleStageInputs(ShuffleInputProvenance.Estimated)) + val (properties, comments) = autoTuner.getRecommendedProperties(showOnlyUpdatedProps = + QualificationAutoTunerRunner.filterByUpdatedPropsEnabled) + val autoTunerOutput = Profiler.getAutoTunerResultsAsString(properties, comments) + assertExpectedLinesExist(Seq("--conf spark.sql.shuffle.partitions=8000"), autoTunerOutput) + assert(comments.map(_.comment).count(_.contains("could not be measured")) == 1) + } + + test("test AutoTuner for Qualification provider reuses the existing SQL plan analyzer") { + val hadoopConf = RapidsToolsConfUtil.newHadoopConf() + val (_, allEventLogs) = EventLogPathProcessor.processAllPaths( + None, None, List(s"$qualLogDir/nds_q86_test"), hadoopConf) + val app = QualificationAppInfo.createApp(allEventLogs.head, hadoopConf, + new PluginTypeChecker(), reportSqlLevel = false, mlOpsEnabled = false, + penalizeTransitions = true, PlatformFactory.createInstance()) match { + case Right(a) => a + case Left(_) => fail("could not build the qualification application") + } + val sqlAnalyzer = AppSQLPlanAnalyzer(app) + val rawAggMetrics = QualSparkMetricsAggregator.getAggRawMetrics(app, 1, Some(sqlAnalyzer)) + + val withAnalyzer = + new QualAppSummaryInfoProvider(app, None, rawAggMetrics, Seq.empty, Some(sqlAnalyzer)) + val analysis = withAnalyzer.getShuffleStageInputAnalysis + // The provider must hand back the analyzer's own cached analysis, not a fresh traversal. + assert(analysis eq sqlAnalyzer.shuffleStageInputAnalysis) + assert(analysis.isComplete) + assert(analysis.provenance == ShuffleInputProvenance.Estimated) + assert(analysis.records.nonEmpty) + + // Without an analyzer there is no evidence at all, which must fail closed. + val withoutAnalyzer = + new QualAppSummaryInfoProvider(app, None, rawAggMetrics, Seq.empty, None) + assert(!withoutAnalyzer.getShuffleStageInputAnalysis.analyzed) + } + test("Qualification PySpark evidence atomically rebalances executor heap") { val sourceProps = mutable.LinkedHashMap[String, String]( "spark.executor.cores" -> "8",