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

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


The following commit(s) were added to refs/heads/main by this push:
     new a3335bee7d [CELEBORN-2467] Fix silent shuffle data loss when a 
replica's commit times out on OSS and HDFS storage
a3335bee7d is described below

commit a3335bee7dc49a47ad1a8214f19460b5777d98bc
Author: ewoodbury <[email protected]>
AuthorDate: Wed Oct 21 10:54:41 2026 +0800

    [CELEBORN-2467] Fix silent shuffle data loss when a replica's commit times 
out on OSS and HDFS storage
    
    ### What changes were proposed in this pull request?
    
    Move DFS replica dedup from the worker (at close) to the driver (after 
commit):
    
    - Worker: stop deleting the local copy when the peer's `.success` exists, 
so both copies are always reported. Remove the unused peer-path helpers.
    - Worker: fail the commit if a DFS copy's data stream or index can't be 
closed. Write `.success` last; failing to create it is only logged.
    - Worker: count a non-empty file with no storage info as failed. This is a 
safety check; DFS writers no longer return null here.
    - Driver: when both copies of a location commit on HDFS, publish only the 
primary and release the replica with `DestroyWorkerSlots` on a separate pool. 
Applies to reduce and map partitions.
    - Worker: if `DestroyWorkerSlots` can't delete an HDFS file, keep tracking 
it and reply `PARTIAL_SUCCESS`; the driver retries, and shuffle expiry removes 
it otherwise.
    
    ### Why are the changes needed?
    
    With DFS storage and `celeborn.client.push.replicate.enabled=true`, it's 
possible to return shuffle data which is missing a small percent of records in 
some edge cases (anecdotally ~0.01% to 0.5%).
    
    The worker that closes a partition second sees the peer's `.success`, 
deletes its own copy, and leaves the id out of its `CommitFiles` response. If 
the peer's commit then times out, the peer reports every id failed. 
`checkDataLost` only flags ids that failed on both sides, so the partition gets 
no location and reducers read nothing. The same happens when a copy whose flush 
failed still writes `.success`.
    
    `.success` only shows that the peer wrote its file, not that its commit 
succeeded. Only the driver sees both results. The dedup came from ISSUE-440 
(#466).
    
    The docs recommend replication off with HDFS, but users enable it, S3/OSS 
are supported since 0.6, and the loss is silent 
(`celeborn.client.shuffle.integrityCheck.enabled` is off by default).
    
    Notes:
    - Storage: HDFS still keeps one copy after commit. With `MEMORY,HDFS`, the 
replica files are deleted, but the worker keeps a stale entry until shuffle 
expiry (CELEBORN-2485). S3/OSS are unchanged, since the old delete only removed 
an empty placeholder object. Each replica now also writes an index and 
`.success`.
    - Reads: after an HDFS replica is released, losing the primary worker means 
a stage rerun, as today. S3/OSS can now fail over to the replica.
    - Upgrade workers first. New workers with an old client keep both HDFS 
copies until shuffle expiry.
    - CELEBORN-2341 touches the same timeout path but does not fix this.
    
    ### Does this PR resolve a correctness bug?
    
    - [x] Yes
    
    Fixes silent shuffle data loss with DFS storage and replication.
    
    ### Does this PR introduce _any_ user-facing change?
    
    - [ ] Yes
    
    No new configuration, API, or protocol change. Storage after commit is the 
same as today.
    
    ### How was this patch tested?
    
    Tested end-to-end on a real Celeborn cluster on top of S3-compatible object 
store, with a Spark 4.0 job where I forced a timeout and confirmed the result 
was not missing data.
    
    CI tests:
    - `HdfsReplicaReleaseClusterSuite` (mini cluster, HDFS simulated with 
`RawLocalFileSystem`):
      - Primary commit times out: the partition reads back intact from the 
replica. Fails on `main` with an empty file group.
      - Both copies commit: the replica is unpublished and deleted. Also passes 
on `main` (storage parity).
    - `TierWriterSuite`: keeps its copy when the peer `.success` exists; index 
failure fails close; `.success` failure does not; the original flush error is 
kept. The first three fail on `main`.
    - `HdfsReplicaReleaseSuite`: replica selection across storage types and 
failure cases.
    - `StorageManagerSuite`: failed HDFS deletes keep the file tracked across 
retries and during shuffle expiry.
    - Client and worker module tests, `spotless:check`, and Scala 2.13 
test-compile. `DataPushQueueSuiteJ` and `ChunkFetchIntegrationSuiteJ` were 
excluded; they also fail on `main` locally.
    
    Closes #3879 from ewoodbury/ewoodbury/celeborn-2467.
    
    Authored-by: ewoodbury <[email protected]>
    Signed-off-by: Nicholas Jiang <[email protected]>
---
 .../org/apache/celeborn/client/CommitManager.scala |   3 +-
 .../apache/celeborn/client/LifecycleManager.scala  |  20 +++
 .../celeborn/client/commit/CommitHandler.scala     |  72 +++++++-
 .../client/commit/MapPartitionCommitHandler.scala  |  12 +-
 .../commit/ReducePartitionCommitHandler.scala      |   7 +-
 .../client/commit/HdfsReplicaReleaseSuite.scala    |  96 ++++++++++
 .../commit/MapPartitionCommitHandlerTest.scala     |   3 +-
 .../apache/celeborn/common/meta/DiskFileInfo.java  |   4 -
 .../org/apache/celeborn/common/util/Utils.scala    |   8 -
 .../deploy/worker/storage/PartitionDataWriter.java |   4 +-
 .../service/deploy/worker/Controller.scala         |   5 +-
 .../service/deploy/worker/storage/FlushTask.scala  |  16 +-
 .../deploy/worker/storage/StorageManager.scala     |  36 +++-
 .../service/deploy/worker/storage/TierWriter.scala |  74 ++++----
 .../cluster/HdfsReplicaReleaseClusterSuite.scala   | 197 +++++++++++++++++++++
 .../worker/storage/StorageManagerSuite.scala       |  72 +++++++-
 .../deploy/worker/storage/TierWriterSuite.scala    | 162 ++++++++++++++++-
 17 files changed, 717 insertions(+), 74 deletions(-)

diff --git 
a/client/src/main/scala/org/apache/celeborn/client/CommitManager.scala 
b/client/src/main/scala/org/apache/celeborn/client/CommitManager.scala
index 64a7c95e9d..d75e1d1db4 100644
--- a/client/src/main/scala/org/apache/celeborn/client/CommitManager.scala
+++ b/client/src/main/scala/org/apache/celeborn/client/CommitManager.scala
@@ -324,7 +324,8 @@ class CommitManager(appUniqueId: String, val conf: 
CelebornConf, lifecycleManage
               committedPartitionInfo,
               lifecycleManager.workerStatusTracker,
               lifecycleManager.rpcSharedThreadPool,
-              commitRetryScheduler)
+              commitRetryScheduler,
+              lifecycleManager)
           case _ => throw new UnsupportedOperationException(
               s"Unexpected ShufflePartitionType for CommitManager: 
$partitionType")
         }
diff --git 
a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala 
b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
index 3b4ea83ef7..facf767b01 100644
--- a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
+++ b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
@@ -33,6 +33,7 @@ import scala.collection.mutable.ArrayBuffer
 import scala.concurrent.{ExecutionContext, Future}
 import scala.concurrent.duration.Duration
 import scala.util.Random
+import scala.util.control.NonFatal
 
 import com.google.common.annotations.VisibleForTesting
 import com.google.common.cache.{Cache, CacheBuilder}
@@ -182,6 +183,8 @@ class LifecycleManager(val appUniqueId: String, val conf: 
CelebornConf) extends
       conf.clientRpcSharedThreads,
       30)
   val ec = ExecutionContext.fromExecutor(rpcSharedThreadPool)
