[SPARK-40480][SHUFFLE] Remove push-based shuffle data after query finished #37922

wankunde · 2022-09-18T14:25:52Z

What changes were proposed in this pull request?

Cleanup merged shuffle data files after query finished.

The main changes are:

After push merge service received FinalizeShuffleMerge RPC from driver, it will mark the MergePartitionsInfo as finilized instead of remove it.
Delete the shuffle merge files if partition.mapTracker is empty.
When spark driver begin to cleanup shuffle, if push-based shuffle is enabled, spark driver will send RemoveShuffleMerge RPC to all merger locations to delete all merge data files.

Why are the changes needed?

There will be too many merged shuffle data files for long running spark applications.

Does this PR introduce any user-facing change?

No

How was this patch tested?

Local cluster test.

mridulm · 2022-09-18T18:10:13Z

+CC @akpatnam25, @otterc - since you were planning to work on this already.

mridulm

Thanks for working on this @wankunde !

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/BlockStoreClient.java

...work-shuffle/src/main/java/org/apache/spark/network/shuffle/protocol/RemoveShuffleMerge.java

mridulm · 2022-09-18T18:16:32Z

.../network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java

@@ -393,6 +393,20 @@ public void applicationRemoved(String appId, boolean cleanupLocalDirs) {
    }
  }

