From a9361cff9706f4dfe4888bd5effde3681968666c Mon Sep 17 00:00:00 2001 From: Peter Rudenko Date: Fri, 7 Aug 2020 15:04:35 +0300 Subject: [PATCH] Performance benchmark --- pom.xml | 19 ++ .../shuffle/ucx/UcxShuffleTransport.scala | 4 +- .../spark/shuffle/ucx/UcxWorkerWrapper.scala | 62 ++--- .../perf/UcxShuffleTransportPerfTool.scala | 224 ++++++++++++++++++ .../ucx/rpc/GlobalWorkerRpcThread.scala | 2 +- .../spark/shuffle/ucx/GpuMemoryPool.scala | 30 +++ .../ucx/UcxShuffleTransportTestSuite.scala | 2 +- 7 files changed, 312 insertions(+), 31 deletions(-) create mode 100755 src/main/scala/org/apache/spark/shuffle/ucx/perf/UcxShuffleTransportPerfTool.scala create mode 100755 src/test/scala/org/apache/spark/shuffle/ucx/GpuMemoryPool.scala diff --git a/pom.xml b/pom.xml index d9330461..389cd465 100755 --- a/pom.xml +++ b/pom.xml @@ -35,6 +35,7 @@ See file LICENSE for terms. 2.12.12 2.12 1.10.0-SNAPSHOT + 0.16 @@ -55,6 +56,12 @@ See file LICENSE for terms. 3.2.1 test + + ai.rapids + cudf + ${cudf.version} + test + @@ -137,6 +144,18 @@ See file LICENSE for terms. + + org.apache.maven.plugins + maven-jar-plugin + 3.2.0 + + + + test-jar + + + + diff --git a/src/main/scala/org/apache/spark/shuffle/ucx/UcxShuffleTransport.scala b/src/main/scala/org/apache/spark/shuffle/ucx/UcxShuffleTransport.scala index 10aab2f2..8d150b9e 100755 --- a/src/main/scala/org/apache/spark/shuffle/ucx/UcxShuffleTransport.scala +++ b/src/main/scala/org/apache/spark/shuffle/ucx/UcxShuffleTransport.scala @@ -153,7 +153,7 @@ class UcxShuffleTransport(var ucxShuffleConf: UcxShuffleConf = null, val executo private[ucx] def handlePrefetchRequest(workerId: String, workerAddress: ByteBuffer, blockIds: Seq[BlockId]) { - logInfo(s"Prefetching blocks: ${blockIds.mkString(",")}") + logDebug(s"Prefetching blocks: ${blockIds.mkString(",")}") clientConnections.getOrElseUpdate(workerId, globalWorker.newEndpoint(new UcpEndpointParams().setUcpAddress(workerAddress)) ) @@ -184,7 +184,7 @@ class UcxShuffleTransport(var ucxShuffleConf: UcxShuffleConf = null, val executo lock.lock() val blockMemory = block.getMemoryBlock - logInfo(s"Sending $blockId of size ${blockMemory.size} to tag: $tag") + logDebug(s"Sending $blockId of size ${blockMemory.size} to tag: $tag") ep.sendTaggedNonBlocking(blockMemory.address, blockMemory.size, tag, new UcxCallback { override def onSuccess(request: UcpRequest): Unit = { if (block.isInstanceOf[UcxPinnedBlock]) { diff --git a/src/main/scala/org/apache/spark/shuffle/ucx/UcxWorkerWrapper.scala b/src/main/scala/org/apache/spark/shuffle/ucx/UcxWorkerWrapper.scala index 2b55eb06..d4337e71 100755 --- a/src/main/scala/org/apache/spark/shuffle/ucx/UcxWorkerWrapper.scala +++ b/src/main/scala/org/apache/spark/shuffle/ucx/UcxWorkerWrapper.scala @@ -5,6 +5,7 @@ package org.apache.spark.shuffle.ucx import java.io.{Closeable, ObjectOutputStream} +import java.nio.BufferOverflowException import java.util.concurrent.ThreadLocalRandom import scala.collection.mutable @@ -91,7 +92,7 @@ class UcxWorkerWrapper(worker: UcpWorker, transport: UcxShuffleTransport, ucxCon logInfo(s"Worker from thread ${Thread.currentThread().getName} connecting to $executorId") val endpointParams = new UcpEndpointParams() .setUcpAddress(workerAdresses.get(executorId)) - worker.newEndpoint(endpointParams) + worker.newEndpoint(endpointParams) }) } @@ -115,16 +116,16 @@ class UcxWorkerWrapper(worker: UcpWorker, transport: UcxShuffleTransport, ucxCon ep.sendTaggedNonBlocking(mem.address, transport.ucxShuffleConf.rpcMessageSize, UcxRpcMessages.PREFETCH_TAG, new UcxCallback() { - override def onSuccess(request: UcpRequest): Unit = { - logTrace(s"Sent prefetch ${blockIds.length} blocks to $executorId") - memoryPool.put(mem) - } - }) + override def onSuccess(request: UcpRequest): Unit = { + logTrace(s"Sent prefetch ${blockIds.length} blocks to $executorId") + memoryPool.put(mem) + } + }) } private[ucx] def fetchBlocksByBlockIds(executorId: String, blockIds: Seq[BlockId], - resultBuffer: Seq[MemoryBlock], - callbacks: Seq[OperationCallback]): Seq[Request] = { + resultBuffer: Seq[MemoryBlock], + callbacks: Seq[OperationCallback]): Seq[Request] = { val ep = getConnection(executorId) val mem = memoryPool.get(transport.ucxShuffleConf.rpcMessageSize) val buffer = UcxUtils.getByteBufferView(mem.address, @@ -136,13 +137,20 @@ class UcxWorkerWrapper(worker: UcpWorker, transport: UcxShuffleTransport, ucxCon Utils.tryWithResource(new ByteBufferBackedOutputStream(buffer)) { bos => val out = new ObjectOutputStream(bos) - out.writeObject(message) - out.flush() - out.close() + try { + out.writeObject(message) + out.flush() + out.close() + } catch { + case _: BufferOverflowException => + throw new UcxException(s"Prefetch blocks message size > " + + s"${transport.ucxShuffleConf.RPC_MESSAGE_SIZE.key}:${transport.ucxShuffleConf.rpcMessageSize}") + case ex: Exception => throw new UcxException(ex.getMessage) + } } val tag = ThreadLocalRandom.current().nextLong(Long.MinValue, 0) - logInfo(s"Sending message to $executorId to fetch ${blockIds.length} blocks on tag $tag") + logTrace(s"Sending message to $executorId to fetch ${blockIds.length} blocks on tag $tag") ep.sendTaggedNonBlocking(mem.address, transport.ucxShuffleConf.rpcMessageSize, tag, new UcxCallback() { override def onSuccess(request: UcpRequest): Unit = { @@ -179,7 +187,7 @@ class UcxWorkerWrapper(worker: UcpWorker, transport: UcxShuffleTransport, ucxCon } private[ucx] def fetchBlockByBlockId(executorId: String, blockId: BlockId, - resultBuffer: MemoryBlock, cb: OperationCallback): UcxRequest = { + resultBuffer: MemoryBlock, cb: OperationCallback): UcxRequest = { val stats = new UcxStats() val ep = getConnection(executorId) val mem = memoryPool.get(transport.ucxShuffleConf.rpcMessageSize) @@ -206,28 +214,28 @@ class UcxWorkerWrapper(worker: UcpWorker, transport: UcxShuffleTransport, ucxCon override def onSuccess(request: UcpRequest): Unit = { memoryPool.put(mem) } - }) + }) val result = new UcxSuccessOperationResult(stats) val request = worker.recvTaggedNonBlocking(resultBuffer.address, resultBuffer.size, tag, -1L, new UcxCallback () { - override def onError(ucsStatus: Int, errorMsg: String): Unit = { - logError(s"Failed to receive blockId $blockId on tag: $tag, from executorId: $executorId " + - s" of size: ${resultBuffer.size}: $errorMsg") - if (cb != null ) { - cb.onComplete(new UcxFailureOperationResult(errorMsg)) + override def onError(ucsStatus: Int, errorMsg: String): Unit = { + logError(s"Failed to receive blockId $blockId on tag: $tag, from executorId: $executorId " + + s" of size: ${resultBuffer.size}: $errorMsg") + if (cb != null ) { + cb.onComplete(new UcxFailureOperationResult(errorMsg)) + } } - } - override def onSuccess(request: UcpRequest): Unit = { - stats.endTime = System.nanoTime() - stats.receiveSize = request.getRecvSize - if (cb != null) { - cb.onComplete(result) + override def onSuccess(request: UcpRequest): Unit = { + stats.endTime = System.nanoTime() + stats.receiveSize = request.getRecvSize + if (cb != null) { + cb.onComplete(result) + } } - } - }) + }) new UcxRequest(request, stats) } diff --git a/src/main/scala/org/apache/spark/shuffle/ucx/perf/UcxShuffleTransportPerfTool.scala b/src/main/scala/org/apache/spark/shuffle/ucx/perf/UcxShuffleTransportPerfTool.scala new file mode 100755 index 00000000..039b3219 --- /dev/null +++ b/src/main/scala/org/apache/spark/shuffle/ucx/perf/UcxShuffleTransportPerfTool.scala @@ -0,0 +1,224 @@ +/* +* Copyright (C) Mellanox Technologies Ltd. 2020. ALL RIGHTS RESERVED. +* See file LICENSE for terms. +*/ +package org.apache.spark.shuffle.ucx.perf + +import java.net.{InetAddress, InetSocketAddress, ServerSocket, Socket} +import java.nio.ByteBuffer +import java.util.concurrent.atomic.{AtomicInteger, AtomicLong} +import java.util.concurrent.{CountDownLatch, Executors, TimeUnit} + +import org.apache.commons.cli.{GnuParser, HelpFormatter, Options} +import org.apache.spark.SparkConf +import org.apache.spark.shuffle.ucx._ +import org.apache.spark.shuffle.ucx.memory.MemoryPool +import org.apache.spark.util.Utils + +object UcxShuffleTransportPerfTool { + private val HELP_OPTION = "h" + private val ADDRESS_OPTION = "a" + private val NUM_BLOCKS_OPTION = "n" + private val SIZE_OPTION = "s" + private val PORT_OPTION = "p" + private val ITER_OPTION = "i" + private val MEMORY_TYPE_OPTION = "m" + private val NUM_THREADS_OPTION = "t" + + private val ucxShuffleConf = new UcxShuffleConf(new SparkConf()) + private val transport = new UcxShuffleTransport(ucxShuffleConf, "e") + private val workerAddress = transport.init() + private var memoryPool: MemoryPool = transport.memoryPool + + case class TestBlockId(id: Int) extends BlockId + + case class PerfOptions(remoteAddress: InetSocketAddress, numBlocks: Int, blockSize: Long, + serverPort: Int, numIterations: Int, numThreads: Int) + + private def initOptions(): Options = { + val options = new Options() + options.addOption(HELP_OPTION, "help", false, + "display help message") + options.addOption(ADDRESS_OPTION, "address", true, + "address of the remote host") + options.addOption(NUM_BLOCKS_OPTION, "num-blocks", true, + "number of blocks to transfer. Default: 1") + options.addOption(SIZE_OPTION, "block-size", true, + "size of block to transfer. Default: 4m") + options.addOption(PORT_OPTION, "server-port", true, + "server port. Default: 12345") + options.addOption(ITER_OPTION, "num-iterations", true, + "number of iterations. Default: 5") + options.addOption(NUM_THREADS_OPTION, "num-threads", true, + "number of threads. Default: 1") + options.addOption(MEMORY_TYPE_OPTION, "memory-type", true, + "memory type: host (default), cuda") + } + + private def parseOptions(args: Array[String]): PerfOptions = { + val parser = new GnuParser() + val options = initOptions() + val cmd = parser.parse(options, args) + + if (cmd.hasOption(HELP_OPTION)) { + new HelpFormatter().printHelp("UcxShufflePerfTool", options) + System.exit(0) + } + + val inetAddress = if (cmd.hasOption(ADDRESS_OPTION)) { + val Array(host, port) = cmd.getOptionValue(ADDRESS_OPTION).split(":") + new InetSocketAddress(host, Integer.parseInt(port)) + } else { + null + } + + val serverPort = Integer.parseInt(cmd.getOptionValue(PORT_OPTION, "12345")) + + val numIterations = Integer.parseInt(cmd.getOptionValue(ITER_OPTION, "5")) + + val threadsNumber = Integer.parseInt(cmd.getOptionValue(NUM_THREADS_OPTION, "1")) + + if (cmd.hasOption(MEMORY_TYPE_OPTION) && cmd.getOptionValue(MEMORY_TYPE_OPTION) == "cuda") { + val className = "org.apache.spark.shuffle.ucx.GpuMemoryPool" + val cls = Utils.classForName(className) + memoryPool = cls.getConstructor().newInstance().asInstanceOf[MemoryPool] + } + + PerfOptions(inetAddress, + Integer.parseInt(cmd.getOptionValue(NUM_BLOCKS_OPTION, "1")), + Utils.byteStringAsBytes(cmd.getOptionValue(SIZE_OPTION, "4m")), + serverPort, numIterations, threadsNumber) + } + + private def startServer(perfOptions: PerfOptions): Unit = { + val blocks: Seq[Block] = (0 until perfOptions.numBlocks).map { _ => + val block = memoryPool.get(perfOptions.blockSize) + new Block { + override def getMemoryBlock: MemoryBlock = + block + } + } + + val blockIds = (0 until perfOptions.numBlocks).map(i => TestBlockId(i)) + blockIds.zip(blocks).foreach { + case (blockId, block) => transport.register(blockId, block) + } + + val serverSocket = new ServerSocket(perfOptions.serverPort) + + println(s"Waiting for connections on " + + s"${InetAddress.getLocalHost.getHostName}:${perfOptions.serverPort} ") + + val clientSocket = serverSocket.accept() + val out = clientSocket.getOutputStream + val in = clientSocket.getInputStream + + val buf = ByteBuffer.allocate(workerAddress.capacity()) + buf.put(workerAddress) + buf.flip() + + out.write(buf.array()) + out.flush() + + println(s"Sending worker address to ${clientSocket.getInetAddress}") + + buf.flip() + + in.read(buf.array()) + clientSocket.close() + serverSocket.close() + + blocks.foreach(block => memoryPool.put(block.getMemoryBlock)) + blockIds.foreach(transport.unregister) + transport.close() + } + + private def startClient(perfOptions: PerfOptions): Unit = { + val socket = new Socket(perfOptions.remoteAddress.getHostName, + perfOptions.remoteAddress.getPort) + + val buf = new Array[Byte](4096) + val readSize = socket.getInputStream.read(buf) + val executorId = "1" + val workerAddress = ByteBuffer.allocateDirect(readSize) + + workerAddress.put(buf, 0, readSize) + println("Received worker address") + + transport.addExecutor(executorId, workerAddress) + + val resultSize = perfOptions.numBlocks * perfOptions.blockSize + val resultMemory = memoryPool.get(resultSize) + + val blockIds = (0 until perfOptions.numBlocks).map(i => TestBlockId(i)) + + val threadPool = Executors.newFixedThreadPool(perfOptions.numThreads) + + for (i <- 1 to perfOptions.numIterations) { + val elapsedTime = new AtomicLong(0) + val countDownLatch = new CountDownLatch(perfOptions.numThreads) + + for (_ <- 0 until perfOptions.numThreads) { + threadPool.execute(() => { + val completed = new AtomicInteger(0) + + val mem = new Array[MemoryBlock](perfOptions.numBlocks) + val callbacks = new Array[OperationCallback](perfOptions.numBlocks) + + for (j <- 0 until perfOptions.numBlocks) { + mem(j) = MemoryBlock(resultMemory.address + j * perfOptions.blockSize, perfOptions.blockSize) + callbacks(j) = (result: OperationResult) => { + elapsedTime.addAndGet(result.getStats.get.getElapsedTimeNs) + completed.incrementAndGet() + } + } + + transport.fetchBlocksByBlockIds(executorId, blockIds, mem, callbacks) + + while (completed.get() != perfOptions.numBlocks) { + transport.progress() + } + countDownLatch.countDown() + }) + } + + countDownLatch.await() + + val totalTime = if (elapsedTime.get() < TimeUnit.MILLISECONDS.toNanos(1)) { + s"$elapsedTime ns" + } else { + s"${TimeUnit.NANOSECONDS.toMillis(elapsedTime.get())} ms" + } + val throughput: Double = (resultSize * perfOptions.numThreads / 1024.0D / 1024.0D / 1024.0D) / + (elapsedTime.get() / 1e9D) + + println(f"${s"[$i/${perfOptions.numIterations}]"}%12s" + + s" numBlocks: ${perfOptions.numBlocks}" + + s" numThreads: ${perfOptions.numThreads}" + + s" size: ${Utils.bytesToString(perfOptions.blockSize)}," + + s" total size: ${Utils.bytesToString(resultSize * perfOptions.numThreads)}," + + f" time: $totalTime%3s" + + f" throughput: $throughput%.5f GB/s") + } + + val out = socket.getOutputStream + out.write(buf) + out.flush() + out.close() + socket.close() + + memoryPool.put(resultMemory) + transport.close() + } + + def main(args: Array[String]): Unit = { + val perfOptions = parseOptions(args) + + if (perfOptions.remoteAddress == null) { + startServer(perfOptions) + } else { + startClient(perfOptions) + } + } + +} diff --git a/src/main/scala/org/apache/spark/shuffle/ucx/rpc/GlobalWorkerRpcThread.scala b/src/main/scala/org/apache/spark/shuffle/ucx/rpc/GlobalWorkerRpcThread.scala index c65dba23..efff6d76 100755 --- a/src/main/scala/org/apache/spark/shuffle/ucx/rpc/GlobalWorkerRpcThread.scala +++ b/src/main/scala/org/apache/spark/shuffle/ucx/rpc/GlobalWorkerRpcThread.scala @@ -21,7 +21,7 @@ class GlobalWorkerRpcThread(globalWorker: UcpWorker, memPool: MemoryPool, setDaemon(true) setName("Ucx Shuffle Transport Progress Thread") - + override def run(): Unit = { val numRecvs = transport.ucxShuffleConf.recvQueueSize val msgSize = transport.ucxShuffleConf.rpcMessageSize diff --git a/src/test/scala/org/apache/spark/shuffle/ucx/GpuMemoryPool.scala b/src/test/scala/org/apache/spark/shuffle/ucx/GpuMemoryPool.scala new file mode 100755 index 00000000..6d3e633a --- /dev/null +++ b/src/test/scala/org/apache/spark/shuffle/ucx/GpuMemoryPool.scala @@ -0,0 +1,30 @@ +/* +* Copyright (C) Mellanox Technologies Ltd. 2020. ALL RIGHTS RESERVED. +* See file LICENSE for terms. +*/ +package org.apache.spark.shuffle.ucx + +import ai.rapids.cudf.DeviceMemoryBuffer +import org.apache.spark.shuffle.ucx.memory.MemoryPool + + +/** + * Test GPU mempool to run with [[ UcxShuffleTransportPerfTool ]] + */ +class GpuMemoryPool extends MemoryPool { + + class GpuMemoryBlock(val deviceBuffer: DeviceMemoryBuffer, + override val address: Long, override val size: Long) + extends MemoryBlock(address, size, isHostMemory = false) + + override def get(size: Long): MemoryBlock = { + val deviceBuffer = DeviceMemoryBuffer.allocate(size) + new GpuMemoryBlock(deviceBuffer, deviceBuffer.getAddress, size) + } + + override def put(mem: MemoryBlock): Unit = { + mem.asInstanceOf[GpuMemoryBlock].deviceBuffer.close() + } + + override def close(): Unit = ??? +} diff --git a/src/test/scala/org/apache/spark/shuffle/ucx/UcxShuffleTransportTestSuite.scala b/src/test/scala/org/apache/spark/shuffle/ucx/UcxShuffleTransportTestSuite.scala index f391f3cd..7997042a 100755 --- a/src/test/scala/org/apache/spark/shuffle/ucx/UcxShuffleTransportTestSuite.scala +++ b/src/test/scala/org/apache/spark/shuffle/ucx/UcxShuffleTransportTestSuite.scala @@ -1,5 +1,5 @@ /* -* Copyright (C) Mellanox Technologies Ltd. 2019. ALL RIGHTS RESERVED. +* Copyright (C) Mellanox Technologies Ltd. 2020. ALL RIGHTS RESERVED. * See file LICENSE for terms. */ package org.apache.spark.shuffle.ucx