Skip to content

[SPARK-34840][SHUFFLE] Fixes cases of corruption in merged shuffle … #31934

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
otterc wants to merge 3 commits into from
Closed

[SPARK-34840][SHUFFLE] Fixes cases of corruption in merged shuffle … #31934

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
9382fdc
[SPARK-34840] Fixes cases of corruption in merged shuffle blocks
otterc Mar 22, 2021
c33f961
empty commit by Min Shen
Victsm Mar 22, 2021
7c38487
Removed unused import
otterc Mar 24, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 13 additions & 40 deletions ...twork-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,6 @@
import java.util.Arrays;
import java.util.Collection;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentMap;
Expand Down Expand Up @@ -445,9 +444,9 @@ static class PushBlockStreamCallback implements StreamCallbackWithID {
private final AppShufflePartitionInfo partitionInfo;
private int length = 0;
// This indicates that this stream got the opportunity to write the blocks to the merged file.
// Once this is set to true and the stream encounters a failure then it will take necessary
// action to overwrite any partial written data. This is reset to false when the stream
// completes without any failures.
// Once this is set to true and the stream encounters a failure then it will unset the
// currentMapId of the partition so that another stream can start merging the blocks to the
// partition. This is reset to false when the stream completes.
private boolean isWriting = false;
// Use on-heap instead of direct ByteBuffer since these buffers will be GC'ed very quickly
private List<ByteBuffer> deferredBufs;
Expand Down Expand Up @@ -477,16 +476,11 @@ public String getID() {
*/
private void writeBuf(ByteBuffer buf) throws IOException {
while (buf.hasRemaining()) {
if (partitionInfo.isEncounteredFailure()) {
long updatedPos = partitionInfo.getDataFilePos() + length;
logger.debug(
"{} shuffleId {} reduceId {} encountered failure current pos {} updated pos {}",
partitionInfo.appShuffleId.appId, partitionInfo.appShuffleId.shuffleId,
partitionInfo.reduceId, partitionInfo.getDataFilePos(), updatedPos);
length += partitionInfo.dataChannel.write(buf, updatedPos);
} else {
length += partitionInfo.dataChannel.write(buf);
}
long updatedPos = partitionInfo.getDataFilePos() + length;
logger.debug("{} shuffleId {} reduceId {} current pos {} updated pos {}",
partitionInfo.appShuffleId.appId, partitionInfo.appShuffleId.shuffleId,
partitionInfo.reduceId, partitionInfo.getDataFilePos(), updatedPos);
length += partitionInfo.dataChannel.write(buf, updatedPos);
}
}

Expand Down Expand Up @@ -581,7 +575,6 @@ public void onData(String streamId, ByteBuffer buf) throws IOException {
}
// Check whether we can write to disk
if (allowedToWrite()) {
isWriting = true;
// Identify duplicate block generated by speculative tasks. We respond success to
// the client in cases of duplicate even though no data is written.
if (isDuplicateBlock()) {
Expand All @@ -598,6 +591,7 @@ public void onData(String streamId, ByteBuffer buf) throws IOException {

// If we got here, it's safe to write the block data to the merged shuffle file. We
// first write any deferred block.
isWriting = true;
try {
if (deferredBufs != null && !deferredBufs.isEmpty()) {
writeDeferredBufs();
Expand All @@ -609,16 +603,6 @@ public void onData(String streamId, ByteBuffer buf) throws IOException {
// back to the client so the block could be retried.
throw ioe;
}
// If we got here, it means we successfully write the current chunk of block to merged
// shuffle file. If we encountered failure while writing the previous block, we should
// reset the file channel position and the status of partitionInfo to indicate that we
// have recovered from previous disk write failure. However, we do not update the
// position tracked by partitionInfo here. That is only updated while the entire block
// is successfully written to merged shuffle file.
if (partitionInfo.isEncounteredFailure()) {
partitionInfo.dataChannel.position(partitionInfo.getDataFilePos() + length);
partitionInfo.setEncounteredFailure(false);
}
} else {
logger.trace("{} shuffleId {} reduceId {} onData deferred",
partitionInfo.appShuffleId.appId, partitionInfo.appShuffleId.shuffleId,
Expand All @@ -639,7 +623,7 @@ public void onData(String streamId, ByteBuffer buf) throws IOException {
// written to disk due to this reason. We thus decide to optimize for server
// throughput and memory usage.
if (deferredBufs == null) {
deferredBufs = new LinkedList<>();
deferredBufs = new ArrayList<>();
}
// Write the buffer to the in-memory deferred cache. Since buf is a slice of a larger
// byte buffer, we cache only the relevant bytes not the entire large buffer to save
Expand Down Expand Up @@ -670,7 +654,6 @@ public void onComplete(String streamId) throws IOException {
}
// Check if we can commit this block
if (allowedToWrite()) {
isWriting = true;
// Identify duplicate block generated by speculative tasks. We respond success to
// the client in cases of duplicate even though no data is written.
if (isDuplicateBlock()) {
Expand All @@ -681,6 +664,7 @@ public void onComplete(String streamId) throws IOException {
try {
if (deferredBufs != null && !deferredBufs.isEmpty()) {
abortIfNecessary();
isWriting = true;
writeDeferredBufs();
}
} catch (IOException ioe) {
Expand Down Expand Up @@ -738,14 +722,14 @@ public void onFailure(String streamId, Throwable throwable) throws IOException {
Map<Integer, AppShufflePartitionInfo> shufflePartitions =
mergeManager.partitions.get(partitionInfo.appShuffleId);
if (shufflePartitions != null && shufflePartitions.containsKey(partitionInfo.reduceId)) {
logger.debug("{} shuffleId {} reduceId {} set encountered failure",
logger.debug("{} shuffleId {} reduceId {} encountered failure",
partitionInfo.appShuffleId.appId, partitionInfo.appShuffleId.shuffleId,
partitionInfo.reduceId);
partitionInfo.setCurrentMapIndex(-1);
partitionInfo.setEncounteredFailure(true);
}
}
}
isWriting = false;
Copy link
Member

@Ngone51 Ngone51 Mar 24, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Move this into the if condition scope?

Copy link
Contributor Author

@otterc otterc Mar 24, 2021 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can move this to if scope and that would not change the behavior or cause any issues. The only reason I had it outside because it was consistent with where this flag is unset in onComplete. I understand that is a very trivial cosmetic reason so can move this.

Copy link
Member

@Ngone51 Ngone51 Mar 25, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok, keeping it consistent sounds fine. we can leave it as it is since it's trivial.

}

@VisibleForTesting
Expand Down Expand Up @@ -802,8 +786,6 @@ public static class AppShufflePartitionInfo {
public FileChannel dataChannel;
// Location offset of the last successfully merged block for this shuffle partition
private long dataFilePos;
// Indicating whether failure was encountered when merging the previous block
private boolean encounteredFailure;
// Track the map index whose block is being merged for this shuffle partition
private int currentMapIndex;
// Bitmap tracking which mapper's blocks have been merged for this shuffle partition
Expand Down Expand Up @@ -836,7 +818,6 @@ public static class AppShufflePartitionInfo {
// Writing 0 offset so that we can reuse ShuffleIndexInformation.getIndex()
updateChunkInfo(0L, -1);
this.dataFilePos = 0;
this.encounteredFailure = false;
this.mapTracker = new RoaringBitmap();
this.chunkTracker = new RoaringBitmap();
}
Expand All @@ -851,14 +832,6 @@ public void setDataFilePos(long dataFilePos) {
this.dataFilePos = dataFilePos;
}

boolean isEncounteredFailure() {
return encounteredFailure;
}

void setEncounteredFailure(boolean encounteredFailure) {
this.encounteredFailure = encounteredFailure;
}

int getCurrentMapIndex() {
return currentMapIndex;
}
Expand Down
87 changes: 84 additions & 3 deletions ...-shuffle/src/test/java/org/apache/spark/network/shuffle/RemoteBlockPushResolverSuite.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
import java.nio.file.Paths;
import java.util.Arrays;
import java.util.concurrent.Semaphore;
import java.util.concurrent.ThreadLocalRandom;

import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableMap;
Expand Down Expand Up @@ -292,18 +293,32 @@ public void testTooLateArrival() throws IOException {
@Test
public void testIncompleteStreamsAreOverwritten() throws IOException {
registerExecutor(TEST_APP, prepareLocalDirs(localDirs));
byte[] expectedBytes = new byte[4];
ThreadLocalRandom.current().nextBytes(expectedBytes);

StreamCallbackWithID stream1 =
pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 0, 0, 0));
stream1.onData(stream1.getID(), ByteBuffer.wrap(new byte[4]));
byte[] data = new byte[10];
ThreadLocalRandom.current().nextBytes(data);
stream1.onData(stream1.getID(), ByteBuffer.wrap(data));
// There is a failure
stream1.onFailure(stream1.getID(), new RuntimeException("forced error"));
StreamCallbackWithID stream2 =
pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 1, 0, 0));
stream2.onData(stream2.getID(), ByteBuffer.wrap(new byte[5]));
ByteBuffer nextBuf= ByteBuffer.wrap(expectedBytes, 0, 2);
stream2.onData(stream2.getID(), nextBuf);
stream2.onComplete(stream2.getID());
StreamCallbackWithID stream3 =
pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 2, 0, 0));
nextBuf = ByteBuffer.wrap(expectedBytes, 2, 2);
stream3.onData(stream3.getID(), nextBuf);
stream3.onComplete(stream3.getID());
pushResolver.finalizeShuffleMerge(new FinalizeShuffleMerge(TEST_APP, 0));
MergedBlockMeta blockMeta = pushResolver.getMergedBlockMeta(TEST_APP, 0, 0);
validateChunks(TEST_APP, 0, 0, blockMeta, new int[]{5}, new int[][]{{1}});
validateChunks(TEST_APP, 0, 0, blockMeta, new int[]{4}, new int[][]{{1, 2}});
FileSegmentManagedBuffer mb =
(FileSegmentManagedBuffer) pushResolver.getMergedBlockData(TEST_APP, 0, 0, 0);
assertArrayEquals(expectedBytes, mb.nioByteBuffer().array());
}

@Test (expected = RuntimeException.class)
Expand Down Expand Up @@ -740,6 +755,72 @@ public void testFailureWhileTruncatingFiles() throws IOException {
validateChunks(TEST_APP, 0, 1, meta, new int[]{5, 3}, new int[][]{{0},{1}});
}

@Test
public void testOnFailureInvokedMoreThanOncePerBlock() throws IOException {
StreamCallbackWithID stream1 =
pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 0, 0, 0));
stream1.onData(stream1.getID(), ByteBuffer.wrap(new byte[2]));
stream1.onFailure(stream1.getID(), new RuntimeException("forced error"));
StreamCallbackWithID stream2 =
pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 1, 0, 0));
stream2.onData(stream2.getID(), ByteBuffer.wrap(new byte[5]));
// On failure on stream1 gets invoked again and should cause no interference
stream1.onFailure(stream1.getID(), new RuntimeException("2nd forced error"));
StreamCallbackWithID stream3 =
pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 3, 0, 0));
// This should be deferred as stream 2 is still the active stream
stream3.onData(stream3.getID(), ByteBuffer.wrap(new byte[2]));
// Stream 2 writes more and completes
stream2.onData(stream2.getID(), ByteBuffer.wrap(new byte[4]));
stream2.onComplete(stream2.getID());
stream3.onComplete(stream3.getID());
pushResolver.finalizeShuffleMerge(new FinalizeShuffleMerge(TEST_APP, 0));
MergedBlockMeta blockMeta = pushResolver.getMergedBlockMeta(TEST_APP, 0, 0);
validateChunks(TEST_APP, 0, 0, blockMeta, new int[] {9, 2}, new int[][] {{1},{3}});
removeApplication(TEST_APP);
}

@Test (expected = RuntimeException.class)
public void testFailureAfterDuplicateBlockDoesNotInterfereActiveStream() throws IOException {
StreamCallbackWithID stream1 =
pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 0, 0, 0));
StreamCallbackWithID stream1Duplicate =
pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 0, 0, 0));
stream1.onData(stream1.getID(), ByteBuffer.wrap(new byte[2]));
stream1.onComplete(stream1.getID());
stream1Duplicate.onData(stream1.getID(), ByteBuffer.wrap(new byte[2]));

