From 72438be2ec8227a89709e61466f458558a91ebf6 Mon Sep 17 00:00:00 2001 From: Partho Sarthi Date: Wed, 5 Aug 2026 18:37:16 -0700 Subject: [PATCH 01/10] Add downward shuffle partition policy configuration and calculator Adds the six user-overridable tuning-config entries that drive the new downward-only shuffle partition pass, plus the pure policy layer that turns raw consumer-stage shuffle input records into either an applied reduction or a typed no-op reason. The calculator is deliberately isolated from AutoTuner state so its arithmetic can be proven on its own: configuration is validated as one fail-closed unit, byte and rung math is overflow-safe, rung generation is guaranteed to progress, and a requirement no representable rung can cover skips rather than recommending a partition count that is too small. Also adds the shared raw record types, which distinguish an application that executed no shuffle from a provider that produced no analysis at all. Signed-off-by: Partho Sarthi Co-Authored-By: Claude Opus 5 (1M context) --- .../resources/bootstrap/tuningConfigs.yaml | 55 +++ .../profiling/ShuffleStageInputMetrics.scala | 147 ++++++ .../DownwardShufflePartitionsPolicy.scala | 429 +++++++++++++++++ .../DownwardShufflePartitionsSuite.scala | 452 ++++++++++++++++++ 4 files changed, 1083 insertions(+) create mode 100644 core/src/main/scala/com/nvidia/spark/rapids/tool/profiling/ShuffleStageInputMetrics.scala create mode 100644 core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/DownwardShufflePartitionsPolicy.scala create mode 100644 core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/DownwardShufflePartitionsSuite.scala diff --git a/core/src/main/resources/bootstrap/tuningConfigs.yaml b/core/src/main/resources/bootstrap/tuningConfigs.yaml index 3d8f3a911..f156dbe86 100644 --- a/core/src/main/resources/bootstrap/tuningConfigs.yaml +++ b/core/src/main/resources/bootstrap/tuningConfigs.yaml @@ -180,6 +180,54 @@ 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. + default: true + 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: DOWNWARD_SHUFFLE_PARTITION_FLOOR + description: >- + Lowest partition count the downward pass can recommend. This is also the first generated + partition rung. + default: 500 + usedBy: spark.sql.shuffle.partitions, spark.sql.adaptive.coalescePartitions.initialPartitionNum + + - name: DOWNWARD_SHUFFLE_RUNG_MULTIPLIER + description: >- + Multiplier used to generate successive partition rungs from the floor + (e.g. a floor of 500 and a multiplier of 2 generates 500, 1000, 2000, ...). + The raw partition requirement is always rounded up to the next rung. + default: 2 + usedBy: spark.sql.shuffle.partitions, spark.sql.adaptive.coalescePartitions.initialPartitionNum + + - name: DOWNWARD_SHUFFLE_MIN_REDUCTION_FACTOR + description: >- + Minimum ratio between the normal recommendation and the downward candidate required before + a reduction is applied. A value of 2.0 means the normal value must be at least twice the + candidate. + default: 2.0 + usedBy: spark.sql.shuffle.partitions, spark.sql.adaptive.coalescePartitions.initialPartitionNum + - name: WORKER_GPU_COUNT description: >- Default number of GPUs per worker node @@ -406,6 +454,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/profiling/ShuffleStageInputMetrics.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/profiling/ShuffleStageInputMetrics.scala new file mode 100644 index 000000000..bbdc422bf --- /dev/null +++ b/core/src/main/scala/com/nvidia/spark/rapids/tool/profiling/ShuffleStageInputMetrics.scala @@ -0,0 +1,147 @@ +/* + * 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 + */ +case class ShuffleStageInputAnalysis( + records: Seq[ShuffleStageInputRecord], + incompleteReasons: Seq[ShuffleStageInputIncompleteReason], + provenance: ShuffleInputProvenance, + analyzed: Boolean = true) { + + 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/DownwardShufflePartitionsPolicy.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/DownwardShufflePartitionsPolicy.scala new file mode 100644 index 000000000..b73c0df3a --- /dev/null +++ b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/DownwardShufflePartitionsPolicy.scala @@ -0,0 +1,429 @@ +/* + * 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 + * @param partitionFloor lowest partition count the pass may recommend; first rung + * @param rungMultiplier ratio between successive generated partition rungs + * @param minReductionFactor required ratio of normal value to candidate before applying + */ +case class DownwardShufflePolicyConfig( + enabled: Boolean, + targetPartitionSizeBytes: Long, + inputSizeFactor: Double, + partitionFloor: Int, + rungMultiplier: Double, + minReductionFactor: 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" + val PARTITION_FLOOR_KEY = "DOWNWARD_SHUFFLE_PARTITION_FLOOR" + val RUNG_MULTIPLIER_KEY = "DOWNWARD_SHUFFLE_RUNG_MULTIPLIER" + val MIN_REDUCTION_FACTOR_KEY = "DOWNWARD_SHUFFLE_MIN_REDUCTION_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, + partitionFloor = 0, + rungMultiplier = 0.0, + minReductionFactor = 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 floor = parsePositiveInt(configProvider, PARTITION_FLOOR_KEY) + val multiplier = + parseDouble(configProvider, RUNG_MULTIPLIER_KEY, min = 1.0, minInclusive = false) + val minReduction = + parseDouble(configProvider, MIN_REDUCTION_FACTOR_KEY, min = 1.0, minInclusive = true) + + val errors = Seq(targetSize, factor, floor, multiplier, minReduction).collect { + case Left(err) => err + } + (targetSize, factor, floor, multiplier, minReduction) match { + case (Right(size), Right(f), Right(fl), Right(m), Right(r)) if errors.isEmpty => + Right(DownwardShufflePolicyConfig( + enabled = true, + targetPartitionSizeBytes = size, + inputSizeFactor = f, + partitionFloor = fl, + rungMultiplier = m, + minReductionFactor = r)) + 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 parsePositiveInt( + configProvider: TuningConfigProvider, key: String): Either[String, Int] = { + rawValue(configProvider, key).flatMap { raw => + Try(raw.trim.toInt).toOption + .filter(_ > 0) + .toRight(s"'$key' must be a positive integer 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") + + case class IncompleteEvidence(summary: String) + extends DownwardShuffleSkipReason( + s"shuffle input evidence is incomplete: $summary") { + override def isWarning: Boolean = true + } + + 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, so no rung can + * cover it. Failing closed here is safer than recommending a rung below the requirement. + */ + case class RequirementOutOfRange(requirement: Long) + extends DownwardShuffleSkipReason( + s"the worst consumer stage requires $requirement partitions, which exceeds the largest" + + s" representable partition count (${Int.MaxValue})") { + override def isWarning: Boolean = true + } + + case class NotDownward(candidate: Int, normalValue: Int) + extends DownwardShuffleSkipReason( + s"candidate $candidate does not lower the current recommendation $normalValue") + + case class BelowReductionThreshold(candidate: Int, normalValue: Int, minFactor: Double) + extends DownwardShuffleSkipReason( + s"candidate $candidate is not at least ${minFactor}x smaller than the current" + + s" recommendation $normalValue") + + 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 partition rung 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 a rung + * @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, + 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 a rung sequence that is guaranteed to make progress. + */ +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 analysis raw consumer-stage shuffle input analysis + */ + def decide( + configResult: Either[Seq[String], DownwardShufflePolicyConfig], + normalValue: 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, analysis) + } + } + + private def decideWithConfig( + config: DownwardShufflePolicyConfig, + normalValue: Int, + analysis: ShuffleStageInputAnalysis): DownwardShuffleDecision = { + if (!analysis.isComplete) { + return DownwardShuffleDecision.Skipped( + DownwardShuffleSkipReason.IncompleteEvidence(analysis.incompleteSummary())) + } + if (analysis.records.isEmpty) { + return DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.NoStageEvidence) + } + + val worst = selectWorstStage(config, analysis.records) + rungAtOrAbove(worst.requirement, config.partitionFloor, config.rungMultiplier) match { + case None => + // No representable rung 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 the rung floor, + // because every generated rung is at least the floor. + DownwardShuffleDecision.Skipped( + DownwardShuffleSkipReason.NotDownward(candidate, normalValue)) + case Some(candidate) + if normalValue.toDouble < config.minReductionFactor * candidate.toDouble => + DownwardShuffleDecision.Skipped( + DownwardShuffleSkipReason.BelowReductionThreshold( + candidate, normalValue, config.minReductionFactor)) + case Some(candidate) => + DownwardShuffleDecision.Applied( + normalValue = normalValue, + selectedValue = candidate, + determiningRecord = worst.record, + estimatedInputBytes = worst.estimatedBytes, + rawRequirement = worst.requirement, + 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 the first generated rung that is at least as large. + * + * Rungs start at the floor and grow by the multiplier, rounding each step up so that a fractional + * multiplier still produces a strictly increasing integer sequence. Each step is additionally + * forced to advance by at least one partition, so a multiplier barely above 1.0 cannot stall. + * + * @return the covering rung, or None when no rung within `Int.MaxValue` covers the requirement + */ + private[tuning] def rungAtOrAbove( + requirement: Long, floor: Int, multiplier: Double): Option[Int] = { + if (requirement > Int.MaxValue.toLong) { + return None + } + var rung = floor.toLong + while (rung < requirement) { + val grown = math.ceil(rung.toDouble * multiplier) + if (grown > Int.MaxValue.toDouble) { + return None + } + val next = math.max(grown.toLong, rung + 1L) + if (next > Int.MaxValue.toLong) { + return None + } + rung = next + } + Some(rung.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( + partitionProperties: Seq[String], + decision: DownwardShuffleDecision.Applied): String = { + val record = decision.determiningRecord + s"${partitionProperties.map(p => s"'$p'").mkString(" and ")} lowered from " + + s"${decision.normalValue} to ${decision.selectedValue} 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.selectedValue}." + } +} 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..bfe121a63 --- /dev/null +++ b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/DownwardShufflePartitionsSuite.scala @@ -0,0 +1,452 @@ +/* + * 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 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 + + private def profProvider( + default: List[TuningConfigEntry] = List.empty): ProfTuningConfigProvider = { + TuningConfigProvider.builder + .withUserProvidedConfig(Some(ToolTestUtils.buildTuningConfigs(default = default))) + .build[ProfTuningConfigProvider] + } + + private def qualProvider( + default: List[TuningConfigEntry] = List.empty, + qualification: List[TuningConfigEntry] = List.empty): QualTuningConfigProvider = { + TuningConfigProvider.builder + .withUserProvidedConfig(Some( + ToolTestUtils.buildTuningConfigs(default = default, qualification = qualification))) + .build[QualTuningConfigProvider] + } + + 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, rungs 500/1000/2000/..., 2x reduction. */ + private val baseConfig = DownwardShufflePolicyConfig( + enabled = true, + targetPartitionSizeBytes = GiB, + inputSizeFactor = 1.0, + partitionFloor = 500, + rungMultiplier = 2.0, + minReductionFactor = 2.0) + + /** 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 + ): DownwardShuffleDecision = { + DownwardShufflePartitionsPolicy.decide(Right(config), normalValue, + analysis(records, provenance)) + } + + // + // Configuration + // + + 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) + assert(config.partitionFloor == 500) + assert(config.rungMultiplier == 2.0) + assert(config.minReductionFactor == 2.0) + } + + 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) + assert(config.partitionFloor == 500) + } + + 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"), + TuningConfigEntry(name = "DOWNWARD_SHUFFLE_PARTITION_FLOOR", default = "128")))) + assert(fromDefault == Right(baseConfig.copy( + targetPartitionSizeBytes = 512L * 1024L * 1024L, partitionFloor = 128))) + + 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_PARTITION_FLOOR", default = "not-a-number"))) + val configResult = DownwardShufflePolicyConfig.fromProvider(provider) + assert(configResult == Right(DownwardShufflePolicyConfig.disabled)) + assert(DownwardShufflePartitionsPolicy.decide(configResult, 4000, + 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"), + "DOWNWARD_SHUFFLE_PARTITION_FLOOR" -> Seq("0", "-500", "1.5", "abc"), + // A multiplier of exactly 1.0 or below cannot generate growing rungs. + "DOWNWARD_SHUFFLE_RUNG_MULTIPLIER" -> Seq("1", "1.0", "0.5", "abc"), + // A reduction factor below 1.0 would allow raising the recommendation. + "DOWNWARD_SHUFFLE_MIN_REDUCTION_FACTOR" -> Seq("0.9", "0", "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", "DOWNWARD_SHUFFLE_RUNG_MULTIPLIER", + "DOWNWARD_SHUFFLE_MIN_REDUCTION_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_PARTITION_FLOOR", default = "-1"), + TuningConfigEntry(name = "DOWNWARD_SHUFFLE_RUNG_MULTIPLIER", default = "0.5"))) + val configResult = DownwardShufflePolicyConfig.fromProvider(provider) + val errors = expectErrors(configResult) + assert(errors.size == 2) + DownwardShufflePartitionsPolicy.decide(configResult, 4000, + 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("rungAtOrAbove generates each rung from the floor and the multiplier") { + val rungs = Seq(1L -> 500, 500L -> 500, 501L -> 1000, 1000L -> 1000, 1001L -> 2000, + 2000L -> 2000, 2001L -> 4000, 4001L -> 8000, 64001L -> 128000) + rungs.foreach { case (requirement, expected) => + assert(DownwardShufflePartitionsPolicy.rungAtOrAbove(requirement, 500, 2.0) == + Some(expected), s"requirement $requirement") + } + } + + test("fractional multipliers round each rung up so the sequence stays integral") { + // 100 -> 150 -> 225 -> 338 (ceil of 337.5) -> 507 + assert(DownwardShufflePartitionsPolicy.rungAtOrAbove(101L, 100, 1.5) == Some(150)) + assert(DownwardShufflePartitionsPolicy.rungAtOrAbove(226L, 100, 1.5) == Some(338)) + assert(DownwardShufflePartitionsPolicy.rungAtOrAbove(339L, 100, 1.5) == Some(507)) + } + + test("rung generation always makes progress") { + // Without the forced +1 step, a multiplier this close to 1.0 would never advance. + assert(DownwardShufflePartitionsPolicy.rungAtOrAbove(505L, 500, 1.0000001) == Some(505)) + assert(DownwardShufflePartitionsPolicy.rungAtOrAbove(505L, 500, 1.0) == Some(505)) + } + + test("a requirement no representable rung can cover fails closed") { + // Beyond the largest partition count Spark can express. + assert(DownwardShufflePartitionsPolicy.rungAtOrAbove(Int.MaxValue.toLong + 1L, 500, 2.0) + .isEmpty) + assert(DownwardShufflePartitionsPolicy.rungAtOrAbove(Long.MaxValue, 500, 2.0).isEmpty) + // Representable, but the doubling sequence overshoots Int.MaxValue before reaching it. + assert(DownwardShufflePartitionsPolicy.rungAtOrAbove(Int.MaxValue.toLong, 500, 2.0).isEmpty) + } + + test("an out-of-range requirement skips instead of recommending an uncovering rung") { + 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") + } + } + + // + // 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 = 2000, records = Seq(joinStage)) match { + case applied: DownwardShuffleDecision.Applied => + assert(applied.rawRequirement == 730L) + assert(applied.selectedValue == 1000) + assert(applied.normalValue == 2000) + 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 = 1000L * GiB) + val qualConfig = baseConfig.copy(inputSizeFactor = 0.8) + decide(4000, Seq(stage), qualConfig, ShuffleInputProvenance.Estimated) match { + case applied: DownwardShuffleDecision.Applied => + assert(applied.estimatedInputBytes == 800L * GiB) + assert(applied.rawRequirement == 800L) + assert(applied.selectedValue == 1000) + 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 = 1500L * 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 == 2000) + 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 target multiple does not round up to an extra rung") { + decide(4000, Seq(record(totalBytes = 1000L * GiB))) match { + case applied: DownwardShuffleDecision.Applied => + assert(applied.rawRequirement == 1000L) + assert(applied.selectedValue == 1000) + 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 rung equals the current value. + assert(decide(1000, Seq(record(totalBytes = 900L * GiB))) == + DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.NotDownward(1000, 1000))) + // The candidate rung is larger than the current value. + assert(decide(300, Seq(record(totalBytes = 900L * GiB))) == + DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.NotDownward(1000, 300))) + } + + test("AE7: a normal value below the rung floor is left alone") { + // Even a tiny stage cannot pull the recommendation under the floor of 500. + assert(decide(200, Seq(record(totalBytes = 1L))) == + DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.NotDownward(500, 200))) + } + + test("AE7: a reduction below the minimum factor is not applied") { + // 1999 is downward from a 1000-partition candidate but is not a 2x reduction. + assert(decide(1999, Seq(record(totalBytes = 900L * GiB))) == + DownwardShuffleDecision.Skipped( + DownwardShuffleSkipReason.BelowReductionThreshold(1000, 1999, 2.0))) + // Exactly 2x is enough. + assert(decide(2000, Seq(record(totalBytes = 900L * GiB))) + .isInstanceOf[DownwardShuffleDecision.Applied]) + } + + test("the calculator never returns a value above the normal one or below the floor") { + val configs = Seq(baseConfig, baseConfig.copy(partitionFloor = 128, rungMultiplier = 1.5), + baseConfig.copy(inputSizeFactor = 0.8, minReductionFactor = 1.0)) + val byteSizes = Seq(0L, 1L, GiB, 37L * GiB, 999L * GiB, 100000L * GiB, Long.MaxValue) + val normalValues = Seq(1, 199, 200, 500, 501, 2000, 200000, Int.MaxValue) + for (config <- configs; bytes <- byteSizes; normalValue <- normalValues) { + decide(normalValue, Seq(record(totalBytes = bytes)), config) match { + case applied: DownwardShuffleDecision.Applied => + assert(applied.selectedValue < normalValue, + s"config=$config bytes=$bytes normal=$normalValue") + assert(applied.selectedValue >= config.partitionFloor, + s"config=$config bytes=$bytes normal=$normalValue") + 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, 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 is incomplete, not empty-but-complete") { + val notAnalyzed = ShuffleStageInputAnalysis.empty(ShuffleInputProvenance.Measured) + assert(!notAnalyzed.isComplete) + DownwardShufflePartitionsPolicy.decide(Right(baseConfig), 4000, notAnalyzed) match { + case DownwardShuffleDecision.Skipped(reason: DownwardShuffleSkipReason.IncompleteEvidence) => + assert(reason.description.contains(ShuffleStageInputAnalysis.notAnalyzedSummary)) + case other => fail(s"expected an incomplete-evidence skip but got $other") + } + } + + 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, 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 every input of the decision") { + val stage = record(sqlId = 2L, stageId = 7, stageAttemptId = 1, + totalBytes = 729L * GiB + 1L, numBranches = 2) + val applied = decide(2000, Seq(stage)).asInstanceOf[DownwardShuffleDecision.Applied] + val comment = DownwardShufflePartitionsPolicy.appliedComment( + Seq("spark.sql.shuffle.partitions", + "spark.sql.adaptive.coalescePartitions.initialPartitionNum"), applied) + Seq("spark.sql.shuffle.partitions", + "spark.sql.adaptive.coalescePartitions.initialPartitionNum", + "2000", "1000", "measured", "SQL 2", "stage 7", "attempt 1", + (729L * GiB + 1L).toString, "2 shuffle branch(es)", "1.0", "730").foreach { fragment => + assert(comment.contains(fragment), s"'$fragment' missing from: $comment") + } + } +} From 27b464ecce111b3ffd9fac073516e2d5f0782b09 Mon Sep 17 00:00:00 2001 From: Partho Sarthi Date: Wed, 5 Aug 2026 19:00:28 -0700 Subject: [PATCH 02/10] Build consumer-stage shuffle input analysis with attempt-scoped evidence Adds the analysis that the downward shuffle-partition pass consumes: for every executed non-broadcast shuffle exchange in a final SQL plan, it resolves the uncompressed 'data size' once and attributes it to each distinct downstream branch, then totals those branches per consumer stage. A reused exchange is one node with several outgoing edges, so reuse multiplicity is preserved. The producing side of an exchange is identified by walking one edge upstream rather than by comparing stage numbers, because an exchange's own stage assignment spans both the write and the read side. Two gaps that would otherwise let unsafe evidence through are closed: - SparkListenerSQLExecutionEnd's errorMessage is now recorded, so a completed SQL execution can be told apart from a failed one. - Task accumulables now carry their stage attempt, which lets GPU spill be attributed to the attempt that reported it instead of leaking from a failed attempt into the later successful one. Signed-off-by: Partho Sarthi Co-Authored-By: Claude Opus 5 (1M context) --- .../tool/analysis/AppSQLPlanAnalyzer.scala | 9 +- .../analysis/ShuffleStageInputAnalyzer.scala | 348 ++++++++++++++++++ .../profiling/ProfileClassWarehouse.scala | 11 + .../sql/rapids/tool/EventProcessorBase.scala | 6 +- .../sql/rapids/tool/store/AccumManager.scala | 57 ++- .../sql/rapids/tool/util/EventUtils.scala | 14 +- .../ShuffleStageInputMetricsSuite.scala | 290 +++++++++++++++ 7 files changed, 730 insertions(+), 5 deletions(-) create mode 100644 core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/ShuffleStageInputAnalyzer.scala create mode 100644 core/src/test/scala/com/nvidia/spark/rapids/tool/analysis/ShuffleStageInputMetricsSuite.scala 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..4afd726e9 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 @@ -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..da7cef569 --- /dev/null +++ b/core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/ShuffleStageInputAnalyzer.scala @@ -0,0 +1,348 @@ +/* + * 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 + +/** + * 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) { + def add(moreBytes: Long): BranchTotals = BranchTotals(bytes + moreBytes, 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. + */ + 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) + } + + /** + * 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) + if (branches.isEmpty) { + reasons += ShuffleStageInputIncompleteReason.UnresolvedConsumerStage( + sqlId, node.id, node.name) + return + } + branches.foreach { sinkId => + resolveConsumerStage(graph, edgeIndex, sinkId, producerStages, sqlStageIds) match { + case Some(consumerStageId) => + val current = branchTotals.getOrElse(consumerStageId, BranchTotals(0L, 0)) + branchTotals(consumerStageId) = current.add(dataSize) + case None => + reasons += ShuffleStageInputIncompleteReason.UnresolvedConsumerStage( + sqlId, node.id, node.name) + } + } + } + + /** + * 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. More than one candidate at the same + * node is ambiguous and fails closed. + */ + private def resolveConsumerStage( + graph: ToolsPlanGraph, + edgeIndex: EdgeIndex, + sinkId: Long, + producerStages: Set[Int], + sqlStageIds: Set[Int]): Option[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) + } + }.distinct + if (candidates.size == 1) { + return Some(candidates.head) + } + if (candidates.size > 1) { + return None + } + visited ++= frontier + frontier = frontier.flatMap(edgeIndex.sinksOf).distinct.filterNot(visited.contains) + depth += 1 + } + None + } + + /** + * 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 fromStages = app.accumManager.accumInfoMap.get(metric.accumulatorId).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 + // Speculative and failed task metrics do not describe the work the recommendation governs. + val tasks = app.taskManager + .getTasks(stageId, attemptId, Some(t => t.successful && !t.speculative)) + .toSeq + val hasTaskSpill = tasks.exists(t => t.memoryBytesSpilled > 0L || t.diskBytesSpilled > 0L) + 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 + + /** + * 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/ProfileClassWarehouse.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/profiling/ProfileClassWarehouse.scala index 73ef3e357..848e0579b 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/org/apache/spark/sql/rapids/tool/EventProcessorBase.scala b/core/src/main/scala/org/apache/spark/sql/rapids/tool/EventProcessorBase.scala index 14fdd4d54..1121bf042 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 @@ -189,6 +189,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) } } @@ -396,11 +399,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..c1874ebc2 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 @@ -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 d770e9e62..c9e1f3083 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..dac29857d --- /dev/null +++ b/core/src/test/scala/com/nvidia/spark/rapids/tool/analysis/ShuffleStageInputMetricsSuite.scala @@ -0,0 +1,290 @@ +/* + * 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 + +/** + * 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 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")) + } + + 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("test resources are available") { + assert(new File(s"$profilingLogDir/rapids_join_eventlog.zstd").exists()) + } +} From 372f3e3d121b9725569230dba372da78c612adf9 Mon Sep 17 00:00:00 2001 From: Partho Sarthi Date: Wed, 5 Aug 2026 19:13:49 -0700 Subject: [PATCH 03/10] Apply the downward shuffle partition pass in the AutoTuner Runs the new policy after the normal job-level and cluster-level passes so it sees the effective shuffle partition recommendation, and wires the shuffle-stage input analysis through both the profiling and the qualification providers using the SQL plan analyzer each tool already built. Every upward decision the normal passes made is now recorded where it is applied, so the downward pass can preserve it exactly instead of trying to infer it. Before writing anything, the pass also checks the affected consumer stages for spill and skew, checks that Databricks automatic shuffle optimization is not still governing partitioning, and checks that every property it must write is actually writable -- so the shuffle and AQE partition properties can never end up disagreeing. Ordinary no-ops stay log-only. Only an applied reduction, invalid policy configuration, or evidence the analysis ran on but could not complete produces a user-facing comment. Signed-off-by: Partho Sarthi Co-Authored-By: Claude Opus 5 (1M context) --- .../tool/AppSummaryInfoBaseProvider.scala | 13 +- .../profiling/ApplicationSummaryInfo.scala | 17 ++ .../spark/rapids/tool/tuning/AutoTuner.scala | 144 +++++++++++ .../DownwardShufflePartitionsPolicy.scala | 15 ++ .../tuning/QualAppSummaryInfoProvider.scala | 20 +- .../tuning/QualificationAutoTunerRunner.scala | 8 +- .../rapids/tool/tuning/TunerContext.scala | 4 +- .../tool/tuning/BaseAutoTunerSuite.scala | 39 ++- .../DownwardShufflePartitionsSuite.scala | 12 +- .../tuning/ProfilingAutoTunerSuiteV2.scala | 242 +++++++++++++++++- .../tuning/QualificationAutoTunerSuite.scala | 76 +++++- 11 files changed, 566 insertions(+), 24 deletions(-) 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 05471f6b9..90653e6b9 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,7 +17,8 @@ package com.nvidia.spark.rapids.tool import com.nvidia.spark.rapids.tool.analysis.AggRawMetricsResult -import com.nvidia.spark.rapids.tool.profiling.{AppInfoColumnarExchangeMetrics, AppInfoJobStageAggMetricsVisitor, AppInfoPropertyGetter, AppInfoReadMetrics, AppInfoSqlTaskAggMetricsVisitor, AppInfoSQLTaskInputSizes, BaseProfilingAppSummaryInfoProvider, DataSourceProfileResult, ProfilerResult, SingleAppSummaryInfoProvider} +import com.nvidia.spark.rapids.tool.analysis.AppSQLPlanAnalyzer +import com.nvidia.spark.rapids.tool.profiling.{AppInfoColumnarExchangeMetrics, AppInfoJobStageAggMetricsVisitor, AppInfoPropertyGetter, AppInfoReadMetrics, AppInfoShuffleStageInputMetrics, AppInfoSqlTaskAggMetricsVisitor, AppInfoSQLTaskInputSizes, BaseProfilingAppSummaryInfoProvider, DataSourceProfileResult, ProfilerResult, SingleAppSummaryInfoProvider} import com.nvidia.spark.rapids.tool.tuning.QualAppSummaryInfoProvider import org.apache.spark.sql.rapids.tool.ToolUtils @@ -32,7 +33,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 @@ -84,12 +86,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/profiling/ApplicationSummaryInfo.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/profiling/ApplicationSummaryInfo.scala index e5b063118..2f51ef725 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 @@ -101,6 +101,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. */ @@ -249,6 +261,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/tuning/AutoTuner.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala index 98c3792e8..b1dd50848 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 @@ -253,6 +253,16 @@ abstract class AutoTuner( // 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]() + + /** 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 @@ -1485,6 +1495,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 " + @@ -1690,6 +1703,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) } } @@ -1725,6 +1739,122 @@ 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. It sizes the worst consumer stage from the total + * uncompressed shuffle input entering it 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 decision = DownwardShufflePartitionsPolicy.decide( + DownwardShufflePolicyConfig.fromProvider(configProvider), + shufflePartitionValue, + appInfoProvider.getShuffleStageInputAnalysis) + 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) + } + } + + /** + * 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 => + val partitionProperties = requiredPartitionProperties + partitionProperties.foreach(appendRecommendation(_, applied.selectedValue)) + appendComment( + DownwardShufflePartitionsPolicy.appliedComment(partitionProperties, 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] = { + applyToAllPartitionProperties[String](identity) + } + + /** + * 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. 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") + } + } + // 4. 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. */ @@ -1976,6 +2106,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 { @@ -2252,6 +2385,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 { @@ -2466,6 +2600,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 index b73c0df3a..0a68841e5 100644 --- 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 @@ -168,12 +168,24 @@ 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") @@ -281,6 +293,9 @@ object DownwardShufflePartitionsPolicy { config: DownwardShufflePolicyConfig, normalValue: Int, analysis: ShuffleStageInputAnalysis): DownwardShuffleDecision = { + if (!analysis.analyzed) { + return DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.NoAnalysisAvailable) + } if (!analysis.isComplete) { return DownwardShuffleDecision.Skipped( DownwardShuffleSkipReason.IncompleteEvidence(analysis.incompleteSummary())) 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 f69ffd298..b5a12a2bc 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 +import com.nvidia.spark.rapids.tool.analysis.{AggRawMetricsResult, AppSQLPlanAnalyzer} +import com.nvidia.spark.rapids.tool.profiling.{DataSourceProfileResult, ShuffleInputProvenance, ShuffleStageInputAnalysis} import org.apache.spark.internal.Logging import org.apache.spark.sql.rapids.tool.qualification.{QualificationAppInfo, QualificationSummaryInfo} @@ -30,12 +30,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. @@ -145,4 +149,14 @@ class QualAppSummaryInfoProvider( override def getClassPathEntries: Map[String, String] = { 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)) + } } 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..b17670657 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 @@ -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..aa2b650c7 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 @@ -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/test/scala/com/nvidia/spark/rapids/tool/tuning/BaseAutoTunerSuite.scala b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/BaseAutoTunerSuite.scala index d76ca0b38..39f005249 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 @@ -51,7 +51,9 @@ class AppInfoProviderMockTest(val maxInput: Double, val shuffleSkewStages: Set[Long], val scanStagesWithGpuOomSet: Set[Long], val gpuShuffleStagesWithContainerOomSet: Set[Long], - val maxColumnarExchangeDataSizeBytes: Option[Long] = None) + val maxColumnarExchangeDataSizeBytes: Option[Long] = None, + val shuffleStageInputAnalysis: ShuffleStageInputAnalysis = + ShuffleStageInputAnalysis.empty(ShuffleInputProvenance.Measured)) extends BaseProfilingAppSummaryInfoProvider { override def isAppInfoAvailable = true override def getMaxInput: Double = maxInput @@ -72,6 +74,7 @@ class AppInfoProviderMockTest(val maxInput: Double, override def scanStagesWithGpuOom: Set[Long] = scanStagesWithGpuOomSet override def gpuShuffleStagesWithContainerOom: Set[Long] = gpuShuffleStagesWithContainerOomSet override def getMaxColumnarExchangeDataSizeBytes: Option[Long] = maxColumnarExchangeDataSizeBytes + override def getShuffleStageInputAnalysis: ShuffleStageInputAnalysis = shuffleStageInputAnalysis /** * Sets the spark master property in the properties map. @@ -137,11 +140,41 @@ abstract class BaseAutoTunerSuite extends AnyFunSuite with BeforeAndAfterEach shuffleSkewStages: Set[Long] = Set(), scanStagesWithGpuOom: Set[Long] = Set(), gpuShuffleStagesWithContainerOom: Set[Long] = Set(), - maxColumnarExchangeDataSizeBytes: Option[Long] = None): AppInfoProviderMockTest = { + maxColumnarExchangeDataSizeBytes: Option[Long] = None, + shuffleStageInputAnalysis: ShuffleStageInputAnalysis = + ShuffleStageInputAnalysis.empty(ShuffleInputProvenance.Measured)): AppInfoProviderMockTest = { new AppInfoProviderMockTest(maxInput, spilledMetrics, jvmGCFractions, propsFromLog, sparkVersion, rapidsJars, distinctLocationPct, redundantReadSize, meanInput, meanShuffleRead, shuffleStagesWithPosSpilling, shuffleSkewStages, scanStagesWithGpuOom, - gpuShuffleStagesWithContainerOom, maxColumnarExchangeDataSizeBytes) + gpuShuffleStagesWithContainerOom, maxColumnarExchangeDataSizeBytes, 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): 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) + } + + /** 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 index bfe121a63..628b94dbb 100644 --- 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 @@ -401,14 +401,14 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { } } - test("an analysis that was never produced is incomplete, not empty-but-complete") { + test("an analysis that was never produced fails closed without a user comment") { val notAnalyzed = ShuffleStageInputAnalysis.empty(ShuffleInputProvenance.Measured) assert(!notAnalyzed.isComplete) - DownwardShufflePartitionsPolicy.decide(Right(baseConfig), 4000, notAnalyzed) match { - case DownwardShuffleDecision.Skipped(reason: DownwardShuffleSkipReason.IncompleteEvidence) => - assert(reason.description.contains(ShuffleStageInputAnalysis.notAnalyzedSummary)) - case other => fail(s"expected an incomplete-evidence skip but got $other") - } + val decision = DownwardShufflePartitionsPolicy.decide(Right(baseConfig), 4000, 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") { 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 16bd2ecba..7b9325cc5 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 @@ -19,8 +19,8 @@ package com.nvidia.spark.rapids.tool.tuning 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 -import com.nvidia.spark.rapids.tool.tuning.config.{ConfTypeEnum, TuningConfigEntry, TuningEntryDefinition} +import com.nvidia.spark.rapids.tool.profiling.{Profiler, ShuffleStageInputAnalysis} +import com.nvidia.spark.rapids.tool.tuning.config.{ConfTypeEnum, TuningConfigEntry, TuningConfiguration, TuningEntryDefinition} import org.apache.spark.network.util.ByteUnit import org.apache.spark.sql.{SparkSession, TrampolineUtil} @@ -2459,4 +2459,242 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { // scalastyle:on line.size.limit 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" + + /** + * 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" -> "1", + "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()) + } + + /** + * 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(), + maxColumnarExchangeDataSizeBytes: Option[Long] = None, + userProvidedTuningConfigs: Option[TuningConfiguration] = None, + platformName: String = PlatformNames.DATAPROC + ): (Seq[TuningEntryTrait], Seq[String]) = { + 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, + shuffleStagesWithPosSpilling = shuffleStagesWithPosSpilling, + gpuShuffleStagesWithContainerOom = gpuShuffleStagesWithContainerOom, + maxColumnarExchangeDataSizeBytes = maxColumnarExchangeDataSizeBytes, + shuffleStageInputAnalysis = shuffleStageInputAnalysis) + val platform = PlatformFactory.createInstance(platformName) + configureEventLogClusterInfoForTest(platform, numCores = 16, numWorkers = 1, 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 the 1000 rung, which is a 8x reduction from 8000. + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).contains("1000")) + assert(recommendedValue(properties, AQE_INITIAL_PARTITION_NUM_KEY).contains("1000")) + val appliedComment = comments.filter(_.contains("lowered from 8000 to 1000")) + 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").foreach { fragment => + assert(appliedComment.head.contains(fragment), + s"'$fragment' missing from: ${appliedComment.head}") + } + } + + 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("1000")) + 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("1000")) + assert(recommendedValue(properties, AQE_INITIAL_PARTITION_NUM_KEY).isEmpty) + assert(comments.exists(_.contains("lowered from 8000 to 1000"))) + } + + 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 invalidConfigs = ToolTestUtils.buildTuningConfigs( + default = List(TuningConfigEntry(name = "DOWNWARD_SHUFFLE_RUNG_MULTIPLIER", default = "0"))) + val (properties, comments) = runDownwardPass(worstStageNeeding900Partitions, + userProvidedTuningConfigs = Some(invalidConfigs)) + 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 disabledConfigs = ToolTestUtils.buildTuningConfigs( + default = List(TuningConfigEntry(name = "DOWNWARD_SHUFFLE_ENABLED", default = "false"))) + val (properties, comments) = runDownwardPass(worstStageNeeding900Partitions, + userProvidedTuningConfigs = Some(disabledConfigs)) + 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 respects a sub-threshold reduction") { + // 1500 partitions rounds up to the 2000 rung, which is not a 2x reduction from 3000. + val (properties, comments) = runDownwardPass( + completeShuffleStageInputs(Seq(4 -> 1500L * GiB)), + sourceProps = downwardPassSourceProps(Map(SHUFFLE_PARTITIONS_KEY -> "3000"))) + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).forall(_ == "3000")) + assert(!comments.exists(_.contains("lowered from"))) + } + + test("Downward pass never lowers below the configured floor") { + val (properties, comments) = runDownwardPass(completeShuffleStageInputs(Seq(4 -> 1L))) + // Even a one-byte stage cannot pull the recommendation under the 500 rung floor. + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).contains("500")) + assert(comments.exists(_.contains("lowered from 8000 to 500"))) + } + + 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 = 1, gpuCount = 1, + sparkProperties = sourceProps.toMap) + val autoTuner = buildAutoTunerForTests(infoProvider, platform) + 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 = 1, gpuCount = 1, + sparkProperties = sourceProps.toMap) + val autoTuner = buildAutoTunerForTests(infoProvider, platform) + val (properties, comments) = autoTuner.getRecommendedProperties() + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).forall(_ == "8000")) + assert(!comments.map(_.comment).exists(_.contains("lowered from"))) + } } 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 6f7bb3615..da841a7af 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 @@ -21,9 +21,9 @@ 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 +import com.nvidia.spark.rapids.tool.profiling.{Profiler, ShuffleInputProvenance, ShuffleStageInputAnalysis} import com.nvidia.spark.rapids.tool.qualification.{QualificationArgs, QualificationMain} -import com.nvidia.spark.rapids.tool.tuning.config.{CategoryEnum, ConfTypeEnum, LevelEnum, TuningConfigEntry, TuningEntryDefinition} +import com.nvidia.spark.rapids.tool.tuning.config.{CategoryEnum, ConfTypeEnum, LevelEnum, TuningConfigEntry, TuningConfiguration, TuningEntryDefinition} import com.nvidia.spark.rapids.tool.views.CLUSTER_INFORMATION_LABEL import com.nvidia.spark.rapids.tool.views.qualification.QualReportGenConfProvider import org.scalatest.exceptions.TestFailedException @@ -2268,4 +2268,76 @@ class QualificationAutoTunerSuite extends BaseAutoTunerSuite { s"AutoTuner should produce recommendations for GPU device '$gpuName'") } } + + // + // 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. + */ + private def buildDownwardPassAutoTuner( + shuffleStageInputAnalysis: ShuffleStageInputAnalysis, + userProvidedTuningConfigs: Option[TuningConfiguration] = None): AutoTuner = { + 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) + } + + 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 and rounds up to the 1000 rung. + 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=1000", + "--conf spark.sql.adaptive.coalescePartitions.initialPartitionNum=1000"), + autoTunerOutput) + // The comment must say the input was estimated, not measured, for a CPU event log. + val applied = comments.map(_.comment).filter(_.contains("lowered from 8000 to 1000")) + 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")) + } + + test("test AutoTuner for Qualification honours a custom input size factor") { + // A factor of 0.4 estimates 400 GiB, which needs 400 partitions and lands on the 500 rung. + val userConfigs = ToolTestUtils.buildTuningConfigs( + qualification = List( + TuningConfigEntry(name = "DOWNWARD_SHUFFLE_INPUT_SIZE_FACTOR", default = "0.4"))) + val autoTuner = + buildDownwardPassAutoTuner(cpuStageInput(1000L * QUAL_GiB), Some(userConfigs)) + val (properties, comments) = autoTuner.getRecommendedProperties(showOnlyUpdatedProps = + QualificationAutoTunerRunner.filterByUpdatedPropsEnabled) + val autoTunerOutput = Profiler.getAutoTunerResultsAsString(properties, comments) + assertExpectedLinesExist(Seq("--conf spark.sql.shuffle.partitions=500"), autoTunerOutput) + assert(comments.map(_.comment).exists(_.contains("input size factor 0.4"))) + } + + 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) + } } From 3360cab74a4e0d1912cb6a7b688b72da85053b0a Mon Sep 17 00:00:00 2001 From: Partho Sarthi Date: Wed, 5 Aug 2026 20:17:38 -0700 Subject: [PATCH 04/10] Cover the downward shuffle pass with real event-log fixtures Adds end-to-end assertions over existing GPU and CPU event logs so the metric-to-record path is proven on real plan graphs rather than only on constructed records: exact per-consumer-stage totals for a GPU join and a CPU query, a multi-branch CPU AQE stage whose real spill evidence reaches the record, a log with no terminal SQL end event, and a log with several stage attempts. Also proves the qualification provider hands back the SQL analyzer's own cached analysis instead of traversing the plans a second time, and fails closed when no analyzer is supplied. Signed-off-by: Partho Sarthi Co-Authored-By: Claude Opus 5 (1M context) --- .../ShuffleStageInputMetricsSuite.scala | 57 +++++++++++++++++++ .../tuning/QualificationAutoTunerSuite.scala | 36 +++++++++++- 2 files changed, 90 insertions(+), 3 deletions(-) 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 index dac29857d..dc7b44dac 100644 --- 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 @@ -49,6 +49,8 @@ class ShuffleStageInputMetricsSuite extends AnyFunSuite with Logging { 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 @@ -284,6 +286,61 @@ class ShuffleStageInputMetricsSuite extends AnyFunSuite with Logging { } } + 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("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/QualificationAutoTunerSuite.scala b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/QualificationAutoTunerSuite.scala index da841a7af..6d14fbd51 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.{DynamicAllocationInfo, EventLogPathProcessor, GpuTypes, PlatformFactory, PlatformNames, ToolTestUtils} +import com.nvidia.spark.rapids.tool.analysis.{AppSQLPlanAnalyzer, QualSparkMetricsAggregator} import com.nvidia.spark.rapids.tool.profiling.{Profiler, ShuffleInputProvenance, ShuffleStageInputAnalysis} -import com.nvidia.spark.rapids.tool.qualification.{QualificationArgs, QualificationMain} +import com.nvidia.spark.rapids.tool.qualification.{PluginTypeChecker, QualificationArgs, QualificationMain} import com.nvidia.spark.rapids.tool.tuning.config.{CategoryEnum, ConfTypeEnum, LevelEnum, TuningConfigEntry, TuningConfiguration, TuningEntryDefinition} import com.nvidia.spark.rapids.tool.views.CLUSTER_INFORMATION_LABEL import com.nvidia.spark.rapids.tool.views.qualification.QualReportGenConfProvider @@ -32,7 +33,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 @@ -2340,4 +2342,32 @@ class QualificationAutoTunerSuite extends BaseAutoTunerSuite { 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) + } } From 660ddca0b6d7c223d3e8dd9335ab81cead8fc80a Mon Sep 17 00:00:00 2001 From: Partho Sarthi Date: Thu, 6 Aug 2026 04:11:12 -0700 Subject: [PATCH 05/10] Resolve consumer stages for AQE-split and dead-end exchange branches Real event logs showed the consumer-stage mapping failing on two common shapes, which disabled the downward pass for 17 of 25 qualification applications in a representative run. An exchange branch can reach several stages at once: AQE splits one logical shuffle read across multiple query stages, and the single consuming operator is assigned to all of them. Requiring exactly one stage rejected every such branch. The full exchange size is now attributed to each of those stages instead. That deliberately overstates each split stage, since AQE divided the data between them, but overstating can only raise the partition requirement and make a reduction less likely -- understating it is the unsafe direction. Near the root of a plan the downstream nodes often carry no metrics and therefore no stage assignment, so walking the graph dead-ends and the topmost exchange looked unmappable. Its reading stage is recovered from its own assignment instead, which spans both sides of the shuffle, by removing the writing side. Signed-off-by: Partho Sarthi Co-Authored-By: Claude Opus 5 (1M context) --- .../analysis/ShuffleStageInputAnalyzer.scala | 72 +++++++++++++------ .../ShuffleStageInputMetricsSuite.scala | 46 ++++++++++++ 2 files changed, 97 insertions(+), 21 deletions(-) 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 index da7cef569..38bfb980c 100644 --- 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 @@ -189,38 +189,71 @@ class ShuffleStageInputAnalyzer(app: AppBase) extends Logging { branchTotals: mutable.LinkedHashMap[Int, BranchTotals], reasons: mutable.ArrayBuffer[ShuffleStageInputIncompleteReason]): Unit = { val branches = edgeIndex.sinksOf(node.id) - if (branches.isEmpty) { - reasons += ShuffleStageInputIncompleteReason.UnresolvedConsumerStage( - sqlId, node.id, node.name) - return + 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)) } - branches.foreach { sinkId => - resolveConsumerStage(graph, edgeIndex, sinkId, producerStages, sqlStageIds) match { - case Some(consumerStageId) => + // 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) - case None => - reasons += ShuffleStageInputIncompleteReason.UnresolvedConsumerStage( - sqlId, node.id, node.name) + } } } } + /** + * 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. More than one candidate at the same - * node is ambiguous and fails closed. + * 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. + * + * @return the consumer stages of this branch, or an empty set when none could be resolved */ - private def resolveConsumerStage( + private def resolveConsumerStages( graph: ToolsPlanGraph, edgeIndex: EdgeIndex, sinkId: Long, producerStages: Set[Int], - sqlStageIds: Set[Int]): Option[Int] = { + sqlStageIds: Set[Int]): Set[Int] = { val visited = mutable.HashSet.empty[Long] var frontier = List(sinkId) var depth = 0 @@ -231,18 +264,15 @@ class ShuffleStageInputAnalyzer(app: AppBase) extends Logging { !producerStages.contains(stageId) && sqlStageIds.contains(stageId) && raw.contains(stageId) } - }.distinct - if (candidates.size == 1) { - return Some(candidates.head) - } - if (candidates.size > 1) { - return None + }.toSet + if (candidates.nonEmpty) { + return candidates } visited ++= frontier frontier = frontier.flatMap(edgeIndex.sinksOf).distinct.filterNot(visited.contains) depth += 1 } - None + Set.empty } /** 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 index dc7b44dac..82dff20d8 100644 --- 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 @@ -341,6 +341,52 @@ class ShuffleStageInputMetricsSuite extends AnyFunSuite with Logging { } } + 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()) } From d520a5d9fe508553ab72c77b6a74c70972a85090 Mon Sep 17 00:00:00 2001 From: Partho Sarthi Date: Thu, 6 Aug 2026 15:41:49 -0700 Subject: [PATCH 06/10] Size the downward shuffle pass in whole cluster waves The generated partition rung ladder had no relationship to the cluster being recommended for. Anchored at 500, it recommended 500 partitions on a 125-executor by 16-core cluster for a job whose worst stage carried about 2 GB, leaving three quarters of the cluster idle for that stage. The candidate is now quantized to whole execution waves of the recommended cluster: take the larger of the slot count and the worst-stage requirement, then round up to the next whole multiple of the slot count. A slot count is the recommended executor count times a per-executor multiplier, selected by the new DOWNWARD_SHUFFLE_SLOT_BASIS entry so the GPU-concurrency alternative can be measured without a code change. The executor count is read from the recommendation map only, never through getPropertyValue, which falls back to source properties and would yield the source CPU executor count on Databricks; the cluster record is the fallback when no recommendation exists. The floor and rung-multiplier entries are retired. The minimum-reduction factor survives as an escape hatch but defaults to 1.0, since the wave quantum already prevents trivial reductions on realistic clusters. An unknown cluster shape is an ordinary logged no-op with no user comment, which is the common qualification case. The pass now ships disabled so that enabling wave-based sizing is an explicit per-run opt-in while it is being evaluated. No gate is loosened: the spill and skew check, the atomic property write, every upward-safety floor, and the v1 consumer-stage mapping are unchanged. A property's "was not set" comment is now emitted at most once, because the pass recommending a value for a key a later time would otherwise repeat it. Signed-off-by: Partho Sarthi Co-Authored-By: Claude Opus 5 --- .../resources/bootstrap/tuningConfigs.yaml | 27 +- .../spark/rapids/tool/tuning/AutoTuner.scala | 66 ++++- .../DownwardShufflePartitionsPolicy.scala | 147 +++++---- .../DownwardShufflePartitionsSuite.scala | 278 ++++++++++++------ .../tuning/ProfilingAutoTunerSuiteV2.scala | 162 +++++++--- .../tuning/QualificationAutoTunerSuite.scala | 43 ++- 6 files changed, 503 insertions(+), 220 deletions(-) diff --git a/core/src/main/resources/bootstrap/tuningConfigs.yaml b/core/src/main/resources/bootstrap/tuningConfigs.yaml index f156dbe86..73c6dd790 100644 --- a/core/src/main/resources/bootstrap/tuningConfigs.yaml +++ b/core/src/main/resources/bootstrap/tuningConfigs.yaml @@ -187,7 +187,8 @@ default: 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. - default: true + Disabled by default; opt in per run while the wave-based sizing is being evaluated. + default: false usedBy: spark.sql.shuffle.partitions, spark.sql.adaptive.coalescePartitions.initialPartitionNum - name: DOWNWARD_SHUFFLE_TARGET_PARTITION_SIZE @@ -205,27 +206,23 @@ default: default: 1.0 usedBy: spark.sql.shuffle.partitions, spark.sql.adaptive.coalescePartitions.initialPartitionNum - - name: DOWNWARD_SHUFFLE_PARTITION_FLOOR - description: >- - Lowest partition count the downward pass can recommend. This is also the first generated - partition rung. - default: 500 - usedBy: spark.sql.shuffle.partitions, spark.sql.adaptive.coalescePartitions.initialPartitionNum - - - name: DOWNWARD_SHUFFLE_RUNG_MULTIPLIER + - name: DOWNWARD_SHUFFLE_SLOT_BASIS description: >- - Multiplier used to generate successive partition rungs from the floor - (e.g. a floor of 500 and a multiplier of 2 generates 500, 1000, 2000, ...). - The raw partition requirement is always rounded up to the next rung. - default: 2 + Unit used to size one execution wave of the recommended cluster. The slot count is the + recommended executor count multiplied by this per-executor value, and the downward candidate + is always rounded up to a whole multiple of it. 'cores' uses the recommended + spark.executor.cores; 'concurrentGpuTasks' uses the recommended + spark.rapids.sql.concurrentGpuTasks. + default: cores usedBy: spark.sql.shuffle.partitions, spark.sql.adaptive.coalescePartitions.initialPartitionNum - name: DOWNWARD_SHUFFLE_MIN_REDUCTION_FACTOR description: >- Minimum ratio between the normal recommendation and the downward candidate required before a reduction is applied. A value of 2.0 means the normal value must be at least twice the - candidate. - default: 2.0 + candidate. The default of 1.0 imposes no threshold, because wave quantization already + prevents trivial reductions on realistic clusters. + default: 1.0 usedBy: spark.sql.shuffle.partitions, spark.sql.adaptive.coalescePartitions.initialPartitionNum - name: WORKER_GPU_COUNT 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 b1dd50848..62a28372e 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 @@ -18,6 +18,7 @@ package com.nvidia.spark.rapids.tool.tuning import scala.beans.BeanProperty import scala.collection.mutable +import scala.util.Try import scala.util.control.NonFatal import scala.util.matching.Regex @@ -250,6 +251,9 @@ 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]() @@ -464,13 +468,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) + } } /** @@ -1743,8 +1753,9 @@ 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. It sizes the worst consumer stage from the total - * uncompressed shuffle input entering it and lowers the recommendation only when the reduction + * 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 @@ -1752,9 +1763,13 @@ abstract class AutoTuner( * partition size. */ private def recommendDownwardShufflePartitions(): Unit = { + val configResult = DownwardShufflePolicyConfig.fromProvider(configProvider) + val slotCount = configResult.toOption.filter(_.enabled) + .flatMap(config => downwardShuffleSlotCount(config.slotBasis)) val decision = DownwardShufflePartitionsPolicy.decide( - DownwardShufflePolicyConfig.fromProvider(configProvider), + configResult, shufflePartitionValue, + slotCount, appInfoProvider.getShuffleStageInputAnalysis) decision match { case DownwardShuffleDecision.InvalidConfig(errors) => @@ -1769,6 +1784,43 @@ abstract class AutoTuner( } } + /** + * 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(basis: DownwardShuffleSlotBasis): Option[Int] = { + platform.recommendedClusterInfo.flatMap { clusterInfo => + val executors = + recommendedIntValue("spark.executor.instances").getOrElse(clusterInfo.numExecutors) + val slotsPerExecutor = basis match { + case DownwardShuffleSlotBasis.Cores => Some(clusterInfo.coresPerExecutor) + // Concurrent GPU tasks is a recommendation rather than a field on the cluster record, and + // it is suppressed entirely on platforms whose plugin auto-tunes it. + case DownwardShuffleSlotBasis.ConcurrentGpuTasks => + recommendedIntValue("spark.rapids.sql.concurrentGpuTasks") + } + slotsPerExecutor.filter(_ > 0).filter(_ => executors > 0).flatMap { perExecutor => + val slots = executors.toLong * perExecutor.toLong + if (slots > Int.MaxValue.toLong) None else Some(slots.toInt) + } + } + } + + /** 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) + } + /** * Runs the AutoTuner-owned gates that the pure policy cannot see, then updates every required * partition property or none of them. 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 index 0a68841e5..24afaabf1 100644 --- 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 @@ -24,6 +24,30 @@ 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 +/** + * Basis used to turn the recommended executor count into a cluster task-slot count. + * + * A "slot" is one unit of concurrency the recommended cluster actually has, so the number of slots + * is the number of partitions one execution wave can run. Which unit is right depends on whether + * the CPU task parallelism or the GPU task concurrency is the binding constraint, so the choice is + * exposed as a configuration entry rather than hard-coded. + */ +sealed abstract class DownwardShuffleSlotBasis(val label: String) + +object DownwardShuffleSlotBasis { + /** Slots per executor are the recommended `spark.executor.cores`. */ + case object Cores extends DownwardShuffleSlotBasis("cores") + + /** Slots per executor are the recommended `spark.rapids.sql.concurrentGpuTasks`. */ + case object ConcurrentGpuTasks extends DownwardShuffleSlotBasis("concurrentGpuTasks") + + val values: Seq[DownwardShuffleSlotBasis] = Seq(Cores, ConcurrentGpuTasks) + + def parse(raw: String): Option[DownwardShuffleSlotBasis] = { + values.find(_.label.equalsIgnoreCase(raw.trim)) + } +} + /** * Validated policy inputs of the downward-only shuffle partition pass. * @@ -34,24 +58,21 @@ import org.apache.spark.sql.rapids.tool.util.StringUtils * @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 - * @param partitionFloor lowest partition count the pass may recommend; first rung - * @param rungMultiplier ratio between successive generated partition rungs + * @param slotBasis unit used to size one execution wave of the recommended cluster * @param minReductionFactor required ratio of normal value to candidate before applying */ case class DownwardShufflePolicyConfig( enabled: Boolean, targetPartitionSizeBytes: Long, inputSizeFactor: Double, - partitionFloor: Int, - rungMultiplier: Double, + slotBasis: DownwardShuffleSlotBasis, minReductionFactor: 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" - val PARTITION_FLOOR_KEY = "DOWNWARD_SHUFFLE_PARTITION_FLOOR" - val RUNG_MULTIPLIER_KEY = "DOWNWARD_SHUFFLE_RUNG_MULTIPLIER" + val SLOT_BASIS_KEY = "DOWNWARD_SHUFFLE_SLOT_BASIS" val MIN_REDUCTION_FACTOR_KEY = "DOWNWARD_SHUFFLE_MIN_REDUCTION_FACTOR" /** Config used when the feature is switched off. The remaining fields are never read. */ @@ -59,8 +80,7 @@ object DownwardShufflePolicyConfig { enabled = false, targetPartitionSizeBytes = 0L, inputSizeFactor = 0.0, - partitionFloor = 0, - rungMultiplier = 0.0, + slotBasis = DownwardShuffleSlotBasis.Cores, minReductionFactor = 0.0) /** @@ -85,23 +105,20 @@ object DownwardShufflePolicyConfig { 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 floor = parsePositiveInt(configProvider, PARTITION_FLOOR_KEY) - val multiplier = - parseDouble(configProvider, RUNG_MULTIPLIER_KEY, min = 1.0, minInclusive = false) + val slotBasis = parseSlotBasis(configProvider, SLOT_BASIS_KEY) val minReduction = parseDouble(configProvider, MIN_REDUCTION_FACTOR_KEY, min = 1.0, minInclusive = true) - val errors = Seq(targetSize, factor, floor, multiplier, minReduction).collect { + val errors = Seq(targetSize, factor, slotBasis, minReduction).collect { case Left(err) => err } - (targetSize, factor, floor, multiplier, minReduction) match { - case (Right(size), Right(f), Right(fl), Right(m), Right(r)) if errors.isEmpty => + (targetSize, factor, slotBasis, minReduction) match { + case (Right(size), Right(f), Right(b), Right(r)) if errors.isEmpty => Right(DownwardShufflePolicyConfig( enabled = true, targetPartitionSizeBytes = size, inputSizeFactor = f, - partitionFloor = fl, - rungMultiplier = m, + slotBasis = b, minReductionFactor = r)) case _ => Left(errors) } @@ -132,12 +149,13 @@ object DownwardShufflePolicyConfig { } } - private def parsePositiveInt( - configProvider: TuningConfigProvider, key: String): Either[String, Int] = { + private def parseSlotBasis( + configProvider: TuningConfigProvider, + key: String): Either[String, DownwardShuffleSlotBasis] = { rawValue(configProvider, key).flatMap { raw => - Try(raw.trim.toInt).toOption - .filter(_ > 0) - .toRight(s"'$key' must be a positive integer but was '$raw'") + DownwardShuffleSlotBasis.parse(raw).toRight( + s"'$key' must be one of ${DownwardShuffleSlotBasis.values.map(_.label).mkString(", ")}" + + s" but was '$raw'") } } @@ -190,16 +208,26 @@ object DownwardShuffleSkipReason { extends DownwardShuffleSkipReason("no consumer stage shuffle input was found") /** - * The worst stage needs more partitions than a Spark partition count can express, so no rung can - * cover it. Failing closed here is safer than recommending a rung below the requirement. + * 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 exceeds the largest" + - s" representable partition count (${Int.MaxValue})") { + 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") @@ -234,10 +262,12 @@ object DownwardShuffleDecision { * A reduction that passed every policy gate. * * @param normalValue the effective recommendation produced by normal tuning - * @param selectedValue the partition rung to recommend instead + * @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 a rung + * @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 */ @@ -247,6 +277,8 @@ object DownwardShuffleDecision { determiningRecord: ShuffleStageInputRecord, estimatedInputBytes: Long, rawRequirement: Long, + slotCount: Int, + waveCount: Int, provenance: ShuffleInputProvenance, inputSizeFactor: Double) extends DownwardShuffleDecision @@ -260,7 +292,7 @@ object DownwardShuffleDecision { * * 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 a rung sequence that is guaranteed to make progress. + * and quantizing the result to whole execution waves of the recommended cluster. */ object DownwardShufflePartitionsPolicy { @@ -275,23 +307,26 @@ object DownwardShufflePartitionsPolicy { * * @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, analysis) + 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) @@ -304,16 +339,21 @@ object DownwardShufflePartitionsPolicy { 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) - rungAtOrAbove(worst.requirement, config.partitionFloor, config.rungMultiplier) match { + waveQuantized(worst.requirement, slots) match { case None => - // No representable rung covers the requirement: fail closed rather than recommend a - // partition count that is known to be too small. + // 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 the rung floor, - // because every generated rung is at least the floor. + // 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) @@ -328,6 +368,8 @@ object DownwardShufflePartitionsPolicy { determiningRecord = worst.record, estimatedInputBytes = worst.estimatedBytes, rawRequirement = worst.requirement, + slotCount = slots, + waveCount = candidate / slots, provenance = analysis.provenance, inputSizeFactor = config.inputSizeFactor) } @@ -395,32 +437,24 @@ object DownwardShufflePartitionsPolicy { } /** - * Rounds a raw requirement up to the first generated rung that is at least as large. + * Rounds a raw requirement up to a whole number of execution waves of the recommended cluster. * - * Rungs start at the floor and grow by the multiplier, rounding each step up so that a fractional - * multiplier still produces a strictly increasing integer sequence. Each step is additionally - * forced to advance by at least one partition, so a multiplier barely above 1.0 cannot stall. + * 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 covering rung, or None when no rung within `Int.MaxValue` covers the requirement + * @return the wave-quantized candidate, or None when no whole-wave count within `Int.MaxValue` + * covers the requirement */ - private[tuning] def rungAtOrAbove( - requirement: Long, floor: Int, multiplier: Double): Option[Int] = { - if (requirement > Int.MaxValue.toLong) { + private[tuning] def waveQuantized(requirement: Long, slots: Int): Option[Int] = { + if (slots <= 0 || requirement > Int.MaxValue.toLong) { return None } - var rung = floor.toLong - while (rung < requirement) { - val grown = math.ceil(rung.toDouble * multiplier) - if (grown > Int.MaxValue.toDouble) { - return None - } - val next = math.max(grown.toLong, rung + 1L) - if (next > Int.MaxValue.toLong) { - return None - } - rung = next - } - Some(rung.toInt) + // 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) } /** @@ -439,6 +473,7 @@ object DownwardShufflePartitionsPolicy { s"branch(es), input size factor ${decision.inputSizeFactor}, " + s"${decision.estimatedInputBytes} estimated bytes, " + s"raw requirement ${decision.rawRequirement} partitions rounded up to " + - s"${decision.selectedValue}." + s"${decision.waveCount} execution wave(s) of ${decision.slotCount} cluster task slots " + + s"(${decision.selectedValue} partitions)." } } 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 index 628b94dbb..46e445ca7 100644 --- 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 @@ -16,6 +16,8 @@ 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} @@ -30,10 +32,22 @@ 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 = default))) + .withUserProvidedConfig(Some(ToolTestUtils.buildTuningConfigs(default = withOptIn(default)))) .build[ProfTuningConfigProvider] } @@ -41,11 +55,16 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { default: List[TuningConfigEntry] = List.empty, qualification: List[TuningConfigEntry] = List.empty): QualTuningConfigProvider = { TuningConfigProvider.builder - .withUserProvidedConfig(Some( - ToolTestUtils.buildTuningConfigs(default = default, qualification = qualification))) + .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, @@ -66,14 +85,16 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { ShuffleStageInputAnalysis(records, Seq.empty, provenance) } - /** Config used by the arithmetic tests: 1 GiB target, rungs 500/1000/2000/..., 2x reduction. */ + /** 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, - partitionFloor = 500, - rungMultiplier = 2.0, - minReductionFactor = 2.0) + slotBasis = DownwardShuffleSlotBasis.Cores, + minReductionFactor = 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( @@ -92,9 +113,10 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { normalValue: Int, records: Seq[ShuffleStageInputRecord], config: DownwardShufflePolicyConfig = baseConfig, - provenance: ShuffleInputProvenance = ShuffleInputProvenance.Measured + provenance: ShuffleInputProvenance = ShuffleInputProvenance.Measured, + slotCount: Option[Int] = Some(slots) ): DownwardShuffleDecision = { - DownwardShufflePartitionsPolicy.decide(Right(config), normalValue, + DownwardShufflePartitionsPolicy.decide(Right(config), normalValue, slotCount, analysis(records, provenance)) } @@ -102,14 +124,29 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { // 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) - assert(config.partitionFloor == 500) - assert(config.rungMultiplier == 2.0) - assert(config.minReductionFactor == 2.0) + assert(config.slotBasis == DownwardShuffleSlotBasis.Cores) + // Wave quantization already prevents trivial reductions, so no extra threshold is imposed. + assert(config.minReductionFactor == 1.0) + } + + test("the retired rung entries are gone and their removal does not break config loading") { + Seq("DOWNWARD_SHUFFLE_PARTITION_FLOOR", "DOWNWARD_SHUFFLE_RUNG_MULTIPLIER").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") { @@ -118,16 +155,17 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { assert(config.inputSizeFactor == 0.8) // Everything else still comes from the shared defaults. assert(config.targetPartitionSizeBytes == GiB) - assert(config.partitionFloor == 500) + assert(config.slotBasis == DownwardShuffleSlotBasis.Cores) } 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"), - TuningConfigEntry(name = "DOWNWARD_SHUFFLE_PARTITION_FLOOR", default = "128")))) + TuningConfigEntry(name = "DOWNWARD_SHUFFLE_SLOT_BASIS", default = "concurrentGpuTasks")))) assert(fromDefault == Right(baseConfig.copy( - targetPartitionSizeBytes = 512L * 1024L * 1024L, partitionFloor = 128))) + targetPartitionSizeBytes = 512L * 1024L * 1024L, + slotBasis = DownwardShuffleSlotBasis.ConcurrentGpuTasks))) val fromToolSection = DownwardShufflePolicyConfig.fromProvider(qualProvider( qualification = List( @@ -135,14 +173,30 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { assert(fromToolSection == Right(baseConfig.copy(inputSizeFactor = 0.5))) } + test("the slot basis accepts only a known label, case-insensitively") { + Seq("cores" -> DownwardShuffleSlotBasis.Cores, + "CORES" -> DownwardShuffleSlotBasis.Cores, + " concurrentgputasks " -> DownwardShuffleSlotBasis.ConcurrentGpuTasks + ).foreach { case (raw, expected) => + val result = DownwardShufflePolicyConfig.fromProvider(profProvider( + default = List(TuningConfigEntry(name = "DOWNWARD_SHUFFLE_SLOT_BASIS", default = raw)))) + assert(result.exists(_.slotBasis == expected), s"'$raw' should parse as $expected") + } + Seq("gpus", "tasks", "1").foreach { raw => + val result = DownwardShufflePolicyConfig.fromProvider(profProvider( + default = List(TuningConfigEntry(name = "DOWNWARD_SHUFFLE_SLOT_BASIS", default = raw)))) + assert(result.isLeft, s"'$raw' should not parse as a slot basis") + } + } + 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_PARTITION_FLOOR", default = "not-a-number"))) + TuningConfigEntry(name = "DOWNWARD_SHUFFLE_SLOT_BASIS", default = "not-a-basis"))) val configResult = DownwardShufflePolicyConfig.fromProvider(provider) assert(configResult == Right(DownwardShufflePolicyConfig.disabled)) - assert(DownwardShufflePartitionsPolicy.decide(configResult, 4000, + assert(DownwardShufflePartitionsPolicy.decide(configResult, 4000, Some(slots), analysis(Seq(record(totalBytes = GiB)))) == DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.Disabled)) } @@ -163,9 +217,6 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { val invalidByKey = Seq( "DOWNWARD_SHUFFLE_TARGET_PARTITION_SIZE" -> Seq("0", "-1g", "abc"), "DOWNWARD_SHUFFLE_INPUT_SIZE_FACTOR" -> Seq("0", "-0.5", "abc"), - "DOWNWARD_SHUFFLE_PARTITION_FLOOR" -> Seq("0", "-500", "1.5", "abc"), - // A multiplier of exactly 1.0 or below cannot generate growing rungs. - "DOWNWARD_SHUFFLE_RUNG_MULTIPLIER" -> Seq("1", "1.0", "0.5", "abc"), // A reduction factor below 1.0 would allow raising the recommendation. "DOWNWARD_SHUFFLE_MIN_REDUCTION_FACTOR" -> Seq("0.9", "0", "abc")) @@ -181,7 +232,7 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { } test("NaN and infinite values are rejected") { - Seq("DOWNWARD_SHUFFLE_INPUT_SIZE_FACTOR", "DOWNWARD_SHUFFLE_RUNG_MULTIPLIER", + Seq("DOWNWARD_SHUFFLE_INPUT_SIZE_FACTOR", "DOWNWARD_SHUFFLE_MIN_REDUCTION_FACTOR").foreach { key => Seq("NaN", "Infinity", "-Infinity").foreach { value => val result = DownwardShufflePolicyConfig.fromProvider(profProvider( @@ -193,12 +244,12 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { test("all configuration errors are reported as one fail-closed decision") { val provider = profProvider(default = List( - TuningConfigEntry(name = "DOWNWARD_SHUFFLE_PARTITION_FLOOR", default = "-1"), - TuningConfigEntry(name = "DOWNWARD_SHUFFLE_RUNG_MULTIPLIER", default = "0.5"))) + TuningConfigEntry(name = "DOWNWARD_SHUFFLE_SLOT_BASIS", default = "not-a-basis"), + TuningConfigEntry(name = "DOWNWARD_SHUFFLE_MIN_REDUCTION_FACTOR", default = "0.5"))) val configResult = DownwardShufflePolicyConfig.fromProvider(provider) val errors = expectErrors(configResult) assert(errors.size == 2) - DownwardShufflePartitionsPolicy.decide(configResult, 4000, + 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") @@ -226,44 +277,64 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { assert(DownwardShufflePartitionsPolicy.partitionRequirement(1000L * GiB + 1L, GiB) == 1001L) } - test("rungAtOrAbove generates each rung from the floor and the multiplier") { - val rungs = Seq(1L -> 500, 500L -> 500, 501L -> 1000, 1000L -> 1000, 1001L -> 2000, - 2000L -> 2000, 2001L -> 4000, 4001L -> 8000, 64001L -> 128000) - rungs.foreach { case (requirement, expected) => - assert(DownwardShufflePartitionsPolicy.rungAtOrAbove(requirement, 500, 2.0) == - Some(expected), s"requirement $requirement") + 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("fractional multipliers round each rung up so the sequence stays integral") { - // 100 -> 150 -> 225 -> 338 (ceil of 337.5) -> 507 - assert(DownwardShufflePartitionsPolicy.rungAtOrAbove(101L, 100, 1.5) == Some(150)) - assert(DownwardShufflePartitionsPolicy.rungAtOrAbove(226L, 100, 1.5) == Some(338)) - assert(DownwardShufflePartitionsPolicy.rungAtOrAbove(339L, 100, 1.5) == Some(507)) - } - - test("rung generation always makes progress") { - // Without the forced +1 step, a multiplier this close to 1.0 would never advance. - assert(DownwardShufflePartitionsPolicy.rungAtOrAbove(505L, 500, 1.0000001) == Some(505)) - assert(DownwardShufflePartitionsPolicy.rungAtOrAbove(505L, 500, 1.0) == Some(505)) + 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 rung can cover fails closed") { + test("a requirement no representable wave count can cover fails closed") { // Beyond the largest partition count Spark can express. - assert(DownwardShufflePartitionsPolicy.rungAtOrAbove(Int.MaxValue.toLong + 1L, 500, 2.0) - .isEmpty) - assert(DownwardShufflePartitionsPolicy.rungAtOrAbove(Long.MaxValue, 500, 2.0).isEmpty) - // Representable, but the doubling sequence overshoots Int.MaxValue before reaching it. - assert(DownwardShufflePartitionsPolicy.rungAtOrAbove(Int.MaxValue.toLong, 500, 2.0).isEmpty) + 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 rung") { + 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) } // @@ -273,11 +344,14 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { 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 = 2000, records = Seq(joinStage)) match { + decide(normalValue = 8000, records = Seq(joinStage)) match { case applied: DownwardShuffleDecision.Applied => assert(applied.rawRequirement == 730L) - assert(applied.selectedValue == 1000) - assert(applied.normalValue == 2000) + // 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") @@ -285,13 +359,15 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { } test("the qualification factor lowers the requirement of the same stage") { - val stage = record(totalBytes = 1000L * GiB) + val stage = record(totalBytes = 4000L * GiB) val qualConfig = baseConfig.copy(inputSizeFactor = 0.8) - decide(4000, Seq(stage), qualConfig, ShuffleInputProvenance.Estimated) match { + decide(8000, Seq(stage), qualConfig, ShuffleInputProvenance.Estimated) match { case applied: DownwardShuffleDecision.Applied => - assert(applied.estimatedInputBytes == 800L * GiB) - assert(applied.rawRequirement == 800L) - assert(applied.selectedValue == 1000) + 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") @@ -300,12 +376,12 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { 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 = 1500L * 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 == 2000) + assert(applied.selectedValue == 4000) case other => fail(s"expected an applied reduction but got $other") } } @@ -333,52 +409,62 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { } } - test("an exact target multiple does not round up to an extra rung") { - decide(4000, Seq(record(totalBytes = 1000L * GiB))) match { + 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 == 1000L) - assert(applied.selectedValue == 1000) + 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 rung equals the current value. - assert(decide(1000, Seq(record(totalBytes = 900L * GiB))) == - DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.NotDownward(1000, 1000))) - // The candidate rung is larger than the current value. + // 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(1000, 300))) + DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.NotDownward(2000, 300))) } - test("AE7: a normal value below the rung floor is left alone") { - // Even a tiny stage cannot pull the recommendation under the floor of 500. + 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(500, 200))) + DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.NotDownward(2000, 200))) } - test("AE7: a reduction below the minimum factor is not applied") { - // 1999 is downward from a 1000-partition candidate but is not a 2x reduction. - assert(decide(1999, Seq(record(totalBytes = 900L * GiB))) == + test("the default reduction factor imposes no threshold but an override still blocks") { + // A 3000 -> 2000 reduction is not 2x, so v1's default would have blocked it. + 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") + } + val strictConfig = baseConfig.copy(minReductionFactor = 2.0) + assert(decide(3000, Seq(record(totalBytes = 900L * GiB)), strictConfig) == DownwardShuffleDecision.Skipped( - DownwardShuffleSkipReason.BelowReductionThreshold(1000, 1999, 2.0))) - // Exactly 2x is enough. - assert(decide(2000, Seq(record(totalBytes = 900L * GiB))) + DownwardShuffleSkipReason.BelowReductionThreshold(2000, 3000, 2.0))) + // Exactly 2x still clears the explicit override. + assert(decide(4000, Seq(record(totalBytes = 900L * GiB)), strictConfig) .isInstanceOf[DownwardShuffleDecision.Applied]) } - test("the calculator never returns a value above the normal one or below the floor") { - val configs = Seq(baseConfig, baseConfig.copy(partitionFloor = 128, rungMultiplier = 1.5), - baseConfig.copy(inputSizeFactor = 0.8, minReductionFactor = 1.0)) + 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, minReductionFactor = 2.0)) 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; normalValue <- normalValues) { - decide(normalValue, Seq(record(totalBytes = bytes)), config) match { + for (config <- configs; bytes <- byteSizes; slotCount <- slotCounts; + normalValue <- normalValues) { + decide(normalValue, Seq(record(totalBytes = bytes)), config, slotCount = Some(slotCount)) + match { case applied: DownwardShuffleDecision.Applied => - assert(applied.selectedValue < normalValue, - s"config=$config bytes=$bytes normal=$normalValue") - assert(applied.selectedValue >= config.partitionFloor, - s"config=$config bytes=$bytes normal=$normalValue") + 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 } } @@ -393,7 +479,8 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { Seq(record(totalBytes = 900L * GiB)), Seq(ShuffleStageInputIncompleteReason.MissingExchangeMetric(0L, 3L, "Exchange")), ShuffleInputProvenance.Measured) - DownwardShufflePartitionsPolicy.decide(Right(baseConfig), 4000, incomplete) match { + DownwardShufflePartitionsPolicy.decide(Right(baseConfig), 4000, Some(slots), + incomplete) match { case DownwardShuffleDecision.Skipped(reason: DownwardShuffleSkipReason.IncompleteEvidence) => assert(reason.isWarning) assert(reason.description.contains("data size")) @@ -404,7 +491,8 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { 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, notAnalyzed) + 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. @@ -414,7 +502,8 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { 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, noShuffle) + val decision = + DownwardShufflePartitionsPolicy.decide(Right(baseConfig), 4000, Some(slots), noShuffle) assert(decision == DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.NoStageEvidence)) assert(!DownwardShuffleSkipReason.NoStageEvidence.isWarning) } @@ -435,18 +524,21 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { // Diagnostics // - test("the applied comment names every input of the decision") { + test("the applied comment names the slot count and wave arithmetic, not rung rounding") { val stage = record(sqlId = 2L, stageId = 7, stageAttemptId = 1, - totalBytes = 729L * GiB + 1L, numBranches = 2) - val applied = decide(2000, Seq(stage)).asInstanceOf[DownwardShuffleDecision.Applied] + 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", "spark.sql.adaptive.coalescePartitions.initialPartitionNum"), applied) Seq("spark.sql.shuffle.partitions", "spark.sql.adaptive.coalescePartitions.initialPartitionNum", - "2000", "1000", "measured", "SQL 2", "stage 7", "attempt 1", - (729L * GiB + 1L).toString, "2 shuffle branch(es)", "1.0", "730").foreach { fragment => + "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") } } 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 7b9325cc5..876c4d4dc 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 @@ -2470,6 +2470,13 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { "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. @@ -2478,7 +2485,7 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { extra: Map[String, String] = Map.empty): mutable.Map[String, String] = { val base = mutable.LinkedHashMap[String, String]( "spark.executor.cores" -> "16", - "spark.executor.instances" -> "1", + "spark.executor.instances" -> DOWNWARD_PASS_WORKERS.toString, "spark.executor.memory" -> "80g", "spark.executor.resource.gpu.amount" -> "1", "spark.sql.adaptive.enabled" -> "true", @@ -2496,6 +2503,19 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { 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. @@ -2506,9 +2526,13 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { shuffleStagesWithPosSpilling: Set[Long] = Set(), gpuShuffleStagesWithContainerOom: Set[Long] = Set(), maxColumnarExchangeDataSizeBytes: Option[Long] = None, - userProvidedTuningConfigs: Option[TuningConfiguration] = None, - platformName: String = PlatformNames.DATAPROC + 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), @@ -2522,8 +2546,10 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { maxColumnarExchangeDataSizeBytes = maxColumnarExchangeDataSizeBytes, shuffleStageInputAnalysis = shuffleStageInputAnalysis) val platform = PlatformFactory.createInstance(platformName) - configureEventLogClusterInfoForTest(platform, numCores = 16, numWorkers = 1, gpuCount = 1, - sparkProperties = sourceProps.toMap) + 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() @@ -2537,24 +2563,33 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { test("Downward pass lowers both partition properties atomically when AQE coalescing is on") { val (properties, comments) = runDownwardPass(worstStageNeeding900Partitions) - // 900 partitions rounds up to the 1000 rung, which is a 8x reduction from 8000. - assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).contains("1000")) - assert(recommendedValue(properties, AQE_INITIAL_PARTITION_NUM_KEY).contains("1000")) - val appliedComment = comments.filter(_.contains("lowered from 8000 to 1000")) + // 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").foreach { fragment => - assert(appliedComment.head.contains(fragment), - s"'$fragment' missing from: ${appliedComment.head}") - } + "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("1000")) + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).contains("1200")) assert(recommendedValue(properties, ADVISORY_PARTITION_SIZE_KEY).forall(_ == advisorySize)) } @@ -2562,9 +2597,9 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { val (properties, comments) = runDownwardPass(worstStageNeeding900Partitions, sourceProps = downwardPassSourceProps( Map("spark.sql.adaptive.coalescePartitions.enabled" -> "false"))) - assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).contains("1000")) + 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 1000"))) + assert(comments.exists(_.contains("lowered from 8000 to 1200"))) } test("Downward pass preserves an increase made for shuffle spilling") { @@ -2613,39 +2648,62 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { } test("Downward pass warns once and changes nothing when its configuration is invalid") { - val invalidConfigs = ToolTestUtils.buildTuningConfigs( - default = List(TuningConfigEntry(name = "DOWNWARD_SHUFFLE_RUNG_MULTIPLIER", default = "0"))) val (properties, comments) = runDownwardPass(worstStageNeeding900Partitions, - userProvidedTuningConfigs = Some(invalidConfigs)) + extraTuningConfigs = List( + TuningConfigEntry(name = "DOWNWARD_SHUFFLE_SLOT_BASIS", 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 disabledConfigs = ToolTestUtils.buildTuningConfigs( - default = List(TuningConfigEntry(name = "DOWNWARD_SHUFFLE_ENABLED", default = "false"))) val (properties, comments) = runDownwardPass(worstStageNeeding900Partitions, - userProvidedTuningConfigs = Some(disabledConfigs)) + 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 respects a sub-threshold reduction") { - // 1500 partitions rounds up to the 2000 rung, which is not a 2x reduction from 3000. + test("Downward pass respects a configured minimum reduction factor") { + // 1500 partitions rounds up to 4 waves (1600), which is not a 2x reduction from 3000. val (properties, comments) = runDownwardPass( completeShuffleStageInputs(Seq(4 -> 1500L * GiB)), - sourceProps = downwardPassSourceProps(Map(SHUFFLE_PARTITIONS_KEY -> "3000"))) + sourceProps = downwardPassSourceProps(Map(SHUFFLE_PARTITIONS_KEY -> "3000")), + extraTuningConfigs = List( + TuningConfigEntry(name = "DOWNWARD_SHUFFLE_MIN_REDUCTION_FACTOR", default = "2.0"))) assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).forall(_ == "3000")) assert(!comments.exists(_.contains("lowered from"))) + // The default of 1.0 imposes no threshold, so the same reduction applies without the override. + val (defaultProps, defaultComments) = runDownwardPass( + completeShuffleStageInputs(Seq(4 -> 1500L * GiB)), + sourceProps = downwardPassSourceProps(Map(SHUFFLE_PARTITIONS_KEY -> "3000"))) + assert(recommendedValue(defaultProps, SHUFFLE_PARTITIONS_KEY).contains("1600")) + assert(defaultComments.exists(_.contains("lowered from 3000 to 1600"))) } - test("Downward pass never lowers below the configured floor") { + test("Downward pass never lowers below a single execution wave") { val (properties, comments) = runDownwardPass(completeShuffleStageInputs(Seq(4 -> 1L))) - // Even a one-byte stage cannot pull the recommendation under the 500 rung floor. - assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).contains("500")) - assert(comments.exists(_.contains("lowered from 8000 to 500"))) + // 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 slot count follows the configured slot basis") { + // The GPU-concurrency basis uses the recommended concurrentGpuTasks instead of the cores. + val (properties, comments) = runDownwardPass(worstStageNeeding900Partitions, + extraTuningConfigs = List(TuningConfigEntry(name = "DOWNWARD_SHUFFLE_SLOT_BASIS", + default = "concurrentGpuTasks"))) + 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") + val expected = math.ceil(900.0 / gpuSlots).toInt * gpuSlots + assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).contains(expected.toString)) + assert(comments.exists(_.contains(s"$gpuSlots cluster task slots"))) } test("Downward pass changes neither property when one of them is enforced") { @@ -2662,9 +2720,10 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { meanShuffleRead = 1.0E9, shuffleStageInputAnalysis = worstStageNeeding900Partitions) val platform = PlatformFactory.createInstance(PlatformNames.DATAPROC, Some(targetClusterInfo)) - configureEventLogClusterInfoForTest(platform, numCores = 16, numWorkers = 1, gpuCount = 1, - sparkProperties = sourceProps.toMap) - val autoTuner = buildAutoTunerForTests(infoProvider, platform) + 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")) @@ -2690,11 +2749,44 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { shuffleStageInputAnalysis = worstStageNeeding900Partitions) val platform = PlatformFactory.createInstance(PlatformNames.DATABRICKS_AWS, Some(targetClusterInfo)) - configureEventLogClusterInfoForTest(platform, numCores = 16, numWorkers = 1, gpuCount = 1, - sparkProperties = sourceProps.toMap) - val autoTuner = buildAutoTunerForTests(infoProvider, platform) + 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"))) + } } 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 6d14fbd51..1bc5e1814 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 @@ -24,7 +24,7 @@ import com.nvidia.spark.rapids.tool.{DynamicAllocationInfo, EventLogPathProcesso import com.nvidia.spark.rapids.tool.analysis.{AppSQLPlanAnalyzer, QualSparkMetricsAggregator} import com.nvidia.spark.rapids.tool.profiling.{Profiler, ShuffleInputProvenance, ShuffleStageInputAnalysis} import com.nvidia.spark.rapids.tool.qualification.{PluginTypeChecker, QualificationArgs, QualificationMain} -import com.nvidia.spark.rapids.tool.tuning.config.{CategoryEnum, ConfTypeEnum, LevelEnum, TuningConfigEntry, TuningConfiguration, TuningEntryDefinition} +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 import org.scalatest.exceptions.TestFailedException @@ -2280,10 +2280,18 @@ class QualificationAutoTunerSuite extends BaseAutoTunerSuite { /** * 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, - userProvidedTuningConfigs: Option[TuningConfiguration] = None): AutoTuner = { + 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", @@ -2301,36 +2309,43 @@ class QualificationAutoTunerSuite extends BaseAutoTunerSuite { 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 and rounds up to the 1000 rung. + // 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=1000", - "--conf spark.sql.adaptive.coalescePartitions.initialPartitionNum=1000"), + 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. - val applied = comments.map(_.comment).filter(_.contains("lowered from 8000 to 1000")) + // 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 lands on the 500 rung. - val userConfigs = ToolTestUtils.buildTuningConfigs( - qualification = List( + // 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 autoTuner = - buildDownwardPassAutoTuner(cpuStageInput(1000L * QUAL_GiB), Some(userConfigs)) val (properties, comments) = autoTuner.getRecommendedProperties(showOnlyUpdatedProps = QualificationAutoTunerRunner.filterByUpdatedPropsEnabled) val autoTunerOutput = Profiler.getAutoTunerResultsAsString(properties, comments) - assertExpectedLinesExist(Seq("--conf spark.sql.shuffle.partitions=500"), autoTunerOutput) - assert(comments.map(_.comment).exists(_.contains("input size factor 0.4"))) + 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") { From 49469707254014cda64c871cfcbcc5c1d8d0dce8 Mon Sep 17 00:00:00 2001 From: Partho Sarthi Date: Thu, 6 Aug 2026 20:38:45 -0700 Subject: [PATCH 07/10] Stop the downward shuffle pass on failed or OOM applications Adopts the two application-wide gates the reference implementation of the v2 heuristic carries and this pass did not: a run that failed a stage, or that hit an out-of-memory failure, is not evidence to size a global partition reduction from, even when the consumer stages it analyzed look healthy on their own. The failed-stage signal rides on the shuffle-stage input analysis so both tools see it. The OOM signal is only available on the profiling provider, so it is an overridable hook that reports false elsewhere rather than a base-class call that would not compile. Also stops the pass introducing an AQE partition property the application never carried. In practice the AQE pass has already recommended one by the time this pass runs, so the guard is inert; skipping the write when a value does exist would leave the higher normal value in place and defeat the reduction. Signed-off-by: Partho Sarthi Co-Authored-By: Claude Opus 5 (1M context) --- .../analysis/ShuffleStageInputAnalyzer.scala | 3 +- .../profiling/ShuffleStageInputMetrics.scala | 7 +++- .../spark/rapids/tool/tuning/AutoTuner.scala | 33 +++++++++++++++-- .../DownwardShufflePartitionsPolicy.scala | 10 ++++++ .../tool/tuning/BaseAutoTunerSuite.scala | 6 ++-- .../tuning/ProfilingAutoTunerSuiteV2.scala | 36 ++++++++++++++++++- 6 files changed, 87 insertions(+), 8 deletions(-) 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 index 38bfb980c..827afc2e0 100644 --- 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 @@ -82,7 +82,8 @@ class ShuffleStageInputAnalyzer(app: AppBase) extends Logging { reasons += ShuffleStageInputIncompleteReason.IncompleteSqlExecution(planModel.id) } } - ShuffleStageInputAnalysis(records.toSeq, reasons.toSeq, provenance) + ShuffleStageInputAnalysis(records.toSeq, reasons.toSeq, provenance, + appHasFailedStage = app.stageManager.getFailedStages.nonEmpty) } /** 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 index bbdc422bf..09bd925e3 100644 --- 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 @@ -109,12 +109,17 @@ object ShuffleStageInputIncompleteReason { * 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) { + analyzed: Boolean = true, + appHasFailedStage: Boolean = false) { def isComplete: Boolean = analyzed && incompleteReasons.isEmpty 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 62a28372e..a0212586d 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 @@ -263,6 +263,12 @@ abstract class AutoTuner( 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 @@ -1843,7 +1849,12 @@ abstract class AutoTuner( * active AQE partition property must move with it. */ private def requiredPartitionProperties: Seq[String] = { - applyToAllPartitionProperties[String](identity) + // '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 } /** @@ -1866,7 +1877,18 @@ abstract class AutoTuner( None } } - // 3. Every affected consumer stage must be free of skew and spill. + // 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 => @@ -1875,7 +1897,7 @@ abstract class AutoTuner( DownwardShuffleSkipReason.StagePressure(record.stageId, "shuffle read skew") } } - // 4. Every property this decision must write has to be writable. + // 5. Every property this decision must write has to be writable. .orElse { requiredPartitionProperties.collectFirst { case property if ignoreRecommendation(property) || !isCalculationEnabled(property) => @@ -2426,6 +2448,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 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 index 24afaabf1..d17271da8 100644 --- 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 @@ -237,6 +237,16 @@ object DownwardShuffleSkipReason { s"candidate $candidate is not at least ${minFactor}x smaller than the current" + s" 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") 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 39f005249..565102eca 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 @@ -158,13 +158,15 @@ abstract class BaseAutoTunerSuite extends AnyFunSuite with BeforeAndAfterEach provenance: ShuffleInputProvenance = ShuffleInputProvenance.Measured, numBranches: Int = 1, hasPositiveSpill: Boolean = false, - hasSkew: Boolean = false): ShuffleStageInputAnalysis = { + 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) + ShuffleStageInputAnalysis(records, Seq.empty, provenance, + appHasFailedStage = appHasFailedStage) } /** Builds an analysis that ran but found a gap, which must keep the normal recommendation. */ 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 876c4d4dc..facd80e9b 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 @@ -2525,6 +2525,8 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { 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, @@ -2539,10 +2541,11 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { jvmGCFractions = Seq(0.1, 0.1), propsFromLog = sourceProps, sparkVersion = Some(testSparkVersion), - meanInput = 1.0E9, + meanInput = meanInputOverride.getOrElse(1.0E9), meanShuffleRead = 1.0E9, shuffleStagesWithPosSpilling = shuffleStagesWithPosSpilling, gpuShuffleStagesWithContainerOom = gpuShuffleStagesWithContainerOom, + scanStagesWithGpuOom = scanStagesWithGpuOom, maxColumnarExchangeDataSizeBytes = maxColumnarExchangeDataSizeBytes, shuffleStageInputAnalysis = shuffleStageInputAnalysis) val platform = PlatformFactory.createInstance(platformName) @@ -2789,4 +2792,35 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { 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") + } } From 7fdf6a7d664066aea5bb28b5f001c80a5880b346 Mon Sep 17 00:00:00 2001 From: Partho Sarthi Date: Tue, 25 Aug 2026 00:33:51 -0700 Subject: [PATCH 08/10] Fix scalastyle line lengths and refresh copyright years CI flagged five lines over the 100-character limit and four files whose copyright year had not been refreshed after this branch modified them. Signed-off-by: Partho Sarthi Co-Authored-By: Claude Opus 5 (1M context) --- .../spark/rapids/tool/analysis/AppSQLPlanAnalyzer.scala | 2 +- .../rapids/tool/analysis/ShuffleStageInputAnalyzer.scala | 3 ++- .../rapids/tool/tuning/QualificationAutoTunerRunner.scala | 2 +- .../com/nvidia/spark/rapids/tool/tuning/TunerContext.scala | 2 +- .../apache/spark/sql/rapids/tool/store/AccumManager.scala | 2 +- .../spark/rapids/tool/tuning/BaseAutoTunerSuite.scala | 3 ++- .../rapids/tool/tuning/ProfilingAutoTunerSuiteV2.scala | 3 ++- .../rapids/tool/tuning/QualificationAutoTunerSuite.scala | 6 ++++-- 8 files changed, 14 insertions(+), 9 deletions(-) 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 4afd726e9..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. 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 index 827afc2e0..0e8da6ba0 100644 --- 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 @@ -286,7 +286,8 @@ class ShuffleStageInputAnalyzer(app: AppBase) extends Logging { private def resolveExchangeDataSize( node: SparkPlanGraphNode, producerStages: Set[Int]): Option[Long] = { node.metrics.find(_.name == DATA_SIZE_METRIC).flatMap { metric => - val fromStages = app.accumManager.accumInfoMap.get(metric.accumulatorId).flatMap { accumInfo => + 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. 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 b17670657..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. 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 aa2b650c7..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. 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 c1874ebc2..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. 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 f2751f977..a1aef6032 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 @@ -150,7 +150,8 @@ abstract class BaseAutoTunerSuite extends AnyFunSuite with BeforeAndAfterEach pySparkMemoryEvidence: Seq[PySparkMemoryEvidence] = Seq.empty, hasSqlCache: Boolean = false, shuffleStageInputAnalysis: ShuffleStageInputAnalysis = - ShuffleStageInputAnalysis.empty(ShuffleInputProvenance.Measured)): AppInfoProviderMockTest = { + ShuffleStageInputAnalysis.empty(ShuffleInputProvenance.Measured) + ): AppInfoProviderMockTest = { new AppInfoProviderMockTest(maxInput, spilledMetrics, jvmGCFractions, propsFromLog, sparkVersion, rapidsJars, distinctLocationPct, redundantReadSize, meanInput, meanShuffleRead, shuffleStagesWithPosSpilling, shuffleSkewStages, scanStagesWithGpuOom, 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 1c2359168..e50f0af51 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 @@ -2732,7 +2732,8 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { } } - test("Downward pass keeps the normal recommendation and stays quiet with no recommended cluster") { + 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")) 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 b48e7faa5..c76e84455 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 @@ -2365,7 +2365,8 @@ class QualificationAutoTunerSuite extends BaseAutoTunerSuite { // 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.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")) @@ -2381,7 +2382,8 @@ class QualificationAutoTunerSuite extends BaseAutoTunerSuite { 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.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")) From d2f9c3c58f421cccea8b8efc9328f83920d2d6f5 Mon Sep 17 00:00:00 2001 From: Partho Sarthi Date: Mon, 31 Aug 2026 17:27:48 -0700 Subject: [PATCH 09/10] Address review: drop two sizing configs and widen the spill gate Removes DOWNWARD_SHUFFLE_SLOT_BASIS. The slot count is now always the recommended executor count times coresPerExecutor. GPU task concurrency is auto-tuned by recent plugin versions, and sizing a wave against it would badly under-use a cluster running mixed CPU and GPU stages, so the alternative the entry existed to measure is not one we want to select at runtime. Measuring it on paired runs now needs a code change. Removes DOWNWARD_SHUFFLE_MIN_REDUCTION_FACTOR. It already defaulted to 1.0, so this changes no behaviour. At 1.0 it is inert, and any other value is an arbitrary discontinuity with no principled way to set it; wave quantization is the only size gate worth having. The tests that covered the threshold now pin the opposite property, so a size floor cannot reappear unnoticed. Fixes the spill gate. One filter fed both the byte totals and the spill evidence, so a task that spilled and then failed was invisible to the gate that exists to see exactly that. Totals still require a successful task; the gate now only excludes speculative duplicates. The predicates moved onto the companion object so the asymmetry between them is directly testable. The shuffle-stage input analysis is no longer built when the pass is disabled. It walks every SQL plan, and the pass ships off, so that was the common path paying for an analysis nobody read. Also clamps each partition property against its own current value rather than the effective maximum across them. This is an invariant guard, not a fix: recommendAQEProperties levels the two before this pass runs, so the clamp is inert today and only matters if that ordering changes. The applied comment renders per-property values so it cannot misreport if it ever does. BranchTotals.add saturates instead of wrapping, because a wrapped total becomes small and would understate the requirement. Signed-off-by: Partho Sarthi Co-Authored-By: Claude Opus 5 --- .../resources/bootstrap/tuningConfigs.yaml | 19 ---- .../analysis/ShuffleStageInputAnalyzer.scala | 49 +++++++++-- .../spark/rapids/tool/tuning/AutoTuner.scala | 49 +++++++---- .../DownwardShufflePartitionsPolicy.scala | 85 ++++-------------- .../ShuffleStageInputMetricsSuite.scala | 58 +++++++++++++ .../DownwardShufflePartitionsSuite.scala | 86 ++++++++----------- .../tuning/ProfilingAutoTunerSuiteV2.scala | 31 +++---- 7 files changed, 199 insertions(+), 178 deletions(-) diff --git a/core/src/main/resources/bootstrap/tuningConfigs.yaml b/core/src/main/resources/bootstrap/tuningConfigs.yaml index 5672680d9..9c6bd09ec 100644 --- a/core/src/main/resources/bootstrap/tuningConfigs.yaml +++ b/core/src/main/resources/bootstrap/tuningConfigs.yaml @@ -244,25 +244,6 @@ default: default: 1.0 usedBy: spark.sql.shuffle.partitions, spark.sql.adaptive.coalescePartitions.initialPartitionNum - - name: DOWNWARD_SHUFFLE_SLOT_BASIS - description: >- - Unit used to size one execution wave of the recommended cluster. The slot count is the - recommended executor count multiplied by this per-executor value, and the downward candidate - is always rounded up to a whole multiple of it. 'cores' uses the recommended - spark.executor.cores; 'concurrentGpuTasks' uses the recommended - spark.rapids.sql.concurrentGpuTasks. - default: cores - usedBy: spark.sql.shuffle.partitions, spark.sql.adaptive.coalescePartitions.initialPartitionNum - - - name: DOWNWARD_SHUFFLE_MIN_REDUCTION_FACTOR - description: >- - Minimum ratio between the normal recommendation and the downward candidate required before - a reduction is applied. A value of 2.0 means the normal value must be at least twice the - candidate. The default of 1.0 imposes no threshold, because wave quantization already - prevents trivial reductions on realistic clusters. - 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 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 index 0e8da6ba0..443581bbc 100644 --- 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 @@ -24,7 +24,7 @@ import com.nvidia.spark.rapids.tool.profiling.{ShuffleInputProvenance, ShuffleSt 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 +import org.apache.spark.sql.rapids.tool.store.{StageModel, TaskModel} /** * Builds the raw consumer-stage shuffle input inventory used by the downward shuffle-partition @@ -48,7 +48,13 @@ class ShuffleStageInputAnalyzer(app: AppBase) extends Logging { /** Accumulated bytes and branch count for one (SQL execution, consumer stage). */ private case class BranchTotals(bytes: Long, branches: Int) { - def add(moreBytes: Long): BranchTotals = BranchTotals(bytes + moreBytes, branches + 1) + // 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 = { @@ -61,6 +67,11 @@ class ShuffleStageInputAnalyzer(app: AppBase) extends Logging { * 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: 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) { @@ -247,6 +258,10 @@ class ShuffleStageInputAnalyzer(app: AppBase) extends Logging { * 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: 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( @@ -317,11 +332,10 @@ class ShuffleStageInputAnalyzer(app: AppBase) extends Logging { stageModel: StageModel, totals: BranchTotals): ShuffleStageInputRecord = { val attemptId = stageModel.getAttemptId - // Speculative and failed task metrics do not describe the work the recommendation governs. - val tasks = app.taskManager - .getTasks(stageId, attemptId, Some(t => t.successful && !t.speculative)) - .toSeq - val hasTaskSpill = tasks.exists(t => t.memoryBytesSpilled > 0L || t.diskBytesSpilled > 0L) + 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, @@ -366,6 +380,27 @@ object ShuffleStageInputAnalyzer { /** 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. 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 3f6760ea9..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 @@ -2077,13 +2077,17 @@ abstract class AutoTuner( */ private def recommendDownwardShufflePartitions(): Unit = { val configResult = DownwardShufflePolicyConfig.fromProvider(configProvider) - val slotCount = configResult.toOption.filter(_.enabled) - .flatMap(config => downwardShuffleSlotCount(config.slotBasis)) + 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, - appInfoProvider.getShuffleStageInputAnalysis) + analysis) decision match { case DownwardShuffleDecision.InvalidConfig(errors) => // Fail closed as one decision: no property is touched when any policy input is invalid. @@ -2110,20 +2114,18 @@ abstract class AutoTuner( * @return the slot count, or None when the recommended cluster shape or the per-executor * multiplier cannot be resolved */ - private def downwardShuffleSlotCount(basis: DownwardShuffleSlotBasis): Option[Int] = { + private def downwardShuffleSlotCount: Option[Int] = { platform.recommendedClusterInfo.flatMap { clusterInfo => val executors = recommendedIntValue("spark.executor.instances").getOrElse(clusterInfo.numExecutors) - val slotsPerExecutor = basis match { - case DownwardShuffleSlotBasis.Cores => Some(clusterInfo.coresPerExecutor) - // Concurrent GPU tasks is a recommendation rather than a field on the cluster record, and - // it is suppressed entirely on platforms whose plugin auto-tunes it. - case DownwardShuffleSlotBasis.ConcurrentGpuTasks => - recommendedIntValue("spark.rapids.sql.concurrentGpuTasks") - } - slotsPerExecutor.filter(_ > 0).filter(_ => executors > 0).flatMap { perExecutor => - val slots = executors.toLong * perExecutor.toLong + // 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 } } } @@ -2134,6 +2136,14 @@ abstract class AutoTuner( .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. @@ -2143,10 +2153,17 @@ abstract class AutoTuner( downwardShuffleBlockingReason(applied) match { case Some(reason) => reportDownwardShuffleSkip(reason) case None => - val partitionProperties = requiredPartitionProperties - partitionProperties.foreach(appendRecommendation(_, applied.selectedValue)) + // 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(partitionProperties, applied)) + DownwardShufflePartitionsPolicy.appliedComment(updates, applied)) } } 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 index d17271da8..207f0afa8 100644 --- 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 @@ -24,30 +24,6 @@ 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 -/** - * Basis used to turn the recommended executor count into a cluster task-slot count. - * - * A "slot" is one unit of concurrency the recommended cluster actually has, so the number of slots - * is the number of partitions one execution wave can run. Which unit is right depends on whether - * the CPU task parallelism or the GPU task concurrency is the binding constraint, so the choice is - * exposed as a configuration entry rather than hard-coded. - */ -sealed abstract class DownwardShuffleSlotBasis(val label: String) - -object DownwardShuffleSlotBasis { - /** Slots per executor are the recommended `spark.executor.cores`. */ - case object Cores extends DownwardShuffleSlotBasis("cores") - - /** Slots per executor are the recommended `spark.rapids.sql.concurrentGpuTasks`. */ - case object ConcurrentGpuTasks extends DownwardShuffleSlotBasis("concurrentGpuTasks") - - val values: Seq[DownwardShuffleSlotBasis] = Seq(Cores, ConcurrentGpuTasks) - - def parse(raw: String): Option[DownwardShuffleSlotBasis] = { - values.find(_.label.equalsIgnoreCase(raw.trim)) - } -} - /** * Validated policy inputs of the downward-only shuffle partition pass. * @@ -58,30 +34,22 @@ object DownwardShuffleSlotBasis { * @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 - * @param slotBasis unit used to size one execution wave of the recommended cluster - * @param minReductionFactor required ratio of normal value to candidate before applying */ case class DownwardShufflePolicyConfig( enabled: Boolean, targetPartitionSizeBytes: Long, - inputSizeFactor: Double, - slotBasis: DownwardShuffleSlotBasis, - minReductionFactor: Double) + 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" - val SLOT_BASIS_KEY = "DOWNWARD_SHUFFLE_SLOT_BASIS" - val MIN_REDUCTION_FACTOR_KEY = "DOWNWARD_SHUFFLE_MIN_REDUCTION_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, - slotBasis = DownwardShuffleSlotBasis.Cores, - minReductionFactor = 0.0) + inputSizeFactor = 0.0) /** * Reads and validates every policy entry from the tuning-config provider. @@ -105,21 +73,16 @@ object DownwardShufflePolicyConfig { 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 slotBasis = parseSlotBasis(configProvider, SLOT_BASIS_KEY) - val minReduction = - parseDouble(configProvider, MIN_REDUCTION_FACTOR_KEY, min = 1.0, minInclusive = true) - val errors = Seq(targetSize, factor, slotBasis, minReduction).collect { + val errors = Seq(targetSize, factor).collect { case Left(err) => err } - (targetSize, factor, slotBasis, minReduction) match { - case (Right(size), Right(f), Right(b), Right(r)) if errors.isEmpty => + (targetSize, factor) match { + case (Right(size), Right(f)) if errors.isEmpty => Right(DownwardShufflePolicyConfig( enabled = true, targetPartitionSizeBytes = size, - inputSizeFactor = f, - slotBasis = b, - minReductionFactor = r)) + inputSizeFactor = f)) case _ => Left(errors) } } @@ -149,16 +112,6 @@ object DownwardShufflePolicyConfig { } } - private def parseSlotBasis( - configProvider: TuningConfigProvider, - key: String): Either[String, DownwardShuffleSlotBasis] = { - rawValue(configProvider, key).flatMap { raw => - DownwardShuffleSlotBasis.parse(raw).toRight( - s"'$key' must be one of ${DownwardShuffleSlotBasis.values.map(_.label).mkString(", ")}" + - s" but was '$raw'") - } - } - private def parseDouble( configProvider: TuningConfigProvider, key: String, @@ -232,11 +185,6 @@ object DownwardShuffleSkipReason { extends DownwardShuffleSkipReason( s"candidate $candidate does not lower the current recommendation $normalValue") - case class BelowReductionThreshold(candidate: Int, normalValue: Int, minFactor: Double) - extends DownwardShuffleSkipReason( - s"candidate $candidate is not at least ${minFactor}x smaller than the current" + - s" 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. @@ -366,11 +314,6 @@ object DownwardShufflePartitionsPolicy { // every candidate is at least the slot count. DownwardShuffleDecision.Skipped( DownwardShuffleSkipReason.NotDownward(candidate, normalValue)) - case Some(candidate) - if normalValue.toDouble < config.minReductionFactor * candidate.toDouble => - DownwardShuffleDecision.Skipped( - DownwardShuffleSkipReason.BelowReductionThreshold( - candidate, normalValue, config.minReductionFactor)) case Some(candidate) => DownwardShuffleDecision.Applied( normalValue = normalValue, @@ -472,11 +415,21 @@ object DownwardShufflePartitionsPolicy { * decision so the recommendation can be audited without re-running the tool. */ def appliedComment( - partitionProperties: Seq[String], + partitionUpdates: Seq[(String, Int)], decision: DownwardShuffleDecision.Applied): String = { val record = decision.determiningRecord - s"${partitionProperties.map(p => s"'$p'").mkString(" and ")} lowered from " + - s"${decision.normalValue} to ${decision.selectedValue} based on the " + + // 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 " + 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 index 82dff20d8..130b6f2b6 100644 --- 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 @@ -28,6 +28,7 @@ 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. @@ -242,6 +243,63 @@ class ShuffleStageInputMetricsSuite extends AnyFunSuite with Logging { 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. 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 index 46e445ca7..94d42f6dd 100644 --- 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 @@ -89,9 +89,7 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { private val baseConfig = DownwardShufflePolicyConfig( enabled = true, targetPartitionSizeBytes = GiB, - inputSizeFactor = 1.0, - slotBasis = DownwardShuffleSlotBasis.Cores, - minReductionFactor = 1.0) + inputSizeFactor = 1.0) /** Slot count of the cluster the arithmetic tests recommend for: 125 executors x 16 cores. */ private val slots = 2000 @@ -134,13 +132,11 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { assert(config.enabled) assert(config.targetPartitionSizeBytes == GiB) assert(config.inputSizeFactor == 1.0) - assert(config.slotBasis == DownwardShuffleSlotBasis.Cores) - // Wave quantization already prevents trivial reductions, so no extra threshold is imposed. - assert(config.minReductionFactor == 1.0) } - test("the retired rung entries are gone and their removal does not break config loading") { - Seq("DOWNWARD_SHUFFLE_PARTITION_FLOOR", "DOWNWARD_SHUFFLE_RUNG_MULTIPLIER").foreach { key => + 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") } @@ -155,17 +151,14 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { assert(config.inputSizeFactor == 0.8) // Everything else still comes from the shared defaults. assert(config.targetPartitionSizeBytes == GiB) - assert(config.slotBasis == DownwardShuffleSlotBasis.Cores) } 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"), - TuningConfigEntry(name = "DOWNWARD_SHUFFLE_SLOT_BASIS", default = "concurrentGpuTasks")))) + TuningConfigEntry(name = "DOWNWARD_SHUFFLE_TARGET_PARTITION_SIZE", default = "512m")))) assert(fromDefault == Right(baseConfig.copy( - targetPartitionSizeBytes = 512L * 1024L * 1024L, - slotBasis = DownwardShuffleSlotBasis.ConcurrentGpuTasks))) + targetPartitionSizeBytes = 512L * 1024L * 1024L))) val fromToolSection = DownwardShufflePolicyConfig.fromProvider(qualProvider( qualification = List( @@ -173,27 +166,11 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { assert(fromToolSection == Right(baseConfig.copy(inputSizeFactor = 0.5))) } - test("the slot basis accepts only a known label, case-insensitively") { - Seq("cores" -> DownwardShuffleSlotBasis.Cores, - "CORES" -> DownwardShuffleSlotBasis.Cores, - " concurrentgputasks " -> DownwardShuffleSlotBasis.ConcurrentGpuTasks - ).foreach { case (raw, expected) => - val result = DownwardShufflePolicyConfig.fromProvider(profProvider( - default = List(TuningConfigEntry(name = "DOWNWARD_SHUFFLE_SLOT_BASIS", default = raw)))) - assert(result.exists(_.slotBasis == expected), s"'$raw' should parse as $expected") - } - Seq("gpus", "tasks", "1").foreach { raw => - val result = DownwardShufflePolicyConfig.fromProvider(profProvider( - default = List(TuningConfigEntry(name = "DOWNWARD_SHUFFLE_SLOT_BASIS", default = raw)))) - assert(result.isLeft, s"'$raw' should not parse as a slot basis") - } - } - 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_SLOT_BASIS", default = "not-a-basis"))) + 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), @@ -216,9 +193,7 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { 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"), - // A reduction factor below 1.0 would allow raising the recommendation. - "DOWNWARD_SHUFFLE_MIN_REDUCTION_FACTOR" -> Seq("0.9", "0", "abc")) + "DOWNWARD_SHUFFLE_INPUT_SIZE_FACTOR" -> Seq("0", "-0.5", "abc")) invalidByKey.foreach { case (key, values) => values.foreach { value => @@ -232,8 +207,7 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { } test("NaN and infinite values are rejected") { - Seq("DOWNWARD_SHUFFLE_INPUT_SIZE_FACTOR", - "DOWNWARD_SHUFFLE_MIN_REDUCTION_FACTOR").foreach { key => + 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)))) @@ -244,8 +218,8 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { test("all configuration errors are reported as one fail-closed decision") { val provider = profProvider(default = List( - TuningConfigEntry(name = "DOWNWARD_SHUFFLE_SLOT_BASIS", default = "not-a-basis"), - TuningConfigEntry(name = "DOWNWARD_SHUFFLE_MIN_REDUCTION_FACTOR", default = "0.5"))) + 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) @@ -434,24 +408,22 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { DownwardShuffleDecision.Skipped(DownwardShuffleSkipReason.NotDownward(2000, 200))) } - test("the default reduction factor imposes no threshold but an override still blocks") { - // A 3000 -> 2000 reduction is not 2x, so v1's default would have blocked it. + 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") } - val strictConfig = baseConfig.copy(minReductionFactor = 2.0) - assert(decide(3000, Seq(record(totalBytes = 900L * GiB)), strictConfig) == - DownwardShuffleDecision.Skipped( - DownwardShuffleSkipReason.BelowReductionThreshold(2000, 3000, 2.0))) - // Exactly 2x still clears the explicit override. - assert(decide(4000, Seq(record(totalBytes = 900L * GiB)), strictConfig) - .isInstanceOf[DownwardShuffleDecision.Applied]) + // 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, minReductionFactor = 2.0)) + 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) @@ -530,8 +502,8 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { 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", - "spark.sql.adaptive.coalescePartitions.initialPartitionNum"), applied) + 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", @@ -541,4 +513,20 @@ class DownwardShufflePartitionsSuite extends AnyFunSuite { } 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/ProfilingAutoTunerSuiteV2.scala b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuiteV2.scala index a87091784..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 @@ -2855,7 +2855,7 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { test("Downward pass warns once and changes nothing when its configuration is invalid") { val (properties, comments) = runDownwardPass(worstStageNeeding900Partitions, extraTuningConfigs = List( - TuningConfigEntry(name = "DOWNWARD_SHUFFLE_SLOT_BASIS", default = "0"))) + 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") @@ -2869,21 +2869,14 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { assert(!comments.exists(c => c.contains("lowered from") || c.contains("downward"))) } - test("Downward pass respects a configured minimum reduction factor") { - // 1500 partitions rounds up to 4 waves (1600), which is not a 2x reduction from 3000. + 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")), - extraTuningConfigs = List( - TuningConfigEntry(name = "DOWNWARD_SHUFFLE_MIN_REDUCTION_FACTOR", default = "2.0"))) - assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).forall(_ == "3000")) - assert(!comments.exists(_.contains("lowered from"))) - // The default of 1.0 imposes no threshold, so the same reduction applies without the override. - val (defaultProps, defaultComments) = runDownwardPass( completeShuffleStageInputs(Seq(4 -> 1500L * GiB)), sourceProps = downwardPassSourceProps(Map(SHUFFLE_PARTITIONS_KEY -> "3000"))) - assert(recommendedValue(defaultProps, SHUFFLE_PARTITIONS_KEY).contains("1600")) - assert(defaultComments.exists(_.contains("lowered from 3000 to 1600"))) + 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") { @@ -2895,20 +2888,16 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { assert(comments.exists(_.contains("1 execution wave(s)"))) } - test("Downward pass slot count follows the configured slot basis") { - // The GPU-concurrency basis uses the recommended concurrentGpuTasks instead of the cores. - val (properties, comments) = runDownwardPass(worstStageNeeding900Partitions, - extraTuningConfigs = List(TuningConfigEntry(name = "DOWNWARD_SHUFFLE_SLOT_BASIS", - default = "concurrentGpuTasks"))) + 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") - val expected = math.ceil(900.0 / gpuSlots).toInt * gpuSlots - assert(recommendedValue(properties, SHUFFLE_PARTITIONS_KEY).contains(expected.toString)) - assert(comments.exists(_.contains(s"$gpuSlots cluster task slots"))) + 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") { From 2dbb1d45f20a5bf3d04a2b5cae89eef7926fdfd4 Mon Sep 17 00:00:00 2001 From: Partho Sarthi Date: Tue, 1 Sep 2026 16:02:45 -0700 Subject: [PATCH 10/10] Record the enable criterion and link the deferred analyzer TODOs DOWNWARD_SHUFFLE_ENABLED's description now says what has to be settled before it defaults to true. It is 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. That reconciliation is the gate, and leaving it as tribal knowledge invites the flag being flipped for the wrong reason. The two analyzer TODOs now reference #2133, which describes both: sizing a CPU exchange inside a GPU run, blocked on the input-size factor being application-level, and apportioning an AQE-split exchange by per-stage shuffle read metrics rather than duplicating it. Both are conservative in the safe direction, so neither blocks shipping the pass enabled, which is why they are tracked apart from #2128. Signed-off-by: Partho Sarthi Co-Authored-By: Claude Opus 5 --- core/src/main/resources/bootstrap/tuningConfigs.yaml | 5 ++++- .../tool/analysis/ShuffleStageInputAnalyzer.scala | 12 ++++++------ 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/core/src/main/resources/bootstrap/tuningConfigs.yaml b/core/src/main/resources/bootstrap/tuningConfigs.yaml index 9c6bd09ec..1e22b2ceb 100644 --- a/core/src/main/resources/bootstrap/tuningConfigs.yaml +++ b/core/src/main/resources/bootstrap/tuningConfigs.yaml @@ -225,7 +225,10 @@ default: 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; opt in per run while the wave-based sizing is being evaluated. + 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 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 index 443581bbc..7416e01ab 100644 --- 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 @@ -68,10 +68,10 @@ class ShuffleStageInputAnalyzer(app: AppBase) extends Logging { * 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: 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. + * 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) { @@ -258,8 +258,8 @@ class ShuffleStageInputAnalyzer(app: AppBase) extends Logging { * 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: 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 + * 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