[CELEBORN-2211] Avoid allocating additional buffers When HdfsFlushTask writes data

Author: xy2953396112

common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala

@@ -955,6 +955,7 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable with Logging with Se
   def clientShuffleDynamicResourceFactor: Double = get(CLIENT_SHUFFLE_DYNAMIC_RESOURCE_FACTOR)
   def appHeartbeatTimeoutMs: Long = get(APPLICATION_HEARTBEAT_TIMEOUT)
   def dfsExpireDirsTimeoutMS: Long = get(DFS_EXPIRE_DIRS_TIMEOUT)
+  def reuseHdfsOuputSteamEnabled: Boolean = get(REUSE_HDFS_OUTPUT_STREAM_ENABLED)
   def appHeartbeatIntervalMs: Long = get(APPLICATION_HEARTBEAT_INTERVAL)
   def applicationUnregisterEnabled: Boolean = get(APPLICATION_UNREGISTER_ENABLED)
 
@@ -2488,6 +2489,14 @@ object CelebornConf extends Logging {
       .timeConf(TimeUnit.MILLISECONDS)
       .createWithDefaultString("1h")
 
+  val REUSE_HDFS_OUTPUT_STREAM_ENABLED: ConfigEntry[Boolean] =
+    buildConf("celeborn.worker.reuse.hdfs.outputStream.enabled")
+      .categories("worker")
+      .version("0.7.0")
+      .doc("Whether to enable reuse output stream on hdfs.")
+      .booleanConf
+      .createWithDefault(false)
+
   val WORKER_HEARTBEAT_TIMEOUT: ConfigEntry[Long] =
     buildConf("celeborn.master.heartbeat.worker.timeout")
       .withAlternative("celeborn.worker.heartbeat.timeout")

docs/configuration/worker.md

@@ -177,6 +177,7 @@ license: |
 | celeborn.worker.replicate.port | 0 | false | Server port for Worker to receive replicate data request from other Workers. | 0.2.0 |  | 
 | celeborn.worker.replicate.randomConnection.enabled | true | false | Whether worker will create random connection to peer when replicate data. When false, worker tend to reuse the same cached TransportClient to a specific replicate worker; when true, worker tend to use different cached TransportClient. Netty will use the same thread to serve the same connection, so with more connections replicate server can leverage more netty threads | 0.2.1 |  | 
 | celeborn.worker.replicate.threads | 64 | false | Thread number of worker to replicate shuffle data. | 0.2.0 |  | 
+| celeborn.worker.reuse.hdfs.outputStream.enabled | false | false | Whether to enable reuse output stream on hdfs. | 0.7.0 |  | 
 | celeborn.worker.rpc.port | 0 | false | Server port for Worker to receive RPC request. | 0.2.0 |  | 
 | celeborn.worker.shuffle.partitionSplit.enabled | true | false | enable the partition split on worker side | 0.3.0 | celeborn.worker.partition.split.enabled | 
 | celeborn.worker.shuffle.partitionSplit.max | 2g | false | Specify the maximum partition size for splitting, and ensure that individual partition files are always smaller than this limit. | 0.3.0 |  | 

worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/FlushTask.scala

@@ -21,7 +21,7 @@ import java.io.{ByteArrayInputStream, Closeable, IOException}
 import java.nio.channels.FileChannel
 
 import io.netty.buffer.{ByteBufUtil, CompositeByteBuf}
-import org.apache.hadoop.fs.Path
+import org.apache.hadoop.fs.{FSDataOutputStream, Path}
 
 import org.apache.celeborn.common.internal.Logging
 import org.apache.celeborn.common.metrics.source.AbstractSource
@@ -98,18 +98,28 @@ abstract private[worker] class DfsFlushTask(
 
 private[worker] class HdfsFlushTask(
     buffer: CompositeByteBuf,
+    hdfsStream: FSDataOutputStream,
     val path: Path,
     notifier: FlushNotifier,
     keepBuffer: Boolean,
     source: AbstractSource) extends DfsFlushTask(buffer, notifier, keepBuffer, source) {
   override def flush(copyBytes: Array[Byte]): Unit = {
     val readableBytes = buffer.readableBytes()
-    val hadoopFs = StorageManager.hadoopFs.get(Type.HDFS)
-    val hdfsStream = hadoopFs.append(path, 256 * 1024)
-    flush(hdfsStream) {
-      hdfsStream.write(convertBufferToBytes(buffer, copyBytes, readableBytes))
-      source.incCounter(WorkerSource.HDFS_FLUSH_COUNT)
-      source.incCounter(WorkerSource.HDFS_FLUSH_SIZE, readableBytes)
+    if (hdfsStream != null) {
+      // TODO : If the FSDataOutputStream supports concurrent writes, the lock can be removed.
+      hdfsStream.synchronized {
+        hdfsStream.write(convertBufferToBytes(buffer, copyBytes, readableBytes))
+        source.incCounter(WorkerSource.HDFS_FLUSH_COUNT)
+        source.incCounter(WorkerSource.HDFS_FLUSH_SIZE, readableBytes)
+      }
+    } else {
+      val hadoopFs = StorageManager.hadoopFs.get(Type.HDFS)
+      val hdfsStream = hadoopFs.append(path, 256 * 1024)
+      flush(hdfsStream) {
+        hdfsStream.write(convertBufferToBytes(buffer, copyBytes, readableBytes))
+        source.incCounter(WorkerSource.HDFS_FLUSH_COUNT)
+        source.incCounter(WorkerSource.HDFS_FLUSH_SIZE, readableBytes)
+      }
     }
   }
 }

worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/TierWriter.scala

@@ -26,7 +26,7 @@ import java.util.concurrent.atomic.AtomicInteger
 import scala.collection.JavaConverters.asScalaBufferConverter
 
 import io.netty.buffer.{ByteBuf, CompositeByteBuf}
-import org.apache.hadoop.fs.FileSystem
+import org.apache.hadoop.fs.{FileSystem, FSDataOutputStream}
 
 import org.apache.celeborn.common.CelebornConf
 import org.apache.celeborn.common.exception.AlreadyClosedException
@@ -532,6 +532,8 @@ class DfsTierWriter(
   private val flushWorkerIndex: Int = flusher.getWorkerIndex
   val hadoopFs: FileSystem = StorageManager.hadoopFs.get(storageType)
   var deleted = false
+  private var hdfsStream: FSDataOutputStream = null
+  private val reuseHdfsOutputStreamEnabled = conf.reuseHdfsOuputSteamEnabled
   private var s3MultipartUploadHandler: MultipartUploadHandler = _
   private var ossMultipartUploadHandler: MultipartUploadHandler = _
   var partNumber: Int = 1
@@ -546,7 +548,11 @@ class DfsTierWriter(
     }
 
   try {
-    hadoopFs.create(dfsFileInfo.getDfsPath, true).close()
+    if (reuseHdfsOutputStreamEnabled) {
+      hdfsStream = hadoopFs.create(dfsFileInfo.getDfsPath, true)
+    } else {
+      hadoopFs.create(dfsFileInfo.getDfsPath, true).close()
+    }
     hadoopFs.setReplication(dfsFileInfo.getDfsPath, conf.workerDfsReplicationFactor.toShort);
     if (dfsFileInfo.isS3) {
       val uri = hadoopFs.getUri
@@ -590,7 +596,14 @@ class DfsTierWriter(
         case ex: InterruptedException =>
           throw new RuntimeException(ex)
       }
-      hadoopFs.create(dfsFileInfo.getDfsPath, true).close()
+      if (reuseHdfsOutputStreamEnabled) {
+        if (hdfsStream != null) {
+          hdfsStream.close()
+        }
+        hdfsStream = hadoopFs.create(dfsFileInfo.getDfsPath, true)
+      } else {
+        hadoopFs.create(dfsFileInfo.getDfsPath, true).close()
+      }
   }
 
   storageManager.registerDiskFilePartitionWriter(
@@ -605,7 +618,7 @@ class DfsTierWriter(
   override def genFlushTask(finalFlush: Boolean, keepBuffer: Boolean): FlushTask = {
     notifier.numPendingFlushes.incrementAndGet()
     if (dfsFileInfo.isHdfs) {
-      new HdfsFlushTask(flushBuffer, dfsFileInfo.getDfsPath(), notifier, true, source)
+      new HdfsFlushTask(flushBuffer, hdfsStream, dfsFileInfo.getDfsPath(), notifier, true, source)
     } else if (dfsFileInfo.isOSS) {
       val flushTask = new OssFlushTask(
         flushBuffer,
@@ -659,6 +672,9 @@ class DfsTierWriter(
   }
 
   override def closeStreams(): Unit = {
+    if (hdfsStream != null) {
+      hdfsStream.close()
+    }
     if (hadoopFs.exists(dfsFileInfo.getDfsPeerWriterSuccessPath)) {
       hadoopFs.delete(dfsFileInfo.getDfsPath, false)
       deleted = true