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

Author: xy2953396112

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

@@ -954,6 +954,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 enableReuseHdfsOuputSteam: Boolean = get(ENABLE_REUSE_HDFS_OUTPUT_STREAM)
   def appHeartbeatIntervalMs: Long = get(APPLICATION_HEARTBEAT_INTERVAL)
   def applicationUnregisterEnabled: Boolean = get(APPLICATION_UNREGISTER_ENABLED)
 
@@ -2484,6 +2485,14 @@ object CelebornConf extends Logging {
       .timeConf(TimeUnit.MILLISECONDS)
       .createWithDefaultString("1h")
 
+  val ENABLE_REUSE_HDFS_OUTPUT_STREAM: ConfigEntry[Boolean] =
+    buildConf("celeborn.worker.hdfs.reuseOutputStream.enable")
+      .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

@@ -109,6 +109,7 @@ license: |
 | celeborn.worker.graceful.shutdown.saveCommittedFileInfo.sync | false | false | Whether to call sync method to save committed file infos into DB to handle OS crash. | 0.3.1 |  | 
 | celeborn.worker.graceful.shutdown.timeout | 600s | false | The worker's graceful shutdown timeout time. | 0.2.0 |  | 
 | celeborn.worker.hdfs.replication.factor | 2 | false | HDFS replication factor for shuffle files. | 0.7.0 |  | 
+| celeborn.worker.hdfs.reuseOutputStream.enable | false | false | Whether to enable reuse output stream on hdfs. | 0.7.0 |  | 
 | celeborn.worker.http.auth.administers |  | false | A comma-separated list of users who have admin privileges, Note, when celeborn.worker.http.auth.supportedSchemes is not set, everyone is treated as administrator. | 0.6.0 |  | 
 | celeborn.worker.http.auth.basic.provider | org.apache.celeborn.common.authentication.AnonymousAuthenticationProviderImpl | false | User-defined password authentication implementation of org.apache.celeborn.common.authentication.PasswdAuthenticationProvider | 0.6.0 |  | 
 | celeborn.worker.http.auth.bearer.provider | org.apache.celeborn.common.authentication.AnonymousAuthenticationProviderImpl | false | User-defined token authentication implementation of org.apache.celeborn.common.authentication.TokenAuthenticationProvider | 0.6.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
@@ -522,6 +522,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 enableReuseHdfsOutputStream = conf.enableReuseHdfsOuputSteam
   private var s3MultipartUploadHandler: MultipartUploadHandler = _
   private var ossMultipartUploadHandler: MultipartUploadHandler = _
   var partNumber: Int = 1
@@ -536,7 +538,11 @@ class DfsTierWriter(
     }
 
   try {
-    hadoopFs.create(dfsFileInfo.getDfsPath, true).close()
+    if (enableReuseHdfsOutputStream) {
+      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
@@ -580,7 +586,11 @@ class DfsTierWriter(
         case ex: InterruptedException =>
           throw new RuntimeException(ex)
       }
-      hadoopFs.create(dfsFileInfo.getDfsPath, true).close()
+      if (enableReuseHdfsOutputStream) {
+        hdfsStream = hadoopFs.create(dfsFileInfo.getDfsPath, true)
+      } else {
+        hadoopFs.create(dfsFileInfo.getDfsPath, true).close()
+      }
   }
 
   storageManager.registerDiskFilePartitionWriter(
@@ -595,7 +605,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,
@@ -649,6 +659,9 @@ class DfsTierWriter(
   }
 
   override def closeStreams(): Unit = {
+    if (hdfsStream != null) {
+      hdfsStream.close()
+    }
     if (hadoopFs.exists(dfsFileInfo.getDfsPeerWriterSuccessPath)) {
       hadoopFs.delete(dfsFileInfo.getDfsPath, false)
       deleted = true