[CELEBORN-894] Add checksum for shuffle data #2979

jiang13021 · 2024-12-08T06:00:14Z

What changes were proposed in this pull request?

Added a checksum mechanism to the shuffle data header.
Implemented checksum verification while writing data in the worker.
Integrated checksum validation during data reading in CelebornInputStreamImpl.
Added a PushDataHeaderUtils to manage the shuffle data header.

Why are the changes needed?

Sometimes, the transmitted data may change, leading to errors. Therefore, it is necessary to add a checksum to detect this exception. I only added a checksum to the header because I believe that in a production environment, the data body is mostly compressed, and during compression, a checksum is typically generated. Therefore, in scenarios where compression is enabled, adding a checksum to the header is enough.

Does this PR introduce any user-facing change?

No, this change is compatible.

How was this patch tested?

unit test: org.apache.celeborn.service.deploy.cluster.PushDataWithChecksumSuite

codenohup

Hi, @jiang13021
Thanks for your contribution!
I have two questions, PTAL.

codenohup · 2024-12-09T06:21:50Z

...common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java

@@ -79,6 +76,7 @@ public class FlinkShuffleClientImpl extends ShuffleClientImpl {
  private ConcurrentHashMap<String, TransportClient> currentClient =
      JavaUtils.newConcurrentHashMap();
  private long driverTimestamp;
+  private final int BATCH_HEADER_SIZE = 4 * 4;


Does this feature currently support only Spark? Will flink be supported in the future?

Yes. Currently, this feature does not support Flink, but I will submit another PRs in the future to add support. The current changes should be compatible with Flink.

codenohup · 2024-12-09T06:24:46Z

worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionDataWriter.java

@@ -293,14 +293,15 @@ public void flush(boolean finalFlush, boolean fromEvict) throws IOException {
          // read flush buffer to generate correct chunk offsets
          // data header layout (mapId, attemptId, nextBatchId, length)
          if (numBytes > chunkSize) {
-            ByteBuffer headerBuf = ByteBuffer.allocate(16);
+            ByteBuffer headerBuf = ByteBuffer.allocate(PushDataHeaderUtils.BATCH_HEADER_SIZE);


If a worker receives a data buffer from the an old version client or not supported engine client, will this modification yield accurate results?

If a worker receives a data buffer from an old version client, we will check the checksum flag in the highest bit of batchId for compatibility. Similarly, if the data comes from an unsupported engine's client but follows the old version data header layout, it will also be compatible.

...common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java

FMX · 2024-12-09T08:04:23Z

...common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java

@@ -79,6 +76,7 @@ public class FlinkShuffleClientImpl extends ShuffleClientImpl {
  private ConcurrentHashMap<String, TransportClient> currentClient =
      JavaUtils.newConcurrentHashMap();
  private long driverTimestamp;
+  private final int BATCH_HEADER_SIZE = 4 * 4;


FMX · 2024-12-09T08:25:31Z

common/src/main/java/org/apache/celeborn/common/util/PushDataHeaderUtils.java

+  }
+
+  public static int getLength(byte[] data) {
+    return Platform.getInt(data, LENGTH_OFFSET) - 4;


This will cause errors if an old client reads from Celeborn workers with this feature.

+1, we should think about how maintain compatibility with older versions.

@FMX @zwangsheng Thank you for ~~reviving~~ reviewing this, we will check the checksum flag in the highest bit of batchId for compatibility. I have update the code here, please take another look.

FMX · 2024-12-09T08:28:18Z

common/src/main/java/org/apache/celeborn/common/util/PushDataHeaderUtils.java

+  public static final int LENGTH_OFFSET = Platform.BYTE_ARRAY_OFFSET + 12;
+  public static final int CHECKSUM_OFFSET = Platform.BYTE_ARRAY_OFFSET + 16;
+  public static final int POSITIVE_MASK = 0x7FFFFFFF;
+  public static final int HIGHEST_1_BIT_FLAG_MASK = 0x80000000;


I agree with this design that the batchId can not be a negative number. But this class needs to be compatible with the old clients. During cluster upgrade, there will be a moment when old clients are talking to new servers.

Worker will check the checksum flag in the highest bit of batchId for compatibility.

FMX · 2024-12-09T08:32:25Z

common/src/main/java/org/apache/celeborn/common/util/PushDataHeaderUtils.java

+  public static int computeHeaderChecksum32(byte[] data) {
+    assert data.length >= BATCH_HEADER_SIZE_WITHOUT_CHECKSUM;
+    CRC32 crc32 = new CRC32();
+    crc32.update(data, 0, BATCH_HEADER_SIZE_WITHOUT_CHECKSUM);


Although 16-byte calculation for CRC can be trivial, there can be enormous push data structures to handle for a Celeborn Worker. I think it would be better to add a switch to let the users enable or disable this feature. For some users, the never meet any issues about this, they can just disable this feature to save the CPU.

I have received some reports from users telling me that Celeborn workers are consuming too much CPU during rush hours. If this feature is on by default, this will surely get things worse.
You can make this config a client-side control config and pass it in the ReserveSlots request.

OK, I will add a config this week.

@FMX Hi, I have added a config to disable header checksum, PTAL

CRC32 is not the fastest choice, I haven't gone through the whole design, and not sure if the algo can be configured like Spark's impl.
apache/spark#47929 apache/spark#49258

@pan3793 Thank you for your advice. However, we will verify the checksum in the worker. Therefore, if we have multiple checksum algorithms to choose from, we should consider adding a tag or field in the pushData/pushMerged request. I would prefer to do that at another PR, WDYT?

cc @mridulm, do you have any suggestions here?

@jiang13021 Seems we only have one bit to mark data chunk includes checksum or not, if so, I would suggest choosing the most efficient algo, namely CRC32C, which is the fastest one that ships in JDK 9+, and port the Hadoop PureJavaCrc32C for compatible with JDK 8

...er/src/test/scala/org/apache/celeborn/service/deploy/cluster/PushDataWithChecksumSuite.scala

zwangsheng · 2024-12-09T08:50:34Z

common/src/main/java/org/apache/celeborn/common/util/PushDataHeaderUtils.java

+  }
+
+  public static int getLength(byte[] data) {
+    return Platform.getInt(data, LENGTH_OFFSET) - 4;


+1, we should think about how maintain compatibility with older versions.

RexXiong · 2024-12-13T08:00:45Z

worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala

@@ -1470,7 +1470,24 @@ class PushDataHandler(val workerSource: WorkerSource) extends BaseMessageHandler
        shuffleKey: String,
        index: Int): Unit = {
      try {
-        fileWriter.write(body)
+        val header = new Array[Byte](PushDataHeaderUtils.BATCH_HEADER_SIZE_WITHOUT_CHECKSUM)


We could consider implementing a new method for verifying data checksums. If a checksum validation fails, we would throw a CelebornChecksumException. This approach would easier for ut testing. Additionally, we can handle this exception within the catch block alongside other exceptions.

Fixed, thank you.

common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala

FMX

LGTM.

codenohup reviewed Dec 9, 2024

View reviewed changes

FMX reviewed Dec 9, 2024

View reviewed changes

zwangsheng reviewed Dec 9, 2024

View reviewed changes

RexXiong reviewed Dec 13, 2024

View reviewed changes

pan3793 reviewed Dec 18, 2024

View reviewed changes

common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala Outdated Show resolved Hide resolved

pan3793 reviewed Dec 18, 2024

View reviewed changes

common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala Outdated Show resolved Hide resolved

[CELEBORN-894] Add config to disable checksum

f719b49

jiang13021 force-pushed the celeborn_894_checksum branch from 2eef55a to f719b49 Compare December 18, 2024 11:03

fix

a43e6a7

FMX approved these changes Dec 24, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[CELEBORN-894] Add checksum for shuffle data #2979

[CELEBORN-894] Add checksum for shuffle data #2979

jiang13021 commented Dec 8, 2024

codenohup left a comment

codenohup Dec 9, 2024

jiang13021 Dec 9, 2024

codenohup Dec 9, 2024

jiang13021 Dec 9, 2024

FMX Dec 9, 2024

FMX Dec 9, 2024

zwangsheng Dec 9, 2024

jiang13021 Dec 9, 2024 •

edited

Loading

FMX Dec 9, 2024

jiang13021 Dec 9, 2024

FMX Dec 9, 2024

FMX Dec 9, 2024

jiang13021 Dec 9, 2024

jiang13021 Dec 18, 2024

pan3793 Dec 18, 2024 •

edited

Loading

jiang13021 Dec 19, 2024

pan3793 Dec 24, 2024

pan3793 Dec 24, 2024

zwangsheng Dec 9, 2024

RexXiong Dec 13, 2024

jiang13021 Dec 18, 2024

FMX left a comment

[CELEBORN-894] Add checksum for shuffle data #2979

Are you sure you want to change the base?

[CELEBORN-894] Add checksum for shuffle data #2979

Conversation

jiang13021 commented Dec 8, 2024

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

codenohup left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jiang13021 Dec 9, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pan3793 Dec 18, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

FMX left a comment

Choose a reason for hiding this comment

jiang13021 Dec 9, 2024 •

edited

Loading

pan3793 Dec 18, 2024 •

edited

Loading