Demo of implementation to the local directory

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

@@ -34,6 +34,7 @@
 import org.slf4j.LoggerFactory;
 
 import org.apache.spark.network.buffer.ManagedBuffer;
+import org.apache.spark.network.client.StreamCallbackWithID;
 import org.apache.spark.network.client.RpcResponseCallback;
 import org.apache.spark.network.client.TransportClient;
 import org.apache.spark.network.server.OneForOneStreamManager;
@@ -81,6 +82,22 @@ public void receive(TransportClient client, ByteBuffer message, RpcResponseCallb
     handleMessage(msgObj, client, callback);
   }
 
+  @Override
+  public StreamCallbackWithID receiveStream(
+          TransportClient client,
+          ByteBuffer messageHeader,
+          RpcResponseCallback callback) {
+    BlockTransferMessage header = BlockTransferMessage.Decoder.fromByteBuffer(messageHeader);
+    return handleStream(header, client, callback);
+  }
+
+  protected StreamCallbackWithID handleStream(
+          BlockTransferMessage header,
+          TransportClient client,
+          RpcResponseCallback callback) {
+      throw new UnsupportedOperationException("Unexpected message header: " + header);
+  }
+
   protected void handleMessage(
       BlockTransferMessage msgObj,
       TransportClient client,

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

@@ -19,13 +19,13 @@
 
 import java.io.*;
 import java.nio.charset.StandardCharsets;
+import java.util.regex.Pattern;
 import java.util.*;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
 import java.util.regex.Matcher;
-import java.util.regex.Pattern;
 
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;

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

@@ -140,7 +140,8 @@ public void registerWithShuffleServer(
       ExecutorShuffleInfo executorInfo) throws IOException, InterruptedException {
     checkInit();
     try (TransportClient client = clientFactory.createUnmanagedClient(host, port)) {
-      ByteBuffer registerMessage = new RegisterExecutor(appId, execId, executorInfo).toByteBuffer();
+      ByteBuffer registerMessage =
+              new RegisterExecutor(appId, execId, executorInfo).toByteBuffer();
       client.sendRpcSync(registerMessage, registrationTimeoutMs);
     }
   }

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

@@ -0,0 +1,150 @@
+package org.apache.spark.network.shuffle;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.FileChannel;
+import java.nio.channels.WritableByteChannel;
+import java.nio.file.StandardOpenOption;
+
+import org.apache.spark.network.client.StreamCallbackWithID;
+
+public class FileWriterStreamCallback implements StreamCallbackWithID {
+
+  private static final Logger logger = LoggerFactory.getLogger(FileWriterStreamCallback.class);
+
+  public enum FileType {
+    DATA("shuffle-data"),
+    INDEX("shuffle-index");
+
+    private final String typeString;
+
+    FileType(String typeString) {
+      this.typeString = typeString;
+    }
+
+    @Override
+    public String toString() {
+      return typeString;
+    }
+  }
+
+  private final String appId;
+  private final int shuffleId;
+  private final int mapId;
+  private final File file;
+  private final FileType fileType;
+  private WritableByteChannel fileOutputChannel = null;
+
+  public FileWriterStreamCallback(
+      String appId,
+      int shuffleId,
+      int mapId,
+      File file,
+      FileWriterStreamCallback.FileType fileType) {
+    this.appId = appId;
+    this.shuffleId = shuffleId;
+    this.mapId = mapId;
+    this.file = file;
+    this.fileType = fileType;
+  }
+
+  public void open() {
+    logger.info(
+      "Opening {} for remote writing. File type: {}", file.getAbsolutePath(), fileType);
+    if (fileOutputChannel != null) {
+      throw new IllegalStateException(
+        String.format(
+          "File %s for is already open for writing (type: %s).",
+          file.getAbsolutePath(),
+          fileType));
+    }
+    if (!file.exists()) {
+      try {
+        if (!file.getParentFile().isDirectory() && !file.getParentFile().mkdirs()) {
+          throw new IOException(
+              String.format(
+                  "Failed to create shuffle file directory at"
+                      + file.getParentFile().getAbsolutePath() + "(type: %s).", fileType));
+        }
+
+        if (!file.createNewFile()) {
+          throw new IOException(
+              String.format(
+                  "Failed to create shuffle file (type: %s).", fileType));
+        }
+      } catch (IOException e) {
+        throw new RuntimeException(
+            String.format(
+                "Failed to create shuffle file at %s for backup (type: %s).",
+                file.getAbsolutePath(),
+                fileType),
+            e);
+      }
+    }
+    try {
+      // TODO encryption
+      fileOutputChannel = FileChannel.open(file.toPath(), StandardOpenOption.APPEND);
+    } catch (IOException e) {
+      throw new RuntimeException(
+          String.format(
+              "Failed to find file for writing at %s (type: %s).",
+              file.getAbsolutePath(),
+              fileType),
+          e);
+    }
+  }
+
+  @Override
+  public String getID() {
+    return String.format("%s-%d-%d-%s",
+      appId,
+      shuffleId,
+      mapId,
+      fileType);
+  }
+
+  @Override
+  public void onData(String streamId, ByteBuffer buf) throws IOException {
+    verifyShuffleFileOpenForWriting();
+    while (buf.hasRemaining()) {
+      fileOutputChannel.write(buf);
+    }
+  }
+
+  @Override
+  public void onComplete(String streamId) throws IOException {
+    logger.info(
+      "Finished writing {}. File type: {}", file.getAbsolutePath(), fileType);
+    fileOutputChannel.close();
+  }
+
+  @Override
+  public void onFailure(String streamId, Throwable cause) throws IOException {
+    logger.warn("Failed to write shuffle file at {} (type: %s).",
+      file.getAbsolutePath(),
+      fileType,
+      cause);
+    fileOutputChannel.close();
+    // TODO delete parent dirs too
+    if (!file.delete()) {
+      logger.warn(
+        "Failed to delete incomplete remote shuffle file at %s (type: %s)",
+        file.getAbsolutePath(),
+        fileType);
+    }
+  }
+
+  private void verifyShuffleFileOpenForWriting() {
+    if (fileOutputChannel == null) {
+      throw new RuntimeException(
+          String.format(
+            "Shuffle file at %s not open for writing (type: %s).",
+            file.getAbsolutePath(),
+            fileType));
+    }
+  }
+}

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

@@ -0,0 +1,125 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.network.shuffle.k8s;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+
+import com.google.common.util.concurrent.ThreadFactoryBuilder;
+import org.apache.spark.network.shuffle.protocol.ShuffleServiceHeartbeat;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.spark.network.client.RpcResponseCallback;
+import org.apache.spark.network.client.TransportClient;
+import org.apache.spark.network.sasl.SecretKeyHolder;
+import org.apache.spark.network.shuffle.ExternalShuffleClient;
+import org.apache.spark.network.shuffle.protocol.RegisterDriver;
+import org.apache.spark.network.util.TransportConf;
+
+/**
+ * A client for talking to the external shuffle service in Kubernetes coarse-grained mode.
+ *
+ * This is used by the Spark driver to register with each external shuffle service on the cluster.
+ * The reason why the driver has to talk to the service is for cleaning up shuffle files reliably
+ * after the application exits. Kubernetes does not provide a great alternative to do this, so Spark
+ * has to detect this itself.
+ */
+public class KubernetesExternalShuffleClient extends ExternalShuffleClient {
+  private static final Logger logger =
+    LoggerFactory.getLogger(KubernetesExternalShuffleClient.class);
+
+  private final ScheduledExecutorService heartbeaterThread =
+      Executors.newSingleThreadScheduledExecutor(
+        new ThreadFactoryBuilder()
+          .setDaemon(true)
+          .setNameFormat("kubernetes-external-shuffle-client-heartbeater")
+          .build());
+
+  /**
+   * Creates a Kubernetes external shuffle client that wraps the {@link ExternalShuffleClient}.
+   * Please refer to docs on {@link ExternalShuffleClient} for more information.
+   */
+  public KubernetesExternalShuffleClient(
+      TransportConf conf,
+      SecretKeyHolder secretKeyHolder,
+      boolean authEnabled,
+      long registrationTimeoutMs) {
+    super(conf, secretKeyHolder, authEnabled, registrationTimeoutMs);
+  }
+
+  public void registerDriverWithShuffleService(
+          String host,
+          int port,
+          long heartbeatTimeoutMs,
+          long heartbeatIntervalMs) throws IOException, InterruptedException {
+
+    checkInit();
+    ByteBuffer registerDriver = new RegisterDriver(appId, heartbeatTimeoutMs).toByteBuffer();
+    logger.info("Registering with external shuffle service at " + host + ":" + port);
+    TransportClient client = clientFactory.createClient(host, port);
+    client.sendRpc(registerDriver, new RegisterDriverCallback(client, heartbeatIntervalMs));
+  }
+
+  private class RegisterDriverCallback implements RpcResponseCallback {
+    private final TransportClient client;
+    private final long heartbeatIntervalMs;
+
+    private RegisterDriverCallback(TransportClient client, long heartbeatIntervalMs) {
+      this.client = client;
+      this.heartbeatIntervalMs = heartbeatIntervalMs;
+    }
+
+    @Override
+    public void onSuccess(ByteBuffer response) {
+      heartbeaterThread.scheduleAtFixedRate(
+          new Heartbeater(client), 0, heartbeatIntervalMs, TimeUnit.MILLISECONDS);
+      logger.info("Successfully registered app " + appId + " with external shuffle service.");
+    }
+
+    @Override
+    public void onFailure(Throwable e) {
+      logger.warn("Unable to register app " + appId + " with external shuffle service. " +
+          "Please manually remove shuffle data after driver exit. Error: " + e);
+    }
+  }
+
+  @Override
+  public void close() {
+    heartbeaterThread.shutdownNow();
+    super.close();
+  }
+
+  private class Heartbeater implements Runnable {
+
+    private final TransportClient client;
+
+    private Heartbeater(TransportClient client) {
+      this.client = client;
+    }
+
+    @Override
+    public void run() {
+      // TODO: Stop sending heartbeats if the shuffle service has lost the app due to timeout
+      client.send(new ShuffleServiceHeartbeat(appId).toByteBuffer());
+    }
+  }
+}

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

@@ -24,15 +24,15 @@
 import java.util.concurrent.TimeUnit;
 
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
-import org.apache.spark.network.shuffle.protocol.mesos.ShuffleServiceHeartbeat;
+import org.apache.spark.network.shuffle.protocol.ShuffleServiceHeartbeat;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import org.apache.spark.network.client.RpcResponseCallback;
 import org.apache.spark.network.client.TransportClient;
 import org.apache.spark.network.sasl.SecretKeyHolder;
 import org.apache.spark.network.shuffle.ExternalShuffleClient;
-import org.apache.spark.network.shuffle.protocol.mesos.RegisterDriver;
+import org.apache.spark.network.shuffle.protocol.RegisterDriver;
 import org.apache.spark.network.util.TransportConf;
 
 /**

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

@@ -23,8 +23,6 @@
 import io.netty.buffer.Unpooled;
 
 import org.apache.spark.network.protocol.Encodable;
-import org.apache.spark.network.shuffle.protocol.mesos.RegisterDriver;
-import org.apache.spark.network.shuffle.protocol.mesos.ShuffleServiceHeartbeat;
 
 /**
  * Messages handled by the {@link org.apache.spark.network.shuffle.ExternalShuffleBlockHandler}, or
@@ -42,7 +40,8 @@ public abstract class BlockTransferMessage implements Encodable {
   /** Preceding every serialized message is its type, which allows us to deserialize it. */
   public enum Type {
     OPEN_BLOCKS(0), UPLOAD_BLOCK(1), REGISTER_EXECUTOR(2), STREAM_HANDLE(3), REGISTER_DRIVER(4),
-    HEARTBEAT(5), UPLOAD_BLOCK_STREAM(6);
+    HEARTBEAT(5), UPLOAD_BLOCK_STREAM(6), UPLOAD_SHUFFLE_PARTITION_STREAM(7),
+    REGISTER_SHUFFLE_INDEX(8), OPEN_SHUFFLE_PARTITION(9), UPLOAD_SHUFFLE_INDEX(10);
 
     private final byte id;
 
@@ -68,6 +67,10 @@ public static BlockTransferMessage fromByteBuffer(ByteBuffer msg) {
         case 4: return RegisterDriver.decode(buf);
         case 5: return ShuffleServiceHeartbeat.decode(buf);
         case 6: return UploadBlockStream.decode(buf);
+        case 7: return UploadShufflePartitionStream.decode(buf);
+        case 8: return RegisterShuffleIndex.decode(buf);
+        case 9: return OpenShufflePartition.decode(buf);
+        case 10: return UploadShuffleIndex.decode(buf);
         default: throw new IllegalArgumentException("Unknown message type: " + type);
       }
     }