[SPARK-54663][CORE] Computes RowBasedChecksum in ShuffleWriters #50230
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,114 @@ | ||
| /* | ||
| * 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.spark.shuffle.checksum | ||
|
|
||
| import java.io.ObjectOutputStream | ||
| import java.util.zip.Checksum | ||
|
|
||
| import scala.util.control.NonFatal | ||
|
|
||
| import org.apache.spark.internal.Logging | ||
| import org.apache.spark.network.shuffle.checksum.ShuffleChecksumHelper | ||
| import org.apache.spark.util.MyByteArrayOutputStream | ||
|
|
||
| /** | ||
| * A class for computing checksum for input (key, value) pairs. The checksum is independent of | ||
| * the order of the input (key, value) pairs. It is done by computing a checksum for each row | ||
| * first, and then computing the XOR for all the row checksums. | ||
| */ | ||
| abstract class RowBasedChecksum() extends Serializable with Logging { | ||
| private var hasError: Boolean = false | ||
| private var checksumValue: Long = 0 | ||
| /** Returns the checksum value computed. Tt returns the default checksum value (0) if there | ||
| * are any errors encountered during the checksum computation. | ||
| */ | ||
| def getValue: Long = { | ||
| if (!hasError) checksumValue else 0 | ||
| } | ||
|
|
||
| /** Updates the row-based checksum with the given (key, value) pair */ | ||
| def update(key: Any, value: Any): Unit = { | ||
cloud-fan marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| if (!hasError) { | ||
| try { | ||
| val rowChecksumValue = calculateRowChecksum(key, value) | ||
| checksumValue = checksumValue ^ rowChecksumValue | ||
|
||
| } catch { | ||
| case NonFatal(e) => | ||
| logError("Checksum computation encountered error: ", e) | ||
| hasError = true | ||
| } | ||
| } | ||
| } | ||
|
|
||
| /** Computes and returns the checksum value for the given (key, value) pair */ | ||
| protected def calculateRowChecksum(key: Any, value: Any): Long | ||
| } | ||
|
|
||
| /** | ||
| * A Concrete implementation of RowBasedChecksum. The checksum for each row is | ||
| * computed by first converting the (key, value) pair to byte array using OutputStreams, | ||
| * and then computing the checksum for the byte array. | ||
| * Note that this checksum computation is very expensive, and it is used only in tests | ||
| * in the core component. A much cheaper implementation of RowBasedChecksum is in | ||
| * UnsafeRowChecksum. | ||
| * | ||
| * @param checksumAlgorithm the algorithm used for computing checksum. | ||
| */ | ||
| class OutputStreamRowBasedChecksum(checksumAlgorithm: String) | ||
|
||
| extends RowBasedChecksum() { | ||
|
|
||
| private val DEFAULT_INITIAL_SER_BUFFER_SIZE = 32 * 1024 | ||
|
|
||
| @transient private lazy val serBuffer = | ||
| new MyByteArrayOutputStream(DEFAULT_INITIAL_SER_BUFFER_SIZE) | ||
| @transient private lazy val objOut = new ObjectOutputStream(serBuffer) | ||
mridulm marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
|
|
||
| @transient | ||
| protected lazy val checksum: Checksum = | ||
| ShuffleChecksumHelper.getChecksumByAlgorithm(checksumAlgorithm) | ||
|
|
||
| override protected def calculateRowChecksum(key: Any, value: Any): Long = { | ||
| assert(checksum != null, "Checksum is null") | ||
|
|
||
| // Converts the (key, value) pair into byte array. | ||
| objOut.reset() | ||
| serBuffer.reset() | ||
| objOut.writeObject((key, value)) | ||
| objOut.flush() | ||
| serBuffer.flush() | ||
|
|
||
| // Computes and returns the checksum for the byte array. | ||
| checksum.reset() | ||
| checksum.update(serBuffer.getBuf, 0, serBuffer.size()) | ||
| checksum.getValue | ||
| } | ||
| } | ||
|
|
||
| object RowBasedChecksum { | ||
| def createPartitionRowBasedChecksums( | ||
|
||
| numPartitions: Int, | ||
| checksumAlgorithm: String): Array[RowBasedChecksum] = { | ||
| Array.tabulate(numPartitions)(_ => new OutputStreamRowBasedChecksum(checksumAlgorithm)) | ||
| } | ||
|
|
||
| def getAggregatedChecksumValue(rowBasedChecksums: Array[RowBasedChecksum]): Long = { | ||
| Option(rowBasedChecksums) | ||
| .map(_.foldLeft(0L)((acc, c) => acc * 31L + c.getValue)) | ||
| .getOrElse(0L) | ||
| } | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,26 @@ | ||
| /* | ||
| * 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.spark.util; | ||
|
|
||
| import java.io.ByteArrayOutputStream; | ||
|
|
||
| /** Subclass of ByteArrayOutputStream that exposes `buf` directly. */ | ||
| public final class MyByteArrayOutputStream extends ByteArrayOutputStream { | ||
|
||
| public MyByteArrayOutputStream(int size) { super(size); } | ||
| public byte[] getBuf() { return buf; } | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -23,7 +23,7 @@ import java.util.concurrent.{ConcurrentHashMap, LinkedBlockingQueue, ThreadPoolE | |
| import java.util.concurrent.locks.ReentrantReadWriteLock | ||
|
|
||
| import scala.collection | ||
| import scala.collection.mutable.{HashMap, ListBuffer, Map, Set} | ||
| import scala.concurrent.{ExecutionContext, Future} | ||
| import scala.concurrent.duration.Duration | ||
| import scala.jdk.CollectionConverters._ | ||
|
|
@@ -99,6 +99,11 @@ private class ShuffleStatus( | |
| */ | ||
| val mapStatusesDeleted = new Array[MapStatus](numPartitions) | ||
|
|
||
| /** | ||
| * Keep the indices of the Map tasks whose checksums are different across retries. | ||
| */ | ||
| private[this] val checksumMismatchIndices : Set[Int] = Set() | ||
|
|
||
| /** | ||
| * MergeStatus for each shuffle partition when push-based shuffle is enabled. The index of the | ||
| * array is the shuffle partition id (reduce id). Each value in the array is the MergeStatus for | ||
|
|
@@ -169,6 +174,12 @@ private class ShuffleStatus( | |
| } else { | ||
| mapIdToMapIndex.remove(currentMapStatus.mapId) | ||
| } | ||
|
|
||
| val preStatus = | ||
| if (mapStatuses(mapIndex) != null) mapStatuses(mapIndex) else mapStatusesDeleted(mapIndex) | ||
| if (preStatus != null && preStatus.checksumValue != status.checksumValue) { | ||
| checksumMismatchIndices.add(mapIndex) | ||
| } | ||
|
Contributor
There was a problem hiding this comment. There are three main cases here:
For the latter two, we dont need to track it in
Contributor
Author
There was a problem hiding this comment. Yes, for case 1, we need to track the mismatches. The usage of checksumMismatchIndices is that (in the next PR) we will rollback the downstream stages, if we detect checksum mismatches for its upstream stages. For case 2, if downstream stages have not consumed output, which means they have not started. In this case, the rollback is a no-op, and it doesn't hurt to record the mismatches here. For case 3, I think we need to record the mismatches. Assuming a situation where all partitions of a stage have finished, while some speculative tasks are still running. As all outputs have been produced, the downstream stage can start and read from the data. Later, some speculative tasks finish, and new mapStatus will override the old mapStatus with new data location. For the downstream stage, the not yet started tasks or retried tasks would read from the new data, while the finished and running tasks would read from the old data, resulting in inconsistency.
Contributor
There was a problem hiding this comment.
It is unclear how
That is fair, this is indeed possible. |
||
| mapStatuses(mapIndex) = status | ||
| mapIdToMapIndex(status.mapId) = mapIndex | ||
| } | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Let's update the classdoc. We now also leverage the sum to handle duplicated values better.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
updated