Support external shuffle service

pom.xml

@@ -45,6 +45,8 @@ See file LICENSE for terms.
               <excludes>
                 <exclude>**/spark_2_1/**</exclude>
                 <exclude>**/spark_3_0/**</exclude>
+                <exclude>org/apache/spark/network/**</exclude>
+                <exclude>org/apache/spark/shuffle/ucx/external/server/**</exclude>
               </excludes>
             </configuration>
           </plugin>
@@ -71,6 +73,8 @@ See file LICENSE for terms.
               <excludes>
                 <exclude>**/spark_2_1/**</exclude>
                 <exclude>**/spark_2_4/**</exclude>
+                <exclude>org/apache/spark/network/**</exclude>
+                <exclude>org/apache/spark/shuffle/ucx/external/server/**</exclude>
               </excludes>
             </configuration>
           </plugin>
@@ -83,6 +87,93 @@ See file LICENSE for terms.
         <sonar.exclusions>**/spark_2_1/**, **/spark_2_4/**</sonar.exclusions>
       </properties>
     </profile>
+    <profile>
+      <id>sparkess-2.4</id>
+      <build>
+        <plugins>
+          <plugin>
+            <groupId>net.alchim31.maven</groupId>
+            <artifactId>scala-maven-plugin</artifactId>
+            <configuration>
+              <includes>
+                <include>org/apache/spark/network/**</include>
+                <include>org/apache/spark/shuffle/ucx/external/**</include>
+                <include>org/apache/spark/shuffle/ucx/memory/UcxLimitedMemPool.scala</include>
+                <include>org/apache/spark/shuffle/ucx/ShuffleTransport.scala</include>
+                <include>org/apache/spark/shuffle/utils/**</include>
+              </includes>
+              <excludes>
+                <exclude>**/spark_3_0/**</exclude>
+                <exclude>org/apache/spark/shuffle/ucx/external/client/**</exclude>
+              </excludes>
+            </configuration>
+          </plugin>
+        </plugins>
+      </build>
+      <properties>
+        <spark.version>2.4.0</spark.version>
+        <sonar.exclusions>**/spark_3_0/**, **/spark_2_1/**</sonar.exclusions>
+        <scala.version>2.11.12</scala.version>
+        <scala.compat.version>2.11</scala.compat.version>
+        <fasterxml.jackson.version>2.6.7</fasterxml.jackson.version>
+      </properties>
+      <dependencies>
+        <dependency>
+          <groupId>org.scala-lang</groupId>
+          <artifactId>scala-library</artifactId>
+          <version>${scala.version}</version>
+        </dependency>
+        <dependency>
+          <groupId>com.fasterxml.jackson.core</groupId>
+          <artifactId>jackson-databind</artifactId>
+          <version>${fasterxml.jackson.version}</version>
+        </dependency>
+      </dependencies>
+    </profile>
+    <profile>
+      <id>sparkess-3.0</id>
+      <build>
+        <plugins>
+          <plugin>
+            <groupId>net.alchim31.maven</groupId>
+            <artifactId>scala-maven-plugin</artifactId>
+            <configuration>
+              <includes>
+                <include>org/apache/spark/network/**</include>
+                <include>org/apache/spark/shuffle/ucx/external/**</include>
+                <include>org/apache/spark/shuffle/ucx/memory/UcxLimitedMemPool.scala</include>
+                <include>org/apache/spark/shuffle/ucx/ShuffleTransport.scala</include>
+                <include>org/apache/spark/shuffle/utils/**</include>
+              </includes>
+              <excludes>
+                <exclude>**/spark_2_4/**</exclude>
+                <exclude>org/apache/spark/shuffle/ucx/external/client/**</exclude>
+              </excludes>
+            </configuration>
+          </plugin>
+        </plugins>
+      </build>
+      <properties>
+        <spark.version>3.0.1</spark.version>
+        <scala.version>2.12.10</scala.version>
+        <scala.compat.version>2.12</scala.compat.version>
+        <sonar.exclusions>**/spark_2_1/**, **/spark_2_4/**</sonar.exclusions>
+        <fasterxml.jackson.version>2.10.0</fasterxml.jackson.version>
+      </properties>
+      <dependencies>
+        <dependency>
+          <groupId>org.scala-lang</groupId>
+          <artifactId>scala-library</artifactId>
+          <version>${scala.version}</version>
+        </dependency>
+        <dependency>
+          <groupId>com.fasterxml.jackson.core</groupId>
+          <artifactId>jackson-databind</artifactId>
+          <version>${fasterxml.jackson.version}</version>
+        </dependency>
+      </dependencies>
+    </profile>
+
   </profiles>
 
 
@@ -93,10 +184,16 @@ See file LICENSE for terms.
       <version>${spark.version}</version>
       <scope>provided</scope>
     </dependency>
+    <dependency>
+      <groupId>org.apache.spark</groupId>
+      <artifactId>spark-network-yarn_${scala.compat.version}</artifactId>
+      <version>${spark.version}</version>
+      <scope>provided</scope>
+    </dependency>
     <dependency>
       <groupId>org.openucx</groupId>
       <artifactId>jucx</artifactId>
-      <version>1.13.1</version>
+      <version>1.16.0</version>
     </dependency>
   </dependencies>
 

src/main/scala/org/apache/spark/network/shuffle/ExternalUcxShuffleBlockResolver.scala

@@ -0,0 +1,91 @@
+package org.apache.spark.network.shuffle
+
+import java.io.File
+import java.nio.charset.StandardCharsets
+import java.lang.reflect.{Method, Field}
+
+import scala.collection.mutable
+
+import com.fasterxml.jackson.databind.ObjectMapper
+
+import org.apache.hadoop.metrics2.lib.DefaultMetricsSystem;
+import org.apache.hadoop.metrics2.impl.MetricsSystemImpl;
+
+import org.apache.spark.network.shuffle.protocol.ExecutorShuffleInfo
+import org.apache.spark.network.util.TransportConf
+import org.apache.spark.network.shuffle.ExternalShuffleBlockResolver.AppExecId
+
+import org.apache.spark.shuffle.utils.UcxLogging
+import org.apache.spark.shuffle.ucx.ExternalUcxServerTransport
+
+class ExternalUcxShuffleBlockResolver(conf: TransportConf, registeredExecutorFile: File)
+  extends ExternalShuffleBlockResolver(conf, registeredExecutorFile) with UcxLogging {
+  private[spark] final val APP_KEY_PREFIX = "AppExecShuffleInfo";
+  private[spark] final val ucxMapper = new ObjectMapper
+  private[spark] var dbAppExecKeyMethod: Method = _
+  private[spark] val knownManagers = mutable.Set(
+    "org.apache.spark.shuffle.sort.SortShuffleManager",
+    "org.apache.spark.shuffle.unsafe.UnsafeShuffleManager",
+    "org.apache.spark.shuffle.ExternalUcxShuffleManager")
+  private[spark] var ucxTransport: ExternalUcxServerTransport = _
+
+  // init()
+
+  private[spark] def dbAppExecKey(appExecId: AppExecId): Array[Byte] = {
+    // we stick a common prefix on all the keys so we can find them in the DB
+    val appExecJson = ucxMapper.writeValueAsString(appExecId);
+    val key = (APP_KEY_PREFIX + ";" + appExecJson);
+    key.getBytes(StandardCharsets.UTF_8);
+  }
+
+  // def init(): Unit = {
+  //   val clazz = Class.forName("org.apache.spark.network.shuffle.ExternalShuffleBlockResolver")
+  //   try {
+  //     dbAppExecKeyMethod = clazz.getDeclaredMethod("dbAppExecKey", classOf[AppExecId])
+  //     dbAppExecKeyMethod.setAccessible(true)
+  //   } catch {
+  //     case e: Exception => {
+  //       logError(s"Get dbAppExecKey from ExternalUcxShuffleBlockResolver failed: $e")
+  //     }
+  //   }
+  // }
+
+  // def dbAppExecKey(fullId: AppExecId): Array[Byte] = {
+  //   dbAppExecKeyMethod.invoke(this, fullId).asInstanceOf[Array[Byte]]
+  // }
+
+  def setTransport(transport: ExternalUcxServerTransport): Unit = {
+    ucxTransport = transport
+  }
+  /** Registers a new Executor with all the configuration we need to find its shuffle files. */
+  override def registerExecutor(
+      appId: String,
+      execId: String,
+      executorInfo: ExecutorShuffleInfo): Unit = {
+    val fullId = new AppExecId(appId, execId)
+    logInfo(s"Registered executor ${fullId} with ${executorInfo}")
+    if (!knownManagers.contains(executorInfo.shuffleManager)) {
+      throw new UnsupportedOperationException(
+        "Unsupported shuffle manager of executor: " + executorInfo)
+    }
+    try {
+      if (db != null) {
+        val key = dbAppExecKey(fullId)
+        val value = ucxMapper.writeValueAsString(executorInfo).getBytes(StandardCharsets.UTF_8)
+        db.put(key, value)
+      }
+      executors.put(fullId, executorInfo)
+    } catch {
+      case e: Exception => logError("Error saving registered executors", e)
+    }
+  }
+
+  override def applicationRemoved(appId: String, cleanupLocalDirs: Boolean): Unit = {
+    super.applicationRemoved(appId, cleanupLocalDirs)
+    ucxTransport.applicationRemoved(appId)
+  }
+  override def executorRemoved(executorId: String, appId: String): Unit = {
+    super.executorRemoved(executorId, appId)
+    ucxTransport.executorRemoved(executorId, appId)
+  }
+}

