clear

Author: turboFei

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

@@ -363,7 +363,8 @@ public static void addSparkListener(SparkListener listener) {
   @VisibleForTesting
   public static AtomicInteger getReducerFileGroupResponseBroadcastNum = new AtomicInteger();
 
-  protected static Map<Integer, Tuple2<Broadcast<GetReducerFileGroupResponse>, byte[]>>
+  @VisibleForTesting
+  public static Map<Integer, Tuple2<Broadcast<GetReducerFileGroupResponse>, byte[]>>
       getReducerFileGroupResponseBroadcasts = JavaUtils.newConcurrentHashMap();
 
   public static byte[] serializeGetReducerFileGroupResponse(

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

@@ -480,7 +480,8 @@ public static void addSparkListener(SparkListener listener) {
   @VisibleForTesting
   public static AtomicInteger getReducerFileGroupResponseBroadcastNum = new AtomicInteger();
 
-  protected static Map<Integer, Tuple2<Broadcast<GetReducerFileGroupResponse>, byte[]>>
+  @VisibleForTesting
+  public static Map<Integer, Tuple2<Broadcast<GetReducerFileGroupResponse>, byte[]>>
       getReducerFileGroupResponseBroadcasts = JavaUtils.newConcurrentHashMap();
 
   public static byte[] serializeGetReducerFileGroupResponse(

tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/CelebornHashSuite.scala

@@ -68,6 +68,7 @@ class CelebornHashSuite extends AnyFunSuite
   }
 
   test("celeborn spark integration test - GetReducerFileGroupResponse broadcast") {
+    SparkUtils.getReducerFileGroupResponseBroadcasts.clear()
     SparkUtils.getReducerFileGroupResponseBroadcastNum.set(0)
     val sparkConf = new SparkConf().setAppName("celeborn-demo").setMaster("local[2]")
       .set(

tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/CelebornSortSuite.scala

@@ -69,6 +69,7 @@ class CelebornSortSuite extends AnyFunSuite
   }
 
   test("celeborn spark integration test - GetReducerFileGroupResponse broadcast") {
+    SparkUtils.getReducerFileGroupResponseBroadcasts.clear()
     SparkUtils.getReducerFileGroupResponseBroadcastNum.set(0)
     val sparkConf = new SparkConf().setAppName("celeborn-demo").setMaster("local[2]")
       .set(s"spark.${CelebornConf.CLIENT_PUSH_SORT_RANDOMIZE_PARTITION_ENABLED.key}", "false")

tests/spark-it/src/test/scala/org/apache/spark/shuffle/celeborn/SparkUtilsSuite.scala

@@ -203,6 +203,7 @@ class SparkUtilsSuite extends AnyFunSuite
       assert(!broadcast.isValid)
     } finally {
       sparkSession.stop()
+      SparkUtils.getReducerFileGroupResponseBroadcasts.clear()
     }
   }
 }