This is an automated email from the ASF dual-hosted git repository.
zaynt4606 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 5099daad77 [CELEBORN-2490] Make memory file storage configurable for
replica partitions
5099daad77 is described below
commit 5099daad77b0c8aed6f5a744a96eef4aa70a4f70
Author: Venkata Krishnan Sowrirajan <[email protected]>
AuthorDate: Thu Oct 22 10:19:34 2026 +0800
[CELEBORN-2490] Make memory file storage configurable for replica partitions
### What changes were proposed in this pull request?
Add a worker config `celeborn.worker.memoryFileStorage.replica.enabled`
(default `true`) to preserve existing behavior. When explicitly disabled,
`StoragePolicy` skips `MEMORY` for replica partitions and falls through to the
next eligible storage type in
`celeborn.worker.storage.storagePolicy.createFilePolicy`. The setting also
applies when a custom `celeborn.worker.storage.storagePolicy.evictPolicy`
targets MEMORY for replicas. Primary partitions are unaffected.
JIRA: [CELEBORN-2490](https://issues.apache.org/jira/browse/CELEBORN-2490)
### Why are the changes needed?
When memory file storage and replication are enabled, both primary and
replica partitions can consume the memory file storage budget, which also
contributes to push-pause thresholds. Replica data is normally read only when
fetching from the primary fails. Allowing operators to keep replicas out of
memory leaves more memory available for primary data and can reduce eviction
and push-pause pressure, at the cost of additional non-memory storage I/O and
potentially higher replica-fetch la [...]
### Does this PR resolve a correctness bug?
- [ ] Yes
### Does this PR introduce _any_ user-facing change?
- [x] Yes
Adds `celeborn.worker.memoryFileStorage.replica.enabled`, defaulting to
`true` for backward compatibility. Existing clusters retain their current
storage-placement behavior after upgrade. Set it to `false` to opt out of
memory file storage for replica partitions, including eviction policies that
target MEMORY.
### How was this patch tested?
- `StoragePolicyReplicaSuite` uses explicit memory/local-disk availability
and covers primary behavior, replicas with the default and explicit true/false
settings, and custom eviction policies targeting MEMORY.
- `build/mvn -B -q test -pl worker -am
-DwildcardSuites=org.apache.celeborn.service.deploy.worker.storage.storagePolicy,org.apache.celeborn.ConfigurationSuite
-Dtest=none` (JDK 17): passed, 13 storage-policy tests and 8
configuration-documentation tests.
- `build/mvn -B -q spotless:check -pl common,worker`: passed.
Closes #3859 from venkata91/CELEBORN-replica-skip-memory-tier.
Authored-by: Venkata Krishnan Sowrirajan <[email protected]>
Signed-off-by: zhengtao <[email protected]>
Co-authored-by: Copilot <[email protected]>
Co-Authored-By: Codex <[email protected]>
AI-Model: gpt-6
AI-Contributed/Feature: 0/27
AI-Contributed/UT: 0/137
---
.../org/apache/celeborn/common/CelebornConf.scala | 16 +++
docs/configuration/worker.md | 1 +
.../deploy/worker/storage/StoragePolicy.scala | 10 +-
.../storagePolicy/StoragePolicyReplicaSuite.scala | 137 +++++++++++++++++++++
4 files changed, 162 insertions(+), 2 deletions(-)
diff --git
a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
index 32f120d4a1..5f535fb344 100644
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@@ -1477,6 +1477,8 @@ class CelebornConf(loadDefaults: Boolean) extends
Cloneable with Logging with Se
get(WORKER_MEMORY_FILE_STORAGE_EVICT_AGGRESSIVE_MODE_ENABLED)
def workerMemoryFileStorageEvictRatio: Double =
get(WORKER_MEMORY_FILE_STORAGE_EVICT_RATIO)
+ def workerMemoryFileStorageReplicaEnabled: Boolean =
+ get(WORKER_MEMORY_FILE_STORAGE_REPLICA_ENABLED)
def workerPushDataMergeBufferEnabled: Boolean =
get(WORKER_PUSH_DATA_MERGE_BUFFER_ENABLED)
def workerDirectMemoryRatioToMergeBuffer: Double =
get(WORKER_DIRECT_MEMORY_RATIO_TO_MERGE_BUFFER)
@@ -4442,6 +4444,20 @@ object CelebornConf extends Logging {
.doubleConf
.createWithDefault(0.5)
+ val WORKER_MEMORY_FILE_STORAGE_REPLICA_ENABLED: ConfigEntry[Boolean] =
+ buildConf("celeborn.worker.memoryFileStorage.replica.enabled")
+ .categories("worker")
+ .doc("Whether replica partitions can be stored in memory file storage.
Enabled by default " +
+ "to preserve existing behavior. When disabled, replica partitions skip
MEMORY and " +
+ "fall through to the next eligible storage type in " +
+ "`celeborn.worker.storage.storagePolicy.createFilePolicy`, leaving
memory for " +
+ "primary partitions. This also applies when " +
+ "`celeborn.worker.storage.storagePolicy.evictPolicy` targets MEMORY
for replicas. " +
+ "Primary partitions are unaffected.")
+ .version("1.0.0")
+ .booleanConf
+ .createWithDefault(true)
+
val WORKER_PUSH_DATA_MERGE_BUFFER_ENABLED: ConfigEntry[Boolean] =
buildConf("celeborn.worker.pushdata.mergeBuffer.enabled")
.categories("worker")
diff --git a/docs/configuration/worker.md b/docs/configuration/worker.md
index dd0aa15a55..51bd1bdfe7 100644
--- a/docs/configuration/worker.md
+++ b/docs/configuration/worker.md
@@ -149,6 +149,7 @@ license: |
| celeborn.worker.memoryFileStorage.evict.aggressiveMode.enabled | false |
false | If this set to true, memory shuffle files will be evicted when worker
is in PAUSED state. If the worker's offheap memory is not ample, set this to
true and decrease `celeborn.worker.directMemoryRatioForMemoryFileStorage` will
be helpful. | 0.5.1 | |
| celeborn.worker.memoryFileStorage.evict.ratio | 0.5 | false | If memory
shuffle storage usage rate is above this config, the memory storage shuffle
files will evict to free memory. | 0.5.1 | |
| celeborn.worker.memoryFileStorage.maxFileSize | 8MB | false | Max size for a
memory storage file. It must be less than 2GB. | 0.5.0 | |
+| celeborn.worker.memoryFileStorage.replica.enabled | true | false | Whether
replica partitions can be stored in memory file storage. Enabled by default to
preserve existing behavior. When disabled, replica partitions skip MEMORY and
fall through to the next eligible storage type in
`celeborn.worker.storage.storagePolicy.createFilePolicy`, leaving memory for
primary partitions. This also applies when
`celeborn.worker.storage.storagePolicy.evictPolicy` targets MEMORY for
replicas. Primary [...]
| celeborn.worker.monitor.disk.check.interval | 30s | false | Intervals
between device monitor to check disk. | 0.3.0 |
celeborn.worker.monitor.disk.checkInterval |
| celeborn.worker.monitor.disk.check.timeout | 30s | false | Timeout time for
worker check device status. | 0.3.0 | celeborn.worker.disk.check.timeout |
| celeborn.worker.monitor.disk.checklist | readwrite,diskusage | false |
Monitor type for disk, available items are: iohang, readwrite and diskusage. |
0.2.0 | |
diff --git
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StoragePolicy.scala
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StoragePolicy.scala
index c7d8b4d807..7dbcee46ee 100644
---
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StoragePolicy.scala
+++
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StoragePolicy.scala
@@ -24,7 +24,7 @@ import
org.apache.celeborn.common.exception.CelebornIOException
import org.apache.celeborn.common.internal.Logging
import org.apache.celeborn.common.meta.{DiskFileInfo, FileInfo}
import org.apache.celeborn.common.metrics.source.AbstractSource
-import org.apache.celeborn.common.protocol.{PartitionType, StorageInfo}
+import org.apache.celeborn.common.protocol.{PartitionLocation, PartitionType,
StorageInfo}
import org.apache.celeborn.service.deploy.worker.memory.MemoryManager
class StoragePolicy(conf: CelebornConf, storageManager: StorageManager,
source: AbstractSource)
@@ -32,6 +32,8 @@ class StoragePolicy(conf: CelebornConf, storageManager:
StorageManager, source:
private val createFileOrder: Option[List[String]] =
conf.workerStoragePolicyCreateFilePolicy
private val evictFileOrder: Option[Map[String, List[String]]] =
conf.workerStoragePolicyEvictFilePolicy
+ private val memoryFileStorageReplicaEnabled: Boolean =
+ conf.workerMemoryFileStorageReplicaEnabled
def getEvictedFileWriter(
celebornFile: TierWriterBase,
@@ -102,7 +104,11 @@ class StoragePolicy(conf: CelebornConf, storageManager:
StorageManager, source:
try {
storageInfoType match {
case StorageInfo.Type.MEMORY =>
- if (location.getStorageInfo.memoryAvailable() &&
MemoryManager.instance().memoryFileStorageAvailable()) {
+ if (!memoryFileStorageReplicaEnabled
+ && location.getMode == PartitionLocation.Mode.REPLICA) {
+ logDebug(s"Skip memory file for replica
${partitionDataWriterContext.getShuffleKey}
${partitionDataWriterContext.getPartitionLocation.getFileName}")
+ null
+ } else if (location.getStorageInfo.memoryAvailable() &&
MemoryManager.instance().memoryFileStorageAvailable()) {
logDebug(s"Create memory file for
${partitionDataWriterContext.getShuffleKey}
${partitionDataWriterContext.getPartitionLocation.getFileName}")
val memoryFileInfo = storageManager.createMemoryFileInfo(
partitionDataWriterContext.getAppId,
diff --git
a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/storagePolicy/StoragePolicyReplicaSuite.scala
b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/storagePolicy/StoragePolicyReplicaSuite.scala
new file mode 100644
index 0000000000..77bb47524c
--- /dev/null
+++
b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/storagePolicy/StoragePolicyReplicaSuite.scala
@@ -0,0 +1,137 @@
+/*
+ * 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.worker.storage.storagePolicy
+
+import java.io.File
+import java.util.concurrent.atomic.AtomicInteger
+
+import io.netty.buffer.UnpooledByteBufAllocator
+import org.mockito.ArgumentMatchers.any
+import org.mockito.Mockito
+import org.mockito.MockitoSugar.{mock, when}
+
+import org.apache.celeborn.CelebornFunSuite
+import org.apache.celeborn.common.CelebornConf
+import org.apache.celeborn.common.meta.{DiskFileInfo, MemoryFileInfo}
+import org.apache.celeborn.common.metrics.source.AbstractSource
+import org.apache.celeborn.common.protocol.{PartitionLocation, PartitionType,
StorageInfo}
+import org.apache.celeborn.service.deploy.worker.memory.MemoryManager
+import org.apache.celeborn.service.deploy.worker.storage._
+
+class StoragePolicyReplicaSuite extends CelebornFunSuite {
+ val mockedStorageManager: StorageManager = mock[StorageManager]
+ val mockedSource: AbstractSource = mock[AbstractSource]
+
+ when(
+ mockedStorageManager.createMemoryFileInfo(any(), any(), any(), any(),
any(), any())).thenAnswer(
+ mock[MemoryFileInfo])
+
when(mockedStorageManager.storageBufferAllocator).thenAnswer(UnpooledByteBufAllocator.DEFAULT)
+ when(mockedStorageManager.localOrDfsStorageAvailable).thenAnswer(true)
+
+ val mockedDiskFile: DiskFileInfo = mock[DiskFileInfo]
+ when(mockedDiskFile.getStorageType).thenAnswer(StorageInfo.Type.SSD)
+ when(
+ mockedStorageManager.createDiskFile(
+ any(),
+ any(),
+ any(),
+ any(),
+ any(),
+ any(),
+ any(),
+ any())).thenAnswer((mock[LocalFlusher], mockedDiskFile, mock[File]))
+
+ private def createWriter(
+ mode: PartitionLocation.Mode,
+ replicaEnabled: Option[Boolean],
+ evict: Boolean = false): TierWriterBase = {
+ val context = mock[PartitionDataWriterContext]
+ val location = new PartitionLocation(1, 1, "h1", 1, 2, 3, 4, mode)
+ location.setStorageInfo(new StorageInfo(
+ "mountpoint",
+ StorageInfo.Type.SSD,
+ StorageInfo.MEMORY_MASK | StorageInfo.LOCAL_DISK_MASK))
+ when(context.getPartitionLocation).thenAnswer(location)
+ when(context.getPartitionType).thenAnswer(PartitionType.REDUCE)
+ val conf = new CelebornConf()
+ conf.set("celeborn.worker.storage.storagePolicy.createFilePolicy",
"MEMORY,SSD,HDD,HDFS,OSS,S3")
+ conf.set(CelebornConf.WORKER_STORAGE_EVICT_POLICY.key, "SSD,MEMORY,HDD")
+ replicaEnabled.foreach(v =>
+ conf.set(CelebornConf.WORKER_MEMORY_FILE_STORAGE_REPLICA_ENABLED.key,
v.toString))
+ val storagePolicy = new StoragePolicy(conf, mockedStorageManager,
mockedSource)
+ val memoryManager = mock[MemoryManager]
+ when(memoryManager.memoryFileStorageAvailable()).thenReturn(true)
+ val memoryManagerStatic = Mockito.mockStatic(classOf[MemoryManager])
+ try {
+ memoryManagerStatic.when[MemoryManager](() =>
MemoryManager.instance()).thenReturn(
+ memoryManager)
+ if (evict) {
+ val writer = mock[LocalTierWriter]
+ when(writer.storageType).thenAnswer(StorageInfo.Type.SSD)
+ storagePolicy.getEvictedFileWriter(
+ writer,
+ context,
+ new AtomicInteger(),
+ new FlushNotifier)
+ } else {
+ storagePolicy.createFileWriter(context, new AtomicInteger(), new
FlushNotifier)
+ }
+ } finally {
+ memoryManagerStatic.close()
+ }
+ }
+
+ test("primary partition uses memory file storage") {
+ assert(createWriter(PartitionLocation.Mode.PRIMARY,
None).isInstanceOf[MemoryTierWriter])
+ }
+
+ test("primary partition uses memory file storage when replica memory storage
is disabled") {
+ assert(createWriter(PartitionLocation.Mode.PRIMARY,
Some(false)).isInstanceOf[MemoryTierWriter])
+ }
+
+ test("replica partition uses memory file storage by default") {
+ assert(createWriter(PartitionLocation.Mode.REPLICA,
None).isInstanceOf[MemoryTierWriter])
+ }
+
+ test("replica partition skips memory file storage when disabled") {
+ assert(createWriter(PartitionLocation.Mode.REPLICA,
Some(false)).isInstanceOf[LocalTierWriter])
+ }
+
+ test("replica partition uses memory file storage when enabled") {
+ assert(
+ createWriter(PartitionLocation.Mode.REPLICA,
Some(true)).isInstanceOf[MemoryTierWriter])
+ }
+
+ test("replica eviction can target memory file storage by default") {
+ assert(
+ createWriter(PartitionLocation.Mode.REPLICA, None, evict = true)
+ .isInstanceOf[MemoryTierWriter])
+ }
+
+ test("replica eviction skips memory file storage when disabled") {
+ assert(
+ createWriter(PartitionLocation.Mode.REPLICA, Some(false), evict = true)
+ .isInstanceOf[LocalTierWriter])
+ }
+
+ test("primary eviction can target memory when replica memory storage is
disabled") {
+ assert(
+ createWriter(PartitionLocation.Mode.PRIMARY, Some(false), evict = true)
+ .isInstanceOf[MemoryTierWriter])
+ }
+}
