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 f9ecab863a [CELEBORN-2498][TEST] Fix testDataPushQueue starvation on
fast machines
f9ecab863a is described below
commit f9ecab863a4722107e0a51046b1ec931f6eb7913
Author: yew1eb <[email protected]>
AuthorDate: Thu Oct 22 11:05:25 2026 +0800
[CELEBORN-2498][TEST] Fix testDataPushQueue starvation on fast machines
### What changes were proposed in this pull request?
Move the `pushState.addBatch` call in
`DataPushQueueSuiteJ.testDataPushQueue` from the writer-thread loop into the
mocked `pushData`, mirroring `ShuffleClientImpl#pushData`, which registers a
batch as in-flight when it is actually pushed:
```java
Override
protected void pushData(PushTask task) throws IOException {
...
pushState.addBatch(
partitionBatchIdMap.get(partitionId),
task.getSize(),
reducePartitionMap.get(partitionId).hostAndPushPort());
pushState.removeBatch(
partitionBatchIdMap.get(partitionId),
reducePartitionMap.get(partitionId).hostAndPushPort());
}
```
### Why are the changes needed?
The test pre-registers all `numPartitions = 1000000` batches on the writer
thread before enqueueing them. On machines where the writer thread outruns the
push thread, per-worker in-flight counts (limit 2,
`celeborn.client.push.maxReqsInFlightPerWorker`) exceed the cap, and
`DataPushQueue.takePushTasks` degrades to its escape hatch: 1 task per worker
per `takeTaskWaitIntervalMs` (50ms). With 30 workers that is ~500 tasks/s, so
the suite takes ~30+ minutes and looks like a hang.
On CI's slower runners the push thread keeps ahead of the writer, the suite
takes the fast path, and it passes in seconds — which is why this has gone
unnoticed since [CELEBORN-875] raised `numPartitions` to 1,000,000.
Measured locally (JDK 8 + macOS, M-series) before the fix:
- throughput ~500 tasks/s (3.5 min to advance 100k tasks),
- writer steadily 512 tasks (= `celeborn.client.push.queue.capacity`) ahead,
- heap healthy; thread dump shows the test thread parked in
`DataPusher.addTask`'s `idleQueue.poll` while the push thread sleeps in
`takePushTasks`.
### Does this PR resolve a correctness bug?
- [ ] Yes
### Does this PR introduce _any_ user-facing change?
- [ ] Yes
No, test-only change.
### How was this patch tested?
```
./build/sbt "celeborn-client/clean" \
"celeborn-client/testOnly
org.apache.celeborn.client.write.DataPushQueueSuiteJ"
```
Before: ~500 tasks/s, run not finished after 30+ minutes. After: suite
completes in 2.5s / 4.4s across two runs, 2/2 tests passed.
Closes #3890 from yew1eb/CELEBORN-data-push-queue-hang.
Authored-by: yew1eb <[email protected]>
Signed-off-by: zhengtao <[email protected]>
Co-Authored-By: Codex <[email protected]>
AI-Model: gpt-6
AI-Contributed/Feature: 0/0
AI-Contributed/UT: 0/9
---
.../org/apache/celeborn/client/write/DataPushQueueSuiteJ.java | 9 ++++++++-
1 file changed, 8 insertions(+), 1 deletion(-)
diff --git
a/client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuiteJ.java
b/client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuiteJ.java
index 841e04eee0..ecee79387c 100644
---
a/client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuiteJ.java
+++
b/client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuiteJ.java
@@ -115,6 +115,14 @@ public class DataPushQueueSuiteJ {
byte[] buffer = task.getBuffer();
int partitionId = task.getPartitionId();
tarWorkerData.get(partitionId % numWorker).add(bytesToInt(buffer));
+ // Mirrors ShuffleClientImpl#pushData: the batch is registered as
in-flight when it
+ // is actually pushed, not when the writer thread enqueues it.
Pre-registering all
+ // batches on the writer side inflates per-worker in-flight counts
past
+ // celeborn.client.push.maxReqsInFlightPerWorker and starves
takePushTasks.
+ pushState.addBatch(
+ partitionBatchIdMap.get(partitionId),
+ task.getSize(),
+ reducePartitionMap.get(partitionId).hostAndPushPort());
pushState.removeBatch(
partitionBatchIdMap.get(partitionId),
reducePartitionMap.get(partitionId).hostAndPushPort());
@@ -124,7 +132,6 @@ public class DataPushQueueSuiteJ {
for (int i = 0; i < numPartitions; i++) {
byte[] b = intToBytes(workerData.get(i % numWorker).get(i / numWorker));
int batchId = pushState.nextBatchId();
- pushState.addBatch(batchId, b.length,
reducePartitionMap.get(i).hostAndPushPort());
partitionBatchIdMap.put(i, batchId);
dataPusher.addTask(i, b, b.length);
}
