This is an automated email from the ASF dual-hosted git repository.
RexXiong 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 90b9b44eb [CELEBORN-2351] Partition file sorting should only be paused
for PUSH_AND_REPLICATE_PAUSED
90b9b44eb is described below
commit 90b9b44eb18228a580b5fbce9293e2d750f01d2f
Author: Sanskar Modi <[email protected]>
AuthorDate: Thu Jun 11 22:38:03 2026 +0800
[CELEBORN-2351] Partition file sorting should only be paused for
PUSH_AND_REPLICATE_PAUSED
### What changes were proposed in this pull request?
Partition file sorting should only be paused for
`PUSH_AND_REPLICATE_PAUSED`, which represent very high memory pressure and
cause OOM for workers. Sorting should be allowed for `PUSH_PAUSED` state.
### Why are the changes needed?
Currently even for push pause state we stop the sorting for partition
files. If pause is sustained for a longer time then sorting can timeout and
reader waiting for sorting will fail or be delayed.
### Does this PR resolve a correctness bug?
- [ ] Yes
### Does this PR introduce _any_ user-facing change?
- [X] Yes
### How was this patch tested?
Existing UTs
Closes #3720 from s0nskar/sort_memory_ready.
Authored-by: Sanskar Modi <[email protected]>
Signed-off-by: Shuang <[email protected]>
---
.../deploy/worker/memory/MemoryManager.java | 5 +-
.../service/deploy/memory/MemoryManagerSuite.scala | 94 ++++++++++++++++++++++
2 files changed, 97 insertions(+), 2 deletions(-)
diff --git
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/memory/MemoryManager.java
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/memory/MemoryManager.java
index 0bdf20a75..2f73d62c8 100644
---
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/memory/MemoryManager.java
+++
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/memory/MemoryManager.java
@@ -72,7 +72,7 @@ public class MemoryManager {
private final AtomicLong diskBufferCounter = new AtomicLong(0);
private final LongAdder pausePushDataCounter = new LongAdder();
private final LongAdder pausePushDataAndReplicateCounter = new LongAdder();
- public ServingState servingState = ServingState.NONE_PAUSED;
+ public volatile ServingState servingState = ServingState.NONE_PAUSED;
private long pausePushDataStartTime = -1L;
private long pausePushDataTime = 0L;
private long pausePushDataAndReplicateStartTime = -1L;
@@ -435,7 +435,8 @@ public class MemoryManager {
public boolean sortMemoryReady() {
return maxSortMemory == 0
- || (servingState == ServingState.NONE_PAUSED &&
sortMemoryCounter.get() < maxSortMemory);
+ || (servingState != ServingState.PUSH_AND_REPLICATE_PAUSED
+ && sortMemoryCounter.get() < maxSortMemory);
}
public void releaseSortMemory(long size) {
diff --git
a/worker/src/test/scala/org/apache/celeborn/service/deploy/memory/MemoryManagerSuite.scala
b/worker/src/test/scala/org/apache/celeborn/service/deploy/memory/MemoryManagerSuite.scala
index 6362273d1..0b95384bc 100644
---
a/worker/src/test/scala/org/apache/celeborn/service/deploy/memory/MemoryManagerSuite.scala
+++
b/worker/src/test/scala/org/apache/celeborn/service/deploy/memory/MemoryManagerSuite.scala
@@ -346,6 +346,100 @@ class MemoryManagerSuite extends CelebornFunSuite {
MemoryManager.reset()
}
+ test("sortMemoryReady allows sorting in PUSH_PAUSED but blocks in
PUSH_AND_REPLICATE_PAUSED") {
+ val conf = new CelebornConf()
+ // Disable the automatic check thread so we drive state transitions
manually
+ conf.set(CelebornConf.WORKER_DIRECT_MEMORY_CHECK_INTERVAL.key, "300s")
+ conf.set(CelebornConf.WORKER_PINNED_MEMORY_CHECK_INTERVAL.key, "0")
+ conf.set(CelebornConf.WORKER_PINNED_MEMORY_CHECK_ENABLED.key, "false")
+ val memoryManager = MockitoSugar.spy(MemoryManager.initialize(conf))
+ val maxDirectMemory = memoryManager.maxDirectMemory
+ val pushThreshold =
+ (conf.workerDirectMemoryRatioToPauseReceive *
maxDirectMemory).longValue()
+ val replicateThreshold =
+ (conf.workerDirectMemoryRatioToPauseReplicate *
maxDirectMemory).longValue()
+ val maxSortMemory =
+ (conf.workerPartitionSorterDirectMemoryRatioThreshold *
maxDirectMemory).longValue()
+ val sortMemoryCounter = memoryManager.getSortMemoryCounter
+
+ Mockito.when(memoryManager.getNettyPinnedDirectMemory).thenReturn(0L)
+
+ // NONE_PAUSED: sort is allowed
+ Mockito.when(memoryManager.getMemoryUsage).thenReturn(0L)
+ memoryManager.switchServingState()
+ assert(memoryManager.servingState == ServingState.NONE_PAUSED)
+ sortMemoryCounter.set(0)
+ assert(memoryManager.sortMemoryReady())
+
+ // PUSH_PAUSED: sort must be allowed so that fetch reads can proceed while
push is
+ // back-pressured (previously sorting was also blocked in this state).
+ Mockito.when(memoryManager.getMemoryUsage).thenReturn(pushThreshold + 1)
+ memoryManager.switchServingState()
+ assert(memoryManager.servingState == ServingState.PUSH_PAUSED)
+ sortMemoryCounter.set(0)
+ assert(
+ memoryManager.sortMemoryReady(),
+ "sortMemoryReady must return true in PUSH_PAUSED: fetch reads of
already-written " +
+ "data should not be blocked by push back-pressure")
+
+ // PUSH_PAUSED but sort budget exhausted: sort must be blocked
+ sortMemoryCounter.set(maxSortMemory)
+ assert(!memoryManager.sortMemoryReady())
+
+ // PUSH_AND_REPLICATE_PAUSED: sort must be blocked regardless of sort
budget
+ Mockito.when(memoryManager.getMemoryUsage).thenReturn(replicateThreshold +
1)
+ memoryManager.switchServingState()
+ assert(memoryManager.servingState ==
ServingState.PUSH_AND_REPLICATE_PAUSED)
+ sortMemoryCounter.set(0)
+ assert(
+ !memoryManager.sortMemoryReady(),
+ "sortMemoryReady must return false in PUSH_AND_REPLICATE_PAUSED")
+ MemoryManager.reset()
+ }
+
+ test("sortMemoryReady always returns true when sort memory threshold is
disabled") {
+ val conf = new CelebornConf()
+ conf.set(CelebornConf.WORKER_DIRECT_MEMORY_CHECK_INTERVAL.key, "300s")
+ conf.set(CelebornConf.WORKER_PINNED_MEMORY_CHECK_INTERVAL.key, "0")
+ conf.set(CelebornConf.WORKER_PINNED_MEMORY_CHECK_ENABLED.key, "false")
+ conf.set(
+ CelebornConf.WORKER_PARTITION_SORTER_DIRECT_MEMORY_RATIO_THRESHOLD.key,
+ "0")
+ val memoryManager = MockitoSugar.spy(MemoryManager.initialize(conf))
+ val maxDirectMemory = memoryManager.maxDirectMemory
+ val pushThreshold =
+ (conf.workerDirectMemoryRatioToPauseReceive *
maxDirectMemory).longValue()
+ val replicateThreshold =
+ (conf.workerDirectMemoryRatioToPauseReplicate *
maxDirectMemory).longValue()
+
+ Mockito.when(memoryManager.getNettyPinnedDirectMemory).thenReturn(0L)
+
+ // PUSH_PAUSED, threshold=0 means skip all checks
+ Mockito.when(memoryManager.getMemoryUsage).thenReturn(0L)
+ memoryManager.switchServingState()
+ assert(memoryManager.servingState == ServingState.NONE_PAUSED)
+ assert(
+ memoryManager.sortMemoryReady(),
+ "sortMemoryReady must return true when threshold is 0 regardless of
serving state")
+
+ // PUSH_PAUSED, threshold=0 means skip all checks
+ Mockito.when(memoryManager.getMemoryUsage).thenReturn(pushThreshold + 1)
+ memoryManager.switchServingState()
+ assert(memoryManager.servingState == ServingState.PUSH_PAUSED)
+ assert(
+ memoryManager.sortMemoryReady(),
+ "sortMemoryReady must return true when threshold is 0 regardless of
serving state")
+
+ // PUSH_AND_REPLICATE_PAUSED, threshold=0 means skip all checks
+ Mockito.when(memoryManager.getMemoryUsage).thenReturn(replicateThreshold +
1)
+ memoryManager.switchServingState()
+ assert(memoryManager.servingState ==
ServingState.PUSH_AND_REPLICATE_PAUSED)
+ assert(
+ memoryManager.sortMemoryReady(),
+ "sortMemoryReady must return true when threshold is 0 regardless of
serving state")
+ MemoryManager.reset()
+ }
+
class MockMemoryPressureListener(
val belongModuleName: String,
var isPause: Boolean = false) extends MemoryPressureListener {
