[SHUFFLE] [WIP] Prototype: store shuffle file on external storage like S3

assembly/pom.xml

@@ -268,5 +268,17 @@
         </dependency>
       </dependencies>
     </profile>
+
+    <!-- Profile to include external shuffle storage jar file in Spark distribution -->
+    <profile>
+      <id>external-shuffle-storage</id>
+      <dependencies>
+        <dependency>
+          <groupId>org.apache.spark</groupId>
+          <artifactId>external-shuffle-storage_${scala.binary.version}</artifactId>
+          <version>${project.version}</version>
+        </dependency>
+      </dependencies>
+    </profile>
   </profiles>
 </project>

external-shuffle-storage/README.md

@@ -0,0 +1,46 @@
+# External Shuffle Storage
+
+This module provides support to store shuffle files on external shuffle storage like S3. It helps Dynamic
+Allocation on Kubernetes. Spark driver could release idle executors without worrying about losing
+shuffle data because the shuffle data is store on external shuffle storage which are different 
+from executors.
+
+This module implements a new Shuffle Manager named as StarShuffleManager, and copies a lot of codes
+from Spark SortShuffleManager. This is for a quick prototype. We want to use this as an example to discuss
+with Spark community and get feedback. We will work with the community to remove code duplication later
+and make StarShuffleManager more integrated with Spark code.
+
+## How to Build Spark Distribution with StarShuffleManager jar File
+
+Follow [Building Spark](https://spark.apache.org/docs/latest/building-spark.html) instructions,
+with extra `-Pexternal-shuffle-storage` to generate the new shuffle implementation jar file.
+
+Following is one command example to use `dev/make-distribution.sh` under Spark repo root directory:
+
+```
+./dev/make-distribution.sh --name spark-with-external-shuffle-storage --pip --tgz -Phive -Phive-thriftserver -Pkubernetes -Phadoop-3.2 -Phadoop-cloud -Dhadoop.version=3.2.0 -Pexternal-shuffle-storage
+```
+
+If you want to build a Spark docker image, you could unzip the Spark distribution tgz file, and run command like following:
+
+```
+./bin/docker-image-tool.sh -t spark-with-external-shuffle-storage build
+```
+
+This command creates `external-shuffle-storage_xxx.jar` file for StarShuffleManager
+under `jars` directory in the generated Spark distribution. Now you could use this Spark 
+distribution to run your Spark application with external shuffle storage.
+
+## How to Run Spark Application With External Shuffle Storage in Kubernetes
+
+### Run Spark Application With S3 as External Shuffle Storage and Dynamic Allocation
+
+Add configure to your Spark application like following (you need to adjust the values based on your environment):
+
+```
+spark.shuffle.manager=org.apache.spark.shuffle.StarShuffleManager
+spark.shuffle.star.rootDir=s3://my_bucket_name/my_shuffle_folder
+spark.dynamicAllocation.enabled=true
+spark.dynamicAllocation.shuffleTracking.enabled=true
+spark.dynamicAllocation.shuffleTracking.timeout=1
+```

external-shuffle-storage/pom.xml

@@ -0,0 +1,143 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+  ~ 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.
+  -->
+
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <modelVersion>4.0.0</modelVersion>
+    <parent>
+        <groupId>org.apache.spark</groupId>
+        <artifactId>spark-parent_2.12</artifactId>
+        <version>3.3.0-SNAPSHOT</version>
+        <relativePath>../pom.xml</relativePath>
+    </parent>
+
+    <artifactId>external-shuffle-storage_2.12</artifactId>
+    <packaging>jar</packaging>
+    <name>External Shuffle Storage</name>
+    <url>http://spark.apache.org/</url>
+
+    <properties>
+        <sbt.project.name>external-shuffle-storage</sbt.project.name>
+        <build.testJarPhase>none</build.testJarPhase>
+        <build.copyDependenciesPhase>package</build.copyDependenciesPhase>
+        <hadoop.deps.scope>provided</hadoop.deps.scope>
+        <hive.deps.scope>provided</hive.deps.scope>
+        <parquet.deps.scope>provided</parquet.deps.scope>
+    </properties>
+
+    <dependencies>
+        <dependency>
+            <groupId>org.apache.spark</groupId>
+            <artifactId>spark-core_${scala.binary.version}</artifactId>
+            <version>${project.version}</version>
+            <scope>provided</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.spark</groupId>
+            <artifactId>spark-core_${scala.binary.version}</artifactId>
+            <version>${project.version}</version>
+            <type>test-jar</type>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.spark</groupId>
+            <artifactId>spark-sql_${scala.binary.version}</artifactId>
+            <version>${project.version}</version>
+            <scope>provided</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.commons</groupId>
+            <artifactId>commons-math3</artifactId>
+            <scope>provided</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.scalacheck</groupId>
+            <artifactId>scalacheck_${scala.binary.version}</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.scala-lang</groupId>
+            <artifactId>scala-library</artifactId>
+            <scope>provided</scope>
+        </dependency>
+        <dependency>
+            <groupId>io.netty</groupId>
+            <artifactId>netty-all</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.slf4j</groupId>
+            <artifactId>slf4j-api</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.slf4j</groupId>
+            <artifactId>slf4j-log4j12</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.httpcomponents</groupId>
+            <artifactId>httpclient</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>commons-io</groupId>
+            <artifactId>commons-io</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.commons</groupId>
+            <artifactId>commons-lang3</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>com.amazonaws</groupId>
+            <artifactId>aws-java-sdk-s3</artifactId>
+            <version>1.11.975</version>
+            <scope>provided</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.testng</groupId>
+            <artifactId>testng</artifactId>
+            <version>6.14.3</version>
+            <scope>test</scope>
+        </dependency>
+    </dependencies>
+
+    <build>
+        <outputDirectory>target/scala-${scala.binary.version}/classes</outputDirectory>
+        <testOutputDirectory>target/scala-${scala.binary.version}/test-classes</testOutputDirectory>
+        <plugins>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-deploy-plugin</artifactId>
+                <configuration>
+                    <skip>true</skip>
+                </configuration>
+            </plugin>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-install-plugin</artifactId>
+                <configuration>
+                    <skip>true</skip>
+                </configuration>
+            </plugin>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-jar-plugin</artifactId>
+                <configuration>
+                    <outputDirectory>${jars.target.dir}</outputDirectory>
+                </configuration>
+            </plugin>
+        </plugins>
+    </build>
+</project>

external-shuffle-storage/src/main/java/org/apache/spark/starshuffle/ByteBufUtils.java

@@ -0,0 +1,46 @@
+/*
+ * This file is copied from Uber Remote Shuffle Service
+ * (https://github.com/uber/RemoteShuffleService) and modified.
+ *
+ * Licensed 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.starshuffle;
+
+import io.netty.buffer.ByteBuf;
+
+import java.io.IOException;
+import java.io.OutputStream;
+import java.nio.charset.StandardCharsets;
+
+public class ByteBufUtils {
+  public static final void writeLengthAndString(ByteBuf buf, String str) {
+    if (str == null) {
+      buf.writeInt(-1);
+      return;
+    }
+
+    byte[] bytes = str.getBytes(StandardCharsets.UTF_8);
+    buf.writeInt(bytes.length);
+    buf.writeBytes(bytes);
+  }
+
+  public static final String readLengthAndString(ByteBuf buf) {
+    int length = buf.readInt();
+    if (length == -1) {
+      return null;
+    }
+
+    byte[] bytes = new byte[length];
+    buf.readBytes(bytes);
+    return new String(bytes, StandardCharsets.UTF_8);
+  }
+}

external-shuffle-storage/src/main/java/org/apache/spark/starshuffle/StarBlockStoreClient.java

@@ -0,0 +1,134 @@
+/*
+ * 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.starshuffle;
+
+import org.apache.spark.SparkEnv;
+import org.apache.spark.network.buffer.ManagedBuffer;
+import org.apache.spark.network.buffer.NioManagedBuffer;
+import org.apache.spark.network.netty.SparkTransportConf;
+import org.apache.spark.network.shuffle.*;
+import org.apache.spark.network.util.TransportConf;
+import org.apache.spark.storage.BlockId;
+import org.apache.spark.storage.ShuffleBlockId;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import scala.Option;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.ByteBuffer;
+import java.util.concurrent.CompletableFuture;
+
+/**
+ * This class fetches shuffle blocks from external storage like S3
+ */
+public class StarBlockStoreClient extends BlockStoreClient {
+
+    private static final Logger logger = LoggerFactory.getLogger(StarBlockStoreClient.class);
+
+    // Fetch shuffle blocks from external shuffle storage.
+    // The shuffle location is encoded in the host argument. In the future, we should enhance
+    // Spark internal code to support abstraction of shuffle storage location.
+    @Override
+    public void fetchBlocks(String host, int port, String execId, String[] blockIds, BlockFetchingListener listener, DownloadFileManager downloadFileManager) {
+        for (int i = 0; i < blockIds.length; i++) {
+            String blockId = blockIds[i];
+            CompletableFuture.runAsync(() -> fetchBlock(host, execId, blockId, listener, downloadFileManager));
+        }
+    }
+
+    private void fetchBlock(String host, String execId, String blockIdStr, BlockFetchingListener listener, DownloadFileManager downloadFileManager) {
+        BlockId blockId = BlockId.apply(blockIdStr);
+        if (blockId instanceof ShuffleBlockId) {
+            ShuffleBlockId shuffleBlockId = (ShuffleBlockId)blockId;
+            StarMapResultFileInfo mapResultFileInfo = StarMapResultFileInfo.deserializeFromString(host);
+            long offset = 0;
+            for (int i = 0; i < shuffleBlockId.reduceId(); i++) {
+                offset += mapResultFileInfo.getPartitionLengths()[i];
+            }
+            long size = mapResultFileInfo.getPartitionLengths()[shuffleBlockId.reduceId()];
+            StarShuffleFileManager streamProvider = StarUtils.createShuffleFileManager(SparkEnv.get().conf(),
+                    mapResultFileInfo.getLocation());
+            if (downloadFileManager != null) {
+                try (InputStream inputStream = streamProvider.read(mapResultFileInfo.getLocation(), offset, size)) {
+                    TransportConf transportConf = SparkTransportConf.fromSparkConf(
+                            SparkEnv.get().conf(), "starShuffle", 1, Option.empty());
+                    DownloadFile downloadFile = downloadFileManager.createTempFile(transportConf);
+                    downloadFileManager.registerTempFileToClean(downloadFile);
+                    DownloadFileWritableChannel downloadFileWritableChannel = downloadFile.openForWriting();
+
+                    int bufferSize = 64 * 1024;
+                    byte[] bytes = new byte[bufferSize];
+                    int readBytes = 0;
+                    while (readBytes < size) {
+                        int toReadBytes = Math.min((int)size - readBytes, bufferSize);
+                        int n = inputStream.read(bytes, 0, toReadBytes);
+                        if (n == -1) {
+                            throw new RuntimeException(String.format(
+                                    "Failed to read file %s for shuffle block %s, hit end with remaining %s bytes",
+                                    mapResultFileInfo.getLocation(),
+                                    blockId,
+                                    size - readBytes));
+                        }
+                        readBytes += n;
+                        downloadFileWritableChannel.write(ByteBuffer.wrap(bytes, 0, n));
+                    }
+                    ManagedBuffer managedBuffer = downloadFileWritableChannel.closeAndRead();
+                    listener.onBlockFetchSuccess(blockIdStr, managedBuffer);
+                } catch (IOException e) {
+                    throw new RuntimeException(String.format(
+                            "Failed to read file %s for shuffle block %s",
+                            mapResultFileInfo.getLocation(),
+                            blockId));
+                }
+            } else {
+                try (InputStream inputStream = streamProvider.read(mapResultFileInfo.getLocation(), offset, size)) {
+                    ByteBuffer byteBuffer = ByteBuffer.allocate((int)size);
+                    int b = inputStream.read();
+                    while (b != -1) {
+                        byteBuffer.put((byte)b);
+                        if (byteBuffer.position() == size) {
+                            break;
+                        }
+                        b = inputStream.read();
+                    }
+                    byteBuffer.flip();
+                    NioManagedBuffer managedBuffer = new NioManagedBuffer(byteBuffer);
+                    listener.onBlockFetchSuccess(blockIdStr, managedBuffer);
+                } catch (IOException e) {
+                    throw new RuntimeException(String.format(
+                            "Failed to read file %s for shuffle block %s",
+                            mapResultFileInfo.getLocation(),
+                            blockId));
+                }
+            }
+            logger.info("Fetch blocks: {}, {}", host, execId);
+        } else {
+            throw new RuntimeException(String.format(
+                    "%s does not support %s: %s",
+                    this.getClass().getSimpleName(),
+                    blockId.getClass().getSimpleName(),
+                    blockId));
+        }
+    }
+
+    @Override
+    public void close() throws IOException {
+        logger.info("Close");
+    }
+}