[CELEBORN-1663] Only register appShuffleDeterminate if stage using ce…

…leborn for shuffle Only register appShuffleDeterminate if stage using celeborn for shuffle Currently we are passing stage info to lifecyclemanager, eventhough it is not required. NA Existing UTs Closes apache#2832 from s0nskar/fix_register. Authored-by: Sanskar Modi <[email protected]> Signed-off-by: Shuang <[email protected]> (cherry picked from commit 813b45f) Signed-off-by: Shuang <[email protected]>
SteNicholas · Oct 22, 2024 · 8b7b181 · 8b7b181
1 parent bce6166
commit 8b7b181
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/...nt-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java b/...nt-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
@@ -159,10 +159,6 @@ public <K, V, C> ShuffleHandle registerShuffle(
     appUniqueId = SparkUtils.appUniqueId(dependency.rdd().context());
     initializeLifecycleManager();
 
-    lifecycleManager.registerAppShuffleDeterminate(
-        shuffleId,
-        dependency.rdd().getOutputDeterministicLevel() != DeterministicLevel.INDETERMINATE());
-
     if (fallbackPolicyRunner.applyAllFallbackPolicy(
         lifecycleManager, dependency.partitioner().numPartitions())) {
       if (conf.getBoolean("spark.dynamicAllocation.enabled", false)
@@ -178,6 +174,11 @@ public <K, V, C> ShuffleHandle registerShuffle(
       sortShuffleIds.add(shuffleId);
       return sortShuffleManager().registerShuffle(shuffleId, dependency);
     } else {
+      lifecycleManager.registerAppShuffleDeterminate(
+          shuffleId,
+          !DeterministicLevel.INDETERMINATE()
+              .equals(dependency.rdd().getOutputDeterministicLevel()));
+
       return new CelebornShuffleHandle<>(
           appUniqueId,
           lifecycleManager.getHost(),