pass first test

Author: CodingCat

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

@@ -24,6 +24,7 @@
 import org.apache.celeborn.spark.StageDependencyManager;
 import org.apache.spark.*;
 import org.apache.spark.launcher.SparkLauncher;
+import org.apache.spark.listener.ListenerHelper;
 import org.apache.spark.rdd.DeterministicLevel;
 import org.apache.spark.scheduler.RunningStageManager;
 import org.apache.spark.scheduler.RunningStageManagerImpl;
@@ -162,6 +163,7 @@ private void initializeLifecycleManager() {
           }
           if (lifecycleManager.conf().clientShuffleEarlyDeletion()) {
             logger.info("register early deletion callbacks");
+            ListenerHelper.addShuffleStatsTrackingListener();
             lifecycleManager.registerStageToWriteCelebornShuffleCallback(
                 (celebornShuffleId, appShuffleIdentifier) ->
                     SparkUtils.addStageToWriteCelebornShuffleIdDep(

…istner/CelebornShuffleEarlyCleanup.scala …stener/CelebornShuffleEarlyCleanup.scalaclient-spark/spark-3/src/main/scala/org/apache/spark/listner/CelebornShuffleEarlyCleanup.scala renamed to client-spark/spark-3/src/main/scala/org/apache/spark/listener/CelebornShuffleEarlyCleanup.scala client-spark/spark-3/src/main/scala/org/apache/spark/listner/CelebornShuffleEarlyCleanup.scala renamed to client-spark/spark-3/src/main/scala/org/apache/spark/listener/CelebornShuffleEarlyCleanup.scala

@@ -43,7 +43,7 @@ class ShuffleStatsTrackingListener extends SparkListener with Logging {
   override def onStageCompleted(stageCompleted: SparkListenerStageCompleted): Unit = {
     val stageIdentifier = s"${stageCompleted.stageInfo.stageId}-" +
       s"${stageCompleted.stageInfo.attemptNumber()}"
-    logInfo(s"stage $stageIdentifier finished with" +
+    println(s"stage $stageIdentifier finished with" +
       s" ${stageCompleted.stageInfo.taskMetrics.shuffleWriteMetrics.bytesWritten} shuffle bytes")
     val shuffleMgr = SparkEnv.get.shuffleManager.asInstanceOf[SparkShuffleManager]
     if (shuffleMgr.getLifecycleManager.conf.clientShuffleEarlyDeletion) {

…r/CelebornShuffleEarlyCleanupEvent.scala …r/CelebornShuffleEarlyCleanupEvent.scalaclient-spark/spark-3/src/main/scala/org/apache/spark/listner/CelebornShuffleEarlyCleanupEvent.scala renamed to client-spark/spark-3/src/main/scala/org/apache/spark/listener/CelebornShuffleEarlyCleanupEvent.scala client-spark/spark-3/src/main/scala/org/apache/spark/listner/CelebornShuffleEarlyCleanupEvent.scala renamed to client-spark/spark-3/src/main/scala/org/apache/spark/listener/CelebornShuffleEarlyCleanupEvent.scala

@@ -89,11 +89,11 @@ object StorageCheckUtils {
   }
 
   def triggerStorageCheckThread(
-                                 workerDirs: Seq[String],
-                                 shuffleIdShouldNotExist: Seq[Int],
-                                 shuffleIdMustExist: Seq[Int],
-                                 sparkSession: SparkSession,
-                                 forStableStatusChecking: Boolean): CheckingThread = {
+      workerDirs: Seq[String],
+      shuffleIdShouldNotExist: Seq[Int],
+      shuffleIdMustExist: Seq[Int],
+      sparkSession: SparkSession,
+      forStableStatusChecking: Boolean): CheckingThread = {
     val checkingThread =
       if (!forStableStatusChecking) {
         new CheckingThread(workerDirs, shuffleIdShouldNotExist, shuffleIdMustExist, sparkSession)