+  private val replicaReleasePool =
+    
ThreadUtils.newDaemonCachedThreadPool("celeborn-client-lifecycle-manager-replica-release",
 4)
 
   // init driver celeborn LifecycleManager rpc service
   override val rpcEnv: RpcEnv = RpcEnv.create(
@@ -318,6 +321,7 @@ class LifecycleManager(val appUniqueId: String, val conf: 
CelebornConf) extends
     commitManager.stop()
     changePartitionManager.stop()
     releasePartitionManager.stop()
+    ThreadUtils.shutdown(replicaReleasePool)
     heartbeater.stop()
 
     masterClient.close()
@@ -1728,6 +1732,22 @@ class LifecycleManager(val appUniqueId: String, val 
conf: CelebornConf) extends
       var retryTimes: Int,
       var startTime: Long)
 
+  /** Best-effort: a failed or rejected release only leaves the replica until 
shuffle expiry. */
+  private[client] def releaseReplicaSlotsAsync(
+      shuffleId: Int,
+      replicaSlots: WorkerResource): Unit = {
+    try {
+      replicaReleasePool.submit(new Runnable {
+        override def run(): Unit = Utils.tryLogNonFatalError {
+          destroySlotsWithRetry(shuffleId, replicaSlots)
+        }
+      })
+    } catch {
+      case NonFatal(e) =>
+        logWarning(s"Skip releasing committed replicas for shuffle 
$shuffleId.", e)
+    }
+  }
+
   /**
    * For the slots that need to be destroyed, LifecycleManager will ask the 
corresponding worker
    * to destroy related FileWriter.
diff --git 
a/client/src/main/scala/org/apache/celeborn/client/commit/CommitHandler.scala 
b/client/src/main/scala/org/apache/celeborn/client/commit/CommitHandler.scala
index f54c2b990d..7660c375f1 100644
--- 
a/client/src/main/scala/org/apache/celeborn/client/commit/CommitHandler.scala
+++ 
b/client/src/main/scala/org/apache/celeborn/client/commit/CommitHandler.scala
@@ -35,8 +35,8 @@ import org.apache.celeborn.common.{CelebornConf, 
CommitMetadata}
 import org.apache.celeborn.common.internal.Logging
 import org.apache.celeborn.common.meta.{ShufflePartitionLocationInfo, 
WorkerInfo}
 import org.apache.celeborn.common.network.protocol.SerdeVersion
-import org.apache.celeborn.common.protocol.{PartitionLocation, PartitionType}
-import 
org.apache.celeborn.common.protocol.message.ControlMessages.{CommitFiles, 
CommitFilesResponse}
+import org.apache.celeborn.common.protocol.{PartitionLocation, PartitionType, 
StorageInfo}
+import 
org.apache.celeborn.common.protocol.message.ControlMessages.{CommitFiles, 
CommitFilesResponse, WorkerResource}
 import org.apache.celeborn.common.protocol.message.StatusCode
 import org.apache.celeborn.common.rpc.RpcCallContext
 import org.apache.celeborn.common.util.{CollectionUtils, JavaUtils, Utils}
@@ -565,6 +565,42 @@ abstract class CommitHandler(
     }
   }
 
+  /** Like collectResult, but leaves out and releases HDFS replicas whose 
primary also committed. */
+  protected def collectResultAndReleaseHdfsReplicas(
+      shuffleId: Int,
+      shuffleCommittedInfo: ShuffleCommittedInfo,
+      primaryPartitionUniqueIds: util.Iterator[String],
+      replicaPartitionUniqueIds: util.Iterator[String],
+      commitResult: CommitResult,
+      allocatedWorkers: util.Map[String, ShufflePartitionLocationInfo],
+      releaseReplicaSlots: WorkerResource => Unit,
+      isSegmentGranularityVisible: Boolean = false): Unit = {
+    val primaryIds = primaryPartitionUniqueIds.asScala.toSet
+    val replicaIds = replicaPartitionUniqueIds.asScala.toSeq
+    val replicaSlots = CommitHandler.selectHdfsReplicasToRelease(
+      primaryIds,
+      replicaIds,
+      shuffleCommittedInfo.committedPrimaryStorageInfos,
+      shuffleCommittedInfo.committedReplicaStorageInfos,
+      commitResult.replicaPartitionLocationMap,
+      allocatedWorkers)
+    val releasedReplicaIds =
+      
replicaSlots.values().asScala.flatMap(_._2.asScala.map(_.getUniqueId)).toSet
+    collectResult(
+      shuffleId,
+      shuffleCommittedInfo,
+      primaryIds.iterator.asJava,
+      replicaIds.iterator.filterNot(releasedReplicaIds.contains).asJava,
+      commitResult.primaryPartitionLocationMap,
+      commitResult.replicaPartitionLocationMap,
+      isSegmentGranularityVisible)
+    if (!replicaSlots.isEmpty) {
+      logInfo(s"Releasing ${releasedReplicaIds.size} HDFS replicas of shuffle 
$shuffleId " +
+        "whose primary also committed.")
+      releaseReplicaSlots(replicaSlots)
+    }
+  }
+
   def collectResult(
       shuffleId: Int,
       shuffleCommittedInfo: ShuffleCommittedInfo,
@@ -660,3 +696,35 @@ abstract class CommitHandler(
     }
   }
 }
+
+object CommitHandler {
+
+  /** Replicas committed on HDFS alongside an HDFS primary, keyed by the 
allocated worker. */
+  private[client] def selectHdfsReplicasToRelease(
+      committedPrimaryIds: Set[String],
+      committedReplicaIds: Iterable[String],
+      primaryStorageInfos: util.Map[String, StorageInfo],
+      replicaStorageInfos: util.Map[String, StorageInfo],
+      replicaLocations: util.Map[String, PartitionLocation],
+      allocatedWorkers: util.Map[String, ShufflePartitionLocationInfo]): 
WorkerResource = {
+    val resource = new WorkerResource
+    committedReplicaIds.toSet.foreach { id: String =>
+      val replica = replicaLocations.get(id)
+      val allocated =
+        if (replica == null) null else 
allocatedWorkers.get(replica.getWorker.toUniqueId)
+      if (committedPrimaryIds.contains(id) &&
+        isHdfs(primaryStorageInfos.get(id)) && 
isHdfs(replicaStorageInfos.get(id)) &&
+        allocated != null && allocated.workerInfo.endpoint != null) {
+        resource.computeIfAbsent(
+          allocated.workerInfo,
+          (_: WorkerInfo) =>
+            (new util.ArrayList[PartitionLocation](), new 
util.ArrayList[PartitionLocation]()))
+          ._2.add(replica)
+      }
+    }
+    resource
+  }
+
+  private def isHdfs(storageInfo: StorageInfo): Boolean =
+    storageInfo != null && storageInfo.getType == StorageInfo.Type.HDFS
+}
diff --git 
a/client/src/main/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandler.scala
 
b/client/src/main/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandler.scala
index 04701f2f47..e24ab13570 100644
--- 
a/client/src/main/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandler.scala
+++ 
b/client/src/main/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandler.scala
@@ -27,7 +27,7 @@ import scala.collection.mutable
 
 import org.roaringbitmap.RoaringBitmap
 
-import org.apache.celeborn.client.{ShuffleCommittedInfo, WorkerStatusTracker}
+import org.apache.celeborn.client.{LifecycleManager, ShuffleCommittedInfo, 
WorkerStatusTracker}
 import org.apache.celeborn.client.CommitManager.CommittedPartitionInfo
 import org.apache.celeborn.client.LifecycleManager.{ShuffleAllocatedWorkers, 
ShuffleFailedWorkers}
 import org.apache.celeborn.common.{CelebornConf, CommitMetadata}
