ExternalShuffleBlockResolver can handle blockIds w/out stageAttemptId

Author: squito

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

@@ -168,15 +168,14 @@ public void registerExecutor(
    */
   public ManagedBuffer getBlockData(String appId, String execId, String blockId) {
     String[] blockIdParts = blockId.split("_");
-    if (blockIdParts.length != 5) {
+    if (blockIdParts.length < 4 || blockIdParts.length > 5) {
       throw new IllegalArgumentException("Unexpected block id format: " + blockId);
     } else if (!blockIdParts[0].equals("shuffle")) {
       throw new IllegalArgumentException("Expected shuffle block id, got: " + blockId);
     }
     int shuffleId = Integer.parseInt(blockIdParts[1]);
     int mapId = Integer.parseInt(blockIdParts[2]);
     int reduceId = Integer.parseInt(blockIdParts[3]);
-    int stageAttemptId = Integer.parseInt(blockIdParts[4]);
 
     ExecutorShuffleInfo executor = executors.get(new AppExecId(appId, execId));
     if (executor == null) {
@@ -188,7 +187,14 @@ public ManagedBuffer getBlockData(String appId, String execId, String blockId) {
       return getHashBasedShuffleBlockData(executor, blockId);
     } else if ("org.apache.spark.shuffle.sort.SortShuffleManager".equals(executor.shuffleManager)
       || "org.apache.spark.shuffle.unsafe.UnsafeShuffleManager".equals(executor.shuffleManager)) {
-      return getSortBasedShuffleBlockData(executor, shuffleId, mapId, reduceId, stageAttemptId);
+      // for backwards compatibility, we also handle legacy shuffle block ids which don't have
+      // a stageAttemptId
+      String baseFileName = "shuffle_" + shuffleId + "_" + mapId + "_0";
+      if (blockIdParts.length == 5) {
+        int stageAttemptId = Integer.parseInt(blockIdParts[4]);
+        baseFileName = baseFileName + "_" + stageAttemptId;
+      }
+      return getSortBasedShuffleBlockData(executor, baseFileName, reduceId);
     } else {
       throw new UnsupportedOperationException(
         "Unsupported shuffle manager: " + executor.shuffleManager);
@@ -267,8 +273,9 @@ private ManagedBuffer getHashBasedShuffleBlockData(ExecutorShuffleInfo executor,
    * and the block id format is from ShuffleDataBlockId and ShuffleIndexBlockId.
    */
   private ManagedBuffer getSortBasedShuffleBlockData(
-    ExecutorShuffleInfo executor, int shuffleId, int mapId, int reduceId, int stageAttemptId) {
-    String baseFileName = "shuffle_" + shuffleId + "_" + mapId + "_0_" + stageAttemptId;
+      ExecutorShuffleInfo executor,
+      String baseFileName,
+      int reduceId) {
     File indexFile = getFile(executor.localDirs, executor.subDirsPerLocalDir,
       baseFileName + ".index");
 

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

@@ -102,9 +102,16 @@ public void testBadRequests() throws IOException {
       // pass
     }
 
-    // no stageAttemptId
+    // wrong number of parts (note that we allow a missing stageAttemptId)
     try {
-      resolver.getBlockData("app0", "exec1", "shuffle_1_1_0");
+      resolver.getBlockData("app0", "exec1", "shuffle_1_1_0_0_0");
+      fail("Should have failed");
+    } catch (RuntimeException e) {
+      assertTrue("Bad error message: " + e, e.getMessage().contains("Unexpected block id format"));
+    }
+
+    try {
+      resolver.getBlockData("app0", "exec1", "shuffle_1_1");
       fail("Should have failed");
     } catch (RuntimeException e) {
       assertTrue("Bad error message: " + e, e.getMessage().contains("Unexpected block id format"));
@@ -145,6 +152,54 @@ private void testReadBlockData(ExternalShuffleBlockResolver resolver, String blo
     assertEquals(expected, block0);
   }
 
+  @Test
+  public void supportLegacySortShuffleBlockIds() throws IOException {
+    // In Spark 1.6, the stage attempt ID was added to shuffle block ids (SPARK-8029).  However,
+    // during a rolling upgrade, the shuffle service may be restarted with new code but still
+    // need to serve old apps.  So we make sure we can still handle old blocks
+
+    ExternalShuffleBlockResolver resolver = new ExternalShuffleBlockResolver(conf, null);
+    resolver.registerExecutor("app0", "exec0",
+      dataContext.createExecutorInfo("org.apache.spark.shuffle.sort.SortShuffleManager"));
+
+    dataContext.insertLegacySortShuffleData(2, 1,
+      new byte[][]{"legacy".getBytes(), "block".getBytes()});
+
+    testReadBlockData(resolver, "shuffle_2_1_0", "legacy");
+    testReadBlockData(resolver, "shuffle_2_1_1", "block");
+
+    // verify everything still works when we also register some blocks which do have a
+    // stageAttemptId
+    testSortShuffleBlocks();
+
+    testReadBlockData(resolver, "shuffle_2_1_0", "legacy");
+    testReadBlockData(resolver, "shuffle_2_1_1", "block");
+  }
+
+  @Test
+  public void supportLegacyHashShuffleBlockIds() throws IOException {
+    // In Spark 1.6, the stage attempt ID was added to shuffle block ids (SPARK-8029).  However,
+    // during a rolling upgrade, the shuffle service may be restarted with new code but still
+    // need to serve old apps.  So we make sure we can still handle old blocks
+
+    ExternalShuffleBlockResolver resolver = new ExternalShuffleBlockResolver(conf, null);
+    resolver.registerExecutor("app0", "exec0",
+      dataContext.createExecutorInfo("org.apache.spark.shuffle.hash.HashShuffleManager"));
+
+    dataContext.insertLegacyHashShuffleData(2, 0,
+      new byte[][] { "more legacy".getBytes(), "hash".getBytes() } );
+
+    testReadBlockData(resolver, "shuffle_2_0_0", "more legacy");
+    testReadBlockData(resolver, "shuffle_2_0_1", "hash");
+
+    // verify everything still works when we also register some blocks which do have a
+    // stageAttemptId
+    testHashShuffleBlocks();
+
+    testReadBlockData(resolver, "shuffle_2_0_0", "more legacy");
+    testReadBlockData(resolver, "shuffle_2_0_1", "hash");
+  }
+
   @Test
   public void jsonSerializationOfExecutorRegistration() throws IOException {
     ObjectMapper mapper = new ObjectMapper();

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

@@ -57,10 +57,24 @@ public void cleanup() {
   }
 
   /** Creates reducer blocks in a sort-based data format within our local dirs. */
-  public void insertSortShuffleData(int shuffleId, int mapId, int stageAttemptId,
+  public void insertSortShuffleData(
+      int shuffleId,
+      int mapId,
+      int stageAttemptId,
       byte[][] blocks) throws IOException {
     String blockId = "shuffle_" + shuffleId + "_" + mapId + "_0_" + stageAttemptId;
+    insertSortShuffleData(blockId, blocks);
+  }
+
+  public void insertLegacySortShuffleData(
+      int shuffleId,
+      int mapId,
+      byte[][] blocks) throws IOException {
+    String blockId = "shuffle_" + shuffleId + "_" + mapId + "_0";
+    insertSortShuffleData(blockId, blocks);
+  }
 
+  private void insertSortShuffleData(String blockId, byte[][] blocks) throws IOException {
     OutputStream dataStream = new FileOutputStream(
       ExternalShuffleBlockResolver.getFile(localDirs, subDirsPerLocalDir, blockId + ".data"));
     DataOutputStream indexStream = new DataOutputStream(new FileOutputStream(
@@ -78,8 +92,22 @@ public void insertSortShuffleData(int shuffleId, int mapId, int stageAttemptId,
     indexStream.close();
   }
 
+  public void insertLegacyHashShuffleData(
+      int shuffleId,
+      int mapId,
+      byte[][] blocks) throws IOException {
+    for (int i = 0; i < blocks.length; i ++) {
+      String blockId = "shuffle_" + shuffleId + "_" + mapId + "_" + i;
+      Files.write(blocks[i],
+        ExternalShuffleBlockResolver.getFile(localDirs, subDirsPerLocalDir, blockId));
+    }
+  }
+
   /** Creates reducer blocks in a hash-based data format within our local dirs. */
-  public void insertHashShuffleData(int shuffleId, int mapId, int stageAttemptId,
+  public void insertHashShuffleData(
+      int shuffleId,
+      int mapId,
+      int stageAttemptId,
       byte[][] blocks) throws IOException {
     for (int i = 0; i < blocks.length; i ++) {
       String blockId = "shuffle_" + shuffleId + "_" + mapId + "_" + i + "_" + stageAttemptId;