Address comments

Author: aarondav

core/src/main/scala/org/apache/spark/storage/BlockId.scala

@@ -53,6 +53,8 @@ case class RDDBlockId(rddId: Int, splitIndex: Int) extends BlockId {
   def name = "rdd_" + rddId + "_" + splitIndex
 }
 
+// Format of the shuffle block ids (including data and index) should be kept in sync with
+// org.apache.spark.network.shuffle.StandaloneShuffleBlockManager#getBlockData().
 @DeveloperApi
 case class ShuffleBlockId(shuffleId: Int, mapId: Int, reduceId: Int) extends BlockId {
   def name = "shuffle_" + shuffleId + "_" + mapId + "_" + reduceId

core/src/main/scala/org/apache/spark/storage/BlockManager.scala

@@ -35,10 +35,12 @@ import org.apache.spark.io.CompressionCodec
 import org.apache.spark.network._
 import org.apache.spark.network.buffer.{ManagedBuffer, NioManagedBuffer}
 import org.apache.spark.network.netty.{SparkTransportConf, NettyBlockTransferService}
-import org.apache.spark.network.shuffle.{ExecutorShuffleConfig, StandaloneShuffleClient}
+import org.apache.spark.network.shuffle.{ExecutorShuffleInfo, StandaloneShuffleClient}
 import org.apache.spark.network.util.{ConfigProvider, TransportConf}
 import org.apache.spark.serializer.Serializer
 import org.apache.spark.shuffle.ShuffleManager
+import org.apache.spark.shuffle.hash.HashShuffleManager
+import org.apache.spark.shuffle.sort.SortShuffleManager
 import org.apache.spark.util._
 
 private[spark] sealed trait BlockValues
@@ -91,6 +93,14 @@ private[spark] class BlockManager(
   private[spark]
   val externalShuffleServiceEnabled = conf.getBoolean("spark.shuffle.service.enabled", false)
   private val externalShuffleServicePort = conf.getInt("spark.shuffle.service.port", 7337)
+  // Check that we're not using external shuffle service with consolidated shuffle files.
+  if (externalShuffleServiceEnabled
+      && conf.getBoolean("spark.shuffle.consolidateFiles", false)
+      && shuffleManager.isInstanceOf[HashShuffleManager]) {
+    throw new UnsupportedOperationException("Cannot use external shuffle service with consolidated"
+      + " shuffle files in hash-based shuffle. Please disable spark.shuffle.consolidateFiles or "
+      + " switch to sort-based shuffle.")
+  }
 
   val blockManagerId = BlockManagerId(
     executorId, blockTransferService.hostName, blockTransferService.port)
@@ -174,14 +184,33 @@ private[spark] class BlockManager(
 
     // Register Executors' configuration with the local shuffle service, if one should exist.
     if (externalShuffleServiceEnabled && !blockManagerId.isDriver) {
-      logInfo("Registering executor with local external shuffle service.")
-      // Synchronous and will throw an exception if we cannot connect.
-      val shuffleConfig = new ExecutorShuffleConfig(
-          diskBlockManager.localDirs.map(_.toString),
-          diskBlockManager.subDirsPerLocalDir,
-          shuffleManager.getClass.getName)
-      shuffleClient.asInstanceOf[StandaloneShuffleClient].registerWithStandaloneShuffleService(
-        shuffleServerId.host, shuffleServerId.port, shuffleServerId.executorId, shuffleConfig)
+      registerWithExternalShuffleServer()
+    }
+  }
+
+  private def registerWithExternalShuffleServer() {
+    logInfo("Registering executor with local external shuffle service.")
+    val shuffleConfig = new ExecutorShuffleInfo(
+      diskBlockManager.localDirs.map(_.toString),
+      diskBlockManager.subDirsPerLocalDir,
+      shuffleManager.getClass.getName)
+
+    val MAX_ATTEMPTS = 3
+    val SLEEP_TIME_SECS = 5
+
+    for (i <- 1 to MAX_ATTEMPTS) {
+      try {
+        // Synchronous and will throw an exception if we cannot connect.
+        shuffleClient.asInstanceOf[StandaloneShuffleClient].registerWithShuffleServer(
+          shuffleServerId.host, shuffleServerId.port, shuffleServerId.executorId, shuffleConfig)
+        return
+      } catch {
+        case e: Exception if i < MAX_ATTEMPTS =>
+          val attemptsRemaining =
+          logError(s"Failed to connect to external shuffle server, will retry ${MAX_ATTEMPTS - i}}"
+            + s" more times after waiting $SLEEP_TIME_SECS seconds...", e)
+          Thread.sleep(SLEEP_TIME_SECS * 1000)
+      }
     }
   }
 

network/common/src/main/java/org/apache/spark/network/client/TransportClient.java

@@ -34,7 +34,6 @@
 import org.apache.spark.network.protocol.ChunkFetchRequest;
 import org.apache.spark.network.protocol.RpcRequest;
 import org.apache.spark.network.protocol.StreamChunkId;
-import org.apache.spark.network.util.JavaUtils;
 import org.apache.spark.network.util.NettyUtils;
 
 /**

network/shuffle/src/main/java/org/apache/spark/network/shuffle/ExecutorShuffleConfig.java renamed to network/shuffle/src/main/java/org/apache/spark/network/shuffle/ExecutorShuffleInfo.java

@@ -22,16 +22,16 @@
 
 import com.google.common.base.Objects;
 
-/** Contains all configuration necessary for a single Executor to find its shuffle files. */
-public class ExecutorShuffleConfig implements Serializable {
+/** Contains all configuration necessary for locating the shuffle files of an executor. */
+public class ExecutorShuffleInfo implements Serializable {
   /** The base set of local directories that the executor stores its shuffle files in. */
   final String[] localDirs;
   /** Number of subdirectories created within each localDir. */
   final int subDirsPerLocalDir;
   /** Shuffle manager (SortShuffleManager or HashShuffleManager) that the executor is using. */
   final String shuffleManager;
 
-  public ExecutorShuffleConfig(String[] localDirs, int subDirsPerLocalDir, String shuffleManager) {
+  public ExecutorShuffleInfo(String[] localDirs, int subDirsPerLocalDir, String shuffleManager) {
     this.localDirs = localDirs;
     this.subDirsPerLocalDir = subDirsPerLocalDir;
     this.shuffleManager = shuffleManager;
@@ -53,8 +53,8 @@ public String toString() {
 
   @Override
   public boolean equals(Object other) {
-    if (other != null && other instanceof ExecutorShuffleConfig) {
-      ExecutorShuffleConfig o = (ExecutorShuffleConfig) other;
+    if (other != null && other instanceof ExecutorShuffleInfo) {
+      ExecutorShuffleInfo o = (ExecutorShuffleInfo) other;
       return Arrays.equals(localDirs, o.localDirs)
         && Objects.equal(subDirsPerLocalDir, o.subDirsPerLocalDir)
         && Objects.equal(shuffleManager, o.shuffleManager);

network/shuffle/src/main/java/org/apache/spark/network/shuffle/StandaloneShuffleBlockHandler.java

@@ -79,7 +79,7 @@ public void receive(TransportClient client, byte[] message, RpcResponseCallback
 
     } else if (msgObj instanceof RegisterExecutor) {
       RegisterExecutor msg = (RegisterExecutor) msgObj;
-      blockManager.registerExecutor(msg.appId, msg.execId, msg.executorConfig);
+      blockManager.registerExecutor(msg.appId, msg.execId, msg.executorInfo);
       callback.onSuccess(new byte[0]);
 
     } else {

network/shuffle/src/main/java/org/apache/spark/network/shuffle/StandaloneShuffleBlockManager.java

@@ -44,8 +44,8 @@ public class StandaloneShuffleBlockManager {
   private final Logger logger = LoggerFactory.getLogger(StandaloneShuffleBlockManager.class);
 
   // Map from "appId-execId" to the executor's configuration.
-  private final ConcurrentHashMap<String, ExecutorShuffleConfig> executors =
-    new ConcurrentHashMap<String, ExecutorShuffleConfig>();
+  private final ConcurrentHashMap<String, ExecutorShuffleInfo> executors =
+    new ConcurrentHashMap<String, ExecutorShuffleInfo>();
 
   // Returns an id suitable for a single executor within a single application.
   private String getAppExecId(String appId, String execId) {
@@ -56,15 +56,16 @@ private String getAppExecId(String appId, String execId) {
   public void registerExecutor(
       String appId,
       String execId,
-      ExecutorShuffleConfig executorConfig) {
+      ExecutorShuffleInfo executorInfo) {
     String fullId = getAppExecId(appId, execId);
-    executors.put(fullId, executorConfig);
+    logger.info("Registered executor {} with {}", fullId, executorInfo);
+    executors.put(fullId, executorInfo);
   }
 
   /**
    * Obtains a FileSegmentManagedBuffer from a shuffle block id. We expect the blockId has the
-   * format "shuffle_ShuffleId_MapId_ReduceId", and additionally make assumptions about how the
-   * hash and sort based shuffles store their data.
+   * format "shuffle_ShuffleId_MapId_ReduceId" (from ShuffleBlockId), and additionally make
+   * assumptions about how the hash and sort based shuffles store their data.
    */
   public ManagedBuffer getBlockData(String appId, String execId, String blockId) {
     String[] blockIdParts = blockId.split("_");
@@ -77,7 +78,7 @@ public ManagedBuffer getBlockData(String appId, String execId, String blockId) {
     int mapId = Integer.parseInt(blockIdParts[2]);
     int reduceId = Integer.parseInt(blockIdParts[3]);
 
-    ExecutorShuffleConfig executor = executors.get(getAppExecId(appId, execId));
+    ExecutorShuffleInfo executor = executors.get(getAppExecId(appId, execId));
     if (executor == null) {
       throw new RuntimeException(
         String.format("Executor is not registered (appId=%s, execId=%s)", appId, execId));
@@ -98,17 +99,18 @@ public ManagedBuffer getBlockData(String appId, String execId, String blockId) {
    * This logic is from FileShuffleBlockManager.
    */
   // TODO: Support consolidated hash shuffle files
-  private ManagedBuffer getHashBasedShuffleBlockData(ExecutorShuffleConfig executor, String blockId) {
+  private ManagedBuffer getHashBasedShuffleBlockData(ExecutorShuffleInfo executor, String blockId) {
     File shuffleFile = getFile(executor.localDirs, executor.subDirsPerLocalDir, blockId);
     return new FileSegmentManagedBuffer(shuffleFile, 0, shuffleFile.length());
   }
 
   /**
    * Sort-based shuffle data uses an index called "shuffle_ShuffleId_MapId_0.index" into a data file
-   * called "shuffle_ShuffleId_MapId_0.data". This logic is from IndexShuffleBlockManager.
+   * called "shuffle_ShuffleId_MapId_0.data". This logic is from IndexShuffleBlockManager,
+   * and the block id format is from ShuffleDataBlockId and ShuffleIndexBlockId.
    */
   private ManagedBuffer getSortBasedShuffleBlockData(
-    ExecutorShuffleConfig executor, int shuffleId, int mapId, int reduceId) {
+    ExecutorShuffleInfo executor, int shuffleId, int mapId, int reduceId) {
     File indexFile = getFile(executor.localDirs, executor.subDirsPerLocalDir,
       "shuffle_" + shuffleId + "_" + mapId + "_0.index");
 

network/shuffle/src/main/java/org/apache/spark/network/shuffle/StandaloneShuffleClient.java

@@ -67,16 +67,16 @@ public void fetchBlocks(
    * @param host Host of standalone shuffle server.
    * @param port Port of standalone shuffle server.
    * @param execId This Executor's id.
-   * @param executorConfig Contains all config necessary for the service to find our shuffle files.
+   * @param executorInfo Contains all info necessary for the service to find our shuffle files.
    */
-  public void registerWithStandaloneShuffleService(
+  public void registerWithShuffleServer(
       String host,
       int port,
       String execId,
-      ExecutorShuffleConfig executorConfig) {
+      ExecutorShuffleInfo executorInfo) {
     TransportClient client = clientFactory.createClient(host, port);
     byte[] registerExecutorMessage =
-      JavaUtils.serialize(new RegisterExecutor(appId, execId, executorConfig));
+      JavaUtils.serialize(new RegisterExecutor(appId, execId, executorInfo));
     client.sendRpcSync(registerExecutorMessage, 3000);
   }
 }

network/shuffle/src/main/java/org/apache/spark/network/shuffle/StandaloneShuffleMessages.java

@@ -67,28 +67,28 @@ public boolean equals(Object other) {
   public static class RegisterExecutor implements Serializable {
     public final String appId;
     public final String execId;
-    public final ExecutorShuffleConfig executorConfig;
+    public final ExecutorShuffleInfo executorInfo;
 
     public RegisterExecutor(
         String appId,
         String execId,
-        ExecutorShuffleConfig executorConfig) {
+        ExecutorShuffleInfo executorInfo) {
       this.appId = appId;
       this.execId = execId;
-      this.executorConfig = executorConfig;
+      this.executorInfo = executorInfo;
     }
 
     @Override
     public int hashCode() {
-      return Objects.hashCode(appId, execId, executorConfig);
+      return Objects.hashCode(appId, execId, executorInfo);
     }
 
     @Override
     public String toString() {
       return Objects.toStringHelper(this)
         .add("appId", appId)
         .add("execId", execId)
-        .add("executorConfig", executorConfig)
+        .add("executorInfo", executorInfo)
         .toString();
     }
 
@@ -98,7 +98,7 @@ public boolean equals(Object other) {
         RegisterExecutor o = (RegisterExecutor) other;
         return Objects.equal(appId, o.appId)
           && Objects.equal(execId, o.execId)
-          && Objects.equal(executorConfig, o.executorConfig);
+          && Objects.equal(executorInfo, o.executorInfo);
       }
       return false;
     }

network/shuffle/src/test/java/org/apache/spark/network/shuffle/ShuffleMessagesSuite.java

@@ -36,7 +36,7 @@ public void serializeOpenShuffleBlocks() {
 
   @Test
   public void serializeRegisterExecutor() {
-    RegisterExecutor msg = new RegisterExecutor("app-1", "exec-2", new ExecutorShuffleConfig(
+    RegisterExecutor msg = new RegisterExecutor("app-1", "exec-2", new ExecutorShuffleInfo(
       new String[] { "/local1", "/local2" }, 32, "MyShuffleManager"));
     RegisterExecutor msg2 = JavaUtils.deserialize(JavaUtils.serialize(msg));
     assertEquals(msg, msg2);

network/shuffle/src/test/java/org/apache/spark/network/shuffle/StandaloneShuffleBlockHandlerSuite.java

@@ -56,7 +56,7 @@ public void beforeEach() {
   public void testRegisterExecutor() {
     RpcResponseCallback callback = mock(RpcResponseCallback.class);
 
-    ExecutorShuffleConfig config = new ExecutorShuffleConfig(new String[] {"/a", "/b"}, 16, "sort");
+    ExecutorShuffleInfo config = new ExecutorShuffleInfo(new String[] {"/a", "/b"}, 16, "sort");
     byte[] registerMessage = JavaUtils.serialize(
       new RegisterExecutor("app0", "exec1", config));
     handler.receive(client, registerMessage, callback);
@@ -109,7 +109,7 @@ public void testBadMessages() {
     }
 
     byte[] unexpectedMessage = JavaUtils.serialize(
-      new ExecutorShuffleConfig(new String[] {"/a", "/b"}, 16, "sort"));
+      new ExecutorShuffleInfo(new String[] {"/a", "/b"}, 16, "sort"));
     try {
       handler.receive(client, unexpectedMessage, callback);
       fail("Should have thrown");