venkata91 commented on code in PR #3853: URL: https://github.com/apache/celeborn/pull/3853#discussion_r4251108649
########## client-spark/spark-3/src/main/scala/org/apache/spark/sql/celeborn/CelebornShuffleCleanupListener.scala: ########## @@ -0,0 +1,160 @@ +/* + * 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.celeborn + +import java.lang.reflect.Method +import java.util.concurrent.atomic.AtomicReference + +import scala.collection.mutable.ArrayBuffer +import scala.util.Try +import scala.util.control.NonFatal + +import org.apache.spark.internal.Logging +import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent} +import org.apache.spark.sql.catalyst.plans.physical.RangePartitioning +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, QueryStageExec, ShuffleQueryStageExec} +import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec} +import org.apache.spark.sql.execution.ui.SparkListenerSQLExecutionEnd + +import org.apache.celeborn.client.LifecycleManager + +/** + * Unregisters the shuffles of a SQL execution once it completes, so that long-running + * applications (ThriftServer, Kyuubi, spark-sql, spark-connect-server) do not pile up + * shuffle data while waiting for driver GC. Stage rerun must be enabled so that a + * re-executed DataFrame can regenerate its cleaned shuffles. + */ +class CelebornShuffleCleanupListener(lifecycleManager: LifecycleManager) + extends SparkListener with Logging { + + import CelebornShuffleCleanupListener._ + + override def onOtherEvent(event: SparkListenerEvent): Unit = { + event match { + case end: SparkListenerSQLExecutionEnd => + try { + cleanupShufflesOnQueryEnd(end) + } catch { + case NonFatal(t) => + logWarning( + s"Failed to cleanup shuffles on completion of SQL execution ${end.executionId}.", + t) + } + case _ => + } + } + + private def extractShuffleIds(plan: SparkPlan, queryFailed: Boolean): Seq[Int] = { + (plan +: plan.subqueriesAll).flatMap(collectShuffleIds(_, queryFailed)).distinct + } + + // Runs on the listener-bus thread, so it must never trigger computation. Reading the lazy + // ShuffleExchangeExec.shuffleDependency of an exchange that never executed would run its + // child plan and register a new shuffle, so only materialized exchanges may be touched. + private def collectShuffleIds(plan: SparkPlan, queryFailed: Boolean): Seq[Int] = { + val shuffleIds = ArrayBuffer.empty[Int] + + // inAdaptivePlan: inside an AdaptiveSparkPlanExec without having crossed a materialized + // shuffle stage; raw exchanges there may be unexecuted (failed/cancelled AQE queries). + def visit(p: SparkPlan, inAdaptivePlan: Boolean): Unit = p match { + case adaptivePlan: AdaptiveSparkPlanExec => + visit(adaptivePlan.executedPlan, inAdaptivePlan = true) + + case shuffleStage: ShuffleQueryStageExec => + if (isStageMaterialized(shuffleStage)) { + // Non-vanilla ShuffleExchangeLike (e.g. Gluten's columnar exchange) is ignored. + shuffleStage.shuffle match { + case exchange: ShuffleExchangeExec => + shuffleIds += exchange.shuffleDependency.shuffleId + case _ => + } + // A materialized stage only references materialized child stages; keep descending. + visit(shuffleStage.plan, inAdaptivePlan = false) + } + + // Spark 4 wraps the final AQE plan in a ResultQueryStageExec leaf. + case stage: QueryStageExec => + visit(stage.plan, inAdaptivePlan) + + case reused: ReusedExchangeExec => + visit(reused.child, inAdaptivePlan) + + case exchange: ShuffleExchangeExec => + // In a failed query an exchange here may be unexecuted; RangePartitioning ones are + // skipped since reading them could launch a sampling job on this thread. + if (!inAdaptivePlan && + (!queryFailed || !exchange.outputPartitioning.isInstanceOf[RangePartitioning])) { + shuffleIds += exchange.shuffleDependency.shuffleId + } + exchange.children.foreach(visit(_, inAdaptivePlan)) Review Comment: In bd7f9b3, this still reads `shuffleDependency` for a failed non-AQE query when the outer exchange isn't `RangePartitioning`. That lazy val runs `child.execute()`, so a nested range exchange, subquery or broadcast underneath can still start a job. Repro (AQE off): `bad.union(good).collect()`, where each side is `repartitionByRange(4, x).groupBy(x % 10).count()` and `bad` has a throwing UDF. After the query fails, the driver log shows `Got job 2 (onOtherEvent at SparkListenerBus.scala:100)`, coming from `CelebornShuffleCleanupListener.scala:103 -> shuffleDependency$lzycompute -> RangePartitioner.sketch`. The exception then aborts the walk, so `good`'s shuffle isn't cleaned either. Some options: (a) Skip raw exchanges when `executionFailure` is set and leave those shuffles to GC cleanup. This is simplest, and it's what Spark's own query-end cleanup effectively does, since it only uses AQE stage info. (b) Only read the id when the exchange has already executed, e.g. by checking whether `mapOutputStatisticsFuture` / `cachedShuffleRDD` was initialized. This relies on Spark internals that differ across versions. (c) Catch per exchange so that one failure doesn't stop the rest of the walk. This contains the damage, but the job still runs. ########## client-spark/spark-3/src/main/scala/org/apache/spark/sql/celeborn/CelebornShuffleCleanupListener.scala: ########## @@ -0,0 +1,160 @@ +/* + * 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.celeborn + +import java.lang.reflect.Method +import java.util.concurrent.atomic.AtomicReference + +import scala.collection.mutable.ArrayBuffer +import scala.util.Try +import scala.util.control.NonFatal + +import org.apache.spark.internal.Logging +import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent} +import org.apache.spark.sql.catalyst.plans.physical.RangePartitioning +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, QueryStageExec, ShuffleQueryStageExec} +import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec} +import org.apache.spark.sql.execution.ui.SparkListenerSQLExecutionEnd + +import org.apache.celeborn.client.LifecycleManager + +/** + * Unregisters the shuffles of a SQL execution once it completes, so that long-running + * applications (ThriftServer, Kyuubi, spark-sql, spark-connect-server) do not pile up + * shuffle data while waiting for driver GC. Stage rerun must be enabled so that a + * re-executed DataFrame can regenerate its cleaned shuffles. + */ +class CelebornShuffleCleanupListener(lifecycleManager: LifecycleManager) + extends SparkListener with Logging { + + import CelebornShuffleCleanupListener._ + + override def onOtherEvent(event: SparkListenerEvent): Unit = { + event match { + case end: SparkListenerSQLExecutionEnd => + try { + cleanupShufflesOnQueryEnd(end) + } catch { + case NonFatal(t) => + logWarning( + s"Failed to cleanup shuffles on completion of SQL execution ${end.executionId}.", + t) + } + case _ => + } + } + + private def extractShuffleIds(plan: SparkPlan, queryFailed: Boolean): Seq[Int] = { + (plan +: plan.subqueriesAll).flatMap(collectShuffleIds(_, queryFailed)).distinct + } Review Comment: With AQE on, scalar-subquery shuffles aren't found. The root is an `AdaptiveSparkPlanExec`, which is a leaf, so `subqueriesAll` is empty. The subqueries live in expressions of nodes inside `executedPlan`, and `visit` only follows `children`. Repro (AQE on): `SELECT COUNT(*) FROM ta WHERE id > (SELECT MAX(c) FROM (SELECT id % 10 k, COUNT(1) c FROM ta GROUP BY id % 10))` cleans only the main-query shuffle. The two subquery shuffles stay registered. Some options: (a) Also visit `p.subqueries` for each node inside `visit`. That covers both AQE and non-AQE, and the top-level `subqueriesAll` would no longer be needed. (b) Read stages from `AdaptiveSparkPlanExec.context.stageCache`. AQE subqueries share the context, so their stages are in it too (see the next comment). ########## client-spark/spark-3/src/main/scala/org/apache/spark/sql/celeborn/CelebornShuffleCleanupListener.scala: ########## @@ -0,0 +1,160 @@ +/* + * 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.celeborn + +import java.lang.reflect.Method +import java.util.concurrent.atomic.AtomicReference + +import scala.collection.mutable.ArrayBuffer +import scala.util.Try +import scala.util.control.NonFatal + +import org.apache.spark.internal.Logging +import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent} +import org.apache.spark.sql.catalyst.plans.physical.RangePartitioning +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, QueryStageExec, ShuffleQueryStageExec} +import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec} +import org.apache.spark.sql.execution.ui.SparkListenerSQLExecutionEnd + +import org.apache.celeborn.client.LifecycleManager + +/** + * Unregisters the shuffles of a SQL execution once it completes, so that long-running + * applications (ThriftServer, Kyuubi, spark-sql, spark-connect-server) do not pile up + * shuffle data while waiting for driver GC. Stage rerun must be enabled so that a + * re-executed DataFrame can regenerate its cleaned shuffles. + */ +class CelebornShuffleCleanupListener(lifecycleManager: LifecycleManager) + extends SparkListener with Logging { + + import CelebornShuffleCleanupListener._ + + override def onOtherEvent(event: SparkListenerEvent): Unit = { + event match { + case end: SparkListenerSQLExecutionEnd => + try { + cleanupShufflesOnQueryEnd(end) + } catch { + case NonFatal(t) => + logWarning( + s"Failed to cleanup shuffles on completion of SQL execution ${end.executionId}.", + t) + } + case _ => + } + } + + private def extractShuffleIds(plan: SparkPlan, queryFailed: Boolean): Seq[Int] = { + (plan +: plan.subqueriesAll).flatMap(collectShuffleIds(_, queryFailed)).distinct + } + + // Runs on the listener-bus thread, so it must never trigger computation. Reading the lazy + // ShuffleExchangeExec.shuffleDependency of an exchange that never executed would run its + // child plan and register a new shuffle, so only materialized exchanges may be touched. + private def collectShuffleIds(plan: SparkPlan, queryFailed: Boolean): Seq[Int] = { + val shuffleIds = ArrayBuffer.empty[Int] + + // inAdaptivePlan: inside an AdaptiveSparkPlanExec without having crossed a materialized + // shuffle stage; raw exchanges there may be unexecuted (failed/cancelled AQE queries). + def visit(p: SparkPlan, inAdaptivePlan: Boolean): Unit = p match { + case adaptivePlan: AdaptiveSparkPlanExec => + visit(adaptivePlan.executedPlan, inAdaptivePlan = true) + + case shuffleStage: ShuffleQueryStageExec => + if (isStageMaterialized(shuffleStage)) { + // Non-vanilla ShuffleExchangeLike (e.g. Gluten's columnar exchange) is ignored. + shuffleStage.shuffle match { + case exchange: ShuffleExchangeExec => + shuffleIds += exchange.shuffleDependency.shuffleId + case _ => + } + // A materialized stage only references materialized child stages; keep descending. + visit(shuffleStage.plan, inAdaptivePlan = false) + } Review Comment: Two cases with AQE are missed here: - If a stage isn't materialized, we don't descend into it, so its materialized children are skipped. For example, if the UDF fails in the second map stage, the first stage's shuffle stays registered. - Stages that AQE drops from the final plan are never visited. For example, `ta JOIN (SELECT * FROM tb WHERE id < 0) b ON ...` with broadcast off leaves both shuffles registered after the query ends. Some options: (a) Descend into the stage regardless of whether it is materialized, since each stage below is checked on its own. This fixes the first case only. (b) Also collect the materialized `ShuffleQueryStageExec`s from `adaptivePlan.context.stageCache`. That holds every stage the query created; I checked it exists in 3.0 through 4.x. This covers both cases and is close to Spark 4's `context.shuffleIds`. Minor: the `inAdaptivePlan = false` reset on L88 doesn't seem needed. Keeping `true` would guarantee that no raw exchange inside an AQE subtree is ever touched. ########## client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java: ########## @@ -179,6 +180,18 @@ private void initializeLifecycleManager(String appId) { } } + if (celebornConf.clientSparkShuffleCleanupEnabled()) { + if (!celebornConf.clientStageRerunEnabled()) { + throw new IllegalArgumentException( + CelebornConf.CLIENT_STAGE_RERUN_ENABLED().key() + + " has to be " + + "enabled, when " + + CelebornConf.CLIENT_SPARK_SHUFFLE_CLEANUP_ENABLED().key() + + " is set to true"); + } + SparkUtils.addSparkListener(new CelebornShuffleCleanupListener(lifecycleManager)); + } Review Comment: This check runs lazily on the first `registerShuffle`, after `lifecycleManager` has already been assigned. With stageRerun=false, the first query fails, but later queries run normally with cleanup silently off. The `cleanFailedShuffle` and broadcast callbacks below are also never registered. Some options: (a) Move the check to the constructor (driver only, e.g. next to `validateAttemptConfig`), so the SparkContext fails to start. (b) Keep it here, but validate before `lifecycleManager` is assigned, so every query fails consistently. (c) Log a warning and disable cleanup instead of throwing. (a) seems the most predictable to me. The fail-fast test would then intercept around `getOrCreate()`. ########## client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala: ########## @@ -1304,6 +1304,14 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends } } + def isAppShuffleRegistered(appShuffleId: Int, hasMapping: Boolean): Boolean = { + if (hasMapping) { Review Comment: Nit: now that stageRerun is required whenever this is used, `hasMapping` is always true. We could either drop the parameter and check `shuffleIdMapping` only, or keep it if you expect other callers. ########## common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala: ########## @@ -5273,6 +5275,19 @@ object CelebornConf extends Logging { .booleanConf .createWithDefault(true) + val CLIENT_SPARK_SHUFFLE_CLEANUP_ENABLED: ConfigEntry[Boolean] = + buildConf("celeborn.client.spark.shuffleCleanup.enabled") Review Comment: Nit: `shuffleCleanup.enabled` could be read as general shuffle cleanup, which already happens via ContextCleaner. Some options: keep it as is, or use a name that says when it fires, e.g. `celeborn.client.spark.sql.cleanShuffleOnQueryEnd.enabled`. No strong opinion. Separately, it might help to note the limits in the doc: it only applies to SQL/Dataset executions, non-vanilla exchanges such as Gluten aren't covered, and shuffles of failed or cancelled stages fall back to GC cleanup. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
