# 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.

Index: core/src/main/scala/org/apache/spark/ExecutorAllocationManager.scala
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
diff --git a/core/src/main/scala/org/apache/spark/ExecutorAllocationManager.scala b/core/src/main/scala/org/apache/spark/ExecutorAllocationManager.scala
--- a/core/src/main/scala/org/apache/spark/ExecutorAllocationManager.scala	(revision 799aefc7d92de50ca76a0f60639949a7dda8e2cb)
+++ b/core/src/main/scala/org/apache/spark/ExecutorAllocationManager.scala	(revision e88da1325c5a9efebf39fb2ad73ca084ba24d1fc)
@@ -44,8 +44,8 @@
  * executors that could run all current running and pending tasks at once.
  *
  * Increasing the target number of executors happens in response to backlogged tasks waiting to be
- * scheduled. If the scheduler queue is not drained in N seconds, then new executors are added. If
- * the queue persists for another M seconds, then more executors are added and so on. The number
+ * scheduled. If the scheduler queue is not drained in M seconds, then new executors are added. If
+ * the queue persists for another N seconds, then more executors are added and so on. The number
  * added in each round increases exponentially from the previous round until an upper bound has been
  * reached. The upper bound is based both on a configured property and on the current number of
  * running and pending tasks, as described above.
@@ -212,9 +212,13 @@
     }
     // Require external shuffle service for dynamic allocation
     // Otherwise, we may lose shuffle files when killing executors
-    if (!conf.get(config.SHUFFLE_SERVICE_ENABLED) && !testing) {
-      throw new SparkException("Dynamic allocation of executors requires the external " +
-        "shuffle service. You may enable this through spark.shuffle.service.enabled.")
+    if (!conf.get(config.SHUFFLE_SERVICE_ENABLED) &&
+      !Utils.isCelebornEnabled(conf) && !testing) {
+      throw new SparkException("Dynamic allocation of executors requires the external or remote " +
+        "shuffle service. You may enable this through spark.shuffle.service.enabled or " +
+        "set spark.shuffle.manager=org.apache.spark.shuffle.celeborn.SparkShuffleManager.")
     }
     if (tasksPerExecutorForFullParallelism == 0) {
       throw new SparkException("spark.executor.cores must not be < spark.task.cpus.")
Index: core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
--- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala	(revision 799aefc7d92de50ca76a0f60639949a7dda8e2cb)
+++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala	(revision e88da1325c5a9efebf39fb2ad73ca084ba24d1fc)
@@ -1738,7 +1738,7 @@
     // if the cluster manager explicitly tells us that the entire worker was lost, then
     // we know to unregister shuffle output.  (Note that "worker" specifically refers to the process
     // from a Standalone cluster, where the shuffle service lives in the Worker.)
-    val fileLost = workerLost || !env.blockManager.externalShuffleServiceEnabled
+    val fileLost = !Utils.isCelebornEnabled(sc.getConf) && (workerLost || !env.blockManager.externalShuffleServiceEnabled)
     removeExecutorAndUnregisterOutputs(
       execId = execId,
       fileLost = fileLost,
Index: core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala	(revision 799aefc7d92de50ca76a0f60639949a7dda8e2cb)
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala	(revision e88da1325c5a9efebf39fb2ad73ca084ba24d1fc)
@@ -974,8 +974,10 @@
     // and we are not using an external shuffle server which could serve the shuffle outputs.
     // The reason is the next stage wouldn't be able to fetch the data from this dead executor
     // so we would need to rerun these tasks on other executors.
-    if (tasks(0).isInstanceOf[ShuffleMapTask] && !env.blockManager.externalShuffleServiceEnabled
-        && !isZombie) {
+    if (tasks(0).isInstanceOf[ShuffleMapTask]
+      && !env.blockManager.externalShuffleServiceEnabled
+      && !Utils.isCelebornEnabled(conf)
+      && !isZombie) {
       for ((tid, info) <- taskInfos if info.executorId == execId) {
         val index = taskInfos(tid).index
         // We may have a running task whose partition has been marked as successful,
Index: core/src/main/scala/org/apache/spark/util/Utils.scala
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
diff --git a/core/src/main/scala/org/apache/spark/util/Utils.scala b/core/src/main/scala/org/apache/spark/util/Utils.scala
--- a/core/src/main/scala/org/apache/spark/util/Utils.scala	(revision 799aefc7d92de50ca76a0f60639949a7dda8e2cb)
+++ b/core/src/main/scala/org/apache/spark/util/Utils.scala	(revision e88da1325c5a9efebf39fb2ad73ca084ba24d1fc)
@@ -2881,6 +2881,10 @@
     props.asScala.foreach(entry => resultProps.put(entry._1, entry._2))
     resultProps
   }
+
+  def isCelebornEnabled(conf: SparkConf): Boolean =
+    conf.get("spark.shuffle.manager", "sort").contains("celeborn")
+
 }
 
 private[util] object CallerContext extends Logging {
