yew1eb commented on code in PR #3853: URL: https://github.com/apache/celeborn/pull/3853#discussion_r4366204533
########## client-spark/spark-3/src/main/scala/org/apache/spark/sql/celeborn/CelebornShuffleCleanupListener.scala: ########## @@ -0,0 +1,177 @@ +/* + * 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.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. Shuffles of failed or + * cancelled queries (or stages) that were never materialized are left to GC cleanup. + */ +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] = { + collectShuffleIds(plan, 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 could even launch jobs, e.g. range-partition sampling) on this thread, + // so only exchanges/stages known to have executed may be touched. + private def collectShuffleIds(plan: SparkPlan, queryFailed: Boolean): Seq[Int] = { + val shuffleIds = ArrayBuffer.empty[Int] + + // inAdaptivePlan: inside an AdaptiveSparkPlanExec subtree, where raw exchanges may be + // unexecuted; their shuffles are collected from the shared stage cache instead. + def visit(p: SparkPlan, inAdaptivePlan: Boolean): Unit = { + p match { + case adaptivePlan: AdaptiveSparkPlanExec => + // The stage cache is shared by the whole query, including subqueries, and holds + // every stage the query created -- even stages dropped from the final plan (e.g. + // an empty join side) and children of unmaterialized stages. Only materialized + // stages are read, since their shuffleDependency is already computed. + adaptivePlan.context.stageCache.values.foreach { + case stage: ShuffleQueryStageExec if isStageMaterialized(stage) => + // Non-vanilla ShuffleExchangeLike (e.g. Gluten's columnar exchange) is ignored. + stage.shuffle match { + case exchange: ShuffleExchangeExec => + shuffleIds += exchange.shuffleDependency.shuffleId + case _ => + } + case _ => + } + // Keep looking for nested AdaptiveSparkPlanExec nodes (e.g. under a + // DataWritingCommandExec for CTAS, or subquery plans); stages already drained + // from a shared cache are only re-collected as duplicates, deduplicated later. + visit(adaptivePlan.executedPlan, inAdaptivePlan = true) + + // Shuffle stages encountered here were already collected from the stage cache + // above, so there is nothing more to do. This must be matched before the generic + // QueryStageExec case below. + case _: ShuffleQueryStageExec => () Review Comment: Good catch! Went with (a): materialized `ShuffleQueryStageExec`s met in the plan tree are now collected as well, since the stage cache only fills when `spark.sql.exchange.reuse` is on. The walk also descends into stage plans so nested child stages are reached with reuse off; duplicates with the stage cache are removed by `distinct`. Added a reuse=off test (group-by + multi-stage shuffle join); the suite passes on both Spark 3.5 and 4.0. -- 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]
