You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@celeborn.apache.org by zh...@apache.org on 2023/02/04 06:28:03 UTC

[incubator-celeborn] branch main updated (322495b3 -> 8e903840)

This is an automated email from the ASF dual-hosted git repository.

zhouky pushed a change to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-celeborn.git


 discard 322495b3 [CELEBORN-243] [REVERT]fix bug that os's disk usage is low but celeborn thinks that it's high_disk_usage (#1202)
     new 8e903840 [CELEBORN-243][REWORK]fix bug that os's disk usage is low but celeborn thinks that it's high_disk_usage (#1202)

This update added new revisions after undoing existing revisions.
That is to say, some revisions that were in the old version of the
branch are not in the new version.  This situation occurs
when a user --force pushes a change and generates a repository
containing something like this:

 * -- * -- B -- O -- O -- O   (322495b3)
            \
             N -- N -- N   refs/heads/main (8e903840)

You should already have received notification emails for all of the O
revisions, and so the following emails describe only the N revisions
from the common base, B.

Any revisions marked "omit" are not gone; other references still
refer to them.  Any revisions marked "discard" are gone forever.

The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails.  The revisions
listed as "add" were already present in the repository and have only
been added to this reference.


Summary of changes:


[incubator-celeborn] 01/01: [CELEBORN-243][REWORK]fix bug that os's disk usage is low but celeborn thinks that it's high_disk_usage (#1202)

Posted by zh...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

zhouky pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-celeborn.git

commit 8e903840af3c6a780268be463d3a7ad615913366
Author: zhongqiangchen <zh...@alibaba-inc.com>
AuthorDate: Sat Feb 4 13:49:43 2023 +0800

    [CELEBORN-243][REWORK]fix bug that os's disk usage is low but celeborn thinks that it's high_disk_usage (#1202)
---
 .../tests/spark/RssHashCheckDiskSuite.scala        | 101 ++++++++++++++
 .../service/deploy/worker/storage/FileWriter.java  |  13 --
 .../celeborn/service/deploy/worker/Worker.scala    |  23 ----
 .../deploy/worker/storage/StorageManager.scala     |  62 ++++++---
 .../service/deploy/MiniClusterFeature.scala        |   4 +-
 .../deploy/worker/storage/WorkerSuite.scala        | 149 ++++++++-------------
 6 files changed, 206 insertions(+), 146 deletions(-)

diff --git a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RssHashCheckDiskSuite.scala b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RssHashCheckDiskSuite.scala
new file mode 100644
index 00000000..24a91fdf
--- /dev/null
+++ b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RssHashCheckDiskSuite.scala
@@ -0,0 +1,101 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.celeborn.tests.spark
+
+import scala.collection.JavaConverters._
+
+import org.apache.spark.SparkConf
+import org.apache.spark.sql.SparkSession
+import org.scalatest.BeforeAndAfterEach
+import org.scalatest.funsuite.AnyFunSuite
+
+import org.apache.celeborn.client.ShuffleClient
+import org.apache.celeborn.service.deploy.worker.Worker
+
+class RssHashCheckDiskSuite extends AnyFunSuite
+  with SparkTestBase
+  with BeforeAndAfterEach {
+  var workers: collection.Set[Worker] = null
+  override def beforeAll(): Unit = {
+    logInfo("RssHashCheckDiskSuite test initialized , setup rss mini cluster")
+    val masterConfs = Map("celeborn.application.heartbeat.timeout" -> "10s")
+    val workerConfs = Map(
+      "celeborn.worker.storage.dirs" -> "/tmp:capacity=1000",
+      "celeborn.worker.heartbeat.timeout" -> "10s")
+    workers = setUpMiniCluster(masterConfs, workerConfs)._2
+  }
+
+  override def beforeEach(): Unit = {
+    ShuffleClient.reset()
+  }
+
+  override def afterEach(): Unit = {
+    System.gc()
+  }
+
+  test("celeborn spark integration test - hash-checkDiskFull") {
+    val sparkConf = new SparkConf().setAppName("rss-demo").setMaster("local[4]").set(
+      "spark.celeborn.shuffle.expired.checkInterval",
+      "5s")
+    val sparkSession = SparkSession.builder().config(sparkConf).getOrCreate()
+    val combineResult = combine(sparkSession)
+    val groupbyResult = groupBy(sparkSession)
+    val repartitionResult = repartition(sparkSession)
+    val sqlResult = runsql(sparkSession)
+
+    Thread.sleep(3000L)
+    sparkSession.stop()
+
+    val rssSparkSession = SparkSession.builder()
+      .config(updateSparkConf(sparkConf, false)).getOrCreate()
+    val rssCombineResult = combine(rssSparkSession)
+    val rssGroupbyResult = groupBy(rssSparkSession)
+    val rssRepartitionResult = repartition(rssSparkSession)
+    val rssSqlResult = runsql(rssSparkSession)
+
+    assert(combineResult.equals(rssCombineResult))
+    assert(groupbyResult.equals(rssGroupbyResult))
+    assert(repartitionResult.equals(rssRepartitionResult))
+    assert(combineResult.equals(rssCombineResult))
+    assert(sqlResult.equals(rssSqlResult))
+
+    // shuffle key not expired, diskInfo.actualUsableSpace < 0, no space
+    workers.map(worker => {
+      worker.storageManager.disksSnapshot().map(diskInfo => {
+        assert(diskInfo.actualUsableSpace < 0)
+      })
+    })
+
+    rssSparkSession.stop()
+    // wait shuffle key expired
+    Thread.sleep(30 * 1000L)
+    logInfo("after shuffle key expired")
+    // after shuffle key expired, storageManager.workingDirWriters will be empty
+    workers.map(worker => {
+      worker.storageManager.workingDirWriters.values().asScala.map(t => assert(t.size() == 0))
+    })
+
+    // after shuffle key expired, diskInfo.actualUsableSpace will equal capacity=1000
+    workers.map(worker => {
+      worker.storageManager.disksSnapshot().map(diskInfo => {
+        assert(diskInfo.actualUsableSpace == 1000)
+      })
+    })
+
+  }
+}
diff --git a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/FileWriter.java b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/FileWriter.java
index a10eef9a..925dc15a 100644
--- a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/FileWriter.java
+++ b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/FileWriter.java
@@ -21,7 +21,6 @@ import java.io.File;
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.nio.channels.FileChannel;
-import java.util.List;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -74,7 +73,6 @@ public abstract class FileWriter implements DeviceObserver {
   private final PartitionSplitMode splitMode;
   private final PartitionType partitionType;
   private final boolean rangeReadFilter;
-  private Runnable destroyHook;
   protected boolean deleted = false;
   private RoaringBitmap mapIdBitMap = null;
   protected final FlushNotifier notifier = new FlushNotifier();
@@ -304,20 +302,9 @@ public abstract class FileWriter implements DeviceObserver {
       if (!fileInfo.isHdfs()) {
         deviceMonitor.unregisterFileWriter(this);
       }
-      destroyHook.run();
     }
   }
 
-  public void registerDestroyHook(List<FileWriter> fileWriters) {
-    FileWriter thisFileWriter = this;
-    destroyHook =
-        () -> {
-          synchronized (fileWriters) {
-            fileWriters.remove(thisFileWriter);
-          }
-        };
-  }
-
   public IOException getException() {
     if (notifier.hasException()) {
       return notifier.exception.get();
diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
index e4258ddf..59f1d5ea 100644
--- a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
+++ b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
@@ -412,19 +412,6 @@ private[celeborn] class Worker(
   @VisibleForTesting
   def cleanup(expiredShuffleKeys: JHashSet[String]): Unit = synchronized {
     expiredShuffleKeys.asScala.foreach { shuffleKey =>
-      partitionLocationInfo.getAllMasterLocations(shuffleKey).asScala.foreach { partition =>
-        val fileWriter = partition.asInstanceOf[WorkingPartition].getFileWriter
-        fileWriter.destroy(new IOException(
-          s"Destroy FileWriter ${fileWriter} caused by shuffle ${shuffleKey} expired."))
-        removeExpiredWorkingDirWriters(fileWriter)
-      }
-      partitionLocationInfo.getAllSlaveLocations(shuffleKey).asScala.foreach { partition =>
-        val fileWriter = partition.asInstanceOf[WorkingPartition].getFileWriter
-        fileWriter.destroy(new IOException(
-          s"Destroy FileWriter ${fileWriter} caused by shuffle ${shuffleKey} expired."))
-        removeExpiredWorkingDirWriters(fileWriter)
-      }
-
       partitionLocationInfo.removeMasterPartitions(shuffleKey)
       partitionLocationInfo.removeSlavePartitions(shuffleKey)
       shufflePartitionType.remove(shuffleKey)
@@ -439,16 +426,6 @@ private[celeborn] class Worker(
     fetchHandler.cleanupExpiredShuffleKey(expiredShuffleKeys)
   }
 
-  @VisibleForTesting
-  def removeExpiredWorkingDirWriters(fileWriter: FileWriter): Unit = {
-    // filepath is dir/appId/shuffleId/filename
-    val dir = fileWriter.getFile.getParentFile.getParentFile.getParentFile
-    storageManager.workingDirWriters.asScala.get(dir).map(f =>
-      f.synchronized {
-        f.remove(fileWriter)
-      })
-  }
-
   override def getWorkerInfo: String = workerInfo.toString()
 
   override def getThreadDump: String = Utils.getThreadDump()
diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
index fb8e1b3a..8003182d 100644
--- a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
+++ b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
@@ -26,6 +26,7 @@ import java.util.concurrent.atomic.AtomicInteger
 import java.util.function.IntUnaryOperator
 
 import scala.collection.JavaConverters._
+import scala.collection.mutable
 import scala.concurrent.duration._
 
 import com.google.common.util.concurrent.ThreadFactoryBuilder
@@ -50,7 +51,7 @@ import org.apache.celeborn.service.deploy.worker.storage.StorageManager.hadoopFs
 final private[worker] class StorageManager(conf: CelebornConf, workerSource: AbstractSource)
   extends ShuffleRecoverHelper with DeviceObserver with Logging with MemoryPressureListener {
   // mount point -> filewriter
-  val workingDirWriters = new ConcurrentHashMap[File, util.ArrayList[FileWriter]]()
+  val workingDirWriters = new ConcurrentHashMap[File, ConcurrentHashMap[String, FileWriter]]()
 
   val (deviceInfos, diskInfos) = {
     val workingDirInfos =
@@ -127,7 +128,7 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
     logInfo(s"Initialize HDFS support with path ${hdfsDir}")
   }
   val hdfsPermission = FsPermission.createImmutable(755)
-  val hdfsWriters = new util.ArrayList[FileWriter]()
+  val hdfsWriters = new ConcurrentHashMap[String, FileWriter]()
   val (hdfsFlusher, _totalHdfsFlusherThread) =
     if (!hdfsDir.isEmpty) {
       val hdfsConfiguration = new Configuration
@@ -253,8 +254,9 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
     }
 
   private val workingDirWriterListFunc =
-    new java.util.function.Function[File, util.ArrayList[FileWriter]]() {
-      override def apply(t: File): util.ArrayList[FileWriter] = new util.ArrayList[FileWriter]()
+    new java.util.function.Function[File, ConcurrentHashMap[String, FileWriter]]() {
+      override def apply(t: File): ConcurrentHashMap[String, FileWriter] =
+        new ConcurrentHashMap[String, FileWriter]()
     }
 
   @throws[IOException]
@@ -317,10 +319,7 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
         }
 
         fileInfos.computeIfAbsent(shuffleKey, newMapFunc).put(fileName, fileInfo)
-        hdfsWriters.synchronized {
-          hdfsWriters.add(hdfsWriter)
-        }
-        hdfsWriter.registerDestroyHook(hdfsWriters)
+        hdfsWriters.put(fileInfo.getFilePath, hdfsWriter)
         return hdfsWriter
       } else {
         val dir = dirs(getNextIndex() % dirs.size)
@@ -362,11 +361,8 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
                 rangeReadFilter)
           }
           deviceMonitor.registerFileWriter(fileWriter)
-          val list = workingDirWriters.computeIfAbsent(dir, workingDirWriterListFunc)
-          list.synchronized {
-            list.add(fileWriter)
-          }
-          fileWriter.registerDestroyHook(list)
+          val map = workingDirWriters.computeIfAbsent(dir, workingDirWriterListFunc)
+          map.put(fileInfo.getFilePath, fileWriter)
           fileInfos.computeIfAbsent(shuffleKey, newMapFunc).put(fileName, fileInfo)
           location.getStorageInfo.setMountPoint(mountPoint)
           logDebug(s"location $location set disk hint to ${location.getStorageInfo} ")
@@ -424,7 +420,36 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
     expiredShuffleKeys.asScala.foreach { shuffleKey =>
       logInfo(s"Cleanup expired shuffle $shuffleKey.")
       if (fileInfos.containsKey(shuffleKey)) {
-        val hdfsInfos = fileInfos.remove(shuffleKey).asScala.filter(_._2.isHdfs)
+        val removedFileInfos = fileInfos.remove(shuffleKey)
+        var isHdfsExpired = false
+        if (removedFileInfos != null) {
+          removedFileInfos.asScala.foreach {
+            case (_, fileInfo) => {
+              if (fileInfo.isHdfs) {
+                isHdfsExpired = true
+                val hdfsFileWriter = hdfsWriters.get(fileInfo.getFilePath)
+                if (hdfsFileWriter != null) {
+                  hdfsFileWriter.destroy(new IOException(
+                    s"Destroy FileWriter ${hdfsFileWriter} caused by shuffle ${shuffleKey} expired."))
+                }
+                fileInfo.deleteAllFiles(StorageManager.hadoopFs)
+                hdfsWriters.remove(fileInfo.getFilePath)
+              } else {
+                val workingDir =
+                  fileInfo.getFile.getParentFile.getParentFile.getParentFile
+                val writers = workingDirWriters.get(workingDir)
+                if (writers != null) {
+                  val fileWriter = writers.get(fileInfo.getFilePath)
+                  if (fileWriter != null) {
+                    fileWriter.destroy(new IOException(
+                      s"Destroy FileWriter ${fileWriter} caused by shuffle ${shuffleKey} expired."))
+                    writers.remove(fileInfo.getFilePath)
+                  }
+                }
+              }
+            }
+          }
+        }
         val (appId, shuffleId) = Utils.splitShuffleKey(shuffleKey)
         disksSnapshot().filter(_.status != DiskStatus.IO_HANG).foreach { diskInfo =>
           diskInfo.dirs.foreach { dir =>
@@ -432,8 +457,7 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
             deleteDirectory(file, diskOperators.get(diskInfo.mountPoint))
           }
         }
-        hdfsInfos.foreach(item => item._2.deleteAllFiles(StorageManager.hadoopFs))
-        if (!hdfsInfos.isEmpty) {
+        if (isHdfsExpired) {
           try {
             StorageManager.hadoopFs.delete(
               new Path(new Path(hdfsDir, conf.workerWorkingDir), s"$appId/$shuffleId"),
@@ -610,9 +634,10 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
     workingDirWriters.asScala.foreach { case (_, writers) =>
       writers.synchronized {
         // Filter out FileWriter that already has IOException to avoid printing too many error logs
-        allWriters.addAll(writers.asScala.filter(_.getException == null).asJava)
+        allWriters.addAll(writers.values().asScala.filter(_.getException == null).asJavaCollection)
       }
     }
+
     allWriters.asScala.foreach { writer =>
       try {
         writer.flushOnMemoryPressure()
@@ -643,7 +668,7 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
         val writers = workingDirWriters.get(dir)
         if (writers != null) {
           writers.synchronized {
-            writers.asScala.map(_.getFileInfo.getFileLength).sum
+            writers.values.asScala.map(_.getFileInfo.getFileLength).sum
           }
         } else {
           0
@@ -653,6 +678,7 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
         Paths.get(diskInfo.mountPoint)).getUsableSpace
       val workingDirUsableSpace =
         Math.min(diskInfo.configuredUsableSpace - totalUsage, fileSystemReportedUsableSpace)
+      logDebug(s"updateDiskInfos  workingDirUsableSpace:$workingDirUsableSpace filemeta:$fileSystemReportedUsableSpace conf:${diskInfo.configuredUsableSpace} totalUsage:$totalUsage")
       val flushTimeAverage = localFlushers.get(diskInfo.mountPoint).averageFlushTime()
       diskInfo.setUsableSpace(workingDirUsableSpace)
       diskInfo.setFlushTime(flushTimeAverage)
diff --git a/worker/src/test/scala/org/apache/celeborn/service/deploy/MiniClusterFeature.scala b/worker/src/test/scala/org/apache/celeborn/service/deploy/MiniClusterFeature.scala
index 142f750a..da6aa25d 100644
--- a/worker/src/test/scala/org/apache/celeborn/service/deploy/MiniClusterFeature.scala
+++ b/worker/src/test/scala/org/apache/celeborn/service/deploy/MiniClusterFeature.scala
@@ -98,7 +98,7 @@ trait MiniClusterFeature extends Logging {
   def setUpMiniCluster(
       masterConfs: Map[String, String] = null,
       workerConfs: Map[String, String] = null,
-      workerNum: Int = 3): Unit = {
+      workerNum: Int = 3): (Master, collection.Set[Worker]) = {
     val master = createMaster(masterConfs)
     val masterThread = runnerWrap(master.rpcEnv.awaitTermination())
     masterThread.start()
@@ -111,12 +111,12 @@ trait MiniClusterFeature extends Logging {
       workerThread.start()
       workerInfos.put(worker, workerThread)
     }
-
     Thread.sleep(5000L)
 
     workerInfos.foreach {
       case (worker, _) => assert(worker.isRegistered())
     }
+    (master, workerInfos.keySet)
   }
 
   def shutdownMiniCluster(): Unit = {
diff --git a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/WorkerSuite.scala b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/WorkerSuite.scala
index 0939f168..3dc35e59 100644
--- a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/WorkerSuite.scala
+++ b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/WorkerSuite.scala
@@ -20,113 +20,82 @@ package org.apache.celeborn.service.deploy.worker.storage
 import java.io.File
 import java.util
 import java.util.{HashSet => JHashSet}
+import java.util.concurrent.ConcurrentHashMap
+
+import scala.collection.JavaConverters._
 
 import org.junit.Assert
 import org.mockito.MockitoSugar._
 import org.scalatest.funsuite.AnyFunSuite
 
 import org.apache.celeborn.common.CelebornConf
-import org.apache.celeborn.common.meta.FileInfo
-import org.apache.celeborn.common.protocol.{PartitionLocation, PartitionType}
-import org.apache.celeborn.service.deploy.worker.{Worker, WorkerArguments, WorkerSource, WorkingPartition}
+import org.apache.celeborn.common.identity.UserIdentifier
+import org.apache.celeborn.common.protocol.{PartitionLocation, PartitionSplitMode, PartitionType}
+import org.apache.celeborn.service.deploy.worker.{Worker, WorkerArguments}
 
 class WorkerSuite extends AnyFunSuite {
   val conf = new CelebornConf()
   val workerArgs = new WorkerArguments(Array(), conf)
   test("clean up") {
-
+    conf.set("celeborn.worker.storage.dirs", "/tmp")
     val worker = new Worker(conf, workerArgs)
-    val expiredShuffleKeys = new JHashSet[String]()
-    val shuffleKey1 = "s1"
-    val shuffleKey2 = "s2"
-    expiredShuffleKeys.add(shuffleKey1)
-    expiredShuffleKeys.add(shuffleKey2)
-    val pl1 = new PartitionLocation(0, 0, "12", 0, 0, 0, 0, PartitionLocation.Mode.MASTER)
-    val pl2 = new PartitionLocation(0, 0, "12", 0, 0, 0, 0, PartitionLocation.Mode.SLAVE)
-
-    val dir1 = new File("/tmp/work1")
-    val dir2 = new File("/tmp/work2")
-    val dir3 = new File("/tmp/work3")
-
-    val filePath1 = new File(dir1, "/1/1")
-    filePath1.mkdirs()
-    val file1 = new File(filePath1, "/1")
-    file1.createNewFile()
-    val filePath2 = new File(dir2, "/2/2")
-    filePath2.mkdirs()
-    val file2 = new File(filePath2, "/2")
-    file2.createNewFile()
-
-    val filePath3 = new File(dir3, "/3/3")
-    filePath3.mkdirs()
-    val file3 = new File(filePath3, "/3")
-    file3.createNewFile()
 
-    val fw1 = new ReducePartitionFileWriter(
-      new FileInfo(file1.getAbsolutePath, null, null, PartitionType.REDUCE),
-      mock[Flusher],
-      new WorkerSource(conf),
-      conf,
-      mock[DeviceMonitor],
-      0,
-      null,
-      false)
-    fw1.registerDestroyHook(new util.ArrayList(util.Arrays.asList(fw1)))
-    val fw2 = new ReducePartitionFileWriter(
-      new FileInfo(file2.getAbsolutePath, null, null, PartitionType.REDUCE),
-      mock[Flusher],
-      new WorkerSource(conf),
-      conf,
-      mock[DeviceMonitor],
-      0,
-      null,
-      false)
-    fw2.registerDestroyHook(new util.ArrayList(util.Arrays.asList(fw2)))
-    val fw3 = new ReducePartitionFileWriter(
-      new FileInfo(file3.getAbsolutePath, null, null, PartitionType.REDUCE),
-      mock[Flusher],
-      new WorkerSource(conf),
-      conf,
-      mock[DeviceMonitor],
-      0,
-      null,
-      false)
-    fw3.registerDestroyHook(new util.ArrayList(util.Arrays.asList(fw3)))
+    val pl1 = new PartitionLocation(0, 0, "12", 0, 0, 0, 0, PartitionLocation.Mode.MASTER)
+    val pl2 = new PartitionLocation(1, 0, "12", 0, 0, 0, 0, PartitionLocation.Mode.SLAVE)
 
-    val wl1 = new WorkingPartition(pl1, fw1)
-    val wl2 = new WorkingPartition(pl2, fw2)
-    val wl3 = new WorkingPartition(pl2, fw3)
-    worker.partitionLocationInfo.addMasterPartitions(shuffleKey1, util.Arrays.asList(wl1))
-    worker.partitionLocationInfo.addMasterPartitions(shuffleKey1, util.Arrays.asList(wl3))
-    worker.partitionLocationInfo.addSlavePartitions(shuffleKey1, util.Arrays.asList(wl2))
+    worker.storageManager.createWriter(
+      "1",
+      1,
+      pl1,
+      100000,
+      PartitionSplitMode.SOFT,
+      PartitionType.REDUCE,
+      true,
+      new UserIdentifier("1", "2"))
+    worker.storageManager.createWriter(
+      "2",
+      2,
+      pl2,
+      100000,
+      PartitionSplitMode.SOFT,
+      PartitionType.REDUCE,
+      true,
+      new UserIdentifier("1", "2"))
 
-    val fws1 = new util.ArrayList[FileWriter]()
-    fws1.add(fw1)
-    val fws2 = new util.ArrayList[FileWriter]()
-    fws2.add(fw2)
-    worker.storageManager.workingDirWriters.put(dir1, fws1)
-    worker.storageManager.workingDirWriters.put(dir2, fws2)
-    Assert.assertEquals(1, worker.storageManager.workingDirWriters.get(dir1).size())
-    Assert.assertEquals(1, worker.storageManager.workingDirWriters.get(dir2).size())
+    Assert.assertEquals(1, worker.storageManager.workingDirWriters.values().size())
+    val expiredShuffleKeys = new JHashSet[String]()
+    val shuffleKey1 = "1-1"
+    val shuffleKey2 = "2-2"
+    expiredShuffleKeys.add(shuffleKey1)
+    expiredShuffleKeys.add(shuffleKey2)
     worker.cleanup(expiredShuffleKeys)
-    Assert.assertEquals(0, worker.storageManager.workingDirWriters.get(dir1).size())
-    Assert.assertEquals(0, worker.storageManager.workingDirWriters.get(dir2).size())
-
-    deleteFile(dir1)
-    deleteFile(dir2)
-    deleteFile(dir3)
+    worker.storageManager.workingDirWriters.values().asScala.map(t => assert(t.size() == 0))
   }
-  def deleteFile(dir: File): Unit = {
-    val files = dir.listFiles();
-    if (files != null) {
-      files.foreach(file => {
-        if (file.isFile()) {
-          file.delete();
-        } else {
-          deleteFile(file);
-        }
-      })
-      dir.delete();
+
+  test("flush filewriters") {
+    conf.set("celeborn.worker.storage.dirs", "/tmp")
+    val worker = new Worker(conf, workerArgs)
+    val dir = new File("/tmp")
+    val allWriters = new util.HashSet[FileWriter]()
+    val map = new ConcurrentHashMap[String, FileWriter]()
+    worker.storageManager.workingDirWriters.put(dir, map)
+    worker.storageManager.workingDirWriters.asScala.foreach { case (_, writers) =>
+      writers.synchronized {
+        // Filter out FileWriter that already has IOException to avoid printing too many error logs
+        allWriters.addAll(writers.values().asScala.filter(_.getException == null).asJavaCollection)
+      }
+    }
+    Assert.assertEquals(0, allWriters.size())
+
+    val fileWriter = mock[FileWriter]
+    when(fileWriter.getException).thenReturn(null)
+    map.put("1", fileWriter)
+    worker.storageManager.workingDirWriters.asScala.foreach { case (_, writers) =>
+      writers.synchronized {
+        // Filter out FileWriter that already has IOException to avoid printing too many error logs
+        allWriters.addAll(writers.values().asScala.filter(_.getException == null).asJavaCollection)
+      }
     }
+    Assert.assertEquals(1, allWriters.size())
   }
 }