StreamCallbackWithID stream2 =
pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 1, 0, 0));
stream2.onData(stream2.getID(), ByteBuffer.wrap(new byte[5]));
// Should not change the current map id of the reduce partition
stream1Duplicate.onFailure(stream2.getID(), new RuntimeException("forced error"));

StreamCallbackWithID stream3 =
pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 2, 0, 0));
// This should be deferred as stream 2 is still the active stream
stream3.onData(stream3.getID(), ByteBuffer.wrap(new byte[2]));
RuntimeException failedEx = null;
try {
stream3.onComplete(stream3.getID());
} catch (RuntimeException re) {
assertEquals(
"Couldn't find an opportunity to write block shufflePush_0_2_0 to merged shuffle",
re.getMessage());
failedEx = re;
}
// Stream 2 writes more and completes
stream2.onData(stream2.getID(), ByteBuffer.wrap(new byte[4]));
stream2.onComplete(stream2.getID());
pushResolver.finalizeShuffleMerge(new FinalizeShuffleMerge(TEST_APP, 0));
MergedBlockMeta blockMeta = pushResolver.getMergedBlockMeta(TEST_APP, 0, 0);
validateChunks(TEST_APP, 0, 0, blockMeta, new int[] {11}, new int[][] {{0, 1}});
removeApplication(TEST_APP);
if (failedEx != null) {
throw failedEx;
}
}

private void useTestFiles(boolean useTestIndexFile, boolean useTestMetaFile) throws IOException {
pushResolver = new RemoteBlockPushResolver(conf) {
@Override
Expand Down