@@ -58,7 +58,8 @@ class MapPartitionCommitHandler(
     committedPartitionInfo: CommittedPartitionInfo,
     workerStatusTracker: WorkerStatusTracker,
     sharedRpcPool: ThreadPoolExecutor,
-    commitRetryScheduler: ScheduledExecutorService)
+    commitRetryScheduler: ScheduledExecutorService,
+    lifecycleManager: LifecycleManager)
   extends CommitHandler(
     appId,
     conf,
@@ -171,13 +172,14 @@ class MapPartitionCommitHandler(
 
     // collect partition result
     if (!dataLost) {
-      collectResult(
+      collectResultAndReleaseHdfsReplicas(
         shuffleId,
         shuffleCommittedInfo,
         getPartitionUniqueIds(shuffleCommittedInfo.committedPrimaryIds, 
partitionId),
         getPartitionUniqueIds(shuffleCommittedInfo.committedReplicaIds, 
partitionId),
-        parallelCommitResult.primaryPartitionLocationMap,
-        parallelCommitResult.replicaPartitionLocationMap,
+        parallelCommitResult,
+        allocatedWorkers,
+        lifecycleManager.releaseReplicaSlotsAsync(shuffleId, _),
         shuffleIsSegmentGranularityVisible.get(shuffleId))
     }
 
diff --git 
a/client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala
 
b/client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala
index 7477cb9dd3..68aec0c977 100644
--- 
a/client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala
+++ 
b/client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala
@@ -254,13 +254,14 @@ class ReducePartitionCommitHandler(
 
     // collect result
     if (!dataLost) {
-      collectResult(
+      collectResultAndReleaseHdfsReplicas(
         shuffleId,
         shuffleCommittedInfo,
         getPartitionUniqueIds(shuffleCommittedInfo.committedPrimaryIds),
         getPartitionUniqueIds(shuffleCommittedInfo.committedReplicaIds),
-        parallelCommitResult.primaryPartitionLocationMap,
-        parallelCommitResult.replicaPartitionLocationMap)
+        parallelCommitResult,
+        allocatedWorkers,
+        lifecycleManager.releaseReplicaSlotsAsync(shuffleId, _))
     }
 
     (dataLost, parallelCommitResult.commitFilesFailedWorkers)
diff --git 
a/client/src/test/scala/org/apache/celeborn/client/commit/HdfsReplicaReleaseSuite.scala
 
b/client/src/test/scala/org/apache/celeborn/client/commit/HdfsReplicaReleaseSuite.scala
new file mode 100644
index 0000000000..1418e09533
--- /dev/null
+++ 
b/client/src/test/scala/org/apache/celeborn/client/commit/HdfsReplicaReleaseSuite.scala
@@ -0,0 +1,96 @@
+/*
+ * 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.client.commit
+
+import java.util
+
+import scala.collection.JavaConverters._
+
+import org.mockito.Mockito
+
+import org.apache.celeborn.CelebornFunSuite
+import org.apache.celeborn.common.meta.{ShufflePartitionLocationInfo, 
WorkerInfo}
+import org.apache.celeborn.common.protocol.{PartitionLocation, StorageInfo}
+import org.apache.celeborn.common.protocol.StorageInfo.Type
+import org.apache.celeborn.common.rpc.RpcEndpointRef
+
+class HdfsReplicaReleaseSuite extends CelebornFunSuite {
+
+  private val replicaWorker = worker()
+
+  private def worker(withEndpoint: Boolean = true): WorkerInfo = {
+    val w = new WorkerInfo("replica-host", 1, 2, 3, 4)
+    if (withEndpoint) {
+      w.endpoint = Mockito.mock(classOf[RpcEndpointRef])
+    }
+    w
+  }
+
+  private def storageInfos(id: String, t: Option[Type]): util.Map[String, 
StorageInfo] =
+    t.map(storageType => id -> new StorageInfo(storageType, true, 
"")).toMap.asJava
+
+  private case class Case(
+      name: String,
+      primaryType: Option[Type],
+      replicaType: Option[Type],
+      allocatedWorker: Option[WorkerInfo] = Some(replicaWorker),
+      primaryCommitted: Boolean = true,
+      replicaCommitted: Boolean = true,
+      expectReleased: Boolean = false)
+
+  Seq(
+    Case("both HDFS", Some(Type.HDFS), Some(Type.HDFS), expectReleased = true),
+    Case("primary not committed", Some(Type.HDFS), Some(Type.HDFS), 
primaryCommitted = false),
+    Case("replica not committed", Some(Type.HDFS), Some(Type.HDFS), 
replicaCommitted = false),
+    Case("primary MEMORY, replica HDFS", Some(Type.MEMORY), Some(Type.HDFS)),
+    Case("primary HDFS, replica MEMORY", Some(Type.HDFS), Some(Type.MEMORY)),
+    Case("both S3", Some(Type.S3), Some(Type.S3)),
+    Case("both OSS", Some(Type.OSS), Some(Type.OSS)),
+    Case("primary storage info missing", None, Some(Type.HDFS)),
+    Case("replica storage info missing", Some(Type.HDFS), None),
+    Case("replica worker not allocated", Some(Type.HDFS), Some(Type.HDFS), 
None),
+    Case(
+      "replica worker has no endpoint",
+      Some(Type.HDFS),
+      Some(Type.HDFS),
+      Some(worker(withEndpoint = false)))).foreach { c =>
+    test(s"selectHdfsReplicasToRelease: ${c.name}") {
+      val replica =
+        new PartitionLocation(0, 0, "replica-host", 1, 2, 3, 4, 
PartitionLocation.Mode.REPLICA)
+      val id = replica.getUniqueId
+      val workers = c.allocatedWorker.map(w =>
+        w.toUniqueId -> new ShufflePartitionLocationInfo(w)).toMap.asJava
+      val slots = CommitHandler.selectHdfsReplicasToRelease(
+        if (c.primaryCommitted) Set(id) else Set.empty,
+        if (c.replicaCommitted) Seq(id) else Seq.empty,
+        storageInfos(id, c.primaryType),
+        storageInfos(id, c.replicaType),
+        Map(id -> replica).asJava,
+        workers)
+
+      if (c.expectReleased) {
+        assert(slots.size === 1)
+        val (slotWorker, (primaries, replicas)) = slots.asScala.head
+        assert(slotWorker eq replicaWorker)
+        assert(primaries.isEmpty && replicas.asScala === Seq(replica))
+      } else {
+        assert(slots.isEmpty)
+      }
+    }
+  }
+}
diff --git 
a/client/src/test/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandlerTest.scala
 
b/client/src/test/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandlerTest.scala
index 47a485b604..4e1f52d1c6 100644
--- 
a/client/src/test/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandlerTest.scala
+++ 
b/client/src/test/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandlerTest.scala
@@ -63,7 +63,8 @@ class MapPartitionCommitHandlerTest extends CelebornFunSuite {
       new CommittedPartitionInfo(),
       new WorkerStatusTracker(conf, null),
       rpcPool,
-      commitScheduler)
+      commitScheduler,
+      null)
   }
 
   /** Write-side per-subpartition checksum/byte arrays computed over the given 
payloads. */
diff --git 
a/common/src/main/java/org/apache/celeborn/common/meta/DiskFileInfo.java 
b/common/src/main/java/org/apache/celeborn/common/meta/DiskFileInfo.java
index d4571fa4bb..2248863dd5 100644
--- a/common/src/main/java/org/apache/celeborn/common/meta/DiskFileInfo.java
+++ b/common/src/main/java/org/apache/celeborn/common/meta/DiskFileInfo.java
@@ -118,10 +118,6 @@ public class DiskFileInfo extends FileInfo {
     return new Path(Utils.getWriteSuccessFilePath(filePath));
   }
 
-  public Path getDfsPeerWriterSuccessPath() {
-    return new 
Path(Utils.getWriteSuccessFilePath(Utils.getPeerPath(filePath)));
-  }
-
   public void deleteAllFiles(FileSystem dfsFs) {
     if (isDFS()) {
       try {
diff --git a/common/src/main/scala/org/apache/celeborn/common/util/Utils.scala 
b/common/src/main/scala/org/apache/celeborn/common/util/Utils.scala
index 3a57a0900f..1c7a8fb3c2 100644
--- a/common/src/main/scala/org/apache/celeborn/common/util/Utils.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/util/Utils.scala
@@ -1174,14 +1174,6 @@ object Utils extends Logging {
     }
   }
 
-  def getPeerPath(path: String): String = {
-    if (path.endsWith("0")) {
-      path.substring(0, path.length - 1) + "1"
-    } else {
-      path.substring(0, path.length - 1) + "0"
-    }
-  }
-
   val SORTED_SUFFIX = ".sorted"
   val INDEX_SUFFIX = ".index"
   val SUFFIX_HDFS_WRITE_SUCCESS = ".success"
diff --git 
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionDataWriter.java
 
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionDataWriter.java
index 233b7a0efc..08d44edcc1 100644
--- 
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionDataWriter.java
+++ 
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionDataWriter.java
@@ -167,9 +167,7 @@ public class PartitionDataWriter implements DeviceObserver {
     if (fileInfo instanceof DiskFileInfo) {
       DiskFileInfo diskFileInfo = (DiskFileInfo) fileInfo;
       if (diskFileInfo.isDFS()) {
-        if (((DfsTierWriter) currentTierWriter).deleted()) {
-          return null;
-        } else if (diskFileInfo.isS3()) {
+        if (diskFileInfo.isS3()) {
           storageInfo = new StorageInfo(StorageInfo.Type.S3, true, 
diskFileInfo.getFilePath());
         } else if (diskFileInfo.isOSS()) {
           storageInfo = new StorageInfo(StorageInfo.Type.OSS, true, 
diskFileInfo.getFilePath());
diff --git 
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Controller.scala
 
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Controller.scala
index 565acb4418..9a723d6c60 100644
--- 
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Controller.scala
+++ 
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Controller.scala
@@ -407,8 +407,9 @@ private[deploy] class Controller(
                   logDebug(
                     s"Commit file $shuffleKey $uniqueId 
${fileWriter.getFilePath} success fileSize: $bytes")
                   if (fileWriter.getStorageInfo == null) {
-                    // Only HDFS can be null, means that this partition 
location is deleted.
-                    logDebug(s"Location $uniqueId is deleted.")
+                    // A non-empty file must be reported, or the client cannot 
detect its loss.
+                    logError(s"Location $uniqueId has no storage info, treat 
as failed.")
+                    failedIds.add(uniqueId)
                   } else {
                     val storageInfo = fileWriter.getStorageInfo
                     committedStorageInfos.put(uniqueId, storageInfo)
diff --git 
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/FlushTask.scala
 
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/FlushTask.scala
index 67f45e28b4..3119b20e03 100644
--- 
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/FlushTask.scala
+++ 
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/FlushTask.scala
@@ -83,16 +83,20 @@ abstract private[worker] class DfsFlushTask(
     notifier: FlushNotifier,
     keepBuffer: Boolean,
     source: AbstractSource) extends FlushTask(buffer, notifier, keepBuffer, 
source) with Logging {
+  // Closing an HDFS append stream persists the data, so a close failure must 
fail the flush.
   def flush(stream: Closeable)(block: => Unit): Unit = {
     try {
       block
-    } finally {
-      try {
-        stream.close()
-      } catch {
-        case e: IOException => logWarning("Close flush dfs stream failed.", e)
-      }
+    } catch {
+      case t: Throwable =>
+        try {
+          stream.close()
+        } catch {
+          case e: IOException => t.addSuppressed(e)
+        }
+        throw t
     }
+    stream.close()
   }
 }
 
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 99f87cd187..4ef0e6f366 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
@@ -583,6 +583,24 @@ final private[worker] class StorageManager(conf: 
CelebornConf, workerSource: Abs
     null
   }
 
+  // The writer ignores delete failures and is destroyed only once, so check 
and retry here.
+  private def ensureDfsFileDeleted(info: DiskFileInfo): Unit = {
+    val fs = StorageManager.hadoopFs.get(info.getStorageType)
+    val paths = Seq(
+      info.getDfsPath,
+      info.getDfsWriterSuccessPath,
+      info.getDfsIndexPath,
+      info.getDfsSortedPath)
+    if (paths.exists(fs.exists)) {
+      info.deleteAllFiles(fs)
+      val remaining = paths.filter(fs.exists)
+      if (remaining.nonEmpty) {
+        logWarning(s"Failed to delete $remaining, keeping them for retry or 
shuffle expiry.")
+        throw new IOException(s"Failed to delete $remaining")
+      }
+    }
+  }
+
   private def removeFileInfo(shuffleKey: String, fileName: String): FileInfo = 
{
     val memoryShuffleMap = memoryFileInfos.get(shuffleKey)
     if (memoryShuffleMap != null) {
@@ -619,22 +637,32 @@ final private[worker] class StorageManager(conf: 
CelebornConf, workerSource: Abs
   }
 
   def cleanFile(shuffleKey: String, fileName: String): Unit = {
-    cleanFileInternal(shuffleKey, removeFileInfo(shuffleKey, fileName))
+    getFileInfo(shuffleKey, fileName) match {
+      case info: DiskFileInfo if info.isHdfs =>
+        // Untrack only after the delete succeeds, so a retry or shuffle 
expiry can still delete it.
+        cleanFileInternal(shuffleKey, info, failOnHdfsDeleteError = true)
+        removeFileInfo(shuffleKey, fileName)
+      case _ =>
+        cleanFileInternal(shuffleKey, removeFileInfo(shuffleKey, fileName))
+    }
   }
 
-  def cleanFileInternal(shuffleKey: String, fileInfo: FileInfo): Boolean = {
+  def cleanFileInternal(
+      shuffleKey: String,
+      fileInfo: FileInfo,
+      failOnHdfsDeleteError: Boolean = false): Boolean = {
     if (fileInfo == null) return false
     var isDfsExpired = false
     fileInfo match {
       case info: DiskFileInfo =>
         if (info.isHdfs) {
           isDfsExpired = true
-          val hdfsFileWriter = hdfsWriters.get(info.getFilePath)
+          val hdfsFileWriter = hdfsWriters.remove(info.getFilePath)
           if (hdfsFileWriter != null) {
             hdfsFileWriter.destroy(new IOException(
               s"Destroy FileWriter $hdfsFileWriter caused by shuffle 
$shuffleKey expired."))
-            hdfsWriters.remove(info.getFilePath)
           }
+          if (failOnHdfsDeleteError) ensureDfsFileDeleted(info)
         } else if (info.isS3) {
           isDfsExpired = true
           val s3FileWriter = s3Writers.get(info.getFilePath)
diff --git 
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/TierWriter.scala
 
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/TierWriter.scala
index 5c55cb0790..1284013587 100644
--- 
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/TierWriter.scala
+++ 
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/TierWriter.scala
@@ -130,6 +130,7 @@ abstract class TierWriterBase(
   // close and destroy need to be invoked in synchronized blocks
   def close(): Long = {
     ensureNotClosed()
+    var closeStreamsError: IOException = null
     try {
       waitOnNoPending(numPendingWrites, false)
       closed = true
@@ -143,8 +144,14 @@ abstract class TierWriterBase(
       } catch {
         case e: IOException =>
           logWarning(s"close file writer $this failed", e)
+          if (failCommitOnCloseStreamsError) {
+            closeStreamsError = e
+          }
       }
     }
+    if (closeStreamsError != null) {
+      throw closeStreamsError
+    }
     committed = true
     notifyFileCommitted()
 
@@ -194,6 +201,8 @@ abstract class TierWriterBase(
 
   def closeStreams(): Unit
 
+  protected def failCommitOnCloseStreamsError: Boolean = false
+
   def genFlushTask(finalFlush: Boolean, keepBuffer: Boolean): FlushTask
 
   def flush(finalFlush: Boolean, rebuildChunkOffsets: Boolean = false): Unit = 
{
@@ -556,7 +565,6 @@ class DfsTierWriter(
   flusherBufferSize = conf.workerHdfsFlusherBufferSize
   private val flushWorkerIndex: Int = flusher.getWorkerIndex
   val hadoopFs: FileSystem = StorageManager.hadoopFs.get(storageType)
-  var deleted = false
   private var hdfsStream: FSDataOutputStream = null
   private val workerReuseHdfsOutputStreamEnabled = 
conf.workerReuseHdfsOuputSteamEnabled
   private var s3MultipartUploadHandler: MultipartUploadHandler = _
@@ -691,39 +699,36 @@ class DfsTierWriter(
     }
   }
 
+  // A copy whose data stream or index failed to close is unreadable, so it 
must not count as
+  // committed. S3/OSS upload completion failures are unchecked and already 
fail the commit.
+  override protected def failCommitOnCloseStreamsError: Boolean = true
+
+  // .success goes last, after the index and the upload are complete.
   override def closeStreams(): Unit = {
     closeHdfsStream()
-    if (hadoopFs.exists(dfsFileInfo.getDfsPeerWriterSuccessPath)) {
-      hadoopFs.delete(dfsFileInfo.getDfsPath, false)
-      deleted = true
-    } else {
-      val retryCount = conf.workerWriterHdfsCreateAuxiliaryFileMaxRetries
-      val retryWait = conf.workerWriterHdfsCreateAuxiliaryFileRetryWait
+    val retryCount = conf.workerWriterHdfsCreateAuxiliaryFileMaxRetries
+    val retryWait = conf.workerWriterHdfsCreateAuxiliaryFileRetryWait
+    if (dfsFileInfo.isReduceFileMeta) {
       Utils.withRetryOnTimeoutOrIOException(retryCount, retryWait) {
-        hadoopFs.create(dfsFileInfo.getDfsWriterSuccessPath).close()
-      }
-      if (dfsFileInfo.isReduceFileMeta) {
-        Utils.withRetryOnTimeoutOrIOException(retryCount, retryWait) {
-          if (hadoopFs.exists(dfsFileInfo.getDfsIndexPath)) {
-            hadoopFs.delete(dfsFileInfo.getDfsIndexPath, true)
+        if (hadoopFs.exists(dfsFileInfo.getDfsIndexPath)) {
+          hadoopFs.delete(dfsFileInfo.getDfsIndexPath, true)
+        }
+        val indexOutputStream = hadoopFs.create(dfsFileInfo.getDfsIndexPath)
+        hadoopFs.setReplication(
+          dfsFileInfo.getDfsIndexPath,
+          conf.workerDfsReplicationFactor.toShort)
+        val byteStream: ByteArrayOutputStream = new ByteArrayOutputStream()
+        val dataStream = new DataOutputStream(byteStream)
+        try {
+          
dataStream.writeInt(dfsFileInfo.getReduceFileMeta.getChunkOffsets.size)
+          
dfsFileInfo.getReduceFileMeta.getChunkOffsets.asScala.foreach(dataStream.writeLong(_))
+          indexOutputStream.write(byteStream.toByteArray)
+        } finally {
+          if (dataStream != null) {
+            dataStream.close()
           }
-          val indexOutputStream = hadoopFs.create(dfsFileInfo.getDfsIndexPath)
-          hadoopFs.setReplication(
-            dfsFileInfo.getDfsIndexPath,
-            conf.workerDfsReplicationFactor.toShort)
-          val byteStream: ByteArrayOutputStream = new ByteArrayOutputStream()
-          val dataStream = new DataOutputStream(byteStream)
-          try {
-            
dataStream.writeInt(dfsFileInfo.getReduceFileMeta.getChunkOffsets.size)
-            
dfsFileInfo.getReduceFileMeta.getChunkOffsets.asScala.foreach(dataStream.writeLong(_))
-            indexOutputStream.write(byteStream.toByteArray)
-          } finally {
-            if (dataStream != null) {
-              dataStream.close()
-            }
-            if (indexOutputStream != null) {
-              indexOutputStream.close()
-            }
+          if (indexOutputStream != null) {
+            indexOutputStream.close()
           }
         }
       }
@@ -736,6 +741,15 @@ class DfsTierWriter(
       ossMultipartUploadHandler.complete()
       ossMultipartUploadHandler.close()
     }
+    try {
+      Utils.withRetryOnTimeoutOrIOException(retryCount, retryWait) {
+        hadoopFs.create(dfsFileInfo.getDfsWriterSuccessPath).close()
+      }
+    } catch {
+      // Readers do not need the marker, so a missing one must not fail a 
readable copy.
+      case e: IOException =>
+        logWarning(s"Failed to create ${dfsFileInfo.getDfsWriterSuccessPath}", 
e)
+    }
   }
 
   override def notifyFileCommitted(): Unit =
diff --git 
a/worker/src/test/scala/org/apache/celeborn/service/deploy/cluster/HdfsReplicaReleaseClusterSuite.scala
 
b/worker/src/test/scala/org/apache/celeborn/service/deploy/cluster/HdfsReplicaReleaseClusterSuite.scala
new file mode 100644
index 0000000000..cc37db6486
--- /dev/null
+++ 
b/worker/src/test/scala/org/apache/celeborn/service/deploy/cluster/HdfsReplicaReleaseClusterSuite.scala
@@ -0,0 +1,197 @@
+/*
+ * 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.service.deploy.cluster
+
+import java.io.ByteArrayOutputStream
+import java.nio.charset.StandardCharsets
+import java.nio.file.{Files, Path => JPath}
+import java.util.concurrent.{CountDownLatch, TimeUnit}
+
+import scala.collection.JavaConverters._
+
+import org.apache.commons.io.{FileUtils, IOUtils}
+import org.apache.hadoop.fs.{FSDataOutputStream, Path, RawLocalFileSystem}
+import org.apache.hadoop.util.Progressable
+import org.scalatest.BeforeAndAfterAll
+import org.scalatest.concurrent.Eventually.eventually
+import org.scalatest.concurrent.Futures.{interval, timeout}
+import org.scalatest.funsuite.AnyFunSuite
+import org.scalatest.time.SpanSugar._
+
+import org.apache.celeborn.client.{LifecycleManager, ShuffleClientImpl}
+import org.apache.celeborn.client.read.MetricsCallback
+import org.apache.celeborn.common.CelebornConf
+import org.apache.celeborn.common.identity.UserIdentifier
+import org.apache.celeborn.common.protocol.{PartitionLocation, StorageInfo}
+import org.apache.celeborn.common.protocol.message.StatusCode
+import org.apache.celeborn.common.util.Utils
+import org.apache.celeborn.service.deploy.MiniClusterFeature
+
+/**
+ * Local filesystem standing in for HDFS. For [[SlowPrimaryFileSystem.App]] it 
stalls the
+ * primary's index past the worker commit timeout, and holds the replica's 
final flush until the
+ * primary has stalled, so the replica always finishes closing second.
+ */
+class SlowPrimaryFileSystem extends RawLocalFileSystem {
+  import SlowPrimaryFileSystem._
+
+  override def create(
+      f: Path,
+      overwrite: Boolean,
+      bufferSize: Int,
+      replication: Short,
+      blockSize: Long,
+      progress: Progressable): FSDataOutputStream = {
+    if (inApp(f) && f.getName.endsWith(s"-0${Utils.INDEX_SUFFIX}")) {
+      primaryStalled.countDown()
+      Thread.sleep(8000)
+    }
+    super.create(f, overwrite, bufferSize, replication, blockSize, progress)
+  }
+
+  override def append(f: Path, bufferSize: Int, progress: Progressable): 
FSDataOutputStream = {
+    if (inApp(f) && f.getName.endsWith("-1")) primaryStalled.await(30, 
TimeUnit.SECONDS)
+    super.append(f, bufferSize, progress)
+  }
+}
+
+object SlowPrimaryFileSystem {
+  val App = "app-commit-timeout"
+  private val primaryStalled = new CountDownLatch(1)
+
+  private def inApp(f: Path): Boolean = f.toString.contains(s"/$App/")
+}
+
+class HdfsReplicaReleaseClusterSuite extends AnyFunSuite with 
MiniClusterFeature
+  with BeforeAndAfterAll {
+
+  private val hdfsRoot: JPath = Files.createTempDirectory("celeborn-hdfs")
+  private val storageConf = Map(
+    CelebornConf.ACTIVE_STORAGE_TYPES.key -> "HDFS",
+    CelebornConf.HDFS_DIR.key -> hdfsRoot.toUri.toString,
+    // The checksummed local filesystem does not support the appends the HDFS 
flusher uses.
+    "celeborn.hadoop.fs.file.impl" -> classOf[SlowPrimaryFileSystem].getName,
+    "celeborn.hadoop.fs.file.impl.disable.cache" -> "true")
+  private var masterPort = 0
+
+  override def beforeAll(): Unit = {
+    val workerConf = storageConf + 
(CelebornConf.WORKER_SHUFFLE_COMMIT_TIMEOUT.key -> "3s")
+    val (master, _) = setupMiniClusterWithRandomPorts(storageConf, workerConf)
+    masterPort = master.conf.masterPort
+  }
+
+  override def afterAll(): Unit = {
+    shutdownMiniCluster()
+    FileUtils.deleteDirectory(hdfsRoot.toFile)
+  }
+
+  // Shuffle file names end with the location mode: 0 for primary, 1 for 
replica.
+  private def shuffleFiles(app: String, mode: Int): Seq[String] = {
+    val files = Files.walk(hdfsRoot)
+    try {
+      files.iterator().asScala
+        .filter(p => Files.isRegularFile(p) && p.toString.contains(s"/$app/"))
+        .map(_.getFileName.toString)
+        .filter(_.matches(s"\\d+-\\d+-$mode(\\..*)?"))
+        .toList
+    } finally {
+      files.close()
+    }
+  }
+
+  /**
+   * Pushes one payload per partition, ends the map stage, and passes each 
partition's published
+   * locations, expected payload and read payload to `check` while the 
application is still alive.
+   */
+  private def runShuffle(
+      app: String,
+      numPartitions: Int,
+      extraConf: Map[String, String] = Map.empty)(
+      check: (LifecycleManager, Seq[(Set[PartitionLocation], String, String)]) 
=> Unit): Unit = {
+    val clientConf = new CelebornConf()
+      .set(CelebornConf.MASTER_ENDPOINTS.key, s"localhost:$masterPort")
+      .set(CelebornConf.CLIENT_PUSH_REPLICATE_ENABLED.key, "true")
+      .set(CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE.key, "256K")
+    (storageConf ++ extraConf).foreach { case (k, v) => clientConf.set(k, v) }
+    val lifecycleManager = new LifecycleManager(app, clientConf)
+    val shuffleClient = new ShuffleClientImpl(app, clientConf, 
UserIdentifier("mock", "mock"))
+    shuffleClient.setupLifecycleManagerRef(lifecycleManager.self)
+    val metricsCallback = new MetricsCallback {
+      override def incBytesRead(bytesRead: Long): Unit = {}
+      override def incReadTime(time: Long): Unit = {}
+    }
+    try {
+      val payloads = (0 until numPartitions).map(p => s"partition-$p-" * 1024)
+      payloads.zipWithIndex.foreach { case (payload, partitionId) =>
+        val bytes = payload.getBytes(StandardCharsets.UTF_8)
+        shuffleClient.pushData(1, 0, 0, partitionId, bytes, 0, bytes.length, 
1, numPartitions)
+      }
+      assert(shuffleFiles(app, 1).nonEmpty, "replicas should be written before 
commit")
+      shuffleClient.mapperEnd(1, 0, 0, 1, numPartitions)
+      check(
+        lifecycleManager,
+        payloads.zipWithIndex.map { case (payload, partitionId) =>
+          val group = shuffleClient.updateFileGroup(1, 
partitionId).partitionGroups.get(partitionId)
+          val in =
+            shuffleClient.readPartition(1, partitionId, 0, 0, 0, Int.MaxValue, 
metricsCallback)
+          val out = new ByteArrayOutputStream()
+          try IOUtils.copy(in, out)
+          finally in.close()
+          val locations = 
Option(group).map(_.asScala.toSet).getOrElse(Set.empty)
+          (locations, payload, new String(out.toByteArray, 
StandardCharsets.UTF_8))
+        })
+    } finally {
+      shuffleClient.shutdown()
+      lifecycleManager.stop()
+    }
+  }
+
+  test("HDFS replica is left out of the file group and deleted once both 
copies commit") {
+    val app = "app-hdfs-replica-release"
+    runShuffle(app, numPartitions = 2) { (_, results) =>
+      results.foreach { case (locations, expected, read) =>
+        assert(locations.size === 1)
+        val location = locations.head
+        assert(location.getMode == PartitionLocation.Mode.PRIMARY && 
location.getPeer == null)
+        assert(location.getStorageInfo.getType == StorageInfo.Type.HDFS)
+        assert(read === expected)
+      }
+      eventually(timeout(30.seconds), interval(500.millis)) {
+        assert(shuffleFiles(app, 1).isEmpty)
+      }
+      assert(shuffleFiles(app, 0).count(!_.contains(".")) === 2)
+    }
+  }
+
+  test("a primary whose commit times out leaves its partition readable from 
the replica") {
+    val app = SlowPrimaryFileSystem.App
+    // The timed-out worker never answers the first CommitFiles; the retry 
gets its cached failure.
+    val retryQuickly = Map(
+      CelebornConf.CLIENT_RPC_COMMIT_FILES_ASK_TIMEOUT.key -> "5s",
+      CelebornConf.CLIENT_COMMIT_FILE_REQUEST_RETRY_WAIT.key -> "0s")
+    runShuffle(app, numPartitions = 1, retryQuickly) { (lifecycleManager, 
results) =>
+      val commitFailures =
+        
lifecycleManager.workerStatusTracker.excludedWorkers.values().asScala.map(_._1)
+      assert(commitFailures.toSet === Set(StatusCode.COMMIT_FILE_EXCEPTION))
+      val (locations, expected, read) = results.head
+      assert(locations.map(_.getMode) === Set(PartitionLocation.Mode.REPLICA))
+      assert(read === expected)
+      assert(shuffleFiles(app, 1).count(!_.contains(".")) === 1)
+    }
+  }
+}
diff --git 
a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManagerSuite.scala
 
b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManagerSuite.scala
index 6107faf986..0e2a8b4200 100644
--- 
a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManagerSuite.scala
+++ 
b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManagerSuite.scala
@@ -19,16 +19,22 @@ package org.apache.celeborn.service.deploy.worker.storage
 
 import java.{lang, util}
 import java.io.IOException
+import java.nio.file.Files
 
+import scala.collection.JavaConverters._
+
+import org.apache.commons.io.FileUtils
+import org.apache.hadoop.fs.FileSystem
 import org.mockito.{Mockito, MockitoSugar}
 import org.mockito.ArgumentMatchersSugar.any
+import org.mockito.invocation.InvocationOnMock
 import org.mockito.stubbing.Stubber
 
 import org.apache.celeborn.CelebornFunSuite
 import org.apache.celeborn.common.CelebornConf
 import org.apache.celeborn.common.CelebornConf.{WORKER_DISK_RESERVE_SIZE, 
WORKER_GRACEFUL_SHUTDOWN_ENABLED, WORKER_GRACEFUL_SHUTDOWN_RECOVER_PATH, 
WORKER_STORAGE_DIRS}
 import org.apache.celeborn.common.identity.UserIdentifier
-import org.apache.celeborn.common.meta.{DiskInfo, DiskStatus}
+import org.apache.celeborn.common.meta.{DiskFileInfo, DiskInfo, DiskStatus}
 import org.apache.celeborn.common.protocol.{PartitionLocation, PartitionType, 
StorageInfo}
 import org.apache.celeborn.common.util.Utils
 import org.apache.celeborn.service.deploy.worker.WorkerSource
@@ -147,6 +153,70 @@ class StorageManagerSuite extends CelebornFunSuite with 
MockitoHelper {
     }
   }
 
+  test("cleanFile keeps an HDFS file it could not delete for a retry or 
shuffle expiry") {
+    val hdfsDir = Files.createTempDirectory("celeborn-hdfs")
+    val conf = new CelebornConf()
+      .set(CelebornConf.ACTIVE_STORAGE_TYPES.key, "HDFS")
+      .set(CelebornConf.HDFS_DIR.key, hdfsDir.toUri.toString)
+    val hadoopFsBefore = StorageManager.hadoopFs
+    val storageManager = new StorageManager(conf, new WorkerSource(conf))
+    val fs = mock[FileSystem]
+    StorageManager.hadoopFs = Map[StorageInfo.Type, 
FileSystem](StorageInfo.Type.HDFS -> fs).asJava
+    try {
+      val shuffleKey = "app-1"
+      def createReplicaFile(fileName: String): (DiskFileInfo, 
PartitionDataWriter) = {
+        val location =
+          new PartitionLocation(0, 0, "localhost", 0, 0, 0, 0, 
PartitionLocation.Mode.REPLICA)
+        location.setStorageInfo(
+          new StorageInfo("", StorageInfo.Type.HDFS, 
StorageInfo.ALL_TYPES_AVAILABLE_MASK))
+        val fileInfo = storageManager.createDiskFile(
+          location,
+          "app",
+          1,
+          fileName,
+          new UserIdentifier("t1", "u1"),
+          PartitionType.REDUCE,
+          partitionSplitEnabled = false,
+          StorageInfo.Type.HDFS)._2
+        val writer = mock[PartitionDataWriter]
+        storageManager.hdfsWriters.put(fileInfo.getFilePath, writer)
+        (fileInfo, writer)
+      }
+
+      // Every delete fails: the file stays tracked and each retry deletes 
again.
+      val (retried, writer) = createReplicaFile("0-0-1")
+      when(fs.exists(retried.getDfsPath)).thenReturn(true)
+      Seq(1, 2).foreach { attempt =>
+        intercept[IOException](storageManager.cleanFile(shuffleKey, "0-0-1"))
+        assert(storageManager.getFileInfo(shuffleKey, "0-0-1") != null)
+        verify(fs, times(attempt)).delete(retried.getDfsPath, false)
+      }
+      verify(writer).destroy(any[IOException])
+      assert(!storageManager.hdfsWriters.containsKey(retried.getFilePath))
+      // A leftover sidecar file also keeps it tracked.
+      when(fs.exists(retried.getDfsPath)).thenReturn(false)
+      when(fs.exists(retried.getDfsIndexPath)).thenReturn(true)
+      intercept[IOException](storageManager.cleanFile(shuffleKey, "0-0-1"))
+      assert(storageManager.getFileInfo(shuffleKey, "0-0-1") != null)
+      when(fs.exists(retried.getDfsIndexPath)).thenReturn(false)
+      storageManager.cleanFile(shuffleKey, "0-0-1")
+      assert(storageManager.getFileInfo(shuffleKey, "0-0-1") == null)
+
+      // Shuffle expiry during a failed delete still finds the file and 
deletes the directory.
+      val (expired, _) = createReplicaFile("1-0-1")
+      val expiredShuffles = new util.HashSet(Seq(shuffleKey).asJava)
+      when(fs.exists(expired.getDfsPath)).thenAnswer { (_: InvocationOnMock) =>
+        storageManager.cleanupExpiredShuffleKey(expiredShuffles)
+        true
+      }
+      intercept[IOException](storageManager.cleanFile(shuffleKey, "1-0-1"))
+      verify(fs).delete(expired.getDfsPath.getParent, true)
+    } finally {
+      StorageManager.hadoopFs = hadoopFsBefore
+      FileUtils.deleteDirectory(hdfsDir.toFile)
+    }
+  }
+
   private def genPartitionLocation(epoch: Int, offsets: Array[Long]): 
PartitionLocation = {
     val location: PartitionLocation =
       new PartitionLocation(0, epoch, "localhost", 0, 0, 0, 0, 
PartitionLocation.Mode.PRIMARY)
diff --git 
a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/TierWriterSuite.scala
 
b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/TierWriterSuite.scala
index ba40bd67f7..cb57960677 100644
--- 
a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/TierWriterSuite.scala
+++ 
b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/TierWriterSuite.scala
@@ -17,26 +17,34 @@
 
 package org.apache.celeborn.service.deploy.worker.storage
 
-import java.io.IOException
+import java.io.{IOException, OutputStream}
+import java.net.URI
 import java.nio.file.Files
 import java.util.concurrent.atomic.{AtomicInteger, AtomicLong}
 
+import scala.collection.JavaConverters._
+
 import io.netty.buffer.{AbstractByteBufAllocator, ByteBuf, CompositeByteBuf, 
DuplicatedByteBuf, UnpooledByteBufAllocator}
+import org.apache.commons.io.FileUtils
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.fs.{FileSystem, FSDataOutputStream, Path, 
RawLocalFileSystem}
+import org.apache.hadoop.util.Progressable
 import org.mockito.Mockito
 import org.mockito.MockitoSugar.when
-import org.scalatest.BeforeAndAfterEach
+import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach}
 import org.scalatest.funsuite.AnyFunSuite
 
 import org.apache.celeborn.common.CelebornConf
 import org.apache.celeborn.common.exception.AlreadyClosedException
 import org.apache.celeborn.common.identity.UserIdentifier
-import org.apache.celeborn.common.meta.{DiskFileInfo, MemoryFileInfo, 
ReduceFileMeta}
+import org.apache.celeborn.common.meta.{DiskFileInfo, DiskStatus, 
MemoryFileInfo, ReduceFileMeta}
 import org.apache.celeborn.common.network.util.{NettyUtils, TransportConf}
 import org.apache.celeborn.common.protocol.{PartitionLocation, 
PartitionSplitMode, PartitionType, StorageInfo}
+import org.apache.celeborn.common.util.Utils
 import org.apache.celeborn.service.deploy.worker.WorkerSource
 import org.apache.celeborn.service.deploy.worker.memory.MemoryManager
 
-class TierWriterSuite extends AnyFunSuite with BeforeAndAfterEach {
+class TierWriterSuite extends AnyFunSuite with BeforeAndAfterAll with 
BeforeAndAfterEach {
   private class LargeReadableByteBuf(buffer: ByteBuf, virtualSize: Int)
     extends DuplicatedByteBuf(buffer) {
     override def capacity(): Int = virtualSize
@@ -586,4 +594,150 @@ class TierWriterSuite extends AnyFunSuite with 
BeforeAndAfterEach {
       }
     }
   }
+
+  /** A DFS flusher the suite can stop; production flushers live as long as 
the worker. */
+  private class StoppableDfsFlusher(conf: CelebornConf) extends Flusher(
+      new WorkerSource(conf),
+      1,
+      NettyUtils.getByteBufAllocator(new TransportConf("test", conf), null, 
true),
+      256,
+      null,
+      "HDFS",
+      false,
+      conf.workerHdfsFlusherBufferSize) {
+    override def processIOException(e: IOException, deviceErrorType: 
DiskStatus): Unit = {}
+
+    override def getFlushTimeMetric(): String = 
WorkerSource.FLUSH_HDFS_DATA_TIME
+
+    def stop(): Unit = {
+      stopFlag.set(true)
+      workers.foreach(_.shutdownNow())
+      bufferQueue.asScala.foreach(_.release())
+      bufferQueue.clear()
+    }
+  }
+
+  private lazy val dfsFlusher = new StoppableDfsFlusher(new CelebornConf())
+
+  override def afterAll(): Unit = {
+    dfsFlusher.stop()
+    super.afterAll()
+  }
+
+  /** Runs `body` on a DFS writer holding 1024 bytes; creating a file ending 
in `failOn` fails. */
+  private def withDfsTierWriter(failOn: String = null, failAppendClose: 
Boolean = false)(
+      body: (DfsTierWriter, DiskFileInfo, FileSystem) => Unit): Unit = {
+    val conf = new CelebornConf()
+      
.set(CelebornConf.WORKER_WRITER_HDFS_CREATE_AUXILIARY_FILE_MAX_RETRIES.key, "0")
+      .set(CelebornConf.WORKER_REUSE_HDFS_OUTPUT_STREAM_ENABLED.key, "false")
+    val fs = new RawLocalFileSystem() {
+      override def create(
+          f: Path,
+          overwrite: Boolean,
+          bufferSize: Int,
+          replication: Short,
+          blockSize: Long,
+          progress: Progressable): FSDataOutputStream = {
+        if (failOn != null && f.getName.endsWith(failOn)) throw new 
IOException("injected")
+        super.create(f, overwrite, bufferSize, replication, blockSize, 
progress)
+      }
+
+      // Drops the appended bytes and fails on close, like an append that 
never persisted.
+      override def append(f: Path, bufferSize: Int, progress: Progressable): 
FSDataOutputStream =
+        if (!failAppendClose) super.append(f, bufferSize, progress)
+        else new FSDataOutputStream(
+          new OutputStream {
+            override def write(b: Int): Unit = {}
+            override def close(): Unit = throw new IOException("append close 
failed")
+          },
+          null)
+    }
+    fs.initialize(URI.create("file:///"), new Configuration())
+    val hadoopFsBefore = StorageManager.hadoopFs
+    StorageManager.hadoopFs = Map[StorageInfo.Type, 
FileSystem](StorageInfo.Type.HDFS -> fs).asJava
+    val dir = Files.createTempDirectory("celeborn-dfs")
+    try {
+      val userIdentifier = UserIdentifier("`aa`.`bb`")
+      val fileInfo = new DiskFileInfo(
+        userIdentifier,
+        false,
+        new ReduceFileMeta(conf.shuffleChunkSize),
+        dir.resolve("1-0-0").toUri.toString,
+        StorageInfo.Type.HDFS)
+      val writerContext = new PartitionDataWriterContext(
+        1024 * 1024 * 1024,
+        PartitionSplitMode.SOFT,
+        false,
+        new PartitionLocation(1, 0, "host", 1111, 1112, 1113, 1114, 
PartitionLocation.Mode.PRIMARY),
+        "app1-1",
+        1,
+        userIdentifier,
+        PartitionType.REDUCE,
+        false,
+        false)
+      val storageManager: StorageManager = 
Mockito.mock(classOf[StorageManager])
+      MemoryManager.initialize(conf, storageManager, null)
+      val dfsTierWriter = new DfsTierWriter(
+        conf,
+        new ReducePartitionMetaHandler(conf.shuffleRangeReadFilterEnabled, 
fileInfo),
+        new AtomicInteger(),
+        new FlushNotifier(),
+        dfsFlusher,
+        new WorkerSource(conf),
+        fileInfo,
+        StorageInfo.Type.HDFS,
+        writerContext,
+        storageManager)
+      val buf = 
WriterUtils.generateSparkFormatData(UnpooledByteBufAllocator.DEFAULT, 0)
+      dfsTierWriter.numPendingWrites.incrementAndGet()
+      dfsTierWriter.write(buf)
+      buf.release()
+      body(dfsTierWriter, fileInfo, fs)
+    } finally {
+      StorageManager.hadoopFs = hadoopFsBefore
+      fs.close()
+      FileUtils.deleteDirectory(dir.toFile)
+    }
+  }
+
+  test("dfs tier writer keeps its copy when the peer success file exists") {
+    withDfsTierWriter() { (dfsTierWriter, fileInfo, fs) =>
+      val dataPath = fileInfo.getDfsPath.toString
+      fs.create(new Path(dataPath.dropRight(1) + "1" + 
Utils.SUFFIX_HDFS_WRITE_SUCCESS)).close()
+
+      assert(dfsTierWriter.close() === 1024)
+      assert(fs.getFileStatus(fileInfo.getDfsPath).getLen === 1024)
+      assert(fs.exists(fileInfo.getDfsIndexPath))
+      assert(fs.exists(fileInfo.getDfsWriterSuccessPath))
+    }
+  }
+
+  test("dfs tier writer fails close when its index cannot be written") {
+    withDfsTierWriter(failOn = Utils.INDEX_SUFFIX) { (dfsTierWriter, fileInfo, 
fs) =>
+      intercept[IOException](dfsTierWriter.close())
+      assert(!fs.exists(fileInfo.getDfsWriterSuccessPath))
+    }
+  }
+
+  test("dfs tier writer reports the flush failure when its index also fails") {
+    withDfsTierWriter(failOn = Utils.INDEX_SUFFIX) { (dfsTierWriter, _, _) =>
+      dfsTierWriter.notifier.setException(new IOException("flush failed"))
+      assert(intercept[IOException](dfsTierWriter.close()).getMessage === 
"flush failed")
+    }
+  }
+
+  test("dfs tier writer fails close when its append stream cannot be closed") {
+    // Stream reuse is off, so each flush appends and closes its own stream.
+    withDfsTierWriter(failAppendClose = true) { (dfsTierWriter, fileInfo, fs) 
=>
+      assert(intercept[IOException](dfsTierWriter.close()).getMessage === 
"append close failed")
+      assert(fs.getFileStatus(fileInfo.getDfsPath).getLen === 0)
+    }
+  }
+
+  test("dfs tier writer still commits when only the success file cannot be 
written") {
+    withDfsTierWriter(failOn = Utils.SUFFIX_HDFS_WRITE_SUCCESS) { 
(dfsTierWriter, fileInfo, fs) =>
+      assert(dfsTierWriter.close() === 1024)
+      assert(fs.exists(fileInfo.getDfsIndexPath))
+    }
+  }
 }

Reply via email to