src/main/scala/org/apache/spark/network/shuffle/spark_2_4/ExternalUcxShuffleBlockHandler.scala

@@ -0,0 +1,20 @@
+package org.apache.spark.network.shuffle
+
+import java.io.File
+
+import org.apache.spark.network.server.OneForOneStreamManager
+import org.apache.spark.network.util.TransportConf
+
+import org.apache.spark.shuffle.utils.UcxLogging
+import org.apache.spark.shuffle.ucx.ExternalUcxServerTransport
+
+class ExternalUcxShuffleBlockHandler(conf: TransportConf, registeredExecutorFile: File)
+  extends ExternalShuffleBlockHandler(new OneForOneStreamManager(),
+  new ExternalUcxShuffleBlockResolver(conf, registeredExecutorFile)) with UcxLogging {
+    def ucxBlockManager(): ExternalUcxShuffleBlockResolver = {
+      blockManager.asInstanceOf[ExternalUcxShuffleBlockResolver]
+    }
+    def setTransport(transport: ExternalUcxServerTransport): Unit = {
+      ucxBlockManager.setTransport(transport)
+    }
+}

src/main/scala/org/apache/spark/network/shuffle/spark_3_0/ExternalUcxShuffleBlockHandler.scala

@@ -0,0 +1,20 @@
+package org.apache.spark.network.shuffle
+
+import java.io.File
+
+import org.apache.spark.network.server.OneForOneStreamManager
+import org.apache.spark.network.util.TransportConf
+
+import org.apache.spark.shuffle.utils.UcxLogging
+import org.apache.spark.shuffle.ucx.ExternalUcxServerTransport
+
+class ExternalUcxShuffleBlockHandler(conf: TransportConf, registeredExecutorFile: File)
+  extends ExternalBlockHandler(new OneForOneStreamManager(),
+  new ExternalUcxShuffleBlockResolver(conf, registeredExecutorFile)) with UcxLogging {
+    def ucxBlockManager(): ExternalUcxShuffleBlockResolver = {
+      blockManager.asInstanceOf[ExternalUcxShuffleBlockResolver]
+    }
+    def setTransport(transport: ExternalUcxServerTransport): Unit = {
+      ucxBlockManager.setTransport(transport)
+    }
+}