Add config key to enable the "stop NM on failure" behavior.

Author: Marcelo Vanzin

common/network-yarn/src/main/java/org/apache/spark/network/yarn/YarnShuffleService.java

@@ -70,6 +70,11 @@ public class YarnShuffleService extends AuxiliaryService {
 
   private static final String RECOVERY_FILE_NAME = "registeredExecutor.ldb";
 
+  // Whether failure during service initialization should stop the NM.
+  @VisibleForTesting
+  static final String STOP_ON_FAILURE_KEY = "spark.yarn.shuffle.stopOnFailure";
+  private static final boolean DEFAULT_STOP_ON_FAILURE = false;
+
   // An entity that manages the shuffle secret per application
   // This is used only if authentication is enabled
   private ShuffleSecretManager secretManager;
@@ -122,37 +127,47 @@ private boolean isAuthenticationEnabled() {
   protected void serviceInit(Configuration conf) throws Exception {
     _conf = conf;
 
-    // In case this NM was killed while there were running spark applications, we need to restore
-    // lost state for the existing executors.  We look for an existing file in the NM's local dirs.
-    // If we don't find one, then we choose a file to use to save the state next time.  Even if
-    // an application was stopped while the NM was down, we expect yarn to call stopApplication()
-    // when it comes back
-    registeredExecutorFile =
-      new File(getRecoveryPath().toUri().getPath(), RECOVERY_FILE_NAME);
-
-    TransportConf transportConf = new TransportConf("shuffle", new HadoopConfigProvider(conf));
-    blockHandler = new ExternalShuffleBlockHandler(transportConf, registeredExecutorFile);
-
-    // If authentication is enabled, set up the shuffle server to use a
-    // special RPC handler that filters out unauthenticated fetch requests
-    boolean authEnabled = conf.getBoolean(SPARK_AUTHENTICATE_KEY, DEFAULT_SPARK_AUTHENTICATE);
-    List<TransportServerBootstrap> bootstraps = Lists.newArrayList();
-    if (authEnabled) {
-      secretManager = new ShuffleSecretManager();
-      bootstraps.add(new SaslServerBootstrap(transportConf, secretManager));
-    }
+    boolean stopOnFailure = conf.getBoolean(STOP_ON_FAILURE_KEY, DEFAULT_STOP_ON_FAILURE);
 
-    int port = conf.getInt(
-      SPARK_SHUFFLE_SERVICE_PORT_KEY, DEFAULT_SPARK_SHUFFLE_SERVICE_PORT);
-    TransportContext transportContext = new TransportContext(transportConf, blockHandler);
-    shuffleServer = transportContext.createServer(port, bootstraps);
-    // the port should normally be fixed, but for tests its useful to find an open port
-    port = shuffleServer.getPort();
-    boundPort = port;
-    String authEnabledString = authEnabled ? "enabled" : "not enabled";
-    logger.info("Started YARN shuffle service for Spark on port {}. " +
-      "Authentication is {}.  Registered executor file is {}", port, authEnabledString,
-      registeredExecutorFile);
+    try {
+      // In case this NM was killed while there were running spark applications, we need to restore
+      // lost state for the existing executors.  We look for an existing file in the NM's local dirs.
+      // If we don't find one, then we choose a file to use to save the state next time.  Even if
+      // an application was stopped while the NM was down, we expect yarn to call stopApplication()
+      // when it comes back
+      registeredExecutorFile =
+        new File(getRecoveryPath().toUri().getPath(), RECOVERY_FILE_NAME);
+
+      TransportConf transportConf = new TransportConf("shuffle", new HadoopConfigProvider(conf));
+      blockHandler = new ExternalShuffleBlockHandler(transportConf, registeredExecutorFile);
+
+      // If authentication is enabled, set up the shuffle server to use a
+      // special RPC handler that filters out unauthenticated fetch requests
+      boolean authEnabled = conf.getBoolean(SPARK_AUTHENTICATE_KEY, DEFAULT_SPARK_AUTHENTICATE);
+      List<TransportServerBootstrap> bootstraps = Lists.newArrayList();
+      if (authEnabled) {
+        secretManager = new ShuffleSecretManager();
+        bootstraps.add(new SaslServerBootstrap(transportConf, secretManager));
+      }
+
+      int port = conf.getInt(
+        SPARK_SHUFFLE_SERVICE_PORT_KEY, DEFAULT_SPARK_SHUFFLE_SERVICE_PORT);
+      TransportContext transportContext = new TransportContext(transportConf, blockHandler);
+      shuffleServer = transportContext.createServer(port, bootstraps);
+      // the port should normally be fixed, but for tests its useful to find an open port
+      port = shuffleServer.getPort();
+      boundPort = port;
+      String authEnabledString = authEnabled ? "enabled" : "not enabled";
+      logger.info("Started YARN shuffle service for Spark on port {}. " +
+        "Authentication is {}.  Registered executor file is {}", port, authEnabledString,
+        registeredExecutorFile);
+    } catch (Exception e) {
+      if (stopOnFailure) {
+        throw e;
+      } else {
+        noteFailure(e);
+      }
+    }
   }
 
   @Override

yarn/src/test/scala/org/apache/spark/network/yarn/YarnShuffleServiceSuite.scala

@@ -327,6 +327,7 @@ class YarnShuffleServiceSuite extends SparkFunSuite with Matchers with BeforeAnd
     val roDir = Utils.createTempDir()
     Files.setPosixFilePermissions(roDir.toPath(), EnumSet.of(OWNER_READ, OWNER_EXECUTE))
     roConfig.set(YarnConfiguration.NM_LOCAL_DIRS, roDir.getAbsolutePath())
+    roConfig.setBoolean(YarnShuffleService.STOP_ON_FAILURE_KEY, true)
 
     // Try to start the shuffle service, it should fail.
     val service = new YarnShuffleService()