[CELEBORN-1896] delete data from failed to fetch shuffles

client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SparkCommonUtils.java

@@ -52,6 +52,14 @@ public static void validateAttemptConfig(SparkConf conf) throws IllegalArgumentE
     }
   }
 
+  public static String encodeAppShuffleIdentifier(int appShuffleId, TaskContext context) {
+    return appShuffleId + "-" + context.stageId() + "-" + context.stageAttemptNumber();
+  }
+
+  public static String[] decodeAppShuffleIdentifier(String appShuffleIdentifier) {
+    return appShuffleIdentifier.split("-");
+  }
+
   public static int getEncodedAttemptNumber(TaskContext context) {
     return (context.stageAttemptNumber() << 16) | context.attemptNumber();
   }

client-spark/common/src/main/scala/org/apache/celeborn/spark/FailedShuffleCleaner.scala

@@ -0,0 +1,93 @@
+/*
+ * 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.celeborn.spark
+
+import java.util
+import java.util.concurrent.{LinkedBlockingQueue, ScheduledExecutorService, TimeUnit}
+
+import scala.collection.JavaConverters._
+import scala.collection.mutable
+
+import org.apache.spark.shuffle.celeborn.SparkCommonUtils
+
+import org.apache.celeborn.client.LifecycleManager
+import org.apache.celeborn.common.internal.Logging
+import org.apache.celeborn.common.util.ThreadUtils
+
+private[celeborn] class FailedShuffleCleaner(lifecycleManager: LifecycleManager) extends Logging {
+
+  // in celeborn ids
+  private val shufflesToBeCleaned = new LinkedBlockingQueue[Int]()
+  private val cleanedShuffleIds = new mutable.HashSet[Int]
+
+  private lazy val cleanInterval =
+    lifecycleManager.conf.clientFetchCleanFailedShuffleIntervalMS
+
+  // for test
+  def reset(): Unit = {
+    shufflesToBeCleaned.clear()
+    cleanedShuffleIds.clear()
+    if (cleanerThreadPool != null) {
+      cleanerThreadPool.shutdownNow()
+      cleanerThreadPool = null
+    }
+  }
+
+  def addShuffleIdToBeCleaned(appShuffleIdentifier: String): Unit = {
+    val Array(appShuffleId, _, _) = SparkCommonUtils.decodeAppShuffleIdentifier(
+      appShuffleIdentifier)
+    lifecycleManager.getShuffleIdMapping.get(appShuffleId.toInt).foreach {
+      case (_, (celebornShuffleId, _)) => shufflesToBeCleaned.put(celebornShuffleId)
+    }
+  }
+
+  def init(): Unit = {
+    cleanerThreadPool = ThreadUtils.newDaemonSingleThreadScheduledExecutor(
+      "failedShuffleCleanerThreadPool")
+    cleanerThreadPool.scheduleWithFixedDelay(
+      new Runnable {
+        override def run(): Unit = {
+          try {
+            val allShuffleIds = new util.ArrayList[Int]
+            shufflesToBeCleaned.drainTo(allShuffleIds)
+            allShuffleIds.asScala.foreach { shuffleId =>
+              if (!cleanedShuffleIds.contains(shuffleId)) {
+                lifecycleManager.unregisterShuffle(shuffleId)
+                logInfo(
+                  s"sent unregister shuffle request for shuffle $shuffleId (celeborn shuffle id)")
+                cleanedShuffleIds += shuffleId
+              }
+            }
+          } catch {
+            case e: Exception =>
+              logError("unexpected exception in cleaner thread", e)
+          }
+        }
+      },
+      cleanInterval,
+      cleanInterval,
+      TimeUnit.MILLISECONDS)
+  }
+
+  init()
+
+  def removeCleanedShuffleId(celebornShuffleId: Int): Unit = {
+    cleanedShuffleIds.remove(celebornShuffleId)
+  }
+
+  private var cleanerThreadPool: ScheduledExecutorService = _
+}

client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java

@@ -37,6 +37,7 @@
 import org.apache.celeborn.common.CelebornConf;
 import org.apache.celeborn.common.protocol.ShuffleMode;
 import org.apache.celeborn.reflect.DynMethods;
+import org.apache.celeborn.spark.FailedShuffleCleaner;
 
 /**
  * In order to support Spark Stage resubmit with ShuffleReader FetchFails, Celeborn shuffleId has to
@@ -84,6 +85,8 @@ public class SparkShuffleManager implements ShuffleManager {
       ConcurrentHashMap.newKeySet();
   private final CelebornShuffleFallbackPolicyRunner fallbackPolicyRunner;
 
+  private FailedShuffleCleaner failedShuffleCleaner = null;
+
   private long sendBufferPoolCheckInterval;
   private long sendBufferPoolExpireTimeout;
 
@@ -158,6 +161,23 @@ private void initializeLifecycleManager(String appId) {
             }
           }
 
+          if (lifecycleManager.conf().clientFetchCleanFailedShuffle()) {
+            if (!lifecycleManager.conf().clientStageRerunEnabled()) {
+              throw new IllegalArgumentException(
+                  CelebornConf.CLIENT_STAGE_RERUN_ENABLED().key()
+                      + " has to be "
+                      + "enabled, when "
+                      + CelebornConf.CLIENT_FETCH_CLEAN_FAILED_SHUFFLE().key()
+                      + " is set to true");
+            }
+            failedShuffleCleaner = new FailedShuffleCleaner(lifecycleManager);
+            lifecycleManager.registerValidateCelebornShuffleIdForCleanCallback(
+                (appShuffleIdentifier) ->
+                    SparkUtils.addWriterShuffleIdsToBeCleaned(this, appShuffleIdentifier));
+            lifecycleManager.registerUnregisterShuffleCallback(
+                (celebornShuffleId) -> SparkUtils.removeCleanedShuffleId(this, celebornShuffleId));
+          }
+
           if (celebornConf.getReducerFileGroupBroadcastEnabled()) {
             lifecycleManager.registerBroadcastGetReducerFileGroupResponseCallback(
                 (shuffleId, getReducerFileGroupResponse) ->
@@ -249,6 +269,9 @@ public void stop() {
       _sortShuffleManager.stop();
       _sortShuffleManager = null;
     }
+    if (celebornConf.clientFetchCleanFailedShuffle()) {
+      failedShuffleCleaner.reset();
+    }
   }
 
   @Override
@@ -470,4 +493,8 @@ private void checkUserClassPathFirst(ShuffleHandle handle) {
   public LifecycleManager getLifecycleManager() {
     return this.lifecycleManager;
   }
+
+  public FailedShuffleCleaner getFailedShuffleCleaner() {
+    return this.failedShuffleCleaner;
+  }
 }

client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java

@@ -128,17 +128,14 @@ public static String appUniqueId(SparkContext context) {
         .getOrElse(context::applicationId);
   }
 
-  public static String getAppShuffleIdentifier(int appShuffleId, TaskContext context) {
-    return appShuffleId + "-" + context.stageId() + "-" + context.stageAttemptNumber();
-  }
-
   public static int celebornShuffleId(
       ShuffleClient client,
       CelebornShuffleHandle<?, ?, ?> handle,
       TaskContext context,
       Boolean isWriter) {
     if (handle.throwsFetchFailure()) {
-      String appShuffleIdentifier = getAppShuffleIdentifier(handle.shuffleId(), context);
+      String appShuffleIdentifier =
+          SparkCommonUtils.encodeAppShuffleIdentifier(handle.shuffleId(), context);
       Tuple2<Integer, Boolean> res =
           client.getShuffleId(
               handle.shuffleId(),
@@ -327,7 +324,8 @@ public static void addFailureListenerIfBarrierTask(
 
     if (!(taskContext instanceof BarrierTaskContext)) return;
     int appShuffleId = handle.shuffleId();
-    String appShuffleIdentifier = SparkUtils.getAppShuffleIdentifier(appShuffleId, taskContext);
+    String appShuffleIdentifier =
+        SparkCommonUtils.encodeAppShuffleIdentifier(appShuffleId, taskContext);
 
     BarrierTaskContext barrierContext = (BarrierTaskContext) taskContext;
     barrierContext.addTaskFailureListener(
@@ -625,4 +623,14 @@ public static void invalidateSerializedGetReducerFileGroupResponse(Integer shuff
           return null;
         });
   }
+
+  public static void addWriterShuffleIdsToBeCleaned(
+      SparkShuffleManager sparkShuffleManager, String appShuffleIdentifier) {
+    sparkShuffleManager.getFailedShuffleCleaner().addShuffleIdToBeCleaned(appShuffleIdentifier);
+  }
+
+  public static void removeCleanedShuffleId(
+      SparkShuffleManager sparkShuffleManager, int celebornShuffleId) {
+    sparkShuffleManager.getFailedShuffleCleaner().removeCleanedShuffleId(celebornShuffleId);
+  }
 }

client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala

@@ -931,6 +931,7 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
                   logInfo(s"reuse existing shuffleId $id for appShuffleId $appShuffleId appShuffleIdentifier $appShuffleIdentifier")
                   id
                 } else {
+                  // this branch means it is a redo of previous write stage
                   if (isBarrierStage) {
                     // unregister previous shuffle(s) which are still valid
                     val mapUpdates = shuffleIds.filter(_._2._2).map { kv =>
@@ -941,6 +942,8 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
                   }
                   val newShuffleId = shuffleIdGenerator.getAndIncrement()
                   logInfo(s"generate new shuffleId $newShuffleId for appShuffleId $appShuffleId appShuffleIdentifier $appShuffleIdentifier")
+                  validateCelebornShuffleIdForClean.foreach(callback =>
+                    callback.accept(appShuffleIdentifier))
                   shuffleIds.put(appShuffleIdentifier, (newShuffleId, true))
                   newShuffleId
                 }
@@ -954,11 +957,12 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
       } else {
         shuffleIds.values.filter(v => v._2).map(v => v._1).toSeq.reverse.find(
           areAllMapTasksEnd) match {
-          case Some(shuffleId) =>
+          case Some(celebornShuffleId) =>
             val pbGetShuffleIdResponse = {
               logDebug(
-                s"get shuffleId $shuffleId for appShuffleId $appShuffleId appShuffleIdentifier $appShuffleIdentifier isWriter $isWriter")
-              PbGetShuffleIdResponse.newBuilder().setShuffleId(shuffleId).setSuccess(true).build()
+                s"get shuffleId $celebornShuffleId for appShuffleId $appShuffleId appShuffleIdentifier $appShuffleIdentifier isWriter $isWriter")
+              PbGetShuffleIdResponse.newBuilder().setShuffleId(celebornShuffleId).setSuccess(
+                true).build()
             }
             context.reply(pbGetShuffleIdResponse)
           case None =>
@@ -1160,6 +1164,7 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
           shuffleIds.values.map {
             case (shuffleId, _) =>
               unregisterShuffle(shuffleId)
+              unregisterShuffleCallback.foreach(c => c.accept(shuffleId))
           })
       }
     } else {
@@ -1850,6 +1855,19 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
     appShuffleTrackerCallback = Some(callback)
   }
 
+  // expecting celeborn shuffle id and application shuffle identifier
+  @volatile private var validateCelebornShuffleIdForClean: Option[Consumer[String]] =
+    None
+  def registerValidateCelebornShuffleIdForCleanCallback(
+      callback: Consumer[String]): Unit = {
+    validateCelebornShuffleIdForClean = Some(callback)
+  }
+
+  @volatile private var unregisterShuffleCallback: Option[Consumer[Integer]] = None
+  def registerUnregisterShuffleCallback(callback: Consumer[Integer]): Unit = {
+    unregisterShuffleCallback = Some(callback)
+  }
+
   def registerAppShuffleDeterminate(appShuffleId: Int, determinate: Boolean): Unit = {
     appShuffleDeterminateMap.put(appShuffleId, determinate)
   }
@@ -1943,4 +1961,6 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
         }
       })
   }
+
+  def getShuffleIdMapping = shuffleIdMapping
 }

client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala

@@ -70,7 +70,7 @@ class ReducePartitionCommitHandler(
 
   private val getReducerFileGroupRequest =
     JavaUtils.newConcurrentHashMap[Int, util.Set[MultiSerdeVersionRpcContext]]()
-  private val dataLostShuffleSet = ConcurrentHashMap.newKeySet[Int]()
+  private[celeborn] val dataLostShuffleSet = ConcurrentHashMap.newKeySet[Int]()
   private val stageEndShuffleSet = ConcurrentHashMap.newKeySet[Int]()
   private val inProcessStageEndShuffleSet = ConcurrentHashMap.newKeySet[Int]()
   private val shuffleMapperAttempts = JavaUtils.newConcurrentHashMap[Int, Array[Int]]()

common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala

@@ -997,6 +997,9 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable with Logging with Se
 
   def clientFetchMaxRetriesForEachReplica: Int = get(CLIENT_FETCH_MAX_RETRIES_FOR_EACH_REPLICA)
   def clientStageRerunEnabled: Boolean = get(CLIENT_STAGE_RERUN_ENABLED)
+  def clientFetchCleanFailedShuffle: Boolean = get(CLIENT_FETCH_CLEAN_FAILED_SHUFFLE)
+  def clientFetchCleanFailedShuffleIntervalMS: Long =
+    get(CLIENT_FETCH_CLEAN_FAILED_SHUFFLE_INTERVAL)
   def clientFetchExcludeWorkerOnFailureEnabled: Boolean =
     get(CLIENT_FETCH_EXCLUDE_WORKER_ON_FAILURE_ENABLED)
   def clientFetchExcludedWorkerExpireTimeout: Long =
@@ -4813,6 +4816,23 @@ object CelebornConf extends Logging {
       .booleanConf
       .createWithDefault(true)
 
+  val CLIENT_FETCH_CLEAN_FAILED_SHUFFLE: ConfigEntry[Boolean] =
+    buildConf("celeborn.client.spark.fetch.cleanFailedShuffle")
+      .categories("client")
+      .version("0.6.0")
+      .doc("whether to clean those disk space occupied by shuffles which cannot be fetched")
+      .booleanConf
+      .createWithDefault(false)
+
+  val CLIENT_FETCH_CLEAN_FAILED_SHUFFLE_INTERVAL: ConfigEntry[Long] =
+    buildConf("celeborn.client.spark.fetch.cleanFailedShuffleInterval")
+      .categories("client")
+      .version("0.6.0")
+      .doc("the interval to clean the failed-to-fetch shuffle files, only valid when" +
+        s" ${CLIENT_FETCH_CLEAN_FAILED_SHUFFLE.key} is enabled")
+      .timeConf(TimeUnit.MILLISECONDS)
+      .createWithDefaultString("1s")
+
   val CLIENT_FETCH_EXCLUDE_WORKER_ON_FAILURE_ENABLED: ConfigEntry[Boolean] =
     buildConf("celeborn.client.fetch.excludeWorkerOnFailure.enabled")
       .categories("client")

docs/configuration/client.md

@@ -111,6 +111,8 @@ license: |
 | celeborn.client.shuffle.register.filterExcludedWorker.enabled | false | false | Whether to filter excluded worker when register shuffle. | 0.4.0 |  | 
 | celeborn.client.shuffle.reviseLostShuffles.enabled | false | false | Whether to revise lost shuffles. | 0.6.0 |  | 
 | celeborn.client.slot.assign.maxWorkers | 10000 | false | Max workers that slots of one shuffle can be allocated on. Will choose the smaller positive one from Master side and Client side, see `celeborn.master.slot.assign.maxWorkers`. | 0.3.1 |  | 
+| celeborn.client.spark.fetch.cleanFailedShuffle | false | false | whether to clean those disk space occupied by shuffles which cannot be fetched | 0.6.0 |  | 
+| celeborn.client.spark.fetch.cleanFailedShuffleInterval | 1s | false | the interval to clean the failed-to-fetch shuffle files, only valid when celeborn.client.spark.fetch.cleanFailedShuffle is enabled | 0.6.0 |  | 
 | celeborn.client.spark.push.dynamicWriteMode.enabled | false | false | Whether to dynamically switch push write mode based on conditions.If true, shuffle mode will be only determined by partition count | 0.5.0 |  | 
 | celeborn.client.spark.push.dynamicWriteMode.partitionNum.threshold | 2000 | false | Threshold of shuffle partition number for dynamically switching push writer mode. When the shuffle partition number is greater than this value, use the sort-based shuffle writer for memory efficiency; otherwise use the hash-based shuffle writer for speed. This configuration only takes effect when celeborn.client.spark.push.dynamicWriteMode.enabled is true. | 0.5.0 |  | 
 | celeborn.client.spark.push.sort.memory.maxMemoryFactor | 0.4 | false | the max portion of executor memory which can be used for SortBasedWriter buffer (only valid when celeborn.client.spark.push.sort.memory.useAdaptiveThreshold is enabled | 0.5.0 |  | 

pom.xml

@@ -907,7 +907,7 @@
               <log4j.configuration>file:src/test/resources/log4j.properties</log4j.configuration>
               <log4j2.configurationFile>src/test/resources/log4j2-test.xml</log4j2.configurationFile>
               <java.io.tmpdir>${project.build.directory}/tmp</java.io.tmpdir>
-              <spark.driver.memory>1g</spark.driver.memory>
+              <spark.driver.memory>8g</spark.driver.memory>
               <spark.shuffle.sort.io.plugin.class>${spark.shuffle.plugin.class}</spark.shuffle.sort.io.plugin.class>
             </systemProperties>
             <environmentVariables>
@@ -946,7 +946,7 @@
               <log4j.configuration>file:src/test/resources/log4j.properties</log4j.configuration>
               <log4j2.configurationFile>src/test/resources/log4j2-test.xml</log4j2.configurationFile>
               <java.io.tmpdir>${project.build.directory}/tmp</java.io.tmpdir>
-              <spark.driver.memory>1g</spark.driver.memory>
+              <spark.driver.memory>8g</spark.driver.memory>
               <spark.shuffle.sort.io.plugin.class>${spark.shuffle.plugin.class}</spark.shuffle.sort.io.plugin.class>
             </systemProperties>
             <environmentVariables>