+  @Override
+  public void removeShuffleMerge(String appId, int shuffleId) {


When adding support for shuffleMergeId, follow the same pattern as finalizeShuffleMerge - there are a few corner cases here, and that method handles them.

...work-shuffle/src/main/java/org/apache/spark/network/shuffle/protocol/RemoveShuffleMerge.java

core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala

core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala

wankunde · 2022-09-19T05:02:36Z

Hi, @mridulm Thanks for your review.

In my opinion, I do not think we should add ShuffleMergeId in RemoveShuffleMerge.
The push-based shuffle service will only keep the latest shuffleMergeId data for each shuffle, and RemoveShuffleMerge will cleanup its data.
Just like cleaning up normal shuffle data written by mapTasks, these methods only need shuffleId: def removeShuffle(shuffleId: Int)

What do you think?

mridulm · 2022-09-19T05:38:17Z

We should decouple current implementation details when making protocol changes, and make it extensible for future evolution.

In this case though, it is much more straightforward - there is an existing usecase which requires shuffle merge id.
When retrying an indeterminate stage, we should cleanup merged shuffle data for previous stage attempt (in submitMissingTasks, before unregisterAllMapAndMergeOutput) - and given the potential race conditions there, we dont want RemoveShuffleMerge to clean up for the next attempt (when we add support for this).

This specific change can be done in a follow up PR though - I want to get the basic mechanics working in this PR, and ensure the cleanup usecase is handled - before looking at further enhancements.

wankunde · 2022-09-19T06:54:47Z

We should decouple current implementation details when making protocol changes, and make it extensible for future evolution.

In this case though, it is much more straightforward - there is an existing usecase which requires shuffle merge id. When retrying an indeterminate stage, we should cleanup merged shuffle data for previous stage attempt (in submitMissingTasks, before unregisterAllMapAndMergeOutput) - and given the potential race conditions there, we dont want RemoveShuffleMerge to clean up for the next attempt (when we add support for this).

This specific change can be done in a follow up PR though - I want to get the basic mechanics working in this PR, and ensure the cleanup usecase is handled - before looking at further enhancements.

The push-based shuffle service will auto clean up the old shuffle merge data, and the following stage will read the new stage's output, so we don't need send RemoveShuffleMerge RPC for a new ShuffleMerge?
The only scenario I can think of now where a cleanup RPC needs is the spark job completes. Could we think of other scenarios?

wankunde · 2022-09-19T06:56:25Z

retest this please

mridulm · 2022-09-19T17:51:43Z

The push-based shuffle service will auto clean up the old shuffle merge data

Consider the case I mentioned above - stage retry for an INDETERMINATE stage.
We cleanup previous attempts only if it happens to use the same merger in new attempt as well.

The previous attempt's mergers are not reused for the next attempt - and so the previous mergers will continue to hold stale data without cleaning them up - until application terminates.
Note - any merger which happens to be reused in new stage attempt will cleanup - I am referring to those which are not used: old attempt mergerLocs -- new attempts mergerLocs.

AmplabJenkins · 2022-09-19T21:10:27Z

Can one of the admins verify this patch?

wankunde · 2022-09-20T05:59:46Z

Thanks @mridulm , I will the the shuffleMergeId into the protocol later.
This PR also has some other issues, I will fix it soon.

akpatnam25 · 2022-09-21T21:47:57Z

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalBlockHandler.java

+    } else if (msgObj instanceof RemoveShuffleMerge) {
+      RemoveShuffleMerge msg = (RemoveShuffleMerge) msgObj;
+      checkAuth(client, msg.appId);
+      logger.info("Remove shuffle merge data for application %s shuffle %d",


Nit: remove --> removing

...network-shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalBlockStoreClient.java

...network-shuffle/src/main/java/org/apache/spark/network/shuffle/MergedShuffleFileManager.java

core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala

otterc · 2022-10-06T18:06:38Z

Gentle ping @wankunde. Do you think you can update the PR soon? Please let us know if you need any help.

wankunde · 2022-10-09T04:52:53Z

Gentle ping @wankunde. Do you think you can update the PR soon? Please let us know if you need any help.

Sorry for the late reply, I will update this PR today.

wankunde · 2022-10-09T06:39:00Z

Hi, @otterc @mridulm If the spark application stops unexpectedly, there will be some leaked merge files. In our production cluster, yarn nodeManager will cleanup these leaked files if the disk is not full, but not if the disk is full.
Could we change the blockHandler.applicationRemoved(appId, false /* clean up local dirs */) to blockHandler.applicationRemoved(appId, true) in YarnShuffleService to clean up them ?

spark/common/network-yarn/src/main/java/org/apache/spark/network/yarn/YarnShuffleService.java

Line 433 in e669957

blockHandler.applicationRemoved(appId, false /* clean up local dirs */);

mridulm

Added a few comments.
While going over this change, I think I might have hit upon an existing bug as well - so thanks for raising this PR for two reasons now !

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalBlockHandler.java

...network-shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalBlockStoreClient.java

.../network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java

mridulm · 2022-10-09T08:43:44Z

.../network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java

          submitCleanupTask(() ->
-              closeAndDeleteOutdatedPartitions(
-                  appAttemptShuffleMergeId, mergePartitionsInfo.shuffleMergePartitions));
+              closeAndDeleteOutdatedPartitions(appAttemptShuffleMergeId, shuffleMergePartitions));
        } else {


revert this ?

Now in finalizeShuffleMerge() method we will update the AppShuffleMergePartitionsInfo to an empty object, so we can not cleanup the leaked files?

See my note above on how to handle this.

mridulm · 2022-10-09T08:44:28Z

.../network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java

      // Update the DB for the finalized shuffle
      writeAppAttemptShuffleMergeInfoToDB(appAttemptShuffleMergeId);
      // Even when the mergePartitionsInfo is null, we mark the shuffle as finalized but the results
      // sent to the driver will be empty. This can happen when the service didn't receive any
      // blocks for the shuffle yet and the driver didn't wait for enough time to finalize the
      // shuffle.
-      return new AppShuffleMergePartitionsInfo(msg.shuffleMergeId, true);
+      return mergePartitionsInfo;


revert changes to this method ? (assuming the changes I sketched above are fine)

mridulm · 2022-10-09T08:45:19Z

.../network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java

@@ -702,7 +722,8 @@ public MergeStatuses finalizeShuffleMerge(FinalizeShuffleMerge msg) {
                "finalizing shuffle partition {}", msg.appId, msg.appAttemptId, msg.shuffleId,
                msg.shuffleMergeId, partition.reduceId);
          } finally {
-            partition.closeAllFilesAndDeleteIfNeeded(false);
+            Boolean deleteFile = partition.mapTracker.getCardinality() == 0;
+            partition.closeAllFilesAndDeleteIfNeeded(deleteFile);


Won't return the reduceId which partition.mapTracker.getCardinality() == 0, so we can closeAllFilesAndDeleteIfNeeded now.

The actual deletion will be handled by either application termination, or during shuffle cleanup (this PR).
Making this change here might have other issues we will need to think through.

If we do want to do this, can you file a follow up jira and we can investigate it there ?
I want to limit this PR specifically to changes required for this feature.

Please revert this for the time being, and revisit it as a follow up as required.
There are a number of edge cases which needs to be handled in this important PR, and I dont want to add to them with optional changes with minor utility.

mridulm · 2022-10-09T08:45:41Z

.../network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java

@@ -1410,26 +1431,27 @@ public String toString() {
   * required for the shuffles of indeterminate stages.
   */
  public static class AppShuffleMergePartitionsInfo {


Revert changes to this class ? (assuming the changes I sketched above are fine)

.../network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java

core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala

mridulm · 2022-11-17T08:04:28Z

I will try to get to this later this week, do let me know if you are still working on it/have pending comments to address. Thanks

wankunde · 2022-11-21T03:05:25Z

Hi, @mridulm , I've been working on some other issues recently. If @yabola can do all or part of this task in #38560, please go ahead.

mridulm

Took a quick look, and some of the comments from the previous review are not yet addressed.
Do let me know if you have queries or comments on this @wankunde .

Thanks

mridulm · 2022-11-21T06:54:28Z

...network-shuffle/src/main/java/org/apache/spark/network/shuffle/MergedShuffleFileManager.java

+  /**
+   * Remove shuffle merge data files.
+   *
+   * @param removeShuffleMerge Remove shuffle merge RPC


Suggested change

* @param removeShuffleMerge Remove shuffle merge RPC

* @param removeShuffleMerge contains shuffle details (appId, shuffleId, etc) to uniquely identify a shuffle to be removed

mridulm · 2022-11-21T06:55:47Z

...ork-shuffle/src/main/java/org/apache/spark/network/shuffle/NoOpMergedShuffleFileManager.java

+
+  @Override
+  public void removeShuffleMerge(RemoveShuffleMerge removeShuffleMerge) {
+    throw new UnsupportedOperationException("Cannot handle shuffle block merge");


nit:

Suggested change

throw new UnsupportedOperationException("Cannot handle shuffle block merge");

throw new UnsupportedOperationException("Cannot handle merged shuffle remove");

mridulm · 2022-11-21T07:00:06Z

.../network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java

+    }
+
+    appShuffleInfo.shuffles.compute(shuffleId, (shuffleIdKey, partitionsInfo) -> {
+      if (null != partitionsInfo) {


The validation related to shuffleMergeId, from finalizeShuffleMerge, is applicable here depending on shuffleMergeId.
I have given the details here - please refer to it.

wankunde · 2022-12-18T07:35:51Z

Hi, @mridulm @yabola I have updated this PR, could you help to this PR again? Thanks.

wankunde · 2023-01-01T13:28:12Z

@mridulm I rebase the code, could you help to review this PR again? Thanks

mridulm · 2023-01-01T19:48:39Z

I just merged #39316 to unblock this PR, can you update to latest please ? Thx

mridulm

Just couple of minor comments - rest looks fine.
Please ping me when done, I want to try to get this in for 3.4, thanks !

mridulm · 2023-01-07T05:05:20Z

.../network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java

+        File dataFile =
+            appShuffleInfo.getMergedShuffleDataFile(shuffleId, shuffleMergeId, reduceId);
+        dataFile.delete();
+      } catch (Exception e) {


Can we fix this @wankunde ?
Essentially, the changes are:

a) There is no exception thrown in this block - we dont need the try/catch.
b) When the delete fails, we dont need to do the warn (here and below) : this can happen if application exit is racing against remove shuffle - the info at the end of the method will suffice.

mridulm · 2023-01-07T05:08:45Z

...ork-shuffle/src/test/java/org/apache/spark/network/shuffle/RemoteBlockPushResolverSuite.java

+    assertTrue(shuffleInfo.getMergedShuffleMetaFile(3, 1, 0).exists());
+    pushResolver.finalizeShuffleMerge(new FinalizeShuffleMerge(testApp, NO_ATTEMPT_ID, 3, 2));
+    closed.tryAcquire(10, TimeUnit.SECONDS);
+    assertFalse(shuffleInfo.getMergedShuffleMetaFile(3, 1, 0).exists());


Can we also add tests for the additional cases we have ?

removeShuffleMerge for DELETE_ALL_MERGED_SHUFFLE - we can just do what we have for streamCallback1 (Case 1) above, but use DELETE_ALL_MERGED_SHUFFLE for merge id.

additional cases here to make sure of the negative cases as well ? namely:

msg.shuffleMergeId < mergePartitionsInfo.shuffleMergeId

delete for a merge id which does not exist.

mridulm

Looks good to me, thanks for the changes @wankunde !

+CC @otterc, @Ngone51

core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala

...network-shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalBlockStoreClient.java

...ork-shuffle/src/main/java/org/apache/spark/network/shuffle/NoOpMergedShuffleFileManager.java

.../network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java

mridulm

Can we add a test to check whether the message goes to the expected mergers from driver for RemoteShuffleMerge when a shuffle is getting cleaned up ?
While trying to test the proposal here, I realized we are not testing this behavior - even replacing it with Nil always causes all tests to pass !

mridulm · 2023-01-12T11:37:28Z

core/src/test/scala/org/apache/spark/MapOutputTrackerSuite.scala

+      sc.cleaner.get.doCleanupShuffle(SHUFFLE_ID, blocking = true)
+      assert(foundHosts.asScala == mergerLocs.map(_.host).toSet)
+    }
+  }


There is a bug in this test - sorry about that.
Please take a look here: 26b2d18

Backport the code, much thanks

otterc

Just have some minor comments

otterc · 2023-01-12T17:06:02Z

.../network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java

+              msg.appId, msg.appAttemptId, msg.shuffleId, mergePartitionsInfo.shuffleMergeId);
+      AppAttemptShuffleMergeId appAttemptShuffleMergeId = new AppAttemptShuffleMergeId(
+          msg.appId, msg.appAttemptId, msg.shuffleId, shuffleMergeId);
+      if(deleteCurrentMergedShuffle || shuffleMergeId > mergePartitionsInfo.shuffleMergeId) {


Nit: missing a space after if

otterc · 2023-01-12T17:14:42Z

.../network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java

+              deleteMergedFiles(currentAppAttemptShuffleMergeId, appShuffleInfo,
+                  mergePartitionsInfo.getReduceIds(), false));
+        }
+      } else if(shuffleMergeId < mergePartitionsInfo.shuffleMergeId) {


nit: space missing after if. Also the condition here is not necessary. It will always be true (IDE shows that as well).

.../network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java

…shuffle/RemoteBlockPushResolver.java Co-authored-by: otterc <singh.chandni@gmail.com>

wankunde · 2023-01-13T03:46:01Z

@otterc Thanks for your review. Format the code.

mridulm

Added one comment.
Rest looks good to me.

Will it be possible to change that and push an update soon @wankunde ? I will merge it once the tests pass successfully.

...network-shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalBlockStoreClient.java

mridulm · 2023-01-14T03:59:19Z

The tests are failing @wankunde, though I dont think it is due to your PR.
Can you please take a look ? And retrigger it if it is unrelated ? Thanks !

core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala

mridulm · 2023-01-14T10:28:14Z

Merged to master.
Thanks for fixing this @wankunde !
Thanks for the reviews @otterc, @akpatnam25, @yabola :-)

dongjoon-hyun

.

github-actions bot added the CORE label Sep 18, 2022

mridulm reviewed Sep 18, 2022

View reviewed changes

HyukjinKwon changed the title ~~[SPARK-40480][SHUFFLE]Remove push-based shuffle data after query finished~~ [SPARK-40480][SHUFFLE] Remove push-based shuffle data after query finished Sep 19, 2022

wankunde force-pushed the remove_shuffle_merge branch from c32bea1 to 79df01a Compare September 19, 2022 06:59

wankunde changed the title ~~[SPARK-40480][SHUFFLE] Remove push-based shuffle data after query finished~~ [WIP][SPARK-40480][SHUFFLE] Remove push-based shuffle data after query finished Sep 19, 2022

akpatnam25 reviewed Sep 21, 2022

View reviewed changes

...network-shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalBlockStoreClient.java Show resolved Hide resolved

akpatnam25 reviewed Sep 21, 2022

View reviewed changes

...network-shuffle/src/main/java/org/apache/spark/network/shuffle/MergedShuffleFileManager.java Show resolved Hide resolved

weixiuli reviewed Sep 26, 2022

View reviewed changes

core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala Outdated Show resolved Hide resolved

wankunde changed the title ~~[WIP][SPARK-40480][SHUFFLE] Remove push-based shuffle data after query finished~~ [SPARK-40480][SHUFFLE] Remove push-based shuffle data after query finished Oct 9, 2022

mridulm reviewed Oct 9, 2022

View reviewed changes

github-actions bot added the YARN label Oct 12, 2022

mridulm mentioned this pull request Nov 10, 2022

[SPARK-38005][core] Support cleaning up merged shuffle files and state from external shuffle service #38560

Closed

mridulm reviewed Nov 21, 2022

View reviewed changes

wankunde force-pushed the remove_shuffle_merge branch from a3ceecb to ce2d576 Compare November 23, 2022 02:42

wankunde force-pushed the remove_shuffle_merge branch from ce2d576 to 091574b Compare December 18, 2022 03:31

wankunde force-pushed the remove_shuffle_merge branch from 6c409dc to 29f4918 Compare January 1, 2023 02:35

move writeAppAttemptShuffleMergeInfoToDB to main thread

3a7d99c

mridulm reviewed Jan 7, 2023

View reviewed changes

remove try catch for file.delete()

613a99a

mridulm approved these changes Jan 8, 2023

View reviewed changes

otterc reviewed Jan 10, 2023

View reviewed changes

Simplify the code and update some comments

96ba7e1

mridulm requested changes Jan 12, 2023

View reviewed changes

Add UT

e06ef76

mridulm reviewed Jan 12, 2023

View reviewed changes

otterc reviewed Jan 12, 2023

View reviewed changes

wankunde and others added 2 commits January 13, 2023 11:40

Update common/network-shuffle/src/main/java/org/apache/spark/network/…

43085d1

…shuffle/RemoteBlockPushResolver.java Co-authored-by: otterc <singh.chandni@gmail.com>

Format code

371feae

wankunde added 2 commits January 13, 2023 11:49

Add a TODO for RemoveShuffleMerge RPC

6975bcc

Fix bug and backport UT

475f8bd

mridulm approved these changes Jan 13, 2023

View reviewed changes

...network-shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalBlockStoreClient.java Outdated Show resolved Hide resolved

wankunde added 2 commits January 14, 2023 10:11

Change log level to debug

0f1f5eb

Change log level to debug

7284b8f

wankunde force-pushed the remove_shuffle_merge branch from 3cc6896 to 7284b8f Compare January 14, 2023 02:13

mridulm reviewed Jan 14, 2023

View reviewed changes

core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala Outdated Show resolved Hide resolved

Bug fix

3181e64

asfgit closed this in e424499 Jan 14, 2023

dongjoon-hyun reviewed Feb 4, 2023

View reviewed changes

	* @param removeShuffleMerge Remove shuffle merge RPC
	* @param removeShuffleMerge contains shuffle details (appId, shuffleId, etc) to uniquely identify a shuffle to be removed

	throw new UnsupportedOperationException("Cannot handle shuffle block merge");
	throw new UnsupportedOperationException("Cannot handle merged shuffle remove");

[SPARK-40480][SHUFFLE] Remove push-based shuffle data after query finished #37922

[SPARK-40480][SHUFFLE] Remove push-based shuffle data after query finished #37922

Conversation

wankunde commented Sep 18, 2022 • edited

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?

mridulm commented Sep 18, 2022

mridulm left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

wankunde commented Sep 19, 2022

mridulm commented Sep 19, 2022

wankunde commented Sep 19, 2022 • edited

wankunde commented Sep 19, 2022

mridulm commented Sep 19, 2022

AmplabJenkins commented Sep 19, 2022

wankunde commented Sep 20, 2022

Choose a reason for hiding this comment

otterc commented Oct 6, 2022

wankunde commented Oct 9, 2022

wankunde commented Oct 9, 2022 • edited

mridulm left a comment • edited

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 Oct 9, 2022 • edited

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 Oct 9, 2022 • edited

Choose a reason for hiding this comment

mridulm commented Nov 17, 2022

wankunde commented Nov 21, 2022

mridulm 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

wankunde commented Dec 18, 2022

wankunde commented Jan 1, 2023

mridulm commented Jan 1, 2023

mridulm left a 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

mridulm left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

otterc left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

wankunde commented Jan 13, 2023

mridulm left a comment

Choose a reason for hiding this comment

mridulm commented Jan 14, 2023

mridulm commented Jan 14, 2023

dongjoon-hyun left a comment • edited

Choose a reason for hiding this comment

wankunde commented Sep 18, 2022 •

edited

wankunde commented Sep 19, 2022 •

edited

wankunde commented Oct 9, 2022 •

edited

mridulm left a comment •

edited

mridulm Oct 9, 2022 •

edited

mridulm Oct 9, 2022 •

edited

dongjoon-hyun left a comment •

edited