You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@kyuubi.apache.org by ul...@apache.org on 2023/03/30 05:05:10 UTC
[kyuubi] branch master updated: [KYUUBI #4615] Support stage level schedule for final write stage
This is an automated email from the ASF dual-hosted git repository.
ulyssesyou pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/kyuubi.git
The following commit(s) were added to refs/heads/master by this push:
new 9ca00c8aa [KYUUBI #4615] Support stage level schedule for final write stage
9ca00c8aa is described below
commit 9ca00c8aa750c962213c3b27f583f238ef065005
Author: ulysses-you <ul...@gmail.com>
AuthorDate: Thu Mar 30 13:04:51 2023 +0800
[KYUUBI #4615] Support stage level schedule for final write stage
### _Why are the changes needed?_
Add a new rule `InjectCustomResourceProfile` to support custom resource profile for final write stage.
It now supports executor configs:
```
executor core
executor memory
executor memory overhead
executor off heap memory
```
### _How was this patch tested?_
add test and manully test
<img width="778" alt="image" src="https://user-images.githubusercontent.com/12025282/226606147-82a29b8c-1a31-4842-97a7-fe702d80e190.png">
Closes #4615 from ulysses-you/resource-profile.
Closes #4615
852b207cd [ulysses-you] Support stage level schedule for final write stag
Authored-by: ulysses-you <ul...@gmail.com>
Signed-off-by: ulyssesyou <ul...@apache.org>
---
docs/extensions/engines/spark/rules.md | 49 +++++----
.../kyuubi/sql/KyuubiSparkSQLExtension.scala | 4 +-
.../{ => sql}/FinalStageResourceManager.scala | 28 +++---
.../spark/sql/InjectCustomResourceProfile.scala | 60 +++++++++++
.../sql/execution/CustomResourceProfileExec.scala | 112 +++++++++++++++++++++
.../spark/sql/InjectResourceProfileSuite.scala | 79 +++++++++++++++
.../org/apache/kyuubi/sql/KyuubiSQLConf.scala | 40 ++++++++
.../spark/sql/KyuubiSparkSQLExtensionTest.scala | 2 +
8 files changed, 338 insertions(+), 36 deletions(-)
diff --git a/docs/extensions/engines/spark/rules.md b/docs/extensions/engines/spark/rules.md
index f06ecf256..a4bda5d53 100644
--- a/docs/extensions/engines/spark/rules.md
+++ b/docs/extensions/engines/spark/rules.md
@@ -63,26 +63,31 @@ Now, you can enjoy the Kyuubi SQL Extension.
Kyuubi provides some configs to make these feature easy to use.
-| Name | Default Value | Description | Since |
-|---------------------------------------------------------------------|---------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|-------|
-| spark.sql.optimizer.insertRepartitionBeforeWrite.enabled | true | Add repartition node at the top of query plan. An approach of merging small files. | 1.2.0 |
-| spark.sql.optimizer.insertRepartitionNum | none | The partition number if `spark.sql.optimizer.insertRepartitionBeforeWrite.enabled` is enabled. If AQE is disabled, the default value is `spark.sql.shuffle.partitions`. If AQE is enabled, the default value is none that means depend on AQE. | 1.2.0 |
-| spark.sql.optimizer.dynamicPartitionInsertionRepartitionNum | 100 | The partition number of each dynamic partition if `spark.sql.optimizer.insertRepartitionBeforeWrite.enabled` is enabled. We will repartition by dynamic partition columns to reduce the small file but that can cause data skew. This config is to extend the partition of dynamic partition column to avoid skew but may generate some small files. | 1.2.0 |
-| spark.sql.optimizer.forceShuffleBeforeJoin.enabled | false | Ensure shuffle node exists before shuffled join (shj and smj) to make AQE `OptimizeSkewedJoin` works (complex scenario join, multi table join). | 1.2.0 |
-| spark.sql.optimizer.finalStageConfigIsolation.enabled | false | If true, the final stage support use different config with previous stage. The prefix of final stage config key should be `spark.sql.finalStage.`. For example, the raw spark config: `spark.sql.adaptive.advisoryPartitionSizeInBytes`, then the final stage config should be: `spark.sql.finalStage.adaptive.advisoryPartitionSizeInBytes`. | 1.2.0 |
-| spark.sql.analyzer.classification.enabled | false | When true, allows Kyuubi engine to judge this SQL's classification and set `spark.sql.analyzer.classification` back into sessionConf. Through this configuration item, Spark can optimizing configuration dynamic. | 1.4.0 |
-| spark.sql.optimizer.insertZorderBeforeWriting.enabled | true | When true, we will follow target table properties to insert zorder or not. The key properties are: 1) `kyuubi.zorder.enabled`: if this property is true, we will insert zorder before writing data. 2) `kyuubi.zorder.cols`: string split by comma, we will zorder by these cols. | 1.4.0 |
-| spark.sql.optimizer.zorderGlobalSort.enabled | true | When true, we do a global sort using zorder. Note that, it can cause data skew issue if the zorder columns have less cardinality. When false, we only do local sort using zorder. | 1.4.0 |
-| spark.sql.watchdog.maxPartitions | none | Set the max partition number when spark scans a data source. Enable MaxPartitionStrategy by specifying this configuration. Add maxPartitions Strategy to avoid scan excessive partitions on partitioned table, it's optional that works with defined | 1.4.0 |
-| spark.sql.optimizer.dropIgnoreNonExistent | false | When true, do not report an error if DROP DATABASE/TABLE/VIEW/FUNCTION/PARTITION specifies a non-existent database/table/view/function/partition | 1.5.0 |
-| spark.sql.optimizer.rebalanceBeforeZorder.enabled | false | When true, we do a rebalance before zorder in case data skew. Note that, if the insertion is dynamic partition we will use the partition columns to rebalance. Note that, this config only affects with Spark 3.3.x. | 1.6.0 |
-| spark.sql.optimizer.rebalanceZorderColumns.enabled | false | When true and `spark.sql.optimizer.rebalanceBeforeZorder.enabled` is true, we do rebalance before Z-Order. If it's dynamic partition insert, the rebalance expression will include both partition columns and Z-Order columns. Note that, this config only affects with Spark 3.3.x. | 1.6.0 |
-| spark.sql.optimizer.twoPhaseRebalanceBeforeZorder.enabled | false | When true and `spark.sql.optimizer.rebalanceBeforeZorder.enabled` is true, we do two phase rebalance before Z-Order for the dynamic partition write. The first phase rebalance using dynamic partition column; The second phase rebalance using dynamic partition column Z-Order columns. Note that, this config only affects with Spark 3.3.x. | 1.6.0 |
-| spark.sql.optimizer.zorderUsingOriginalOrdering.enabled | false | When true and `spark.sql.optimizer.rebalanceBeforeZorder.enabled` is true, we do sort by the original ordering i.e. lexicographical order. Note that, this config only affects with Spark 3.3.x. | 1.6.0 |
-| spark.sql.optimizer.inferRebalanceAndSortOrders.enabled | false | When ture, infer columns for rebalance and sort orders from original query, e.g. the join keys from join. It can avoid compression ratio regression. | 1.7.0 |
-| spark.sql.optimizer.inferRebalanceAndSortOrdersMaxColumns | 3 | The max columns of inferred columns. | 1.7.0 |
-| spark.sql.optimizer.insertRepartitionBeforeWriteIfNoShuffle.enabled | false | When true, add repartition even if the original plan does not have shuffle. | 1.7.0 |
-| spark.sql.optimizer.finalStageConfigIsolationWriteOnly.enabled | true | When true, only enable final stage isolation for writing. | 1.7.0 |
-| spark.sql.finalWriteStage.eagerlyKillExecutors.enabled | false | When true, eagerly kill redundant executors before running final write stage. | 1.8.0 |
-| spark.sql.finalWriteStage.retainExecutorsFactor | 1.2 | If the target executors * factor < active executors, and target executors * factor > min executors, then inject kill executors or inject custom resource profile. | 1.8.0 |
+| Name | Default Value | Description | Since |
+|---------------------------------------------------------------------|----------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|-------|
+| spark.sql.optimizer.insertRepartitionBeforeWrite.enabled | true | Add repartition node at the top of query plan. An approach of merging small files. | 1.2.0 |
+| spark.sql.optimizer.insertRepartitionNum | none | The partition number if `spark.sql.optimizer.insertRepartitionBeforeWrite.enabled` is enabled. If AQE is disabled, the default value is `spark.sql.shuffle.partitions`. If AQE is enabled, the default value is none that means depend on AQE. | 1.2.0 |
+| spark.sql.optimizer.dynamicPartitionInsertionRepartitionNum | 100 | The partition number of each dynamic partition if `spark.sql.optimizer.insertRepartitionBeforeWrite.enabled` is enabled. We will repartition by dynamic partition columns to reduce the small file but that can cause data skew. This config is to extend the partition of dynamic partition column to avoid skew but may generate some small files. | 1.2.0 |
+| spark.sql.optimizer.forceShuffleBeforeJoin.enabled | false | Ensure shuffle node exists before shuffled join (shj and smj) to make AQE `OptimizeSkewedJoin` works (complex scenario join, multi table join). | 1.2.0 |
+| spark.sql.optimizer.finalStageConfigIsolation.enabled | false | If true, the final stage support use different config with previous stage. The prefix of final stage config key should be `spark.sql.finalStage.`. For example, the raw spark config: `spark.sql.adaptive.advisoryPartitionSizeInBytes`, then the final stage config should be: `spark.sql.finalStage.adaptive.advisoryPartitionSizeInBytes`. | 1.2.0 |
+| spark.sql.analyzer.classification.enabled | false | When true, allows Kyuubi engine to judge this SQL's classification and set `spark.sql.analyzer.classification` back into sessionConf. Through this configuration item, Spark can optimizing configuration dynamic. | 1.4.0 |
+| spark.sql.optimizer.insertZorderBeforeWriting.enabled | true | When true, we will follow target table properties to insert zorder or not. The key properties are: 1) `kyuubi.zorder.enabled`: if this property is true, we will insert zorder before writing data. 2) `kyuubi.zorder.cols`: string split by comma, we will zorder by these cols. | 1.4.0 |
+| spark.sql.optimizer.zorderGlobalSort.enabled | true | When true, we do a global sort using zorder. Note that, it can cause data skew issue if the zorder columns have less cardinality. When false, we only do local sort using zorder. | 1.4.0 |
+| spark.sql.watchdog.maxPartitions | none | Set the max partition number when spark scans a data source. Enable MaxPartitionStrategy by specifying this configuration. Add maxPartitions Strategy to avoid scan excessive partitions on partitioned table, it's optional that works with defined | 1.4.0 |
+| spark.sql.optimizer.dropIgnoreNonExistent | false | When true, do not report an error if DROP DATABASE/TABLE/VIEW/FUNCTION/PARTITION specifies a non-existent database/table/view/function/partition | 1.5.0 |
+| spark.sql.optimizer.rebalanceBeforeZorder.enabled | false | When true, we do a rebalance before zorder in case data skew. Note that, if the insertion is dynamic partition we will use the partition columns to rebalance. Note that, this config only affects with Spark 3.3.x. | 1.6.0 |
+| spark.sql.optimizer.rebalanceZorderColumns.enabled | false | When true and `spark.sql.optimizer.rebalanceBeforeZorder.enabled` is true, we do rebalance before Z-Order. If it's dynamic partition insert, the rebalance expression will include both partition columns and Z-Order columns. Note that, this config only affects with Spark 3.3.x. | 1.6.0 |
+| spark.sql.optimizer.twoPhaseRebalanceBeforeZorder.enabled | false | When true and `spark.sql.optimizer.rebalanceBeforeZorder.enabled` is true, we do two phase rebalance before Z-Order for the dynamic partition write. The first phase rebalance using dynamic partition column; The second phase rebalance using dynamic partition column Z-Order columns. Note that, this config only affects with Spark 3.3.x. | 1.6.0 |
+| spark.sql.optimizer.zorderUsingOriginalOrdering.enabled | false | When true and `spark.sql.optimizer.rebalanceBeforeZorder.enabled` is true, we do sort by the original ordering i.e. lexicographical order. Note that, this config only affects with Spark 3.3.x. | 1.6.0 |
+| spark.sql.optimizer.inferRebalanceAndSortOrders.enabled | false | When ture, infer columns for rebalance and sort orders from original query, e.g. the join keys from join. It can avoid compression ratio regression. | 1.7.0 |
+| spark.sql.optimizer.inferRebalanceAndSortOrdersMaxColumns | 3 | The max columns of inferred columns. | 1.7.0 |
+| spark.sql.optimizer.insertRepartitionBeforeWriteIfNoShuffle.enabled | false | When true, add repartition even if the original plan does not have shuffle. | 1.7.0 |
+| spark.sql.optimizer.finalStageConfigIsolationWriteOnly.enabled | true | When true, only enable final stage isolation for writing. | 1.7.0 |
+| spark.sql.finalWriteStage.eagerlyKillExecutors.enabled | false | When true, eagerly kill redundant executors before running final write stage. | 1.8.0 |
+| spark.sql.finalWriteStage.retainExecutorsFactor | 1.2 | If the target executors * factor < active executors, and target executors * factor > min executors, then inject kill executors or inject custom resource profile. | 1.8.0 |
+| spark.sql.finalWriteStage.resourceIsolation.enabled | false | When true, make final write stage resource isolation using custom RDD resource profile. | 1.2.0 |
+| spark.sql.finalWriteStageExecutorCores | fallback spark.executor.cores | Specify the executor core request for final write stage. It would be passed to the RDD resource profile. | 1.8.0 |
+| spark.sql.finalWriteStageExecutorMemory | fallback spark.executor.memory | Specify the executor on heap memory request for final write stage. It would be passed to the RDD resource profile. | 1.8.0 |
+| spark.sql.finalWriteStageExecutorMemoryOverhead | fallback spark.executor.memoryOverhead | Specify the executor memory overhead request for final write stage. It would be passed to the RDD resource profile. | 1.8.0 |
+| spark.sql.finalWriteStageExecutorOffHeapMemory | NONE | Specify the executor off heap memory request for final write stage. It would be passed to the RDD resource profile. | 1.8.0 |
diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLExtension.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLExtension.scala
index 0d034c26c..0db9b3ab8 100644
--- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLExtension.scala
+++ b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLExtension.scala
@@ -17,8 +17,7 @@
package org.apache.kyuubi.sql
-import org.apache.spark.FinalStageResourceManager
-import org.apache.spark.sql.SparkSessionExtensions
+import org.apache.spark.sql.{FinalStageResourceManager, InjectCustomResourceProfile, SparkSessionExtensions}
import org.apache.kyuubi.sql.watchdog.{ForcedMaxOutputRowsRule, MaxPartitionStrategy}
@@ -42,5 +41,6 @@ class KyuubiSparkSQLExtension extends (SparkSessionExtensions => Unit) {
extensions.injectPlannerStrategy(MaxPartitionStrategy)
extensions.injectQueryStagePrepRule(FinalStageResourceManager)
+ extensions.injectQueryStagePrepRule(InjectCustomResourceProfile)
}
}
diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/FinalStageResourceManager.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/FinalStageResourceManager.scala
similarity index 93%
rename from extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/FinalStageResourceManager.scala
rename to extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/FinalStageResourceManager.scala
index 8eae5d0e2..f568f9df8 100644
--- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/FinalStageResourceManager.scala
+++ b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/FinalStageResourceManager.scala
@@ -15,14 +15,14 @@
* limitations under the License.
*/
-package org.apache.spark
+package org.apache.spark.sql
import scala.annotation.tailrec
import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer
+import org.apache.spark.{ExecutorAllocationClient, MapOutputTrackerMaster, SparkContext, SparkEnv}
import org.apache.spark.scheduler.cluster.CoarseGrainedSchedulerBackend
-import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.execution.{FilterExec, ProjectExec, SortExec, SparkPlan}
import org.apache.spark.sql.execution.adaptive._
@@ -37,7 +37,8 @@ import org.apache.kyuubi.sql.{KyuubiSQLConf, MarkNumOutputColumnsRule}
* It provide a feature:
* 1. Kill redundant executors before running final write stage
*/
-case class FinalStageResourceManager(session: SparkSession) extends Rule[SparkPlan] {
+case class FinalStageResourceManager(session: SparkSession)
+ extends Rule[SparkPlan] with FinalRebalanceStageHelper {
override def apply(plan: SparkPlan): SparkPlan = {
if (!conf.getConf(KyuubiSQLConf.FINAL_WRITE_STAGE_EAGERLY_KILL_EXECUTORS_ENABLED)) {
return plan
@@ -77,12 +78,13 @@ case class FinalStageResourceManager(session: SparkSession) extends Rule[SparkPl
case AQEShuffleReadExec(stage: ShuffleQueryStageExec, partitionSpecs) =>
// The condition whether inject custom resource profile:
// - target executors < active executors
- // - target executors > min executors
+ // - active executors - target executors > min executors
val numActiveExecutors = sc.getExecutorIds().length
val targetCores = partitionSpecs.length
val targetExecutors = (math.ceil(targetCores.toFloat / coresPerExecutor) * factor).toInt
.max(1)
- val hasBenefits = targetExecutors < numActiveExecutors && targetExecutors > minExecutors
+ val hasBenefits = targetExecutors < numActiveExecutors &&
+ (numActiveExecutors - targetExecutors) > minExecutors
if (hasBenefits) {
val shuffleId = stage.plan.asInstanceOf[ShuffleExchangeExec].shuffleDependency.shuffleId
val numReduce = stage.plan.asInstanceOf[ShuffleExchangeExec].numPartitions
@@ -171,7 +173,7 @@ case class FinalStageResourceManager(session: SparkSession) extends Rule[SparkPl
logInfo(s"Request to kill executors, total count ${executorsToKill.size}, " +
s"[${executorsToKill.mkString(", ")}].")
- // Note, `SparkContext#killExecutors` does not allow with ARD enabled,
+ // Note, `SparkContext#killExecutors` does not allow with DRA enabled,
// see `https://github.com/apache/spark/pull/20604`.
// It may cause the status in `ExecutorAllocationManager` inconsistent with
// `CoarseGrainedSchedulerBackend` for a while. But it should be synchronous finally.
@@ -182,8 +184,15 @@ case class FinalStageResourceManager(session: SparkSession) extends Rule[SparkPl
force = false)
}
+ @transient private val queryStageOptimizerRules: Seq[Rule[SparkPlan]] = Seq(
+ OptimizeSkewInRebalancePartitions,
+ CoalesceShufflePartitions(session),
+ OptimizeShuffleWithLocalRead)
+}
+
+trait FinalRebalanceStageHelper {
@tailrec
- private def findFinalRebalanceStage(plan: SparkPlan): Option[ShuffleQueryStageExec] = {
+ final protected def findFinalRebalanceStage(plan: SparkPlan): Option[ShuffleQueryStageExec] = {
plan match {
case p: ProjectExec => findFinalRebalanceStage(p.child)
case f: FilterExec => findFinalRebalanceStage(f.child)
@@ -196,9 +205,4 @@ case class FinalStageResourceManager(session: SparkSession) extends Rule[SparkPl
case _ => None
}
}
-
- @transient private val queryStageOptimizerRules: Seq[Rule[SparkPlan]] = Seq(
- OptimizeSkewInRebalancePartitions,
- CoalesceShufflePartitions(session),
- OptimizeShuffleWithLocalRead)
}
diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/InjectCustomResourceProfile.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/InjectCustomResourceProfile.scala
new file mode 100644
index 000000000..30c042b2a
--- /dev/null
+++ b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/InjectCustomResourceProfile.scala
@@ -0,0 +1,60 @@
+/*
+ * 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.spark.sql
+
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.{CustomResourceProfileExec, SparkPlan}
+import org.apache.spark.sql.execution.adaptive._
+
+import org.apache.kyuubi.sql.{KyuubiSQLConf, MarkNumOutputColumnsRule}
+
+/**
+ * Inject custom resource profile for final write stage, so we can specify custom
+ * executor resource configs.
+ */
+case class InjectCustomResourceProfile(session: SparkSession)
+ extends Rule[SparkPlan] with FinalRebalanceStageHelper {
+ override def apply(plan: SparkPlan): SparkPlan = {
+ if (!conf.getConf(KyuubiSQLConf.FINAL_WRITE_STAGE_RESOURCE_ISOLATION_ENABLED)) {
+ return plan
+ }
+
+ if (!MarkNumOutputColumnsRule.isWrite(session, plan)) {
+ return plan
+ }
+
+ val stage = findFinalRebalanceStage(plan)
+ if (stage.isEmpty) {
+ return plan
+ }
+
+ // TODO: Ideally, We can call `CoarseGrainedSchedulerBackend.requestTotalExecutors` eagerly
+ // to reduce the task submit pending time, but it may lose task locality.
+ //
+ // By default, it would request executors when catch stage submit event.
+ injectCustomResourceProfile(plan, stage.get.id)
+ }
+
+ private def injectCustomResourceProfile(plan: SparkPlan, id: Int): SparkPlan = {
+ plan match {
+ case stage: ShuffleQueryStageExec if stage.id == id =>
+ CustomResourceProfileExec(stage)
+ case _ => plan.mapChildren(child => injectCustomResourceProfile(child, id))
+ }
+ }
+}
diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala
new file mode 100644
index 000000000..3698140fb
--- /dev/null
+++ b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala
@@ -0,0 +1,112 @@
+/*
+ * 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.spark.sql.execution
+
+import org.apache.spark.network.util.{ByteUnit, JavaUtils}
+import org.apache.spark.rdd.RDD
+import org.apache.spark.resource.{ExecutorResourceRequests, ResourceProfileBuilder}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, SortOrder}
+import org.apache.spark.sql.catalyst.plans.physical.Partitioning
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+import org.apache.spark.sql.vectorized.ColumnarBatch
+import org.apache.spark.util.Utils
+
+import org.apache.kyuubi.sql.KyuubiSQLConf._
+
+/**
+ * This node wraps the final executed plan and inject custom resource profile to the RDD.
+ * It assumes that, the produced RDD would create the `ResultStage` in `DAGScheduler`,
+ * so it makes resource isolation between previous and final stage.
+ *
+ * Note that, Spark does not support config `minExecutors` for each resource profile.
+ * Which means, it would retain `minExecutors` for each resource profile.
+ * So, suggest set `spark.dynamicAllocation.minExecutors` to 0 if enable this feature.
+ */
+case class CustomResourceProfileExec(child: SparkPlan) extends UnaryExecNode {
+ override def output: Seq[Attribute] = child.output
+ override def outputPartitioning: Partitioning = child.outputPartitioning
+ override def outputOrdering: Seq[SortOrder] = child.outputOrdering
+ override def supportsColumnar: Boolean = child.supportsColumnar
+ override def supportsRowBased: Boolean = child.supportsRowBased
+ override protected def doCanonicalize(): SparkPlan = child.canonicalized
+
+ private val executorCores = conf.getConf(FINAL_WRITE_STAGE_EXECUTOR_CORES).getOrElse(
+ sparkContext.getConf.getInt("spark.executor.cores", 1))
+ private val executorMemory = conf.getConf(FINAL_WRITE_STAGE_EXECUTOR_MEMORY).getOrElse(
+ sparkContext.getConf.get("spark.executor.memory", "2G"))
+ private val executorMemoryOverhead =
+ conf.getConf(FINAL_WRITE_STAGE_EXECUTOR_MEMORY_OVERHEAD)
+ .getOrElse(sparkContext.getConf.get("spark.executor.memoryOverhead", "1G"))
+ private val executorOffHeapMemory = conf.getConf(FINAL_WRITE_STAGE_EXECUTOR_OFF_HEAP_MEMORY)
+
+ override lazy val metrics: Map[String, SQLMetric] = {
+ val base = Map(
+ "executorCores" -> SQLMetrics.createMetric(sparkContext, "executor cores"),
+ "executorMemory" -> SQLMetrics.createMetric(sparkContext, "executor memory (MiB)"),
+ "executorMemoryOverhead" -> SQLMetrics.createMetric(
+ sparkContext,
+ "executor memory overhead (MiB)"))
+ val addition = executorOffHeapMemory.map(_ =>
+ "executorOffHeapMemory" ->
+ SQLMetrics.createMetric(sparkContext, "executor off heap memory (MiB)")).toMap
+ base ++ addition
+ }
+
+ private def wrapResourceProfile[T](rdd: RDD[T]): RDD[T] = {
+ if (Utils.isTesting) {
+ // do nothing for local testing
+ return rdd
+ }
+
+ metrics("executorCores") += executorCores
+ metrics("executorMemory") += JavaUtils.byteStringAs(executorMemory, ByteUnit.MiB)
+ metrics("executorMemoryOverhead") += JavaUtils.byteStringAs(
+ executorMemoryOverhead,
+ ByteUnit.MiB)
+ executorOffHeapMemory.foreach(m =>
+ metrics("executorOffHeapMemory") += JavaUtils.byteStringAs(m, ByteUnit.MiB))
+
+ val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)
+ SQLMetrics.postDriverMetricUpdates(sparkContext, executionId, metrics.values.toSeq)
+
+ val resourceProfileBuilder = new ResourceProfileBuilder()
+ val executorResourceRequests = new ExecutorResourceRequests()
+ executorResourceRequests.cores(executorCores)
+ executorResourceRequests.memory(executorMemory)
+ executorResourceRequests.memoryOverhead(executorMemoryOverhead)
+ executorOffHeapMemory.foreach(executorResourceRequests.offHeapMemory)
+ resourceProfileBuilder.require(executorResourceRequests)
+ rdd.withResources(resourceProfileBuilder.build())
+ rdd
+ }
+
+ override protected def doExecute(): RDD[InternalRow] = {
+ val rdd = child.execute()
+ wrapResourceProfile(rdd)
+ }
+
+ override protected def doExecuteColumnar(): RDD[ColumnarBatch] = {
+ val rdd = child.executeColumnar()
+ wrapResourceProfile(rdd)
+ }
+
+ override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = {
+ this.copy(child = newChild)
+ }
+}
diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/InjectResourceProfileSuite.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/InjectResourceProfileSuite.scala
new file mode 100644
index 000000000..b0767b187
--- /dev/null
+++ b/extensions/spark/kyuubi-extension-spark-3-3/src/test/scala/org/apache/spark/sql/InjectResourceProfileSuite.scala
@@ -0,0 +1,79 @@
+/*
+ * 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.spark.sql
+
+import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent}
+import org.apache.spark.sql.execution.ui.SparkListenerSQLAdaptiveExecutionUpdate
+
+import org.apache.kyuubi.sql.KyuubiSQLConf
+
+class InjectResourceProfileSuite extends KyuubiSparkSQLExtensionTest {
+ private def checkCustomResourceProfile(sqlString: String, exists: Boolean): Unit = {
+ @volatile var lastEvent: SparkListenerSQLAdaptiveExecutionUpdate = null
+ val listener = new SparkListener {
+ override def onOtherEvent(event: SparkListenerEvent): Unit = {
+ event match {
+ case e: SparkListenerSQLAdaptiveExecutionUpdate => lastEvent = e
+ case _ =>
+ }
+ }
+ }
+
+ spark.sparkContext.addSparkListener(listener)
+ try {
+ sql(sqlString).collect()
+ spark.sparkContext.listenerBus.waitUntilEmpty()
+ assert(lastEvent != null)
+ var current = lastEvent.sparkPlanInfo
+ var shouldStop = false
+ while (!shouldStop) {
+ if (current.nodeName != "CustomResourceProfile") {
+ if (current.children.isEmpty) {
+ assert(!exists)
+ shouldStop = true
+ } else {
+ current = current.children.head
+ }
+ } else {
+ assert(exists)
+ shouldStop = true
+ }
+ }
+ } finally {
+ spark.sparkContext.removeSparkListener(listener)
+ }
+ }
+
+ test("Inject resource profile") {
+ withTable("t") {
+ withSQLConf(
+ "spark.sql.adaptive.forceApply" -> "true",
+ KyuubiSQLConf.FINAL_STAGE_CONFIG_ISOLATION.key -> "true",
+ KyuubiSQLConf.FINAL_WRITE_STAGE_RESOURCE_ISOLATION_ENABLED.key -> "true") {
+
+ sql("CREATE TABLE t (c1 int, c2 string) USING PARQUET")
+
+ checkCustomResourceProfile("INSERT INTO TABLE t VALUES(1, 'a')", false)
+ checkCustomResourceProfile("SELECT 1", false)
+ checkCustomResourceProfile(
+ "INSERT INTO TABLE t SELECT /*+ rebalance */ * FROM VALUES(1, 'a')",
+ true)
+ }
+ }
+ }
+}
diff --git a/extensions/spark/kyuubi-extension-spark-common/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala b/extensions/spark/kyuubi-extension-spark-common/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala
index 15634ee35..4df924b51 100644
--- a/extensions/spark/kyuubi-extension-spark-common/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala
+++ b/extensions/spark/kyuubi-extension-spark-common/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala
@@ -206,4 +206,44 @@ object KyuubiSQLConf {
.doubleConf
.checkValue(_ >= 1, "must be bigger than or equal to 1")
.createWithDefault(1.2)
+
+ val FINAL_WRITE_STAGE_RESOURCE_ISOLATION_ENABLED =
+ buildConf("spark.sql.finalWriteStage.resourceIsolation.enabled")
+ .doc(
+ "When true, make final write stage resource isolation using custom RDD resource profile.")
+ .version("1.8.0")
+ .booleanConf
+ .createWithDefault(false)
+
+ val FINAL_WRITE_STAGE_EXECUTOR_CORES =
+ buildConf("spark.sql.finalWriteStage.executorCores")
+ .doc("Specify the executor core request for final write stage. " +
+ "It would be passed to the RDD resource profile.")
+ .version("1.8.0")
+ .intConf
+ .createOptional
+
+ val FINAL_WRITE_STAGE_EXECUTOR_MEMORY =
+ buildConf("spark.sql.finalWriteStage.executorMemory")
+ .doc("Specify the executor on heap memory request for final write stage. " +
+ "It would be passed to the RDD resource profile.")
+ .version("1.8.0")
+ .stringConf
+ .createOptional
+
+ val FINAL_WRITE_STAGE_EXECUTOR_MEMORY_OVERHEAD =
+ buildConf("spark.sql.finalWriteStage.executorMemoryOverhead")
+ .doc("Specify the executor memory overhead request for final write stage. " +
+ "It would be passed to the RDD resource profile.")
+ .version("1.8.0")
+ .stringConf
+ .createOptional
+
+ val FINAL_WRITE_STAGE_EXECUTOR_OFF_HEAP_MEMORY =
+ buildConf("spark.sql.finalWriteStage.executorOffHeapMemory")
+ .doc("Specify the executor off heap memory request for final write stage. " +
+ "It would be passed to the RDD resource profile.")
+ .version("1.8.0")
+ .stringConf
+ .createOptional
}
diff --git a/extensions/spark/kyuubi-extension-spark-common/src/test/scala/org/apache/spark/sql/KyuubiSparkSQLExtensionTest.scala b/extensions/spark/kyuubi-extension-spark-common/src/test/scala/org/apache/spark/sql/KyuubiSparkSQLExtensionTest.scala
index fd81948c6..e58ac726c 100644
--- a/extensions/spark/kyuubi-extension-spark-common/src/test/scala/org/apache/spark/sql/KyuubiSparkSQLExtensionTest.scala
+++ b/extensions/spark/kyuubi-extension-spark-common/src/test/scala/org/apache/spark/sql/KyuubiSparkSQLExtensionTest.scala
@@ -29,6 +29,8 @@ import org.apache.kyuubi.sql.KyuubiSQLConf
trait KyuubiSparkSQLExtensionTest extends QueryTest
with SQLTestUtils
with AdaptiveSparkPlanHelper {
+ sys.props.put("spark.testing", "1")
+
private var _spark: Option[SparkSession] = None
protected def spark: SparkSession = _spark.getOrElse {
throw new RuntimeException("test spark session don't initial before using it.")