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(
@@ -101,5 +102,7 @@ class CelebornHashSuite extends AnyFunSuite
     assert(SparkUtils.getReducerFileGroupResponseBroadcastNum.get() > 0)
 
     celebornSparkSession.stop()
+    SparkUtils.getReducerFileGroupResponseBroadcasts.clear()
+    SparkUtils.getReducerFileGroupResponseBroadcastNum.set(0)
   }
 }

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")
@@ -104,5 +105,7 @@ class CelebornSortSuite extends AnyFunSuite
     assert(SparkUtils.getReducerFileGroupResponseBroadcastNum.get() > 0)
 
     celebornSparkSession.stop()
+    SparkUtils.getReducerFileGroupResponseBroadcasts.clear()
+    SparkUtils.getReducerFileGroupResponseBroadcastNum.set(0)
   }
 }

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

@@ -173,9 +173,10 @@ class SparkUtilsSuite extends AnyFunSuite
       .getOrCreate()
 
     try {
+      val shuffleId = Integer.MAX_VALUE
       val getReducerFileGroupResponse = GetReducerFileGroupResponse(
         StatusCode.SUCCESS,
-        Map(Integer.valueOf(1) -> Set(new PartitionLocation(
+        Map(Integer.valueOf(shuffleId) -> Set(new PartitionLocation(
           0,
           1,
           "localhost",
@@ -185,24 +186,26 @@ class SparkUtilsSuite extends AnyFunSuite
           4,
           PartitionLocation.Mode.REPLICA)).asJava).asJava,
         Array(1),
-        Set(Integer.valueOf(1)).asJava)
+        Set(Integer.valueOf(shuffleId)).asJava)
 
       val serializedBytes =
-        SparkUtils.serializeGetReducerFileGroupResponse(1, getReducerFileGroupResponse)
+        SparkUtils.serializeGetReducerFileGroupResponse(shuffleId, getReducerFileGroupResponse)
       assert(serializedBytes != null && serializedBytes.length > 0)
       val broadcast = SparkUtils.getReducerFileGroupResponseBroadcasts.values().asScala.head._1
       assert(broadcast.value == getReducerFileGroupResponse)
 
       val deserializedGetReducerFileGroupResponse =
-        SparkUtils.deserializeGetReducerFileGroupResponse(1, serializedBytes)
+        SparkUtils.deserializeGetReducerFileGroupResponse(shuffleId, serializedBytes)
       assert(deserializedGetReducerFileGroupResponse == getReducerFileGroupResponse)
 
       assert(!SparkUtils.getReducerFileGroupResponseBroadcasts.isEmpty)
-      SparkUtils.invalidateSerializedGetReducerFileGroupResponse(1)
+      SparkUtils.invalidateSerializedGetReducerFileGroupResponse(shuffleId)
       assert(SparkUtils.getReducerFileGroupResponseBroadcasts.isEmpty)
       assert(!broadcast.isValid)
     } finally {
       sparkSession.stop()
+      SparkUtils.getReducerFileGroupResponseBroadcasts.clear()
+      SparkUtils.getReducerFileGroupResponseBroadcastNum.set(0)
     }
   }
 }