[WIP][VL] Support setting stage execution mode#12511
Open
marin-ma wants to merge 5 commits into
Open
Conversation
|
Run Gluten Clickhouse CI on x86 |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR introduces an explicit stage execution mode concept (CPU vs GPU) and propagates it across shuffle boundaries so that mapper-side and reducer-side stages can independently select the appropriate shuffle writer/reader and batch-resizing behavior.
Changes:
- Add
StageExecutionModeand propagate it through shuffle read paths (ShuffledColumnarBatchRDD,ColumnarShuffleManager,GlutenShuffleUtils). - Extend
ColumnarShuffleExchangeExecto carry mapper/reducer stage modes and adjust shuffle writer selection accordingly. - Add AQE wrapper
ColumnarAQEShuffleReadExecplus Velox rules (AdjustStageExecutionMode, updates toCudfNodeValidationRule) to set/propagate stage modes.
Reviewed changes
Copilot reviewed 26 out of 26 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| gluten-ut/spark41/src/test/scala/org/apache/spark/sql/GlutenPlanStabilitySuite.scala | Update exchange matching to reflect ColumnarShuffleExchangeExec type. |
| gluten-ut/spark41/src/test/scala/org/apache/spark/sql/GlutenDataFrameWindowFunctionsSuite.scala | Update pattern match for expanded ColumnarShuffleExchangeExec constructor. |
| gluten-ut/spark40/src/test/scala/org/apache/spark/sql/GlutenPlanStabilitySuite.scala | Same as Spark 4.1 suite updates. |
| gluten-ut/spark40/src/test/scala/org/apache/spark/sql/GlutenDataFrameWindowFunctionsSuite.scala | Same as Spark 4.1 suite updates. |
| gluten-ut/spark35/src/test/scala/org/apache/spark/sql/GlutenDataFrameWindowFunctionsSuite.scala | Same pattern match update for Spark 3.5 suite. |
| gluten-ut/spark34/src/test/scala/org/apache/spark/sql/GlutenDataFrameWindowFunctionsSuite.scala | Same pattern match update for Spark 3.4 suite. |
| gluten-ut/spark33/src/test/scala/org/apache/spark/sql/GlutenDataFrameWindowFunctionsSuite.scala | Same pattern match update for Spark 3.3 suite. |
| gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ShuffledColumnarBatchRDD.scala | Add stage execution mode parameter; route to shuffle manager reader with mode. |
| gluten-substrait/src/main/scala/org/apache/spark/sql/execution/GPUColumnarShuffleExchangeExec.scala | Removed legacy GPU-only shuffle exchange node. |
| gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarShuffleExchangeExecBase.scala | Removed base class; logic moved into ColumnarShuffleExchangeExec. |
| gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarShuffleExchangeExec.scala | Add mapper/reducer stage modes, node naming, and writer selection changes. |
| gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala | New AQE shuffle-read wrapper that carries execution mode and forwards to columnar shuffle RDD. |
| gluten-substrait/src/main/scala/org/apache/spark/shuffle/sort/ColumnarShuffleManager.scala | Add getReader(..., executionMode) overload and default CPU behavior. |
| gluten-substrait/src/main/scala/org/apache/spark/shuffle/GlutenShuffleUtils.scala | Pass execution mode through shuffle reader generation; refactor codec validation. |
| gluten-substrait/src/main/scala/org/apache/spark/shuffle/GlutenShuffleReaderWrapper.scala | Extend shuffle-reader parameter bundle with executionMode. |
| gluten-substrait/src/main/scala/org/apache/gluten/execution/StageExecutionMode.scala | New execution-mode abstraction (CPU/GPU + testing mode). |
| gluten-substrait/src/main/scala/org/apache/gluten/execution/ColumnarCollectTailBaseExec.scala | Update shuffle RDD construction for new ShuffledColumnarBatchRDD signature (currently hard-coded CPU). |
| gluten-substrait/src/main/scala/org/apache/gluten/execution/ColumnarCollectLimitBaseExec.scala | Same as tail base exec; currently hard-coded CPU. |
| backends-velox/src/test/scala/org/apache/gluten/execution/MiscOperatorSuite.scala | Adjust assertions to match new VeloxResizeBatchesExec configuration model. |
| backends-velox/src/main/scala/org/apache/spark/sql/execution/AdjustStageExecutionMode.scala | New AQE rule to wrap shuffle reads / tag exchanges and resizers with stage execution mode. |
| backends-velox/src/main/scala/org/apache/gluten/extension/GpuBufferBatchResizeForShuffleInputOutput.scala | Removed older rule in favor of stage-mode-driven resizing logic. |
| backends-velox/src/main/scala/org/apache/gluten/extension/CudfNodeValidationRule.scala | Replace GPU-exchange construction with tagging + stage-mode propagation when AQE is off. |
| backends-velox/src/main/scala/org/apache/gluten/extension/AppendBatchResizeForShuffleInputAndOutput.scala | Simplify resizer insertion now that VeloxResizeBatchesExec reads config internally. |
| backends-velox/src/main/scala/org/apache/gluten/execution/VeloxResizeBatchesExec.scala | Switch resizer implementation based on stage execution mode (GPU vs CPU path). |
| backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala | Wire in AdjustStageExecutionMode rule. |
| backends-velox/src/main/java/org/apache/gluten/utils/GpuBufferColumnarBatchResizer.java | New JNI-backed GPU buffer batch resizer helper for GPU-stage resizing. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+149
to
+158
| def getShuffleWriterType: ShuffleWriterType = { | ||
| mapperStageMode match { | ||
| case Some(GPUStageMode) => | ||
| GpuHashShuffleWriterType | ||
| case _ => | ||
| BackendsApiManager.getSparkPlanExecApiInstance.getShuffleWriterType( | ||
| outputPartitioning, | ||
| output) | ||
| } | ||
| } |
Comment on lines
104
to
108
| ), | ||
| readMetrics | ||
| readMetrics, | ||
| // FIXME: pass proper StageExecutionMode | ||
| CPUStageMode | ||
| ) |
Comment on lines
103
to
107
| ), | ||
| readMetrics | ||
| readMetrics, | ||
| // FIXME: pass proper StageExecutionMode | ||
| CPUStageMode | ||
| ) |
| case h: HashPartitioning => | ||
| val partitions = partitionSpecs.map { | ||
| case CoalescedPartitionSpec(start, end, _) => CoalescedBoundary(start, end) | ||
| // Can not happend due to isCoalescedRead |
|
Run Gluten Clickhouse CI on x86 |
|
Run Gluten Clickhouse CI on x86 |
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 195 out of 195 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (5)
gluten-substrait/src/main/scala/org/apache/gluten/execution/StageExecutionMode.scala:1
MockGPUStageModeuses the sameidvalue asCPUStageMode(0). This makesStageExecutionMode.fromId(0)ambiguous and will always map back toCPUStageMode, soMockGPUStageModecan never round-trip correctly. AssignMockGPUStageModea unique ID (or remove/move it to test scope if it's only intended for tests) and updatefromIdaccordingly.
gluten-substrait/src/main/scala/org/apache/spark/shuffle/sort/ColumnarShuffleManager.scala:1- The Spark
ShuffleManager#getReaderoverride unconditionally forcesCPUStageMode. This contradicts the PR’s goal of allowing reducer stages to run in GPU mode and select a GPU shuffle reader when appropriate. Instead of hardcodingCPUStageMode, derive the execution mode from shuffle metadata (e.g., encode it in the shuffle handle/dependency) or from a task/stage property, and pass the derived value into the mode-aware overload.
gluten-substrait/src/main/scala/org/apache/gluten/execution/ColumnarCollectLimitBaseExec.scala:1 - Leaving a
FIXMEwith a hardcodedCPUStageModewill likely reintroduce mode mismatches for collect-limit paths once GPU stages are enabled downstream. This should be resolved before merging by plumbing the correctStageExecutionModeinto this call site (e.g., from the surrounding plan/stage context), or by redesigning the API so the mode is inferable without ad-hoc literals.
gluten-substrait/src/main/scala/org/apache/gluten/execution/ColumnarCollectTailBaseExec.scala:1 - Same issue as
ColumnarCollectLimitBaseExec: hardcodingCPUStageModewith aFIXMEmakes stage-mode propagation incomplete and risks inconsistent shuffle reader/resizer selection. Please thread the actualStageExecutionModethrough this path (mandatory if this code can run inside a GPU reducer stage).
gluten-substrait/src/main/scala/org/apache/gluten/execution/StageExecutionMode.scala:1 - Consider making
StageExecutionModeasealed traitto ensure exhaustive matching across the codebase (e.g., infromIdand othermatchsites). This is especially useful now that execution mode directly impacts shuffle reader/writer/resizer selection.
|
Run Gluten Clickhouse CI on x86 |
marin-ma
force-pushed
the
add-stage-execution-mode
branch
from
36f3ded to
2872874
Compare
|
Run Gluten Clickhouse CI on x86 |
|
Run Gluten Clickhouse CI on x86 |
Comment on lines
+24
to
+40
| case object CPUStageMode extends StageExecutionMode { | ||
| override def id: Int = 0 | ||
| } | ||
| case object GPUStageMode extends StageExecutionMode { | ||
| override def id: Int = 1 | ||
| } | ||
| case object MockGPUStageMode extends StageExecutionMode { | ||
| override def id: Int = 0 | ||
| } | ||
|
|
||
| object StageExecutionMode { | ||
| def fromId(id: Int): StageExecutionMode = id match { | ||
| case 0 => CPUStageMode | ||
| case 1 => GPUStageMode | ||
| case _ => throw new IllegalArgumentException(s"Unknown StageExecutionMode id: $id") | ||
| } | ||
| } |
Comment on lines
94
to
108
| new ShuffledColumnarBatchRDD( | ||
| BackendsApiManager.getSparkPlanExecApiInstance.genShuffleDependency( | ||
| locallyLimited, | ||
| child.output, | ||
| child.output, | ||
| SinglePartition, | ||
| serializer, | ||
| writeMetrics, | ||
| metrics, | ||
| shuffleWriterType | ||
| ), | ||
| readMetrics | ||
| readMetrics, | ||
| // FIXME: pass proper StageExecutionMode | ||
| CPUStageMode | ||
| ) |
Comment on lines
154
to
178
| def genColumnarShuffleReader[K, C]( | ||
| handle: ShuffleHandle, | ||
| startMapIndex: Int, | ||
| endMapIndex: Int, | ||
| startPartition: Int, | ||
| endPartition: Int, | ||
| context: TaskContext, | ||
| metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = { | ||
| metrics: ShuffleReadMetricsReporter, | ||
| executionMode: StageExecutionMode): ShuffleReader[K, C] = { | ||
| val (blocksByAddress, canEnableBatchFetch) = { | ||
| getReaderParam(handle, startMapIndex, endMapIndex, startPartition, endPartition) | ||
| } | ||
| val shouldBatchFetch = | ||
| canEnableBatchFetch && canUseBatchFetch(startPartition, endPartition, context) | ||
|
|
||
| BackendsApiManager.getSparkPlanExecApiInstance | ||
| .genColumnarShuffleReader( | ||
| GenShuffleReaderParameters( | ||
| handle.asInstanceOf[BaseShuffleHandle[K, _, C]], | ||
| blocksByAddress, | ||
| context, | ||
| metrics, | ||
| shouldBatchFetch)) | ||
| shouldBatchFetch, | ||
| executionMode)) | ||
| .shuffleReader |
Comment on lines
93
to
107
| new ShuffledColumnarBatchRDD( | ||
| BackendsApiManager.getSparkPlanExecApiInstance.genShuffleDependency( | ||
| locallyLimited, | ||
| child.output, | ||
| child.output, | ||
| SinglePartition, | ||
| serializer, | ||
| writeMetrics, | ||
| metrics, | ||
| shuffleWriterType | ||
| ), | ||
| readMetrics | ||
| readMetrics, | ||
| // FIXME: pass proper StageExecutionMode | ||
| CPUStageMode | ||
| ) |
Comment on lines
+30
to
+40
| case object MockGPUStageMode extends StageExecutionMode { | ||
| override def id: Int = 0 | ||
| } | ||
|
|
||
| object StageExecutionMode { | ||
| def fromId(id: Int): StageExecutionMode = id match { | ||
| case 0 => CPUStageMode | ||
| case 1 => GPUStageMode | ||
| case _ => throw new IllegalArgumentException(s"Unknown StageExecutionMode id: $id") | ||
| } | ||
| } |
Comment on lines
104
to
108
| ), | ||
| readMetrics | ||
| readMetrics, | ||
| // FIXME: pass proper StageExecutionMode | ||
| CPUStageMode | ||
| ) |
Comment on lines
103
to
107
| ), | ||
| readMetrics | ||
| readMetrics, | ||
| // FIXME: pass proper StageExecutionMode | ||
| CPUStageMode | ||
| ) |
|
Run Gluten Clickhouse CI on x86 |
|
Run Gluten Clickhouse CI on x86 |
marin-ma
force-pushed
the
add-stage-execution-mode
branch
from
6a5ae36 to
98974fc
Compare
|
Run Gluten Clickhouse CI on x86 |
2 similar comments
|
Run Gluten Clickhouse CI on x86 |
|
Run Gluten Clickhouse CI on x86 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Background
When GPU execution is enabled, a Spark stage can run entirely on the GPU only when every operator in that stage supports cuDF. If any operator falls back, the whole stage must use the CPU execution path.
However, the mapper and reducer sides of a shuffle are separate stages and may use different execution modes. For example:
Previously, this CPU-to-GPU transition could fail. The
ColumnarShuffleExchangeExecwas configured based on the mapper-side plan and was not marked for GPU execution when the scan stage fell back. The downstream reducer stage could nevertheless require the GPU shuffle reader, resulting in inconsistent shuffle execution modes.Changes
This PR introduces explicit execution modes for Spark stages and determines the mapper and reducer modes independently at each shuffle boundary.
For the example above:
CPUStageModeand use the CPU pipeline and CPU shuffle writer.GPUStageModeand uses the GPU shuffle reader and GPU pipeline.The selected execution modes are propagated through the shuffle exchange so that the shuffle writer, shuffle reader, and batch-resizing logic use the correct implementation for their respective stages.