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))
+ }
+ }
}