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 4a9112a59 [CELEBORN-2315] Add iterator fully-consumed validation after
shuffle write
4a9112a59 is described below
commit 4a9112a594ddc3f55a98172e643af1d08d72740d
Author: James Xu <[email protected]>
AuthorDate: Thu Jun 11 14:30:07 2026 +0800
[CELEBORN-2315] Add iterator fully-consumed validation after shuffle write
### What changes were proposed in this pull request?
Adds a post-write safety check to HashBasedShuffleWriter and
SortBasedShuffleWriter: after the write loop completes, verify the input
iterator was fully consumed. If records remain, kill the task with
TaskKilledException. This guards against silent data loss.
### Why are the changes needed?
It could give another layer of correctness guarantee.
### Does this PR resolve a correctness bug?
Enhance correctness guarantee.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
UT.
Closes #3672 from xumingming/iterator-fully-consumed-check.
Authored-by: James Xu <[email protected]>
Signed-off-by: Nicholas Jiang <[email protected]>
---
.../shuffle/celeborn/HashBasedShuffleWriter.java | 43 +++++++++++++++-------
.../shuffle/celeborn/SortBasedShuffleWriter.java | 42 ++++++++++++++-------
.../apache/spark/shuffle/celeborn/SparkUtils.java | 8 ++++
.../celeborn/CelebornShuffleWriterSuiteBase.java | 20 ++++++++++
.../shuffle/celeborn/HashBasedShuffleWriter.java | 42 ++++++++++++++-------
.../shuffle/celeborn/SortBasedShuffleWriter.java | 24 +++++++++---
.../apache/spark/shuffle/celeborn/SparkUtils.java | 15 ++++++++
.../celeborn/CelebornShuffleWriterSuiteBase.java | 23 ++++++++++++
8 files changed, 173 insertions(+), 44 deletions(-)
diff --git
a/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
b/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
index 4d55c14be..628df43c0 100644
---
a/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
+++
b/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
@@ -166,18 +166,8 @@ public class HashBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
public void write(scala.collection.Iterator<Product2<K, V>> records) throws
IOException {
boolean needCleanupPusher = true;
try {
- if (canUseFastWrite()) {
- fastWrite0(records);
- } else if (dep.mapSideCombine()) {
- if (dep.aggregator().isEmpty()) {
- throw new UnsupportedOperationException(
- "When using map side combine, an aggregator must be specified.");
- }
- write0(dep.aggregator().get().combineValuesByKey(records,
taskContext));
- } else {
- write0(records);
- }
- close();
+ boolean iteratorHasNext = doWrite(records);
+ close(iteratorHasNext);
needCleanupPusher = false;
} catch (InterruptedException e) {
TaskInterruptedHelper.throwTaskKillException();
@@ -188,6 +178,26 @@ public class HashBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
}
}
+ boolean doWrite(scala.collection.Iterator<Product2<K, V>> records)
+ throws IOException, InterruptedException {
+ if (canUseFastWrite()) {
+ fastWrite0(records);
+ return records.hasNext();
+ } else if (dep.mapSideCombine()) {
+ if (dep.aggregator().isEmpty()) {
+ throw new UnsupportedOperationException(
+ "When using map side combine, an aggregator must be specified.");
+ }
+ scala.collection.Iterator<?> combinedIterator =
+ dep.aggregator().get().combineValuesByKey(records, taskContext);
+ write0(combinedIterator);
+ return combinedIterator.hasNext();
+ } else {
+ write0(records);
+ return records.hasNext();
+ }
+ }
+
@VisibleForTesting
boolean canUseFastWrite() {
return unsafeRowFastWrite
@@ -331,7 +341,7 @@ public class HashBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
}
}
- private void close() throws IOException, InterruptedException {
+ private void close(boolean iteratorHasNext) throws IOException,
InterruptedException {
// merge and push residual data to reduce network traffic
// NB: since dataPusher thread have no in-flight data at this point,
// we now push merged data by task thread will not introduce any
contention
@@ -363,9 +373,16 @@ public class HashBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
sendBuffers = null;
sendOffsets = null;
+ // The check must come before mapperEnd so a partial map output is never
committed to the
+ // shuffle service.
+ // Check BEFORE releasing buffers and draining the pusher(the outer
finally block will handle
+ // that).
+ SparkUtils.assertIteratorFullyConsumed(iteratorHasNext);
+
long waitStartTime = System.nanoTime();
dataPusher.waitOnTermination();
sendBufferPool.returnPushTaskQueue(dataPusher.getAndResetIdleQueue());
+
shuffleClient.mapperEnd(shuffleId, mapId, encodedAttemptId, numMappers,
numPartitions);
writeMetrics.incWriteTime(System.nanoTime() - waitStartTime);
diff --git
a/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
b/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
index 2ab3139b5..80e4ee275 100644
---
a/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
+++
b/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
@@ -147,18 +147,8 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
public void write(scala.collection.Iterator<Product2<K, V>> records) throws
IOException {
boolean needCleanupPusher = true;
try {
- if (canUseFastWrite()) {
- fastWrite0(records);
- } else if (dep.mapSideCombine()) {
- if (dep.aggregator().isEmpty()) {
- throw new UnsupportedOperationException(
- "When using map side combine, an aggregator must be specified.");
- }
- write0(dep.aggregator().get().combineValuesByKey(records,
taskContext));
- } else {
- write0(records);
- }
- close();
+ boolean iteratorHasNext = doWrite(records);
+ close(iteratorHasNext);
needCleanupPusher = false;
} finally {
if (needCleanupPusher) {
@@ -167,6 +157,25 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
}
}
+ boolean doWrite(scala.collection.Iterator<Product2<K, V>> records) throws
IOException {
+ if (canUseFastWrite()) {
+ fastWrite0(records);
+ return records.hasNext();
+ } else if (dep.mapSideCombine()) {
+ if (dep.aggregator().isEmpty()) {
+ throw new UnsupportedOperationException(
+ "When using map side combine, an aggregator must be specified.");
+ }
+ scala.collection.Iterator<?> combinedIterator =
+ dep.aggregator().get().combineValuesByKey(records, taskContext);
+ write0(combinedIterator);
+ return combinedIterator.hasNext();
+ } else {
+ write0(records);
+ return records.hasNext();
+ }
+ }
+
@VisibleForTesting
boolean canUseFastWrite() {
return unsafeRowFastWrite
@@ -304,10 +313,17 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
}
}
- private void close() throws IOException {
+ private void close(boolean iteratorHasNext) throws IOException {
logger.info("Memory used {}", Utils.bytesToString(pusher.getUsed()));
long pushStartTime = System.nanoTime();
pusher.pushData(false);
+
+ // The check must come BEFORE mapperEnd so a partial map output is never
committed to the
+ // shuffle service.
+ // Check BEFORE pusher.close() so if we throw here the outer finally block
will handle the
+ // closing thing.
+ SparkUtils.assertIteratorFullyConsumed(iteratorHasNext);
+
pusher.close(true);
writeMetrics.incWriteTime(System.nanoTime() - pushStartTime);
diff --git
a/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java
b/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java
index d04f4b0eb..7703a55da 100644
---
a/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java
+++
b/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java
@@ -63,6 +63,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.exception.CelebornIOException;
import org.apache.celeborn.common.exception.CelebornRuntimeException;
import org.apache.celeborn.common.network.protocol.TransportMessage;
import
org.apache.celeborn.common.protocol.message.ControlMessages.GetReducerFileGroupResponse;
@@ -543,4 +544,11 @@ public class SparkUtils {
return null;
});
}
+
+ public static void assertIteratorFullyConsumed(boolean iteratorHasNext)
throws IOException {
+ if (iteratorHasNext) {
+ throw new CelebornIOException(
+ "Shuffle write task finished but iterator was not fully consumed.");
+ }
+ }
}
diff --git
a/client-spark/spark-2/src/test/java/org/apache/spark/shuffle/celeborn/CelebornShuffleWriterSuiteBase.java
b/client-spark/spark-2/src/test/java/org/apache/spark/shuffle/celeborn/CelebornShuffleWriterSuiteBase.java
index cc0224e79..ab0026023 100644
---
a/client-spark/spark-2/src/test/java/org/apache/spark/shuffle/celeborn/CelebornShuffleWriterSuiteBase.java
+++
b/client-spark/spark-2/src/test/java/org/apache/spark/shuffle/celeborn/CelebornShuffleWriterSuiteBase.java
@@ -43,6 +43,7 @@ import org.apache.spark.ShuffleDependency;
import org.apache.spark.SparkConf;
import org.apache.spark.SparkEnv;
import org.apache.spark.TaskContext;
+import org.apache.spark.TaskContext$;
import org.apache.spark.executor.ShuffleWriteMetrics;
import org.apache.spark.executor.TaskMetrics;
import org.apache.spark.memory.TaskMemoryManager;
@@ -145,6 +146,7 @@ public abstract class CelebornShuffleWriterSuiteBase {
Mockito.doReturn(metrics).when(taskContext).taskMetrics();
Mockito.doReturn(taskMemoryManager).when(taskContext).taskMemoryManager();
+ Mockito.doReturn(None$.MODULE$).when(taskContext).getKillReason();
Mockito.doReturn(bmId).when(blockManager).shuffleServerId();
Mockito.doReturn(blockManager).when(env).blockManager();
@@ -214,6 +216,24 @@ public abstract class CelebornShuffleWriterSuiteBase {
check(2 << 30, conf, serializer);
}
+ @Test
+ public void testAssertIteratorFullyConsumed() {
+ SparkUtils.assertIteratorFullyConsumed(false);
+ }
+
+ @Test
+ public void testAssertIteratorFullyConsumedThrows() {
+ TaskContext$.MODULE$.setTaskContext(taskContext);
+ try {
+ SparkUtils.assertIteratorFullyConsumed(true);
+ fail("Expected IOException when iterator is not fully consumed");
+ } catch (IOException e) {
+ assertTrue(e.getMessage().contains("not fully consumed"));
+ } finally {
+ TaskContext$.MODULE$.setTaskContext(null);
+ }
+ }
+
private void check(
final int approximateSize, final CelebornConf conf, final Serializer
serializer)
throws Exception {
diff --git
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
index 8b454ebd5..a36a0932a 100644
---
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
+++
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
@@ -162,18 +162,8 @@ public class HashBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
public void write(scala.collection.Iterator<Product2<K, V>> records) throws
IOException {
boolean needCleanupPusher = true;
try {
- if (canUseFastWrite()) {
- fastWrite0(records);
- } else if (dep.mapSideCombine()) {
- if (dep.aggregator().isEmpty()) {
- throw new UnsupportedOperationException(
- "When using map side combine, an aggregator must be specified.");
- }
- write0(dep.aggregator().get().combineValuesByKey(records,
taskContext));
- } else {
- write0(records);
- }
- close();
+ boolean iteratorHasNext = doWrite(records);
+ close(iteratorHasNext);
needCleanupPusher = false;
} catch (InterruptedException e) {
TaskInterruptedHelper.throwTaskKillException();
@@ -184,6 +174,26 @@ public class HashBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
}
}
+ boolean doWrite(scala.collection.Iterator<Product2<K, V>> records)
+ throws IOException, InterruptedException {
+ if (canUseFastWrite()) {
+ fastWrite0(records);
+ return records.hasNext();
+ } else if (dep.mapSideCombine()) {
+ if (dep.aggregator().isEmpty()) {
+ throw new UnsupportedOperationException(
+ "When using map side combine, an aggregator must be specified.");
+ }
+ scala.collection.Iterator<?> combinedIterator =
+ dep.aggregator().get().combineValuesByKey(records, taskContext);
+ write0(combinedIterator);
+ return combinedIterator.hasNext();
+ } else {
+ write0(records);
+ return records.hasNext();
+ }
+ }
+
@VisibleForTesting
boolean canUseFastWrite() {
boolean keyIsPartitionId = false;
@@ -366,7 +376,7 @@ public class HashBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
}
}
- private void close() throws IOException, InterruptedException {
+ private void close(boolean iteratorHasNext) throws IOException,
InterruptedException {
// Send the remaining data in sendBuffer
long pushMergedDataTime = System.nanoTime();
closeWrite();
@@ -374,6 +384,12 @@ public class HashBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
writeMetrics.incWriteTime(System.nanoTime() - pushMergedDataTime);
updateRecordsWrittenMetrics();
+ // The check must come before mapperEnd so a partial map output is never
committed to the
+ // shuffle service.
+ // Check BEFORE releasing buffers and draining the pusher(the outer
finally block will handle
+ // that).
+ SparkUtils.assertIteratorFullyConsumed(iteratorHasNext);
+
long waitStartTime = System.nanoTime();
dataPusher.waitOnTermination();
sendBufferPool.returnPushTaskQueue(dataPusher.getAndResetIdleQueue());
diff --git
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
index ebc1d44ea..d08b36d19 100644
---
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
+++
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
@@ -212,17 +212,24 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
return peakMemoryUsedBytes;
}
- void doWrite(scala.collection.Iterator<Product2<K, V>> records) throws
IOException {
+ // Returns true if the iterator still has records (i.e., not fully consumed)
+ @VisibleForTesting
+ boolean doWrite(scala.collection.Iterator<Product2<K, V>> records) throws
IOException {
if (canUseFastWrite()) {
fastWrite0(records);
+ return records.hasNext();
} else if (dep.mapSideCombine()) {
if (dep.aggregator().isEmpty()) {
throw new UnsupportedOperationException(
"When using map side combine, an aggregator must be specified.");
}
- write0(dep.aggregator().get().combineValuesByKey(records, taskContext));
+ scala.collection.Iterator<?> combinedIterator =
+ dep.aggregator().get().combineValuesByKey(records, taskContext);
+ write0(combinedIterator);
+ return combinedIterator.hasNext();
} else {
write0(records);
+ return records.hasNext();
}
}
@@ -230,8 +237,8 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
public void write(scala.collection.Iterator<Product2<K, V>> records) throws
IOException {
boolean needCleanupPusher = true;
try {
- doWrite(records);
- close();
+ boolean iteratorHasNext = doWrite(records);
+ close(iteratorHasNext);
needCleanupPusher = false;
} finally {
if (needCleanupPusher) {
@@ -371,10 +378,17 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
}
}
- private void close() throws IOException {
+ private void close(boolean iteratorHasNext) throws IOException {
logger.info("Memory used {}", Utils.bytesToString(pusher.getUsed()));
long pushStartTime = System.nanoTime();
pusher.pushData(false);
+
+ // The check must come BEFORE mapperEnd so a partial map output is never
committed to the
+ // shuffle service.
+ // Check BEFORE pusher.close() so if we throw here the outer finally block
will handle the
+ // closing thing.
+ SparkUtils.assertIteratorFullyConsumed(iteratorHasNext);
+
pusher.close(true);
shuffleClient.pushMergedData(shuffleId, mapId, encodedAttemptId);
diff --git
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java
index 0e68a4b46..696a39087 100644
---
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java
+++
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java
@@ -66,6 +66,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.exception.CelebornIOException;
import org.apache.celeborn.common.exception.CelebornRuntimeException;
import org.apache.celeborn.common.network.protocol.TransportMessage;
import
org.apache.celeborn.common.protocol.message.ControlMessages.GetReducerFileGroupResponse;
@@ -716,4 +717,18 @@ public class SparkUtils {
String master = conf.get("spark.master", "");
return master.equals("local") || master.startsWith("local[");
}
+
+ /**
+ * Asserts that the shuffle writer's iterator has been fully consumed. Only
call this when the
+ * shuffle writer finishes writing records. If records remain in the
iterator, the task will fail
+ * with an IOException.
+ *
+ * @param iteratorHasNext true if the iterator still has records remaining
+ */
+ public static void assertIteratorFullyConsumed(boolean iteratorHasNext)
throws IOException {
+ if (iteratorHasNext) {
+ throw new CelebornIOException(
+ "Shuffle write task finished but iterator was not fully consumed.");
+ }
+ }
}
diff --git
a/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/CelebornShuffleWriterSuiteBase.java
b/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/CelebornShuffleWriterSuiteBase.java
index 8c74e0946..dbc0e8e89 100644
---
a/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/CelebornShuffleWriterSuiteBase.java
+++
b/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/CelebornShuffleWriterSuiteBase.java
@@ -46,6 +46,7 @@ import org.apache.spark.SparkConf;
import org.apache.spark.SparkEnv;
import org.apache.spark.SparkVersionUtil;
import org.apache.spark.TaskContext;
+import org.apache.spark.TaskContext$;
import org.apache.spark.executor.ShuffleWriteMetrics;
import org.apache.spark.executor.TaskMetrics;
import org.apache.spark.memory.TaskMemoryManager;
@@ -148,10 +149,13 @@ public abstract class CelebornShuffleWriterSuiteBase {
Mockito.doReturn(metrics).when(taskContext).taskMetrics();
Mockito.doReturn(taskMemoryManager).when(taskContext).taskMemoryManager();
+ Mockito.doReturn(None$.MODULE$).when(taskContext).getKillReason();
Mockito.doReturn(bmId).when(blockManager).shuffleServerId();
Mockito.doReturn(blockManager).when(env).blockManager();
Mockito.doReturn(sparkConf).when(env).conf();
+ Mockito.doReturn(false).when(dependency).mapSideCombine();
+ Mockito.doReturn(Option.empty()).when(dependency).aggregator();
SparkEnv.set(env);
}
@@ -306,6 +310,25 @@ public abstract class CelebornShuffleWriterSuiteBase {
client.shutdown();
}
+ @Test
+ public void testAssertIteratorFullyConsumed() {
+ // Test that assertIteratorFullyConsumed does not throw when iterator is
empty
+ SparkUtils.assertIteratorFullyConsumed(false);
+ }
+
+ @Test
+ public void testAssertIteratorFullyConsumedThrows() {
+ TaskContext$.MODULE$.setTaskContext(taskContext);
+ try {
+ SparkUtils.assertIteratorFullyConsumed(true);
+ fail("Expected IOException when iterator is not fully consumed");
+ } catch (IOException e) {
+ assertTrue(e.getMessage().contains("not fully consumed"));
+ } finally {
+ TaskContext$.MODULE$.setTaskContext(null);
+ }
+ }
+
private void check(
final int approximateSize, final CelebornConf conf, final Serializer
serializer)
throws Exception {
