This is an automated email from the ASF dual-hosted git repository.
marin-ma pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 65dff04b06 [GLUTEN-11524][VL] Support CPU/GPU hybrid task scheduling
(#12588)
65dff04b06 is described below
commit 65dff04b06f50e5cf90aa1eaf205c9c304ffe81f
Author: Rong Ma <[email protected]>
AuthorDate: Fri Aug 14 10:40:30 2026 +0100
[GLUTEN-11524][VL] Support CPU/GPU hybrid task scheduling (#12588)
---
.../gluten/backendsapi/velox/VeloxRuleApi.scala | 4 +-
.../gluten/extension/CudfNodeValidationRule.scala | 8 +-
.../sql/execution/AdjustStageExecutionMode.scala | 59 ++++++---
.../HybridExecutionResourceProfileSuite.scala | 132 +++++++++++++++++++++
.../gluten/execution/StageExecutionModeSuite.scala | 74 ++++++++++--
docs/Configuration.md | 23 ++--
.../org/apache/gluten/config/GlutenConfig.scala | 56 +++++++++
.../TakeOrderedAndProjectExecTransformer.scala | 4 +-
.../gluten/execution/WholeStageTransformer.scala | 9 +-
.../GlutenAutoAdjustStageResourceProfile.scala | 129 ++++++++++++++------
.../adaptive/ColumnarAQEShuffleReadExec.scala | 2 +-
11 files changed, 417 insertions(+), 83 deletions(-)
diff --git
a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala
b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala
index 6f280757cc..4cfa218097 100644
---
a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala
+++
b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala
@@ -151,10 +151,10 @@ object VeloxRuleApi {
injector.injectFinal(c => RemoveGlutenTableCacheColumnarToRow(c.session))
injector.injectFinal(
c => PreventBatchTypeMismatchInTableCache(c.caller.isCache(),
Set(VeloxBatchType)))
- injector.injectFinal(
- c => GlutenAutoAdjustStageResourceProfile(new GlutenConfig(c.sqlConf),
c.session))
injector.injectFinal(
c => AdjustStageExecutionMode(new GlutenConfig(c.sqlConf), c.session,
c.caller.isAqe()))
+ injector.injectFinal(
+ c => GlutenAutoAdjustStageResourceProfile(new GlutenConfig(c.sqlConf),
c.session))
injector.injectFinal(c => GlutenFallbackReporter(new
GlutenConfig(c.sqlConf), c.session))
injector.injectFinal(_ => RemoveFallbackTagRule())
}
diff --git
a/backends-velox/src/main/scala/org/apache/gluten/extension/CudfNodeValidationRule.scala
b/backends-velox/src/main/scala/org/apache/gluten/extension/CudfNodeValidationRule.scala
index 0a25f2f8b5..ebe3bac28e 100644
---
a/backends-velox/src/main/scala/org/apache/gluten/extension/CudfNodeValidationRule.scala
+++
b/backends-velox/src/main/scala/org/apache/gluten/extension/CudfNodeValidationRule.scala
@@ -58,16 +58,16 @@ object CudfNodeValidationRule {
transformer.foreach {
case _: LeafTransformSupport =>
case t: TransformSupport =>
- t.setTagValue(CudfTag.CudfTag, true)
+ t.setTagValue(CudfTag.CudfValidationTag, true)
case _ =>
}
- transformer.setTagValue(CudfTag.CudfTag, true)
+ transformer.setTagValue(CudfTag.CudfValidationTag, true)
}
} else {
- transformer.setTagValue(CudfTag.CudfTag, !hasLeaf)
+ transformer.setTagValue(CudfTag.CudfValidationTag, !hasLeaf)
}
} else {
- transformer.setTagValue(CudfTag.CudfTag, true)
+ transformer.setTagValue(CudfTag.CudfValidationTag, true)
}
}
}
diff --git
a/backends-velox/src/main/scala/org/apache/spark/sql/execution/AdjustStageExecutionMode.scala
b/backends-velox/src/main/scala/org/apache/spark/sql/execution/AdjustStageExecutionMode.scala
index adddc630eb..1f83ea42ce 100644
---
a/backends-velox/src/main/scala/org/apache/spark/sql/execution/AdjustStageExecutionMode.scala
+++
b/backends-velox/src/main/scala/org/apache/spark/sql/execution/AdjustStageExecutionMode.scala
@@ -27,9 +27,9 @@ import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.trees.TreeNodeTag
import
org.apache.spark.sql.execution.AdjustStageExecutionMode.{adjustExecutionMode,
unsetTag}
import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec,
ColumnarAQEShuffleReadExec, ShuffleQueryStageExec}
+import org.apache.spark.sql.execution.joins.BaseJoinExec
import org.apache.spark.util.SparkTestUtil
-// For ShuffleStage, the resource profile is set to
ColumnarShuffleExchangeExec.inputColumnarRDD.
@Experimental
case class AdjustStageExecutionMode(
glutenConf: GlutenConfig,
@@ -51,30 +51,22 @@ case class AdjustStageExecutionMode(
}
private def adjustExecutionModeForGPU(plan: SparkPlan): SparkPlan = {
- val transformers = plan.collect { case t: WholeStageTransformer => t }
- if (transformers.isEmpty) {
- return plan
- }
- if (transformers.size > 1) {
- logWarning(s"Not offloading GPU because multiple WholeStageTransformer
exist. Remove tags.")
- unsetTag(plan, CudfTag.CudfTag)
+ if (!AdjustStageExecutionMode.offloadCuda(plan, glutenConf)) {
+ unsetTag(plan, CudfTag.CudfValidationTag)
return plan
}
- val transformer = transformers.head
- if (transformer.isCudf) {
- val gpuStageMode = if (SparkTestUtil.isTesting) {
- // Only unset for transformer.
- transformer.unsetTagValue(CudfTag.CudfTag)
- MockGPUStageMode
- } else {
- GPUStageMode
+ val gpuStageMode = if (SparkTestUtil.isTesting) {
+ unsetTag(plan, CudfTag.CudfValidationTag)
+ plan.collect { case t: WholeStageTransformer => t }.foreach {
+ _.setTagValue(CudfTag.CudfTestingTag, true)
}
-
- adjustExecutionMode(plan, gpuStageMode)
+ MockGPUStageMode
} else {
- plan
+ GPUStageMode
}
+
+ adjustExecutionMode(plan, gpuStageMode)
}
}
@@ -113,4 +105,33 @@ object AdjustStageExecutionMode extends Logging {
case _ =>
}
}
+
+ def offloadCuda(plan: SparkPlan, glutenConf: GlutenConfig): Boolean = {
+ if (glutenConf.gpuOnlyOffloadJoinStage) {
+ if (!plan.exists(_.isInstanceOf[BaseJoinExec])) {
+ logWarning(s"Not offloading GPU because missing offload condition.")
+ return false
+ }
+ }
+
+ val transformers = plan.collect { case t: WholeStageTransformer => t }
+
+ if (transformers.isEmpty) {
+ logWarning(s"Not offloading GPU because no WholeStageTransformer.")
+ return false
+ }
+
+ if (transformers.size > 1) {
+ // Do not offload GPU if the whole stage is broken down into multiple
native pipelines.
+ logWarning(s"Not offloading GPU because multiple WholeStageTransformer
exist.")
+ return false
+ }
+
+ if (!transformers.head.offloadCuda) {
+ logWarning(s"Not offloading GPU because WholeStageTransformer is not
tagged cudf.")
+ return false
+ }
+
+ true
+ }
}
diff --git
a/backends-velox/src/test/scala/org/apache/gluten/execution/HybridExecutionResourceProfileSuite.scala
b/backends-velox/src/test/scala/org/apache/gluten/execution/HybridExecutionResourceProfileSuite.scala
new file mode 100644
index 0000000000..3c33f29723
--- /dev/null
+++
b/backends-velox/src/test/scala/org/apache/gluten/execution/HybridExecutionResourceProfileSuite.scala
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.gluten.execution
+
+import org.apache.gluten.config.{GlutenConfig, VeloxConfig}
+
+import org.apache.spark.SparkConf
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.sql.execution.{ApplyResourceProfileExec, SparkPlan}
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * Tests for GlutenAutoAdjustStageResourceProfile when CPU/GPU hybrid
execution is enabled.
+ *
+ * The SparkSession for this suite is created with
ENABLE_HYBRID_EXECUTION=true (a static config),
+ * so all tests run in that context.
+ *
+ * When hybridExecution is enabled together with cudf, the rule inserts an
ApplyResourceProfileExec
+ * node for stages whose WholeStageTransformer is fully cudf-tagged. In Spark
testing mode the
+ * profile is always the default one, but the wrapper node itself must be
present in the plan tree.
+ */
+@Experimental
+class HybridExecutionResourceProfileSuite extends
VeloxWholeStageTransformerSuite {
+
+ override protected val resourcePath: String = "/tpch-data-parquet"
+ override protected val fileFormat: String = "parquet"
+
+ import testImplicits._
+
+ private val tbl = "hybrid_rp_tbl"
+ private val leftTable = "hybrid_rp_left"
+ private val rightTable = "hybrid_rp_right"
+
+ override protected def sparkConf: SparkConf = {
+ super.sparkConf
+ .set("spark.shuffle.manager",
"org.apache.spark.shuffle.sort.ColumnarShuffleManager")
+ .set("spark.sql.shuffle.partitions", "2")
+ .set("spark.sql.adaptive.enabled", "true")
+ .set("spark.memory.offHeap.size", "2g")
+ .set(GlutenConfig.AUTO_ADJUST_STAGE_RESOURCE_PROFILE_ENABLED.key, "true")
+ .set(GlutenConfig.ENABLE_HYBRID_EXECUTION.key, "true")
+ // Custom resource names to verify they are read correctly
+ .set(GlutenConfig.HYBRID_EXECUTION_CPU_RESOURCE_NAME.key, "mycpu")
+ .set(GlutenConfig.HYBRID_EXECUTION_GPU_RESOURCE_NAME.key, "mygpu")
+ .set(GlutenConfig.HYBRID_EXECUTION_GPU_RESOURCE_AMOUNT_PER_TASK.key,
"0.5")
+ .set(VeloxConfig.CUDF_ENABLE_VALIDATION.key, "false")
+ .set(VeloxConfig.CUDF_ENABLE_TABLE_SCAN.key, "false")
+ }
+
+ override def beforeAll(): Unit = {
+ super.beforeAll()
+
+ spark.range(100).selectExpr("cast(id % 5 as int) as k", "id as v").write
+ .mode("overwrite").format("parquet").saveAsTable(tbl)
+
+ Seq((1, "a"), (2, "b"), (3, "c")).toDF("id", "val").write
+ .mode("overwrite").format("parquet").saveAsTable(leftTable)
+
+ Seq((1, "x"), (2, "y"), (4, "z")).toDF("id", "val").write
+ .mode("overwrite").format("parquet").saveAsTable(rightTable)
+ }
+
+ override def afterAll(): Unit = {
+ try {
+ spark.sql(s"DROP TABLE IF EXISTS $tbl")
+ spark.sql(s"DROP TABLE IF EXISTS $leftTable")
+ spark.sql(s"DROP TABLE IF EXISTS $rightTable")
+ } finally {
+ super.afterAll()
+ }
+ }
+
+ private def collectApplyResourceProfileExec(plan: SparkPlan):
Seq[ApplyResourceProfileExec] = {
+ collect(plan) { case a: ApplyResourceProfileExec => a }
+ }
+
+ /**
+ * With hybridExecution=true and cudf enabled,
GlutenAutoAdjustStageResourceProfile should insert
+ * an ApplyResourceProfileExec for a fully-cudf stage (e.g. a simple
aggregation). In testing mode
+ * Spark always falls back to the default resource profile, but the wrapper
node must still appear
+ * in the executed plan.
+ */
+ test("GPU resource profile wrapper is inserted for a fully-cudf aggregate
stage") {
+ withSQLConf(
+ GlutenConfig.COLUMNAR_CUDF_ENABLED.key -> "true",
+ SQLConf.ANSI_ENABLED.key -> "false"
+ ) {
+ val df = sql(s"SELECT k, count(*) FROM $tbl GROUP BY k")
+ df.collect()
+ val plan = df.queryExecution.executedPlan
+ val nodes = collectApplyResourceProfileExec(plan)
+ assert(
+ nodes.nonEmpty,
+ "Expected ApplyResourceProfileExec to be inserted when hybrid
execution is enabled " +
+ "for a fully-cudf-offloaded stage")
+ }
+
+ withSQLConf(
+ GlutenConfig.COLUMNAR_CUDF_ENABLED.key -> "true",
+ GlutenConfig.GPU_ONLY_OFFLOAD_JOIN_STAGE.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ANSI_ENABLED.key -> "false"
+ ) {
+ val df = sql(
+ s"""
+ |SELECT l.id, l.val, r.val
+ |FROM $leftTable l
+ |JOIN $rightTable r ON l.id = r.id
+ |""".stripMargin)
+ df.collect()
+ val plan = df.queryExecution.executedPlan
+ val nodes = collectApplyResourceProfileExec(plan)
+ assert(
+ nodes.nonEmpty,
+ "Expected ApplyResourceProfileExec for a cudf join stage with hybrid
execution enabled")
+ }
+ }
+}
diff --git
a/backends-velox/src/test/scala/org/apache/gluten/execution/StageExecutionModeSuite.scala
b/backends-velox/src/test/scala/org/apache/gluten/execution/StageExecutionModeSuite.scala
index 6efaeae753..5f4f68971e 100644
---
a/backends-velox/src/test/scala/org/apache/gluten/execution/StageExecutionModeSuite.scala
+++
b/backends-velox/src/test/scala/org/apache/gluten/execution/StageExecutionModeSuite.scala
@@ -78,6 +78,16 @@ class StageExecutionModeSuite extends
VeloxWholeStageTransformerSuite {
.set(VeloxConfig.CUDF_ENABLE_TABLE_SCAN.key, "false")
}
+ /**
+ * Collects all ColumnarAQEShuffleReadExec nodes from the executed plan of
`df`, irrespective of
+ * plan nesting.
+ */
+ private def collectColumnarAQEReaders(
+ df: org.apache.spark.sql.DataFrame): Seq[ColumnarAQEShuffleReadExec] = {
+ val plan = getExecutedPlan(df)
+ plan.collect { case r: ColumnarAQEShuffleReadExec => r }
+ }
+
test("CPU shuffle mapper and GPU shuffle reader with AQE") {
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
@@ -99,15 +109,10 @@ class StageExecutionModeSuite extends
VeloxWholeStageTransformerSuite {
Row(1, "left-1", "right-1"),
Row(2, "left-2", "right-2")))
- val plan = getExecutedPlan(df)
-
- val shuffleReaders = plan.collect {
- case reader: ColumnarAQEShuffleReadExec => reader
- }
-
- assert(shuffleReaders.nonEmpty)
+ val readers = collectColumnarAQEReaders(df)
+ assert(readers.nonEmpty, "Expected at least one
ColumnarAQEShuffleReadExec")
- shuffleReaders.foreach {
+ readers.foreach {
reader =>
val canonicalized = reader.canonicalized
// canonicalized plan before applying query stage optimizer rules.
@@ -117,7 +122,7 @@ class StageExecutionModeSuite extends
VeloxWholeStageTransformerSuite {
s"Expected GPU AQE shuffle reader, but got
${reader.executionMode}")
}
- val shuffleStages: Seq[ShuffleQueryStageExec] =
shuffleReaders.map(_.delegate).map {
+ val shuffleStages: Seq[ShuffleQueryStageExec] =
readers.map(_.delegate).map {
case a: AQEShuffleReadExec =>
assert(a.child.isInstanceOf[ShuffleQueryStageExec])
a.child.asInstanceOf[ShuffleQueryStageExec]
@@ -136,4 +141,55 @@ class StageExecutionModeSuite extends
VeloxWholeStageTransformerSuite {
}
}
}
+
+ test("with gpuOnlyOffloadJoinStage=true, join stage is offloaded to GPU") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ GlutenConfig.COLUMNAR_CUDF_ENABLED.key -> "true",
+ GlutenConfig.GPU_ONLY_OFFLOAD_JOIN_STAGE.key -> "true",
+ SQLConf.ANSI_ENABLED.key -> "false"
+ ) {
+ // A join query: the join stage should still be offloaded to GPU.
+ val df = sql(
+ s"""
+ |SELECT l.id, l.left_value, r.right_value
+ |FROM $leftTable l
+ |JOIN $rightTable r
+ | ON l.id = r.id
+ |""".stripMargin)
+
+ checkAnswer(
+ df,
+ Seq(
+ Row(1, "left-1", "right-1"),
+ Row(2, "left-2", "right-2")))
+
+ val readers = collectColumnarAQEReaders(df)
+ assert(readers.nonEmpty, "Expected at least one
ColumnarAQEShuffleReadExec for a join query")
+ readers.foreach {
+ reader =>
+ assert(
+ reader.executionMode == MockGPUStageMode,
+ s"Join stage should be GPU-offloaded, but got
${reader.executionMode}")
+ }
+ }
+ }
+
+ test("with gpuOnlyOffloadJoinStage=true, non-join stage is NOT offloaded to
GPU") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ GlutenConfig.COLUMNAR_CUDF_ENABLED.key -> "true",
+ GlutenConfig.GPU_ONLY_OFFLOAD_JOIN_STAGE.key -> "true",
+ SQLConf.ANSI_ENABLED.key -> "false"
+ ) {
+ // A pure aggregation query (no join): the stage must NOT be
GPU-offloaded.
+ val df = sql(s"SELECT id, count(*) FROM $leftTable GROUP BY id")
+ df.collect()
+
+ val readers = collectColumnarAQEReaders(df)
+ assert(readers.isEmpty, "Expected no ColumnarAQEShuffleReadExec")
+ }
+ }
}
diff --git a/docs/Configuration.md b/docs/Configuration.md
index 83f44ec941..7d995c75e2 100644
--- a/docs/Configuration.md
+++ b/docs/Configuration.md
@@ -152,13 +152,18 @@ nav_order: 15
## Gluten *experimental* configurations
-| Key |
Modifiability | Default |
Description [...]
-|-------------------------------------------------------------------|---------------|---------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
[...]
-| spark.gluten.auto.adjustStageResource.enabled | 🔄
Dynamic | false | Experimental: If enabled, gluten will try to set the
stage resource according to stage execution plan. Only worked when aqe is
enabled at the same time!!
[...]
-| spark.gluten.auto.adjustStageResources.fallenNode.ratio.threshold | 🔄
Dynamic | 0.5 | Experimental: Increase executor heap memory when stage
contains fallen node count exceeds the total node count ratio.
[...]
-| spark.gluten.auto.adjustStageResources.heap.ratio | 🔄
Dynamic | 2.0 | Experimental: Increase executor heap memory when match
adjust stage resource rule.
[...]
-| spark.gluten.auto.adjustStageResources.offheap.ratio | 🔄
Dynamic | 0.5 | Experimental: Decrease executor offheap memory when
match adjust stage resource rule.
[...]
-| spark.gluten.memory.dynamic.offHeap.sizing.enabled | âš“ Static
| false | Experimental: When set to true, the offheap config
(spark.memory.offHeap.size) will be ignored and instead we will consider onheap
and offheap memory in combination, both counting towards the executor memory
config (spark.executor.memory). We will make use of JVM APIs to determine how
much onheap memory is use, alongside tracking offheap allocations made by
Gluten. We will then proceed to enfo [...]
-| spark.gluten.memory.dynamic.offHeap.sizing.memory.fraction | âš“ Static
| 0.6 | Experimental: Determines the memory fraction used to determine
the total memory available for offheap and onheap allocations when the dynamic
offheap sizing feature is enabled. The default is set to match
spark.executor.memoryFraction.
[...]
-| spark.gluten.sql.columnar.cudf | 🔄
Dynamic | false | Enable or disable cudf support. This is an experimental
feature.
[...]
+| Key |
Modifiability | Default |
Description [...]
+|---------------------------------------------------------------------|---------------|---------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
[...]
+| spark.gluten.auto.adjustStageResource.enabled | 🔄
Dynamic | false | Experimental: If enabled, gluten will try to set the
stage resource according to stage execution plan. Only worked when aqe is
enabled at the same time!!
[...]
+| spark.gluten.auto.adjustStageResources.fallenNode.ratio.threshold | 🔄
Dynamic | 0.5 | Experimental: Increase executor heap memory when stage
contains fallen node count exceeds the total node count ratio.
[...]
+| spark.gluten.auto.adjustStageResources.heap.ratio | 🔄
Dynamic | 2.0 | Experimental: Increase executor heap memory when match
adjust stage resource rule.
[...]
+| spark.gluten.auto.adjustStageResources.offheap.ratio | 🔄
Dynamic | 0.5 | Experimental: Decrease executor offheap memory when
match adjust stage resource rule.
[...]
+| spark.gluten.memory.dynamic.offHeap.sizing.enabled | âš“
Static | false | Experimental: When set to true, the offheap config
(spark.memory.offHeap.size) will be ignored and instead we will consider onheap
and offheap memory in combination, both counting towards the executor memory
config (spark.executor.memory). We will make use of JVM APIs to determine how
much onheap memory is use, alongside tracking offheap allocations made by
Gluten. We will then proceed to en [...]
+| spark.gluten.memory.dynamic.offHeap.sizing.memory.fraction | âš“
Static | 0.6 | Experimental: Determines the memory fraction used to
determine the total memory available for offheap and onheap allocations when
the dynamic offheap sizing feature is enabled. The default is set to match
spark.executor.memoryFraction.
[...]
+| spark.gluten.sql.columnar.cudf | 🔄
Dynamic | false | Enable or disable cudf support. This is an experimental
feature.
[...]
+| spark.gluten.sql.columnar.gpu.onlyOffloadJoinStage | 🔄
Dynamic | false | If true, Gluten will only offload join stages to GPU.
Other stages will be executed on CPU.
[...]
+| spark.gluten.sql.columnar.hybridExecution.cpuResource.name | âš“
Static | cpu | The CPU resource name (Spark custom resource). This
must match the resource name configured via spark.executor.resource.<name>.* /
spark.task.resource.<name>.* for CPU-stage scheduling to take effect.
[...]
+| spark.gluten.sql.columnar.hybridExecution.enabled | âš“
Static | false | Enable CPU/GPU hybrid execution. At runtime, the
execution will be scheduled to target nodes based on the selected execution
mode.
[...]
+| spark.gluten.sql.columnar.hybridExecution.gpuResource.amountPerTask | âš“
Static | 0.1 | The GPU resource amount per task. This is used to limit
GPU tasks to target nodes.
[...]
+| spark.gluten.sql.columnar.hybridExecution.gpuResource.name | âš“
Static | gpu | The GPU resource name (Spark custom resource). This
must match the resource name configured via spark.executor.resource.<name>.* /
spark.task.resource.<name>.* for GPU-stage scheduling to take effect.
[...]
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/config/GlutenConfig.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/config/GlutenConfig.scala
index 8ffb6345c3..3bb5b0958f 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/config/GlutenConfig.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/config/GlutenConfig.scala
@@ -396,6 +396,14 @@ class GlutenConfig(conf: SQLConf) extends
GlutenCoreConfig(conf) {
def maxBroadcastTableSize: Long =
JavaUtils.byteStringAsBytes(conf.getConfString(SPARK_MAX_BROADCAST_TABLE_SIZE,
"8GB"))
+
+ def enableHybridExecution: Boolean = getConf(ENABLE_HYBRID_EXECUTION)
+
+ def cpuResourceName: String = getConf(HYBRID_EXECUTION_CPU_RESOURCE_NAME)
+ def gpuResourceName: String = getConf(HYBRID_EXECUTION_GPU_RESOURCE_NAME)
+ def gpuResourceAmountPerTask: Double =
getConf(HYBRID_EXECUTION_GPU_RESOURCE_AMOUNT_PER_TASK)
+
+ def gpuOnlyOffloadJoinStage: Boolean = getConf(GPU_ONLY_OFFLOAD_JOIN_STAGE)
}
object GlutenConfig extends ConfigRegistry {
@@ -1718,4 +1726,52 @@ object GlutenConfig extends ConfigRegistry {
"total size of small files is below this threshold.")
.doubleConf
.createWithDefault(0.5)
+
+ val ENABLE_HYBRID_EXECUTION =
+ buildStaticConf("spark.gluten.sql.columnar.hybridExecution.enabled")
+ .experimental()
+ .doc(
+ "Enable CPU/GPU hybrid execution. At runtime, the execution will be
scheduled to target " +
+ "nodes based on the selected execution mode.")
+ .booleanConf
+ .createWithDefault(false)
+
+ val HYBRID_EXECUTION_CPU_RESOURCE_NAME =
+
buildStaticConf("spark.gluten.sql.columnar.hybridExecution.cpuResource.name")
+ .experimental()
+ .doc(
+ "The CPU resource name (Spark custom resource). " +
+ "This must match the resource name configured via
spark.executor.resource.<name>.* / " +
+ "spark.task.resource.<name>.* for CPU-stage scheduling to take
effect."
+ )
+ .stringConf
+ .createWithDefault("cpu")
+
+ val HYBRID_EXECUTION_GPU_RESOURCE_NAME =
+
buildStaticConf("spark.gluten.sql.columnar.hybridExecution.gpuResource.name")
+ .experimental()
+ .doc(
+ "The GPU resource name (Spark custom resource). " +
+ "This must match the resource name configured via
spark.executor.resource.<name>.* / " +
+ "spark.task.resource.<name>.* for GPU-stage scheduling to take
effect."
+ )
+ .stringConf
+ .createWithDefault("gpu")
+
+ val HYBRID_EXECUTION_GPU_RESOURCE_AMOUNT_PER_TASK =
+
buildStaticConf("spark.gluten.sql.columnar.hybridExecution.gpuResource.amountPerTask")
+ .experimental()
+ .doc(
+ "The GPU resource amount per task. This is used to limit GPU tasks to
target nodes.")
+ .doubleConf
+ .createWithDefault(0.1)
+
+ val GPU_ONLY_OFFLOAD_JOIN_STAGE =
+ buildConf("spark.gluten.sql.columnar.gpu.onlyOffloadJoinStage")
+ .experimental()
+ .doc(
+ "If true, Gluten will only offload join stages to GPU." +
+ " Other stages will be executed on CPU.")
+ .booleanConf
+ .createWithDefault(false)
}
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/TakeOrderedAndProjectExecTransformer.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/TakeOrderedAndProjectExecTransformer.scala
index dc47dbba10..8011dfa7a8 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/TakeOrderedAndProjectExecTransformer.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/TakeOrderedAndProjectExecTransformer.scala
@@ -143,9 +143,9 @@ case class TakeOrderedAndProjectExecTransformer(
} else {
val limitStagePlan =
WholeStageTransformer(limitBeforeShuffle)(transformStageCounter.incrementAndGet())
- val cudfTag = child.getTagValue(CudfTag.CudfTag)
+ val cudfTag = child.getTagValue(CudfTag.CudfValidationTag)
if (cudfTag.isDefined) {
- limitStagePlan.setTagValue(CudfTag.CudfTag, cudfTag.get)
+ limitStagePlan.setTagValue(CudfTag.CudfValidationTag, cudfTag.get)
}
val shuffleExec = ShuffleExchangeExec(SinglePartition, limitStagePlan)
val transformedShuffleExec =
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/WholeStageTransformer.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/WholeStageTransformer.scala
index 71cac1b5f3..d7a3999217 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/WholeStageTransformer.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/WholeStageTransformer.scala
@@ -70,12 +70,12 @@ trait TransformSupport extends ValidatablePlan {
s"${this.getClass.getSimpleName} doesn't support doExecute")
}
- def isCudf: Boolean = getTagValue[Boolean](CudfTag.CudfTag).getOrElse(false)
+ def offloadCuda: Boolean =
getTagValue[Boolean](CudfTag.CudfValidationTag).getOrElse(false)
// Use super.nodeName will cause exception scala 213 Super calls can only
target methods
// for FileSourceScan.
override def nodeName: String =
- if (isCudf) {
+ if (offloadCuda) {
"Cudf" + getClass.getSimpleName.replaceAll("Exec$", "")
} else getClass.getSimpleName
@@ -261,7 +261,7 @@ case class WholeStageTransformer(child: SparkPlan,
materializeInput: Boolean = f
WholeStageTransformContext(
planNode,
substraitContext,
- isCudf,
+ offloadCuda,
!hasNonDeterministicExprInJoinProbe(child))
}
@@ -520,5 +520,6 @@ class ColumnarInputRDDsWrapper(columnarInputRDDs:
Seq[RDD[ColumnarBatch]]) exten
}
object CudfTag {
- val CudfTag = TreeNodeTag[Boolean]("org.apache.gluten.CudfTag")
+ val CudfValidationTag =
TreeNodeTag[Boolean]("org.apache.gluten.CudfTag.CudfValidationTag")
+ val CudfTestingTag =
TreeNodeTag[Boolean]("org.apache.gluten.CudfTag.CudfTestingTag")
}
diff --git
a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/GlutenAutoAdjustStageResourceProfile.scala
b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/GlutenAutoAdjustStageResourceProfile.scala
index 5f80d99dee..8682fa6889 100644
---
a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/GlutenAutoAdjustStageResourceProfile.scala
+++
b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/GlutenAutoAdjustStageResourceProfile.scala
@@ -17,7 +17,7 @@
package org.apache.spark.sql.execution
import org.apache.gluten.config.{GlutenConfig, GlutenCoreConfig}
-import org.apache.gluten.execution.{ColumnarToRowExecBase, GlutenPlan}
+import org.apache.gluten.execution.{ColumnarToRowExecBase, CudfTag,
GlutenPlan, WholeStageTransformer}
import org.apache.gluten.logging.LogLevelUtil
import org.apache.spark.SparkConf
@@ -38,11 +38,20 @@ import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer
/**
- * This rule is used to dynamic adjust stage resource profile for following
purposes:
- * 1. Decrease offheap and increase onheap memory size when whole stage
fallback happened; 2.
- * Increase executor heap memory if stage contains gluten operator and
spark operator at the
- * same time. Note: we don't support set resource profile for final stage
now. Todo: will
- * support it.
+ * This rule dynamically adjusts the resource profile of each AQE query stage.
It handles three
+ * cases:
+ *
+ * 1. CPU/GPU hybrid execution: if every `WholeStageTransformer` in the
stage is fully
+ * cuDF-offloaded, the stage is assigned a GPU resource profile so that
Spark schedules its
+ * tasks on GPU-equipped executors.
+ * 2. Whole-stage fallback: if a stage contains no native Gluten operators
(or only
+ * columnar-to-row conversion nodes), heap memory is increased and
off-heap memory is reduced,
+ * since the stage runs entirely on the JVM.
+ * 3. Partial fallback: if the ratio of fallen (non-Gluten) nodes in a stage
exceeds
+ * `spark.gluten.auto.adjustStageResources.fallenNode.ratio.threshold`,
heap memory is
+ * increased and off-heap memory is decreased proportionally.
+ *
+ * * Note: Case 2 and 3 are not applied to final (non-Exchange) stages yet.
*/
@Experimental
case class GlutenAutoAdjustStageResourceProfile(glutenConf: GlutenConfig,
spark: SparkSession)
@@ -68,6 +77,32 @@ case class GlutenAutoAdjustStageResourceProfile(glutenConf:
GlutenConfig, spark:
ResourceProfile.getOrCreateDefaultProfile(sparkConf),
sparkConf,
isDefaultProfile = true)
+
+ val rpManager = spark.sparkContext.resourceProfileManager
+ val defaultRP = rpManager.defaultResourceProfile
+
+ // initial resource profile config as default resource profile
+ val taskResource = mutable.Map.empty[String, TaskResourceRequest] ++=
defaultRP.taskResources
+ val executorResource =
+ mutable.Map.empty[String, ExecutorResourceRequest] ++=
defaultRP.executorResources
+
+ if (glutenConf.enableColumnarCudf && glutenConf.enableHybridExecution) {
+ val transformers = plan.collect { case t: WholeStageTransformer => t }
+ if (
+ transformers.nonEmpty && transformers.forall {
+ t => t.offloadCuda ||
t.getTagValue(CudfTag.CudfTestingTag).getOrElse(false)
+ }
+ ) {
+ return GlutenResourceProfile.setResourceProfileForGpu(
+ plan,
+ executorResource,
+ taskResource,
+ rpManager,
+ sparkConf,
+ glutenConf)
+ }
+ }
+
if (!plan.isInstanceOf[Exchange]) {
// todo: support set resource profile for final stage
return plan
@@ -78,25 +113,18 @@ case class
GlutenAutoAdjustStageResourceProfile(glutenConf: GlutenConfig, spark:
}
log.info(s"detailPlanNodes ${planNodes.map(_.nodeName).mkString("Array(",
", ", ")")}")
- // one stage is considered as fallback if all node is not GlutenPlan
- // or all GlutenPlan node is C2R node.
- val wholeStageFallback = planNodes
- .filter(_.isInstanceOf[GlutenPlan])
- .count(!_.isInstanceOf[ColumnarToRowExecBase]) == 0
-
- val rpManager = spark.sparkContext.resourceProfileManager
- val defaultRP = rpManager.defaultResourceProfile
-
- // initial resource profile config as default resource profile
- val taskResource = mutable.Map.empty[String, TaskResourceRequest] ++=
defaultRP.taskResources
- val executorResource =
- mutable.Map.empty[String, ExecutorResourceRequest] ++=
defaultRP.executorResources
val memoryRequest = executorResource.get(ResourceProfile.MEMORY)
val offheapRequest = executorResource.get(ResourceProfile.OFFHEAP_MEM)
logInfo(s"default memory request $memoryRequest")
logInfo(s"default offheap request $offheapRequest")
// case 1: whole stage fallback to vanilla spark in such case we increase
the heap
+ //
+ // one stage is considered as fallback if all node is not GlutenPlan
+ // or all GlutenPlan node is C2R node.
+ val wholeStageFallback = planNodes
+ .filter(_.isInstanceOf[GlutenPlan])
+ .count(!_.isInstanceOf[ColumnarToRowExecBase]) == 0
if (wholeStageFallback) {
val newMemoryAmount = memoryRequest.get.amount *
glutenConf.autoAdjustStageRPHeapRatio
val newExecutorMemory =
@@ -107,10 +135,10 @@ case class
GlutenAutoAdjustStageResourceProfile(glutenConf: GlutenConfig, spark:
new ExecutorResourceRequest(ResourceProfile.OFFHEAP_MEM,
offheapRequest.get.amount / 10)
executorResource.put(ResourceProfile.OFFHEAP_MEM, newExecutorOffheap)
- val newRP = new ResourceProfile(executorResource.toMap,
taskResource.toMap)
- return GlutenResourceProfile.applyNewResourceProfileIfPossible(
+ return GlutenResourceProfile.applyNewResourceProfile(
plan,
- newRP,
+ executorResource,
+ taskResource,
rpManager,
sparkConf)
}
@@ -132,10 +160,10 @@ case class
GlutenAutoAdjustStageResourceProfile(glutenConf: GlutenConfig, spark:
new ExecutorResourceRequest(ResourceProfile.OFFHEAP_MEM,
newOffHeapMemoryAmount.toLong)
executorResource.put(ResourceProfile.OFFHEAP_MEM, newExecutorOffheap)
- val newRP = new ResourceProfile(executorResource.toMap,
taskResource.toMap)
- return GlutenResourceProfile.applyNewResourceProfileIfPossible(
+ return GlutenResourceProfile.applyNewResourceProfile(
plan,
- newRP,
+ executorResource,
+ taskResource,
rpManager,
sparkConf)
}
@@ -228,17 +256,52 @@ object GlutenAutoAdjustStageResourceProfile extends
Logging {
(offHeapSize / taskSlots).toString)
}
- def applyNewResourceProfileIfPossible(
+ def applyNewResourceProfile(
plan: SparkPlan,
- rp: ResourceProfile,
+ executorResource: mutable.Map[String, ExecutorResourceRequest],
+ taskResource: mutable.Map[String, TaskResourceRequest],
rpManager: ResourceProfileManager,
sparkConf: SparkConf): SparkPlan = {
- updateResourceSetting(rp, sparkConf)
-
+ val rp = new ResourceProfile(executorResource.toMap, taskResource.toMap)
val finalRP = getFinalResourceProfile(rpManager, rp)
- // Wrap the plan with ApplyResourceProfileExec so that we can apply new
ResourceProfile
- val wrapperPlan = ApplyResourceProfileExec(plan.children.head, finalRP)
- logInfo(s"Apply resource profile $finalRP for plan
${wrapperPlan.nodeName}")
- plan.withNewChildren(IndexedSeq(wrapperPlan))
+ updateResourceSetting(finalRP, sparkConf)
+
+ plan match {
+ case shuffle: Exchange =>
+ logInfo(s"Apply resource profile $finalRP for plan
${shuffle.child.nodeName}")
+ // Wrap the plan with ApplyResourceProfileExec so that we can apply
new ResourceProfile
+ val wrapperPlan = ApplyResourceProfileExec(shuffle.child, finalRP)
+ shuffle.withNewChildren(Seq(wrapperPlan))
+ case other =>
+ logInfo(s"Apply resource profile $finalRP for plan ${other.nodeName}")
+ ApplyResourceProfileExec(other, finalRP)
+ }
+ }
+
+ def setResourceProfileForGpu(
+ plan: SparkPlan,
+ executorResource: mutable.Map[String, ExecutorResourceRequest],
+ taskResource: mutable.Map[String, TaskResourceRequest],
+ rpManager: ResourceProfileManager,
+ sparkConf: SparkConf,
+ glutenConf: GlutenConfig): SparkPlan = {
+ val cpuResourceName = glutenConf.cpuResourceName
+ val gpuResourceName = glutenConf.gpuResourceName
+
+ executorResource.remove(glutenConf.cpuResourceName)
+ taskResource.remove(cpuResourceName)
+
+ executorResource.put(gpuResourceName, new
ExecutorResourceRequest(gpuResourceName, 1))
+ // The gpu task resource limits how many tasks can be launched in one
executor.
+ taskResource.put(
+ gpuResourceName,
+ new TaskResourceRequest(gpuResourceName,
glutenConf.gpuResourceAmountPerTask))
+
+ applyNewResourceProfile(
+ plan,
+ executorResource,
+ taskResource,
+ rpManager,
+ sparkConf)
}
}
diff --git
a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala
b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala
index 02f7517cc6..ba59c2c06a 100644
---
a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala
+++
b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala
@@ -70,7 +70,7 @@ case class ColumnarAQEShuffleReadExec(
delegate match {
case a: AQEShuffleReadExec => a
case s: ShuffleQueryStageExec =>
- // Wrap ShuffleQueryStageExe with dummy PartitionSpecs by creating
CoalescedPartitionSpec
+ // Wrap ShuffleQueryStageExec with dummy PartitionSpecs by creating
CoalescedPartitionSpec
// for each partition.
val partitionSpecs =
Array.tabulate(s.shuffle.numPartitions)(i =>
CoalescedPartitionSpec(i, i + 1))
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
