[CELEBORN-1319] Optimize skew partition logic for Reduce Mode to avoid sorting shuffle files #2373

wangshengjie123 · 2024-03-11T02:46:37Z

What changes were proposed in this pull request?

Add logic to support avoid sorting shuffle files for Reduce mode when optimize skew partitions

Why are the changes needed?

Current logic need sorting shuffle files when read Reduce mode skew partition shuffle files, we found some shuffle sorting timeout and performance issue

Does this PR introduce any user-facing change?

No

How was this patch tested?

Cluster test and uts

client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java

waitinfuture

Thanks @wangshengjie123 for this PR! I left some comments. In addition, is the small change to Spark missing?

client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java

waitinfuture · 2024-03-12T02:13:42Z

client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java

+
+    int step = locations.length / subPartitionSize;
+
+    // if partition location is [1,2,3,4,5,6,7,8,9,10], and skew partition split to 3 task:


Seems the logic should be like this:

// if partition location is [1,2,3,4,5,6,7,8,9,10], and skew partition split to 3 task: // task 0: 1, 4, 7, 10 // task 1: 2, 4, 8 // task 2: 3, 5, 9 for (int i = 0; i < step + 1; i++) { int index = i * step + subPartitionIndex; if (index < locations.length) { result.add(orderedPartitionLocations[index]); } }

If I am not wrong, the idea is to minimize per row size - and so why column 0 goes "down" the array index, while column 1 goes "up" - and keeps alternating - so that as the size keeps increasing, it is more reasonably distributed for each row (essentially a way to approximate multi-way partition problem).

The result would be different for the formulation above @waitinfuture.

For example:

partition sizes: {1000, 1100, 1300, 1400, 2000, 2500, 3000, 10000, 20000, 25000, 28000, 30000}
subPartitionSize == 3
subPartitionIndex == 1

In formulation from PR we have:

task 0: 1000 , 2500 , 3000 , 30000
task 1: 1100 , 2000 , 10000 , 28000
task 2: 1300 , 1400 , 20000 , 25000

So the sizes will be:
task 0: 36500
task 1: 41100
task 2: 47700

As formulated above, we will end up with:

task 0: 1000 , 1400 , 3000 , 25000
task 1: 1100 , 2000 , 10000 , 28000
task 2: 1300 , 2500 , 20000 , 30000

In this case, the sizes will be:
task 0: 30400
task 1: 41100
task 2: 53800

Personally, I would have looked into either largest remainder or knapsack heuristic (given we are sorting anyway).

(Do let me know if I am missing something here @wangshengjie123)

@mridulm Sorry for late reply, your understanding is correct, and i should optimize the logic

Thanks @mridulm for the explanation, I actually didn't get the idea and was thinking the naive way :)

client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java

client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java

client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala

common/src/main/java/org/apache/celeborn/common/write/PushState.java

common/src/main/proto/TransportMessages.proto

common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala

mridulm

Interesting work @wangshengjie123 !

client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java

mridulm · 2024-03-13T04:58:48Z

client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java

+
+    int step = locations.length / subPartitionSize;
+
+    // if partition location is [1,2,3,4,5,6,7,8,9,10], and skew partition split to 3 task:


If I am not wrong, the idea is to minimize per row size - and so why column 0 goes "down" the array index, while column 1 goes "up" - and keeps alternating - so that as the size keeps increasing, it is more reasonably distributed for each row (essentially a way to approximate multi-way partition problem).

The result would be different for the formulation above @waitinfuture.

For example:

partition sizes: {1000, 1100, 1300, 1400, 2000, 2500, 3000, 10000, 20000, 25000, 28000, 30000}
subPartitionSize == 3
subPartitionIndex == 1

In formulation from PR we have:

task 0: 1000 , 2500 , 3000 , 30000
task 1: 1100 , 2000 , 10000 , 28000
task 2: 1300 , 1400 , 20000 , 25000

So the sizes will be:
task 0: 36500
task 1: 41100
task 2: 47700

As formulated above, we will end up with:

task 0: 1000 , 1400 , 3000 , 25000
task 1: 1100 , 2000 , 10000 , 28000
task 2: 1300 , 2500 , 20000 , 30000

In this case, the sizes will be:
task 0: 30400
task 1: 41100
task 2: 53800

Personally, I would have looked into either largest remainder or knapsack heuristic (given we are sorting anyway).

(Do let me know if I am missing something here @wangshengjie123)

client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java

