First cut integration of shuffle service with Yarn aux service

Author: Andrew Or

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

@@ -1,10 +1,16 @@
 
 package org.apache.spark.network.yarn;
 
+import java.lang.Override;
 import java.nio.ByteBuffer;
-import java.sql.Timestamp;
-import java.util.Date;
 
+import org.apache.spark.network.TransportContext;
+import org.apache.spark.network.server.RpcHandler;
+import org.apache.spark.network.shuffle.ExternalShuffleBlockHandler;
+import org.apache.spark.network.util.TransportConf;
+import org.apache.spark.network.util.SystemPropertyConfigProvider;
+
+import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.mapreduce.security.token.JobTokenSecretManager;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
 import org.apache.hadoop.yarn.api.records.ContainerId;
@@ -13,51 +19,68 @@
 import org.apache.hadoop.yarn.server.api.ApplicationTerminationContext;
 import org.apache.hadoop.yarn.server.api.ContainerInitializationContext;
 import org.apache.hadoop.yarn.server.api.ContainerTerminationContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * External shuffle service used by Spark on Yarn.
  */
 public class YarnShuffleService extends AuxiliaryService {
-
+    private final Logger logger = LoggerFactory.getLogger(YarnShuffleService.class);
     private static final JobTokenSecretManager secretManager = new JobTokenSecretManager();
 
+    private static final String SPARK_SHUFFLE_SERVICE_PORT_KEY = "spark.shuffle.service.port";
+    private static final int DEFAULT_SPARK_SHUFFLE_SERVICE_PORT = 7337;
+
     public YarnShuffleService() {
-        super("sparkshuffleservice");
-        log("--- [ Welcome to YarnShuffleService v0.1 ] ---");
+        super("spark_shuffle");
+        logger.info("Initializing Yarn shuffle service for Spark");
+    }
+
+    /**
+     * Start the shuffle server with the given configuration.
+     */
+    @Override
+    protected void serviceInit(Configuration conf) {
+        try {
+            int port = conf.getInt(
+                SPARK_SHUFFLE_SERVICE_PORT_KEY, DEFAULT_SPARK_SHUFFLE_SERVICE_PORT);
+            TransportConf transportConf = new TransportConf(new SystemPropertyConfigProvider());
+            RpcHandler rpcHandler = new ExternalShuffleBlockHandler();
+            TransportContext transportContext = new TransportContext(transportConf, rpcHandler);
+            transportContext.createServer(port);
+        } catch (Exception e) {
+            logger.error("Exception in starting Yarn shuffle service for Spark!", e);
+        }
     }
 
     @Override
     public void initializeApplication(ApplicationInitializationContext context) {
         ApplicationId appId = context.getApplicationId();
-        log("Initializing application " + appId + "!");
+        logger.debug("Initializing application " + appId + "!");
     }
 
     @Override
     public void stopApplication(ApplicationTerminationContext context) {
         ApplicationId appId = context.getApplicationId();
-        log("Stopping application " + appId + "!");
+        logger.debug("Stopping application " + appId + "!");
     }
 
     @Override
     public ByteBuffer getMetaData() {
-        log("Getting meta data");
-        return ByteBuffer.wrap("".getBytes());
+        logger.debug("Getting meta data");
+        return ByteBuffer.allocate(0);
     }
 
     @Override
     public void initializeContainer(ContainerInitializationContext context) {
         ContainerId containerId = context.getContainerId();
-        log("Initializing container " + containerId + "!");
+        logger.debug("Initializing container " + containerId + "!");
     }
 
     @Override
     public void stopContainer(ContainerTerminationContext context) {
         ContainerId containerId = context.getContainerId();
-        log("Stopping container " + containerId + "!");
-    }
-
-    private void log(String msg) {
-        Timestamp timestamp = new Timestamp(new Date().getTime());
-        System.out.println("* org.apache.spark.YarnShuffleService " + timestamp + ": " + msg);
+        logger.debug("Stopping container " + containerId + "!");
     }
 }

yarn/stable/src/main/scala/org/apache/spark/deploy/yarn/ExecutorRunnable.scala

@@ -89,6 +89,8 @@ class ExecutorRunnable(
 
     ctx.setApplicationACLs(YarnSparkHadoopUtil.getApplicationAclsForYarn(securityMgr))
 
+    ctx.setServiceData(Map[String, ByteBuffer]("spark_shuffle" -> ByteBuffer.allocate(0)))
+
     // Send the start request to the ContainerManager
     nmClient.startContainer(container, ctx)
   }