Topic/distribute according to NUMA

pom.xml

@@ -96,7 +96,17 @@ See file LICENSE for terms.
     <dependency>
       <groupId>org.openucx</groupId>
       <artifactId>jucx</artifactId>
-      <version>1.13.1</version>
+      <version>1.16.0</version>
+    </dependency>
+    <dependency>
+      <groupId>org.openucx</groupId>
+      <artifactId>nvkv</artifactId>
+      <version>1.0</version>
+  </dependency>
+    <dependency>
+      <groupId>org.ini4j</groupId>
+      <artifactId>ini4j</artifactId>
+      <version>0.5.4</version>
     </dependency>
   </dependencies>
 

src/main/scala/org/apache/spark/shuffle/compat/spark_2_4/UcxShuffleClient.scala

@@ -3,19 +3,19 @@ package org.apache.spark.shuffle.compat.spark_2_4
 import org.openucx.jucx.UcxUtils
 import org.apache.spark.network.buffer.{ManagedBuffer, NioManagedBuffer}
 import org.apache.spark.network.shuffle.{BlockFetchingListener, DownloadFileManager, ShuffleClient}
-import org.apache.spark.shuffle.ucx.{OperationCallback, OperationResult, UcxShuffleBockId, UcxShuffleTransport}
+import org.apache.spark.shuffle.ucx.{OperationCallback, OperationResult, UcxShuffleBlockId, UcxShuffleTransport}
 import org.apache.spark.shuffle.utils.UnsafeUtils
 import org.apache.spark.storage.{BlockId => SparkBlockId, ShuffleBlockId => SparkShuffleBlockId}
 
 class UcxShuffleClient(val transport: UcxShuffleTransport) extends ShuffleClient{
   override def fetchBlocks(host: String, port: Int, execId: String, blockIds: Array[String],
                            listener: BlockFetchingListener,
                            downloadFileManager: DownloadFileManager): Unit = {
-    val ucxBlockIds = Array.ofDim[UcxShuffleBockId](blockIds.length)
+    val ucxBlockIds = Array.ofDim[UcxShuffleBlockId](blockIds.length)
     val callbacks = Array.ofDim[OperationCallback](blockIds.length)
     for (i <- blockIds.indices) {
       val blockId = SparkBlockId.apply(blockIds(i)).asInstanceOf[SparkShuffleBlockId]
-      ucxBlockIds(i) = UcxShuffleBockId(blockId.shuffleId, blockId.mapId, blockId.reduceId)
+      ucxBlockIds(i) = UcxShuffleBlockId(blockId.shuffleId, blockId.mapId, blockId.reduceId)
       callbacks(i) = (result: OperationResult) => {
         val memBlock = result.getData
         val buffer = UnsafeUtils.getByteBufferView(memBlock.address, memBlock.size.toInt)

src/main/scala/org/apache/spark/shuffle/compat/spark_3_0/UcxLocalDiskShuffleExecutorComponents.scala

@@ -8,10 +8,12 @@ import java.util
 import java.util.Optional
 
 import org.apache.spark.internal.Logging
+import org.apache.spark.shuffle.utils.CommonUtils
 import org.apache.spark.{SparkConf, SparkEnv}
 import org.apache.spark.shuffle.sort.io.{LocalDiskShuffleExecutorComponents, LocalDiskShuffleMapOutputWriter, LocalDiskSingleSpillMapOutputWriter}
 import org.apache.spark.shuffle.UcxShuffleManager
 import org.apache.spark.shuffle.api.{ShuffleMapOutputWriter, SingleSpillShuffleMapOutputWriter}
+import scala.concurrent.duration._
 
 /**
  * Entry point to UCX executor.
@@ -22,14 +24,18 @@ class UcxLocalDiskShuffleExecutorComponents(sparkConf: SparkConf)
   private var blockResolver: UcxShuffleBlockResolver = _
 
   override def initializeExecutor(appId: String, execId: String, extraConfigs: util.Map[String, String]): Unit = {
+    logDebug("UcxLocalDiskShuffleExecutorComponents initializeExecutor appId: " + appId + " execId: " + execId)
     val ucxShuffleManager = SparkEnv.get.shuffleManager.asInstanceOf[UcxShuffleManager]
-    while (ucxShuffleManager.ucxTransport == null) {
-      Thread.sleep(5)
-    }
+    CommonUtils.safePolling(() => {},
+      () => {ucxShuffleManager.ucxTransport == null},
+      10.seconds.fromNow,
+      "Got timeout when polling", Duration(5, "millis"))
+
     blockResolver = ucxShuffleManager.shuffleBlockResolver
   }
 
   override def createMapOutputWriter(shuffleId: Int, mapTaskId: Long, numPartitions: Int): ShuffleMapOutputWriter = {
+    logDebug("UcxLocalDiskShuffleExecutorComponents createMapOutputWriter shuffleId: " + shuffleId + " mapTaskId: " + mapTaskId + " numPartitions: " + numPartitions)
     if (blockResolver == null) {
       throw new IllegalStateException(
         "Executor components must be initialized before getting writers.")
@@ -39,10 +45,14 @@ class UcxLocalDiskShuffleExecutorComponents(sparkConf: SparkConf)
   }
 
   override def createSingleFileMapOutputWriter(shuffleId: Int, mapId: Long): Optional[SingleSpillShuffleMapOutputWriter] = {
+    // Called for each mapper
+    logDebug("UcxLocalDiskShuffleExecutorComponents createSingleFileMapOutputWriter shuffleId: " + shuffleId + " mapId: " + mapId)
     if (blockResolver == null) {
       throw new IllegalStateException(
         "Executor components must be initialized before getting writers.")
     }
+
+    // Need to implement an alternative to LocalDiskSingleSpillMapOutputWriter?
     Optional.of(new LocalDiskSingleSpillMapOutputWriter(shuffleId, mapId, blockResolver))
   }
 

src/main/scala/org/apache/spark/shuffle/compat/spark_3_0/UcxShuffleBlockResolver.scala

@@ -5,9 +5,14 @@
 package org.apache.spark.shuffle.compat.spark_3_0
 
 import java.io.{File, RandomAccessFile}
+import java.nio.ByteBuffer
+
+import org.apache.spark.{TaskContext, SparkEnv}
+import org.apache.spark.storage._
+import org.apache.spark.network.buffer.{NioManagedBuffer, ManagedBuffer}
+import org.apache.spark.shuffle.utils.UnsafeUtils
+import org.apache.spark.shuffle.ucx.{OperationCallback, OperationResult, UcxShuffleTransport, CommonUcxShuffleBlockResolver, CommonUcxShuffleManager}
 
-import org.apache.spark.TaskContext
-import org.apache.spark.shuffle.ucx.{CommonUcxShuffleBlockResolver, CommonUcxShuffleManager}
 
 /**
  * Mapper entry point for UcxShuffle plugin. Performs memory registration
@@ -16,6 +21,8 @@ import org.apache.spark.shuffle.ucx.{CommonUcxShuffleBlockResolver, CommonUcxShu
 class UcxShuffleBlockResolver(ucxShuffleManager: CommonUcxShuffleManager)
   extends CommonUcxShuffleBlockResolver(ucxShuffleManager) {
 
+  val shuffleManager: CommonUcxShuffleManager = ucxShuffleManager
+
 
   override def writeIndexFileAndCommit(shuffleId: ShuffleId, mapId: Long,
                                        lengths: Array[Long], dataTmp: File): Unit = {
@@ -29,4 +36,27 @@ class UcxShuffleBlockResolver(ucxShuffleManager: CommonUcxShuffleManager)
     }
     writeIndexFileAndCommitCommon(shuffleId, partitionId, lengths, new RandomAccessFile(dataFile, "r"))
   }
+
+  override def getBlockData(
+      blockId: BlockId,
+      dirs: Option[Array[String]]): ManagedBuffer = {
+
+    val ucxTransport: UcxShuffleTransport = shuffleManager.ucxTransport
+
+    logInfo("UcxShuffleBlockResolver getBlockData")
+    val (shuffleId, mapId, startReduceId, endReduceId) = blockId match {
+      case id: ShuffleBlockId =>
+        (id.shuffleId, id.mapId, id.reduceId, id.reduceId + 1)
+      case batchId: ShuffleBlockBatchId =>
+        (batchId.shuffleId, batchId.mapId, batchId.startReduceId, batchId.endReduceId)
+      case _ =>
+        throw new IllegalArgumentException("unexpected shuffle block id format: " + blockId)
+     }
+
+    var length = ucxTransport.getNvkvWrapper.getPartitonLength(shuffleId, mapId, startReduceId).toInt
+    var offset = ucxTransport.getNvkvWrapper.getPartitonOffset(shuffleId, mapId, startReduceId)
+    logDebug(s"UcxShuffleBlockResolver - Reading shuffleId $shuffleId mapId $mapId reduceId $startReduceId at offset $offset with length $length from nvkv")
+    var resultBuffer = ucxTransport.getNvkvWrapper.read(length, offset)
+    new NioManagedBuffer(resultBuffer)
+  }
 }

src/main/scala/org/apache/spark/shuffle/compat/spark_3_0/UcxShuffleClient.scala

@@ -4,44 +4,57 @@
 */
 package org.apache.spark.shuffle.compat.spark_3_0
 
+import java.nio.ByteBuffer
 import org.apache.spark.internal.Logging
 import org.apache.spark.network.buffer.{ManagedBuffer, NioManagedBuffer}
 import org.apache.spark.network.shuffle.{BlockFetchingListener, BlockStoreClient, DownloadFileManager}
-import org.apache.spark.shuffle.ucx.{OperationCallback, OperationResult, UcxShuffleBockId, UcxShuffleTransport}
-import org.apache.spark.shuffle.utils.UnsafeUtils
+import org.apache.spark.shuffle.ucx.{OperationCallback, OperationResult, UcxShuffleBlockId, UcxShuffleTransport}
+import org.apache.spark.shuffle.utils.{UnsafeUtils, CommonUtils}
 import org.apache.spark.storage.{BlockId => SparkBlockId, ShuffleBlockId => SparkShuffleBlockId}
+import org.apache.spark.SparkException
 
 class UcxShuffleClient(val transport: UcxShuffleTransport, mapId2PartitionId: Map[Long, Int]) extends BlockStoreClient with Logging {
 
   override def fetchBlocks(host: String, port: Int, execId: String, blockIds: Array[String],
                            listener: BlockFetchingListener,
                            downloadFileManager: DownloadFileManager): Unit = {
-    if (blockIds.length > transport.ucxShuffleConf.maxBlocksPerRequest) {
-      val (b1, b2) = blockIds.splitAt(blockIds.length / 2)
-      fetchBlocks(host, port, execId, b1, listener, downloadFileManager)
-      fetchBlocks(host, port, execId, b2, listener, downloadFileManager)
-      return
-    }
+    //TODO - check if we need to limit max number of request "on the air"
+    // if (blockIds.length > 32) {
+    //   val (b1, b2) = blockIds.splitAt(blockIds.length / 2)
+    //   fetchBlocks(host, port, execId, b1, listener, downloadFileManager)
+    //   fetchBlocks(host, port, execId, b2, listener, downloadFileManager)
+    //   return
+    // }
 
-    val ucxBlockIds = Array.ofDim[UcxShuffleBockId](blockIds.length)
+    val ucxBlockIds = Array.ofDim[UcxShuffleBlockId](blockIds.length)
     val callbacks = Array.ofDim[OperationCallback](blockIds.length)
+    var send = 0
+    var receive = 0
     for (i <- blockIds.indices) {
-      val blockId = SparkBlockId.apply(blockIds(i)).asInstanceOf[SparkShuffleBlockId]
-      ucxBlockIds(i) = UcxShuffleBockId(blockId.shuffleId, mapId2PartitionId(blockId.mapId), blockId.reduceId)
-      callbacks(i) = (result: OperationResult) => {
-        val memBlock = result.getData
-        val buffer = UnsafeUtils.getByteBufferView(memBlock.address, memBlock.size.toInt)
-        listener.onBlockFetchSuccess(blockIds(i), new NioManagedBuffer(buffer) {
-          override def release: ManagedBuffer = {
-            memBlock.close()
-            this
+      send = send + 1
+      SparkBlockId.apply(blockIds(i)) match {
+        case blockId: SparkShuffleBlockId => {
+          ucxBlockIds(i) = UcxShuffleBlockId(blockId.shuffleId, mapId2PartitionId(blockId.mapId), blockId.reduceId)
+          callbacks(i) = (result: OperationResult) => {
+            val memBlock = result.getData
+            val buffer = UnsafeUtils.getByteBufferView(memBlock.address, memBlock.size.toInt)
+            listener.onBlockFetchSuccess(blockIds(i), new NioManagedBuffer(buffer) {
+              override def release: ManagedBuffer = {
+                memBlock.close()
+                this
+              }
+            })
           }
-        })
+          val resultBufferAllocator = (size: Long) => transport.hostBounceBufferMemoryPool.get(size)
+          transport.fetchBlocksByBlockIds(execId.toLong, Array(ucxBlockIds(i)), resultBufferAllocator, 
+            Array(callbacks(i)), () => {receive = receive + 1})
+        }
+        case _ =>
+          throw new SparkException("Unrecognized blockId")
       }
     }
-    val resultBufferAllocator = (size: Long) => transport.hostBounceBufferMemoryPool.get(size)
-    transport.fetchBlocksByBlockIds(execId.toLong, ucxBlockIds, resultBufferAllocator, callbacks)
-    transport.progress()
+
+    CommonUtils.safePolling(() => {transport.progress()}, () => {send != receive})
   }
 
   override def close(): Unit = {

src/main/scala/org/apache/spark/shuffle/compat/spark_3_0/UcxShuffleManager.scala

@@ -7,15 +7,16 @@ package org.apache.spark.shuffle
 import scala.collection.JavaConverters._
 
 import org.apache.spark.shuffle.api.ShuffleExecutorComponents
-import org.apache.spark.shuffle.compat.spark_3_0.{UcxLocalDiskShuffleExecutorComponents, UcxShuffleBlockResolver, UcxShuffleReader}
+import org.apache.spark.shuffle.compat.spark_3_0.{UcxShuffleBlockResolver, UcxShuffleReader}
 import org.apache.spark.shuffle.sort.{SerializedShuffleHandle, SortShuffleWriter, UnsafeShuffleWriter}
-import org.apache.spark.shuffle.ucx.CommonUcxShuffleManager
+import org.apache.spark.shuffle.ucx.{CommonUcxShuffleManager, NvkvShuffleExecutorComponents}
 import org.apache.spark.{SparkConf, SparkEnv, TaskContext}
 
 /**
  * Main entry point of Ucx shuffle plugin. It extends spark's default SortShufflePlugin
  * and injects needed logic in override methods.
  */
+
 class UcxShuffleManager(override val conf: SparkConf, isDriver: Boolean)
   extends CommonUcxShuffleManager(conf, isDriver) {
 
@@ -25,9 +26,11 @@ class UcxShuffleManager(override val conf: SparkConf, isDriver: Boolean)
 
   override def getWriter[K, V](handle: ShuffleHandle, mapId: ReduceId, context: TaskContext,
                                metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] = {
+    logInfo("UcxShuffleManager getWriter")
     val env = SparkEnv.get
     handle match {
       case unsafeShuffleHandle: SerializedShuffleHandle[K@unchecked, V@unchecked] =>
+        logDebug("UcxShuffleManager getWriter unsafeShuffleHandle")
         new UnsafeShuffleWriter(
           env.blockManager,
           context.taskMemoryManager(),
@@ -38,19 +41,22 @@ class UcxShuffleManager(override val conf: SparkConf, isDriver: Boolean)
           metrics,
           shuffleExecutorComponents)
       case other: BaseShuffleHandle[K@unchecked, V@unchecked, _] =>
+        logDebug("UcxShuffleManager getWriter other")
         new SortShuffleWriter(
           shuffleBlockResolver, other, mapId, context, shuffleExecutorComponents)
     }
   }
 
   override def getReader[K, C](handle: ShuffleHandle, startPartition: MapId, endPartition: MapId,
                                context: TaskContext, metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = {
+    logInfo("UcxShuffleManager getReader")
     new UcxShuffleReader(handle.asInstanceOf[BaseShuffleHandle[K,_,C]], startPartition, endPartition,
       context, ucxTransport, readMetrics = metrics, shouldBatchFetch = false)
   }
 
   private def loadShuffleExecutorComponents(conf: SparkConf): ShuffleExecutorComponents = {
-    val executorComponents = new UcxLocalDiskShuffleExecutorComponents(conf)
+    logInfo("UcxShuffleManager loadShuffleExecutorComponents")
+    val executorComponents = new NvkvShuffleExecutorComponents(conf, getTransport)
     val extraConfigs = conf.getAllWithPrefix(ShuffleDataIOUtils.SHUFFLE_SPARK_CONF_PREFIX)
       .toMap
     executorComponents.initializeExecutor(
@@ -59,5 +65,4 @@ class UcxShuffleManager(override val conf: SparkConf, isDriver: Boolean)
       extraConfigs.asJava)
     executorComponents
   }
-
 }

src/main/scala/org/apache/spark/shuffle/compat/spark_3_0/UcxShuffleReader.scala

@@ -24,6 +24,7 @@ import org.apache.spark._
 import org.apache.spark.internal.{Logging, config}
 import org.apache.spark.io.CompressionCodec
 import org.apache.spark.serializer.SerializerManager
+import org.apache.spark.shuffle.utils.CommonUtils
 import org.apache.spark.shuffle.ucx.UcxShuffleTransport
 import org.apache.spark.shuffle.{BaseShuffleHandle, ShuffleReadMetricsReporter, ShuffleReader}
 import org.apache.spark.storage.{BlockId, BlockManager, ShuffleBlockBatchId, ShuffleBlockFetcherIterator, ShuffleBlockId}
@@ -115,10 +116,11 @@ private[spark] class UcxShuffleReader[K, C](handle: BaseShuffleHandle[K, _, C],
     // Do progress if queue is empty before calling next on ShuffleIterator
     val ucxWrappedStream = new Iterator[(BlockId, InputStream)] {
       override def next(): (BlockId, InputStream) = {
+
         val startTime = System.nanoTime()
-        while (resultQueue.isEmpty) {
-          transport.progress()
-        }
+        CommonUtils.safePolling(() => {transport.progress()},
+          () => {resultQueue.isEmpty})
+
         val fetchWaitTime = System.nanoTime() - startTime
         readMetrics.incFetchWaitTime(TimeUnit.NANOSECONDS.toMillis(fetchWaitTime))
         wrappedStreams.next()