client/src/main/scala/org/apache/celeborn/client/CommitManager.scala

client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala

common/src/main/java/org/apache/celeborn/common/write/PushFailedBatch.java

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

cfmcgrady · 2024-03-14T07:28:43Z

Thanks @wangshengjie123 for this PR! I left some comments. In addition, is the small change to Spark missing?

HI, @wangshengjie123
Can you please update the Spark patch? It will help the reviewers understand this PR better. Thanks!

wangshengjie123 · 2024-03-16T07:58:48Z

Thanks @wangshengjie123 for this PR! I left some comments. In addition, is the small change to Spark missing?

HI, @wangshengjie123 Can you please update the Spark patch? It will help the reviewers understand this PR better. Thanks!

Sorry for late reply, the pr will be updated today or tomorrow

codecov · 2024-03-16T08:28:59Z

Codecov Report

Attention: Patch coverage is 1.20482% with 82 lines in your changes are missing coverage. Please review.

Project coverage is 48.51%. Comparing base (12c3779) to head (ef81070).
Report is 12 commits behind head on main.

Files	Patch %	Lines
...born/common/protocol/message/ControlMessages.scala	0.00%	38 Missing ⚠️
.../apache/celeborn/common/write/PushFailedBatch.java	0.00%	24 Missing ⚠️
...org/apache/celeborn/common/util/PbSerDeUtils.scala	0.00%	9 Missing ⚠️
...g/apache/celeborn/common/protocol/StorageInfo.java	0.00%	6 Missing ⚠️
...va/org/apache/celeborn/common/write/PushState.java	16.67%	5 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2373      +/-   ##
==========================================
- Coverage   48.77%   48.51%   -0.26%     
==========================================
  Files         209      210       +1     
  Lines       13109    13186      +77     
  Branches     1134     1139       +5     
==========================================
+ Hits         6393     6396       +3     
- Misses       6294     6368      +74     
  Partials      422      422

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

RexXiong

Thanks @wangshengjie123 nice pr! Another suggestion is better to add UT for this feature.

RexXiong · 2024-03-19T09:14:29Z

client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java

@@ -1393,7 +1414,13 @@ public void onSuccess(ByteBuffer response) {
                    Arrays.toString(partitionIds),
                    groupedBatchId,
                    Arrays.toString(batchIds));
