Merge branch 'upmaster' into guava-cache-to-caffeine

Author: LuciferYang

.github/workflows/build_and_test.yml

@@ -199,7 +199,7 @@ jobs:
     name: "Build modules (${{ format('{0}, {1} job', needs.configure-jobs.outputs.branch, needs.configure-jobs.outputs.type) }}): ${{ matrix.modules }}"
     runs-on: ubuntu-20.04
     container:
-      image: dongjoon/apache-spark-github-action-image:20210602
+      image: dongjoon/apache-spark-github-action-image:20210730
     strategy:
       fail-fast: false
       matrix:
@@ -266,8 +266,6 @@ jobs:
     - name: Run tests
       env: ${{ fromJSON(needs.configure-jobs.outputs.envs) }}
       run: |
-        # TODO(SPARK-36345): Install mlflow>=1.0 and sklearn in Python 3.9 of the base image
-        python3.9 -m pip install 'mlflow>=1.0' sklearn
         # TODO(SPARK-36361): Install coverage in Python 3.9 and PyPy 3 in the base image
         python3.9 -m pip install coverage
         pypy3 -m pip install coverage

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

@@ -206,12 +206,15 @@ public long sendRpc(ByteBuffer message, RpcResponseCallback callback) {
    *
    * @param appId applicationId.
    * @param shuffleId shuffle id.
+   * @param shuffleMergeId shuffleMergeId is used to uniquely identify merging process
+   *                       of shuffle by an indeterminate stage attempt.
    * @param reduceId reduce id.
    * @param callback callback the handle the reply.
    */
   public void sendMergedBlockMetaReq(
       String appId,
       int shuffleId,
+      int shuffleMergeId,
       int reduceId,
       MergedBlockMetaResponseCallback callback) {
     long requestId = requestId();
@@ -222,7 +225,8 @@ public void sendMergedBlockMetaReq(
     handler.addRpcRequest(requestId, callback);
     RpcChannelListener listener = new RpcChannelListener(requestId, callback);
     channel.writeAndFlush(
-      new MergedBlockMetaRequest(requestId, appId, shuffleId, reduceId)).addListener(listener);
+      new MergedBlockMetaRequest(requestId, appId, shuffleId, shuffleMergeId,
+        reduceId)).addListener(listener);
   }
 
   /**

common/network-common/src/main/java/org/apache/spark/network/protocol/MergedBlockMetaRequest.java

@@ -32,13 +32,20 @@ public class MergedBlockMetaRequest extends AbstractMessage implements RequestMe
   public final long requestId;
   public final String appId;
   public final int shuffleId;
+  public final int shuffleMergeId;
   public final int reduceId;
 
-  public MergedBlockMetaRequest(long requestId, String appId, int shuffleId, int reduceId) {
+  public MergedBlockMetaRequest(
+      long requestId,
+      String appId,
+      int shuffleId,
+      int shuffleMergeId,
+      int reduceId) {
     super(null, false);
     this.requestId = requestId;
     this.appId = appId;
     this.shuffleId = shuffleId;
+    this.shuffleMergeId = shuffleMergeId;
     this.reduceId = reduceId;
   }
 
@@ -49,36 +56,39 @@ public Type type() {
 
   @Override
   public int encodedLength() {
-    return 8 + Encoders.Strings.encodedLength(appId) + 4 + 4;
+    return 8 + Encoders.Strings.encodedLength(appId) + 4 + 4 + 4;
   }
 
   @Override
   public void encode(ByteBuf buf) {
     buf.writeLong(requestId);
     Encoders.Strings.encode(buf, appId);
     buf.writeInt(shuffleId);
+    buf.writeInt(shuffleMergeId);
     buf.writeInt(reduceId);
   }
 
   public static MergedBlockMetaRequest decode(ByteBuf buf) {
     long requestId = buf.readLong();
     String appId = Encoders.Strings.decode(buf);
     int shuffleId = buf.readInt();
+    int shuffleMergeId = buf.readInt();
     int reduceId = buf.readInt();
-    return new MergedBlockMetaRequest(requestId, appId, shuffleId, reduceId);
+    return new MergedBlockMetaRequest(requestId, appId, shuffleId, shuffleMergeId, reduceId);
   }
 
   @Override
   public int hashCode() {
-    return Objects.hashCode(requestId, appId, shuffleId, reduceId);
+    return Objects.hashCode(requestId, appId, shuffleId, shuffleMergeId, reduceId);
   }
 
   @Override
   public boolean equals(Object other) {
     if (other instanceof MergedBlockMetaRequest) {
       MergedBlockMetaRequest o = (MergedBlockMetaRequest) other;
-      return requestId == o.requestId && shuffleId == o.shuffleId && reduceId == o.reduceId
-        && Objects.equal(appId, o.appId);
+      return requestId == o.requestId && shuffleId == o.shuffleId &&
+        shuffleMergeId == o.shuffleMergeId && reduceId == o.reduceId &&
+        Objects.equal(appId, o.appId);
     }
     return false;
   }
@@ -89,6 +99,7 @@ public String toString() {
       .append("requestId", requestId)
       .append("appId", appId)
       .append("shuffleId", shuffleId)
+      .append("shuffleMergeId", shuffleMergeId)
       .append("reduceId", reduceId)
       .toString();
   }

common/network-common/src/test/java/org/apache/spark/network/TransportRequestHandlerSuite.java

@@ -152,14 +152,14 @@ public MergedBlockMetaReqHandler getMergedBlockMetaReqHandler() {
     TransportClient reverseClient = mock(TransportClient.class);
     TransportRequestHandler requestHandler = new TransportRequestHandler(channel, reverseClient,
       rpcHandler, 2L, null);
-    MergedBlockMetaRequest validMetaReq = new MergedBlockMetaRequest(19, "app1", 0, 0);
+    MergedBlockMetaRequest validMetaReq = new MergedBlockMetaRequest(19, "app1", 0, 0, 0);
     requestHandler.handle(validMetaReq);
     assertEquals(1, responseAndPromisePairs.size());
     assertTrue(responseAndPromisePairs.get(0).getLeft() instanceof MergedBlockMetaSuccess);
     assertEquals(2,
       ((MergedBlockMetaSuccess) (responseAndPromisePairs.get(0).getLeft())).getNumChunks());
 
-    MergedBlockMetaRequest invalidMetaReq = new MergedBlockMetaRequest(21, "app1", -1, 1);
+    MergedBlockMetaRequest invalidMetaReq = new MergedBlockMetaRequest(21, "app1", -1, 0, 1);
     requestHandler.handle(invalidMetaReq);
     assertEquals(2, responseAndPromisePairs.size());
     assertTrue(responseAndPromisePairs.get(1).getLeft() instanceof RpcFailure);

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/BlockStoreClient.java

@@ -33,9 +33,9 @@
 import org.apache.spark.network.client.RpcResponseCallback;
 import org.apache.spark.network.client.TransportClient;
 import org.apache.spark.network.client.TransportClientFactory;
-import org.apache.spark.network.shuffle.protocol.BlockTransferMessage;
-import org.apache.spark.network.shuffle.protocol.GetLocalDirsForExecutors;
-import org.apache.spark.network.shuffle.protocol.LocalDirsForExecutors;
+import org.apache.spark.network.shuffle.checksum.Cause;
+import org.apache.spark.network.shuffle.protocol.*;
+import org.apache.spark.network.util.TransportConf;
 
 /**
  * Provides an interface for reading both shuffle files and RDD blocks, either from an Executor
@@ -46,6 +46,45 @@ public abstract class BlockStoreClient implements Closeable {
 
   protected volatile TransportClientFactory clientFactory;
   protected String appId;
+  protected TransportConf transportConf;
+
+  /**
+   * Send the diagnosis request for the corrupted shuffle block to the server.
+   *
+   * @param host the host of the remote node.
+   * @param port the port of the remote node.
+   * @param execId the executor id.
+   * @param shuffleId the shuffleId of the corrupted shuffle block
+   * @param mapId the mapId of the corrupted shuffle block
+   * @param reduceId the reduceId of the corrupted shuffle block
+   * @param checksum the shuffle checksum which calculated at client side for the corrupted
+   *                 shuffle block
+   * @return The cause of the shuffle block corruption
+   */
+  public Cause diagnoseCorruption(
+      String host,
+      int port,
+      String execId,
+      int shuffleId,
+      long mapId,
+      int reduceId,
+      long checksum,
+      String algorithm) {
+    try {
+      TransportClient client = clientFactory.createClient(host, port);
+      ByteBuffer response = client.sendRpcSync(
+        new DiagnoseCorruption(appId, execId, shuffleId, mapId, reduceId, checksum, algorithm)
+          .toByteBuffer(),
+        transportConf.connectionTimeoutMs()
+      );
+      CorruptionCause cause =
+        (CorruptionCause) BlockTransferMessage.Decoder.fromByteBuffer(response);
+      return cause.cause;
+    } catch (Exception e) {
+      logger.warn("Failed to get the corruption cause.");
+      return Cause.UNKNOWN_ISSUE;
+    }
+  }
 
   /**
    * Fetch a sequence of blocks from a remote node asynchronously,
@@ -167,6 +206,8 @@ public void pushBlocks(
    * @param host host of shuffle server
    * @param port port of shuffle server.
    * @param shuffleId shuffle ID of the shuffle to be finalized
+   * @param shuffleMergeId shuffleMergeId is used to uniquely identify merging process
+   *                       of shuffle by an indeterminate stage attempt.
    * @param listener the listener to receive MergeStatuses
    *
    * @since 3.1.0
@@ -175,6 +216,7 @@ public void finalizeShuffleMerge(
       String host,
       int port,
       int shuffleId,
+      int shuffleMergeId,
       MergeFinalizerListener listener) {
     throw new UnsupportedOperationException();
   }
@@ -185,6 +227,8 @@ public void finalizeShuffleMerge(
    * @param host the host of the remote node.
    * @param port the port of the remote node.
    * @param shuffleId shuffle id.
+   * @param shuffleMergeId shuffleMergeId is used to uniquely identify merging process
+   *                       of shuffle by an indeterminate stage attempt.
    * @param reduceId reduce id.
    * @param listener the listener to receive chunk counts.
    *
@@ -194,6 +238,7 @@ public void getMergedBlockMeta(
       String host,
       int port,
       int shuffleId,
+      int shuffleMergeId,
       int reduceId,
       MergedBlocksMetaListener listener) {
     throw new UnsupportedOperationException();

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/ErrorHandler.java

@@ -55,12 +55,14 @@ default boolean shouldLogError(Throwable t) {
   class BlockPushErrorHandler implements ErrorHandler {
     /**
      * String constant used for generating exception messages indicating a block to be merged
-     * arrives too late on the server side, and also for later checking such exceptions on the
-     * client side. When we get a block push failure because of the block arrives too late, we
-     * will not retry pushing the block nor log the exception on the client side.
+     * arrives too late or stale block push in the case of indeterminate stage retries on the
+     * server side, and also for later checking such exceptions on the client side. When we get
+     * a block push failure because of the block push being stale or arrives too late, we will
+     * not retry pushing the block nor log the exception on the client side.
      */
-    public static final String TOO_LATE_MESSAGE_SUFFIX =
-      "received after merged shuffle is finalized";
+    public static final String TOO_LATE_OR_STALE_BLOCK_PUSH_MESSAGE_SUFFIX =
+      "received after merged shuffle is finalized or stale block push as shuffle blocks of a"
+        + " higher shuffleMergeId for the shuffle is being pushed";
 
     /**
      * String constant used for generating exception messages indicating the server couldn't
@@ -81,25 +83,54 @@ class BlockPushErrorHandler implements ErrorHandler {
     public static final String IOEXCEPTIONS_EXCEEDED_THRESHOLD_PREFIX =
       "IOExceptions exceeded the threshold";
 
+    /**
+     * String constant used for generating exception messages indicating the server rejecting a
+     * shuffle finalize request since shuffle blocks of a higher shuffleMergeId for a shuffle is
+     * already being pushed. This typically happens in the case of indeterminate stage retries
+     * where if a stage attempt fails then the entirety of the shuffle output needs to be rolled
+     * back. For more details refer SPARK-23243, SPARK-25341 and SPARK-32923.
+     */
+    public static final String STALE_SHUFFLE_FINALIZE_SUFFIX =
+      "stale shuffle finalize request as shuffle blocks of a higher shuffleMergeId for the"
+        + " shuffle is already being pushed";
+
     @Override
     public boolean shouldRetryError(Throwable t) {
       // If it is a connection time-out or a connection closed exception, no need to retry.
       // If it is a FileNotFoundException originating from the client while pushing the shuffle
-      // blocks to the server, even then there is no need to retry. We will still log this exception
-      // once which helps with debugging.
+      // blocks to the server, even then there is no need to retry. We will still log this
+      // exception once which helps with debugging.
       if (t.getCause() != null && (t.getCause() instanceof ConnectException ||
           t.getCause() instanceof FileNotFoundException)) {
         return false;
       }
-      // If the block is too late, there is no need to retry it
-      return !Throwables.getStackTraceAsString(t).contains(TOO_LATE_MESSAGE_SUFFIX);
+
+      String errorStackTrace = Throwables.getStackTraceAsString(t);
+      // If the block is too late or stale block push, there is no need to retry it
+      return !errorStackTrace.contains(TOO_LATE_OR_STALE_BLOCK_PUSH_MESSAGE_SUFFIX);
     }
 
     @Override
     public boolean shouldLogError(Throwable t) {
       String errorStackTrace = Throwables.getStackTraceAsString(t);
-      return !errorStackTrace.contains(BLOCK_APPEND_COLLISION_DETECTED_MSG_PREFIX) &&
-        !errorStackTrace.contains(TOO_LATE_MESSAGE_SUFFIX);
+      return !(errorStackTrace.contains(BLOCK_APPEND_COLLISION_DETECTED_MSG_PREFIX) ||
+        errorStackTrace.contains(TOO_LATE_OR_STALE_BLOCK_PUSH_MESSAGE_SUFFIX));
+    }
+  }
+
+  class BlockFetchErrorHandler implements ErrorHandler {
+    public static final String STALE_SHUFFLE_BLOCK_FETCH =
+      "stale shuffle block fetch request as shuffle blocks of a higher shuffleMergeId for the"
+        + " shuffle is available";
+
+    @Override
+    public boolean shouldRetryError(Throwable t) {
+      return !Throwables.getStackTraceAsString(t).contains(STALE_SHUFFLE_BLOCK_FETCH);
+    }
+
+    @Override
+    public boolean shouldLogError(Throwable t) {
+      return !Throwables.getStackTraceAsString(t).contains(STALE_SHUFFLE_BLOCK_FETCH);
     }
   }
 }