-
+                if (dataPushFailureTrackingEnabled) {


There is no need for HARD_SPLIT to do this. as worker never write the batch when HARD_SPLIT. cc @waitinfuture

I'm not sure if it's possible that the master copy succeeds but the copy fails due to HARD_SPLIT. I will check it again

RexXiong · 2024-03-19T09:46:24Z

client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java

@@ -615,6 +663,17 @@ private boolean fillBuffer() throws IOException {

          // de-duplicate
          if (attemptId == attempts[mapId]) {
+            if (splitSkewPartitionWithoutMapRange) {


We can reuse one PushFailedBatch object and update inner fields to improve memory-efficient.

Better to check failedBatches is empty or not first. May be we never need to check failed batches.

get this

fixed to avid NPE

client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala

RexXiong · 2024-03-19T09:58:01Z

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

@@ -4671,4 +4671,13 @@ object CelebornConf extends Logging {
      .version("0.5.0")
      .intConf
      .createWithDefault(10000)
+
+  val CLIENT_DATA_PUSH_FAILURE_TRACKING_ENABLED: ConfigEntry[Boolean] =


May be we can use another configuration name for enable optimize skew join. The CLIENT_DATA_PUSH_FAILURE_TRACKING_ENABLED doesn't feel so straightforward.

wangshengjie123 · 2024-03-24T09:19:02Z

Thanks @wangshengjie123 nice pr! Another suggestion is better to add UT for this feature.

UTs is doing, test in cluster this week, uts will be submit later

s0nskar · 2024-04-04T08:53:10Z

@wangshengjie123 Is there any doc or ticket explaining this approach? Also for the sort based approach that you mentioned.

s0nskar · 2024-04-04T13:02:29Z

From my understanding, in this PR we're diverting from vanilla spark approach based on mapIndex and just dividing the full partition into multiple sub-partition based on some heuristics. I'm new to Celeborn code, so might be missing something basic but in this PR we're not addressing below issue. If we consider a basic scenario where a partial partition read is happening and we see a FetchFailure.

ShuffleMapStage --> ResultStage

ShuffleMapStage (attempt 0) generated [P0, P1, P2] and P0 is skewed with partition location [0,1,2,3,4,5].
AQE asks for three splits and this PR logic will create three partitions [0, 1], [2, 3], [4, 5]
Now consider is reducer read [0, 1] and [2, 3] and gets FetchFailure while reading [4, 5]
This will trigger a complete mapper stage retry a/c to this doc and will clear the map output corresponding the shuffleID
ShuffleMapStage (attempt 0) will again generate data for P0 at different partition location [a, b, c, d, e, f] and it will get divided like [a, b], [c, d], [e, f]
Now if reader stage is ShuffleMapStage then it will read every sub-partition again but if the reader is ResultStage then it will only read missing partition data which [e, f].

The data generated on location 1 and location a would be different because of other factors like network delay (same thing applies for other locations). Ex – The data that might be present in 1st location in first attempt might be present in 2nd location or any location in different attempt because of the order mapper generated the data and in order server received that data.

This can cause both Data loss and Data duplication, this might be getting addressed in some other place in the codebase that i'm not aware of but i wanted point this problem out.

pan3793 · 2024-04-04T14:00:08Z

@s0nskar Good point, this should be an issue for ResultStage, even though the ShuffleMapStage's output is deterministic.

IIRC, vanilla Spark also has some limitations on stage retry cases for ResultStage when ShuffleMapStage's output is indeterministic, for such cases, we need to fail the job, right?

s0nskar · 2024-04-04T14:13:24Z

@pan3793 This does not become problem if we are maintaining the concept of mapIndex ranges as spark will always read deterministic output for each sub-partition.

As vanilla spark always read deterministic output because of mapIndex range filter, it will not face this issue. In this approach sub-partitions data will be indeterministic across stage attempts. Failing would be only option for such cases until spark start supporting ResultStage rollback.

s0nskar · 2024-04-04T14:26:45Z

Also, I think this issue would not be only limited to ResultStage, this can happen with ShuffleMapStage as well in some complex cases. Consider another scenario –

ShuffleMapStage1 -----> ShuffleMapStage2 ----->

Similar to above example, let's say partition skew P0 generated by ShuffleMapStage1.
ShuffleMapStage2 gets FetchFailure while reading sub-partitions of ShuffleMapStage1.
ShuffleMapStage1 will be recomputed and shuffle outputs will be cleared.
Only missing task of ShuffleMapStage2 will be retries, again causing the same issue.

This is case though, we can rollback the whole lineage till this point instead of failing this job. Similar to what vanilla spark does, what this will be very expensive.

pan3793 · 2024-04-04T16:23:48Z

@s0nskar I see your point. When consuming skew partitions, we should always treat the previous ShuffleMapStage's output as indeterministic under the current approach to avoid correctness issues.

waitinfuture · 2024-04-04T16:42:21Z

Hi @s0nskar , thanks for your point, I think you are correct. Seems this PR conflicts with stage rerun.

we should always treat the previous ShuffleMapStage's output as indeterministic under the current approach to avoid correctness issues.

@pan3793 Is it possible to force make it as indeterministic?

Also, I think Spark doesn't correctly set stage's determinism for some cases, for example a row_number window operator followed by aggregation keyed by the row_number.

cc @mridulm @ErikFang

waitinfuture · 2024-04-04T16:56:40Z

@wangshengjie123 Is there any doc or ticket explaining this approach? Also for the sort based approach that you mentioned.

The sort based approach is roughly like this:

Each sub reducer reads from all partition splits of its partitionId for data within its map range
The first read request will trigger the partition split file to be sorted based on map ids, so each IO will be sequential

s0nskar · 2024-04-04T17:49:59Z

Thanks a lot @waitinfuture for the sort based approach description.

Is it possible to force make it as indeterministic?

IMO this would be very difficult to do it from Celeborn itself but it can be done by putting a patch in the Spark code. ShuffledRowRDD can set Determinacy Level to INDETEMINATE if partial partition reads are happening and Celeborn is getting is used.

cc: @mridulm for viz

pan3793 · 2024-04-04T18:09:09Z

@waitinfuture It seems this PR is getting attention, some discussions happened offline, we'd better update the PR description(or Google Docs) to summarize the whole design and known issues so far

s0nskar · 2024-04-05T04:11:00Z

assets/spark-patch/Celeborn-Optimize-Skew-Partitions-spark3_3.patch

+         }
+-        PartialReducerPartitionSpec(reducerId, startMapIndex, endMapIndex, dataSize)
+        if (splitSkewPartitionWithCeleborn) {
+          PartialReducerPartitionSpec(reducerId, mapStartIndices.length, i, dataSize)


We can maybe add a note here that these dataSize will not be accurate. Even though in the current downstream code, we're only getting the sum of dataSize which should be equal but someone might be using these differently.

mridulm · 2024-04-06T01:48:17Z

It has been a while since I looked at this PR - but as formulated, the split into subranges is deterministic (if it is not, it should be made so).
With that in place, this would not be an issue ...
(I will take a deeper look later next week, but do let me know if I am missing something so that I can add that to my analysis)

waitinfuture · 2024-04-06T03:05:00Z

It has been a while since I looked at this PR - but as formulated, the split into subranges is deterministic (if it is not, it should be made so). With that in place, this would not be an issue ... (I will take a deeper look later next week, but do let me know if I am missing something so that I can add that to my analysis)

the split into subranges is deterministic

The way Celeborn splits partition is not deterministic with stage rerun, for example any push failure will cause split, so I'm afraid this statement does not hold...

mridulm · 2024-04-06T05:14:08Z

Ah, I see what you mean ... PartitionLocation would change between retries.
Yeah, this is a problem then - it will cause data loss. This would be a variant of SPARK-23207

I will need to relook at the PR, and how it interact with Celeborn - but if scenarios directly described in SPARK-23207 (or variants of it) are applicable (and we cant mitigate it), we should not proceed down this path given the correctness implications unfortunately.

mridulm · 2024-04-06T05:27:58Z

+CC @otterc as well.

waitinfuture · 2024-04-06T11:59:22Z

Ah, I see what you mean ... PartitionLocation would change between retries. Yeah, this is a problem then - it will cause data loss. This would be a variant of SPARK-23207

I will need to relook at the PR, and how it interact with Celeborn - but if scenarios directly described in SPARK-23207 (or variants of it) are applicable (and we cant mitigate it), we should not proceed down this path given the correctness implications unfortunately.

Maybe we can remain both this optimization and stage rerun, but only allows one to take effect by checking configs for now. The performance issue this PR solves does happen in production.

This reverts commit 7393d45.

…t.adaptive.optimizeSkewedPartitionRead.enabled`

…n without mapRange

wangshengjie123 changed the title ~~[CELEBORN-1319] Optimize skew partition logic for Reduce Mode to avoid sorting shuffle files~~ [WIP][CELEBORN-1319] Optimize skew partition logic for Reduce Mode to avoid sorting shuffle files Mar 11, 2024

lyy-pineapple reviewed Mar 11, 2024

View reviewed changes

client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java Outdated Show resolved Hide resolved

waitinfuture reviewed Mar 12, 2024

View reviewed changes

mridulm reviewed Mar 13, 2024

View reviewed changes

wangshengjie123 closed this Mar 16, 2024

wangshengjie123 reopened this Mar 16, 2024

wangshengjie123 force-pushed the optimize-skew-partition branch from b3af836 to 599be24 Compare March 16, 2024 08:14

RexXiong reviewed Mar 19, 2024

View reviewed changes

wangshengjie123 force-pushed the optimize-skew-partition branch 2 times, most recently from 8fe3a13 to 12eca26 Compare March 26, 2024 14:12

s0nskar reviewed Apr 5, 2024

View reviewed changes

wangshengjie123 and others added 28 commits December 22, 2024 13:27

remove unused log and refactor code

fe87865

add unit tests and refactor code

0b4eed4

split the skewed partition based on the chunk range

f628aa8

remove CelebornInputStreamSuiteJ

50c7110

refactor code for batch open stream

00b06fb

Revert "refactor code for batch open stream"

a7638c3

This reverts commit 7393d45.

celeborn.client.dataPushFailure.tracking.enabled -> `celeborn.clien…

d79a1b1

…t.adaptive.optimizeSkewedPartitionRead.enabled`

refactor code for batch open stream

10f8a87

add ut and refactor split logic

44079d9

refactor code according to review suggestions

037f32e

add sbt dependency

bdf4782

update spark patch

5fea9ba

address comment

d4bae88

fix

17c1094

add ut

94db413

update

cc99f96

refactor

35fbe50

add CelebornPartitionUtil.java

f728aa0

fix ut for jdk 8

4685946

rebase main and fix npe

f82a25c

fix ut

86eee4c

fix npe when memory storage enabled

3fcda0d

fix code style check error

16a7425

fix comile error

dce8b99

add spark patch

5f2f1ca

fix ut compile error and add license to spark patch

df24a56

update spark patch to abort stage when rerun skew join stage

583cb67

address comment, avoid reading replicate peer when read skew partitio…

57813ab

…n without mapRange

wangshengjie123 force-pushed the optimize-skew-partition branch from 09522d6 to 57813ab Compare December 22, 2024 08:28

update ClientUtils

26a7033

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[CELEBORN-1319] Optimize skew partition logic for Reduce Mode to avoid sorting shuffle files #2373

[CELEBORN-1319] Optimize skew partition logic for Reduce Mode to avoid sorting shuffle files #2373

wangshengjie123 commented Mar 11, 2024

waitinfuture left a comment

waitinfuture Mar 12, 2024

mridulm Mar 13, 2024

wangshengjie123 Mar 13, 2024

waitinfuture Mar 13, 2024

mridulm left a comment

mridulm Mar 13, 2024

cfmcgrady commented Mar 14, 2024

wangshengjie123 commented Mar 16, 2024

codecov bot commented Mar 16, 2024 •

edited

Loading

RexXiong left a comment

RexXiong Mar 19, 2024

wangshengjie123 Mar 24, 2024

RexXiong Mar 19, 2024

wangshengjie123 Mar 24, 2024 •

edited

Loading

RexXiong Mar 19, 2024

wangshengjie123 commented Mar 24, 2024

s0nskar commented Apr 4, 2024 •

edited

Loading

s0nskar commented Apr 4, 2024

pan3793 commented Apr 4, 2024

s0nskar commented Apr 4, 2024

s0nskar commented Apr 4, 2024

pan3793 commented Apr 4, 2024 •

edited

Loading

waitinfuture commented Apr 4, 2024 •

edited

Loading

waitinfuture commented Apr 4, 2024 •

edited

Loading

s0nskar commented Apr 4, 2024

pan3793 commented Apr 4, 2024 •

edited

Loading

s0nskar Apr 5, 2024

mridulm commented Apr 6, 2024 •

edited

Loading

waitinfuture commented Apr 6, 2024

mridulm commented Apr 6, 2024 •

edited

Loading

mridulm commented Apr 6, 2024

waitinfuture commented Apr 6, 2024


		int step = locations.length / subPartitionSize;

		// if partition location is [1,2,3,4,5,6,7,8,9,10], and skew partition split to 3 task:

[CELEBORN-1319] Optimize skew partition logic for Reduce Mode to avoid sorting shuffle files #2373

Are you sure you want to change the base?

[CELEBORN-1319] Optimize skew partition logic for Reduce Mode to avoid sorting shuffle files #2373

Conversation

wangshengjie123 commented Mar 11, 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?

waitinfuture 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

mridulm left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

cfmcgrady commented Mar 14, 2024

wangshengjie123 commented Mar 16, 2024

codecov bot commented Mar 16, 2024 • edited Loading

Codecov Report

RexXiong 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

wangshengjie123 Mar 24, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

wangshengjie123 commented Mar 24, 2024

s0nskar commented Apr 4, 2024 • edited Loading

s0nskar commented Apr 4, 2024

pan3793 commented Apr 4, 2024

s0nskar commented Apr 4, 2024

s0nskar commented Apr 4, 2024

pan3793 commented Apr 4, 2024 • edited Loading

waitinfuture commented Apr 4, 2024 • edited Loading

waitinfuture commented Apr 4, 2024 • edited Loading

s0nskar commented Apr 4, 2024

pan3793 commented Apr 4, 2024 • edited Loading

Choose a reason for hiding this comment

mridulm commented Apr 6, 2024 • edited Loading

waitinfuture commented Apr 6, 2024

mridulm commented Apr 6, 2024 • edited Loading

mridulm commented Apr 6, 2024

waitinfuture commented Apr 6, 2024

codecov bot commented Mar 16, 2024 •

edited

Loading

wangshengjie123 Mar 24, 2024 •

edited

Loading

s0nskar commented Apr 4, 2024 •

edited

Loading

pan3793 commented Apr 4, 2024 •

edited

Loading

waitinfuture commented Apr 4, 2024 •

edited

Loading

waitinfuture commented Apr 4, 2024 •

edited

Loading

pan3793 commented Apr 4, 2024 •

edited

Loading

mridulm commented Apr 6, 2024 •

edited

Loading

mridulm commented Apr 6, 2024 •

edited

Loading