[#1608] feat: Introduce ExpiringClosableSupplier and refactor ShuffleManagerClient creation #1838

xumanbu · 2024-06-26T13:28:55Z

What changes were proposed in this pull request?

(Please outline the changes and how this PR fixes the issue.)

Why are the changes needed?

The createShuffleManagerClient function is defined in both RssShuffleWriter and RssFetchFailedIterator.
The ShuffleManagerClient is created in RssShuffleManager when shuffleManagerRpcServiceEnabled is true. We can reuse the client in RssShuffleReader and RssShuffleWriter.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

Existing UTs.

…d ShuffleWriter

github-actions · 2024-06-26T13:51:45Z

Test Results

2 747 files + 15 2 747 suites +15 9h 48m 23s ⏱️ + 4h 11m 44s
978 tests + 9 977 ✅ + 9 1 💤 ±0 0 ❌ ±0
12 279 runs +135 12 264 ✅ +135 15 💤 ±0 0 ❌ ±0

Results for commit 02ca92c. ± Comparison against base commit 48d23ce.

♻️ This comment has been updated with latest results.

xumanbu · 2024-06-27T02:26:09Z

@zuston @advancedxy PTAL

codecov-commenter · 2024-06-27T09:56:19Z

Codecov Report

Attention: Patch coverage is 41.66667% with 49 lines in your changes missing coverage. Please review.

Project coverage is 52.62%. Comparing base (f8e4329) to head (0bb986c).
Report is 41 commits behind head on master.

Files	Patch %	Lines
...e/spark/shuffle/reader/RssFetchFailedIterator.java	0.00%	14 Missing ⚠️
...uniffle/shuffle/manager/RssShuffleManagerBase.java	0.00%	11 Missing ⚠️
...org/apache/spark/shuffle/RssSparkShuffleUtils.java	0.00%	9 Missing ⚠️
.../shuffle/BlockIdSelfManagedShuffleWriteClient.java	0.00%	5 Missing ⚠️
...le/client/factory/ShuffleManagerClientFactory.java	0.00%	3 Missing ⚠️
...fle/client/impl/grpc/ShuffleManagerGrpcClient.java	50.00%	3 Missing ⚠️
...pache/uniffle/shuffle/RssShuffleClientFactory.java	0.00%	2 Missing ⚠️
...uniffle/common/util/ExpiringCloseableSupplier.java	94.11%	0 Missing and 2 partials ⚠️

Additional details and impacted files

@@              Coverage Diff              @@
##             master    #1838       +/-   ##
=============================================
+ Coverage          0   52.62%   +52.62%     
- Complexity        0     2561     +2561     
=============================================
  Files             0      415      +415     
  Lines             0    18784    +18784     
  Branches          0     1703     +1703     
=============================================
+ Hits              0     9885     +9885     
- Misses            0     8288     +8288     
- Partials          0      611      +611

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

advancedxy · 2024-06-27T11:19:01Z

client-spark/spark2/src/main/java/org/apache/spark/shuffle/RssShuffleManager.java

@@ -446,6 +446,7 @@ public <K, V> ShuffleWriter<K, V> getWriter(
          this,
          sparkConf,
          shuffleWriteClient,
+          shuffleManagerClient,


hmm, I don't think it should be called directly.

A more appreciate way should be defining a new method such as lazyShuffleManagerClient, which will ensure
shuffle manager client is created if null.

Other places should be updated too.

Agree with u. my original idea is same with u, but I discovered it already initialized in the RssShuffleManage construct, so I simply addressed it.

Ok. I will attempt to implement a lazyShuffleManagerClient instead of shuffleManagerClient initialized in new RssShuffleManage at now.

I have refactored the getReader and getWriter constructs, and I understand why the RssShuffleManage should initialize the shuffleManagerClient in the constructor. in case of RssShuffleClientFactory.ExtendWriteClientBuilder need it. so I also make ExtendWriteClientBuilder.shuffleManagerClient is lazy.

@advancedxy @zuston PTALA.

advancedxy

Thanks for your contribution. Left one minor comment.

advancedxy · 2024-06-28T06:49:46Z

client-spark/spark2/src/main/java/org/apache/spark/shuffle/RssShuffleManager.java

@@ -446,6 +442,7 @@ public <K, V> ShuffleWriter<K, V> getWriter(
          this,
          sparkConf,
          shuffleWriteClient,
+          this::getOrCreateShuffleManagerClient,


hmm, maybe I wasn't clear enough and mislead you. It's fine and probably better to accept a ShuffleManagerClient in builders instead of a supplier.

Instead of passing the shuffleManagerClient directly, we probably want to call getOrCreateShuffleManagerClient() instead to make sure the client is actually initialized.

I think the lifecycle of ShuffleManagerClient may be different with the shuffle manager, if this is idle with some times, it should be closed(this is not supported now). If that, we should provide the wrapper to create or get the active ShuffleManagerClient .

I think the lifecycle of ShuffleManagerClient may be different with the shuffle manager, if this is idle with some times, it should be closed(this is not supported now). If that, we should provide the wrapper to create or get the active ShuffleManagerClient .

agree with u. the lifecycle of ShuffleManagerClient is major factor in guiding how to modify

hmm, maybe I wasn't clear enough and mislead you. It's fine and probably better to accept a ShuffleManagerClient in builders instead of a supplier.

Instead of passing the shuffleManagerClient directly, we probably want to call getOrCreateShuffleManagerClient() instead to make sure the client is actually initialized.

Do you might be considering make ShuffleManagerClient as a single object or a LazyShuffleManagerClient impl which logic like a lazy val shuffleManagerClient in Scala ?

Properly not. Maybe like the following code.

class AutoClosableShuffleManagerClient extends ShuffleManagerClient { private ShuffleManagerClient client; private int reference = 0 public get() { reference += 1; if (client == null) { client = new ShuffleManagerClient(); } } public void close() { reference -= 1 if (reference == 0) { client.close() client = null } } }

Maybe this is not right, just my simple thought.

@zuston @advancedxy I have add a AutoCloseWrapper in commit, please take a look and verify it's accuracy?

xumanbu · 2024-06-28T07:41:12Z

client-spark/common/src/main/java/org/apache/uniffle/shuffle/manager/RssShuffleManagerBase.java

+  public void stop() {
+    if (shuffleManagerClient != null) {
+      try {
+        shuffleManagerClient.close();


@zuston I make the lifecycle of ShuffleManagerClient with ShuffleManage, referencee with ShuffleWriteClient.

rickyma · 2024-06-28T09:21:21Z

If we are trying to refactor ShuffleManagerClient, could you also consider to fix #1774 as well? @xumanbu

rssConf in ShuffleManagerGrpcClient is hard-coded.

xumanbu · 2024-07-01T12:44:01Z

If we are trying to refactor ShuffleManagerClient, could you also consider to fix #1774 as well? @xumanbu

rssConf in ShuffleManagerGrpcClient is hard-coded.

ok.

This reverts commit 2906e54.

advancedxy · 2024-07-04T15:14:07Z

common/src/main/java/org/apache/uniffle/common/util/AutoCloseWrapper.java

+
+import org.apache.uniffle.common.exception.RssException;
+
+public class AutoCloseWrapper<T extends Closeable> implements Closeable {


I'm not sure about this, how about we leverage some building blocks in Guava, such as:

CacheBuilder, which closes the the closeable when evicted.

or A sightly modified version of ExpiringMemoizingSupplier which should close the previous closeable(Client) when refreshing to get a new one, we can implement that in Uniffle.

Thank you for your advice, let me reconsider.

@advancedxy I referenced the ExpiringMemoizingSupplier and used ExpireCloseableSupplier instead of AutoCloseWrapper in the last commit. Please check if this meets your expectations.

add some ideas：
ExpireCloseableSupplier is compared to AutoCloseWrapper, the ExpireCloseableSupplier implements the Supplier interface while the AutoCloseWrapper implements Closeable.

I believe AutoCloseWrapper is superior, because it allows for better management of client closure when obtaining a client.

advancedxy · 2024-07-04T15:22:24Z

...rk/common/src/main/java/org/apache/uniffle/shuffle/BlockIdSelfManagedShuffleWriteClient.java

@@ -73,7 +74,11 @@ public void reportShuffleResult(
    RssReportShuffleResultRequest request =
        new RssReportShuffleResultRequest(
            appId, shuffleId, taskAttemptId, partitionToBlockIds, bitmapNum);
-    shuffleManagerClient.reportShuffleResult(request);
+    AutoCloseWrapper.run(


hmmm. This seems a bit of code smell here, I don't think we should introduce such complex construct to simply delay the close of closeable. It brings the burden of using AutoCloseWrapper.run on every invocation.

advancedxy

@xumanbu thanks for your hard work. I think it's in good shape and the right direction.

The biggest issue is how to handle close from external caller for ExpiringCloseableSupplier.

advancedxy · 2024-07-23T02:49:40Z

client-spark/common/src/main/java/org/apache/spark/shuffle/RssSparkShuffleUtils.java

@@ -346,6 +343,7 @@ public static boolean isStageResubmitSupported() {
  }

  public static RssException reportRssFetchFailedException(
+      ExpireCloseableSupplier<ShuffleManagerClient> managerClientSupplier,


I think we can simply use Supplier<ShuffleMangerClient> here?

The delayed close is an implementation detail and doesn't have to be exposed to the API side.

One defeat of this approach is that callers may close the ShuffleManagerClient explicitly, which might break how ExpireCloseableSupplier<ShuffleManagerClient> work: once it's closed by a caller, it cannot be reused again, however, the supplier has no idea it's already been closed.

For that case, I think we should callout in the comment/javadoc to indicate that ShuffleManagerClients are self-closeable.

advancedxy · 2024-07-23T02:49:59Z

client-spark/common/src/main/java/org/apache/spark/shuffle/reader/RssFetchFailedIterator.java

@@ -52,8 +51,7 @@ public static class Builder {
    private int shuffleId;
    private int partitionId;
    private int stageAttemptId;
-    private String reportServerHost;
-    private int reportServerPort;
+    private ExpireCloseableSupplier<ShuffleManagerClient> managerClientSupplier;


client-spark/common/src/main/java/org/apache/spark/shuffle/reader/RssFetchFailedIterator.java

client-spark/common/src/main/java/org/apache/uniffle/shuffle/manager/RssShuffleManagerBase.java

client-spark/spark2/src/main/java/org/apache/spark/shuffle/RssShuffleManager.java

common/src/main/java/org/apache/uniffle/common/util/ExpireCloseableSupplier.java

xumanbu · 2024-07-23T03:42:53Z

@xumanbu thanks for your hard work. I think it's in good shape and the right direction.

The biggest issue is how to handle close from external caller for ExpiringCloseableSupplier.

@advancedxy Very thank you for helping me review the code and for your valuable suggestions. I'm still a bit confused by some comments, I need think about more and optimize the code.

advancedxy · 2024-07-23T09:10:19Z

@xumanbu thanks for your hard work. I think it's in good shape and the right direction.
The biggest issue is how to handle close from external caller for ExpiringCloseableSupplier.

@advancedxy Very thank you for helping me review the code and for your valuable suggestions. I'm still a bit confused by some comments, I need think about more and optimize the code.

You are welcome. And if you have any further questions or problems, feel free to ask here.

common/src/main/java/org/apache/uniffle/common/util/CloseStateful.java

advancedxy · 2024-07-24T10:51:53Z

common/src/test/java/org/apache/uniffle/common/util/ExpiringCloseableSupplierTest.java

+  }
+
+  @Test
+  void testDelegateExtendClose() throws IOException {


Thanks for this test code, looking good.

common/src/test/java/org/apache/uniffle/common/util/ExpiringCloseableSupplierTest.java

client-spark/common/src/main/java/org/apache/uniffle/shuffle/manager/RssShuffleManagerBase.java

common/src/main/java/org/apache/uniffle/common/util/ExpiringCloseableSupplier.java

advancedxy · 2024-07-24T12:30:51Z

@xumanbu other parts lgtm now. The ExpiringCloseableSupplier may still needs some refinements. Please take a look at the PR I raised in your repo.

also cc @zuston to take a look if you are interested.

Refine ExpiringCloseableSupplier

common/src/main/java/org/apache/uniffle/common/util/ExpiringCloseableSupplier.java

advancedxy

LGTM. I will leave this for some more time before others have a chance to take a look at this.

common/src/main/java/org/apache/uniffle/common/util/ExpiringCloseableSupplier.java

rickyma · 2024-07-25T14:41:46Z

client-spark/common/src/main/java/org/apache/uniffle/shuffle/manager/RssShuffleManagerBase.java

-      this.shuffleManagerClient =
-          ShuffleManagerClientFactory.getInstance()
-              .createShuffleManagerClient(ClientType.GRPC, driver, port);
+      long rpcTimeout = rssConf.getLong(RssBaseConf.RSS_CLIENT_TYPE_GRPC_TIMEOUT_MS);


Should we unify rss.client.rpc.timeout.ms and rss.rpc.client.type.grpc.timeout?

Emm, the rss.rpc.client.type.grpc.timeout configuration is only used in ShuffleManageClient, so we can remove it. instead it with rss.client.rpc.timeout.ms

Could you fix this in this PR as well?

How about address this in a follow-up PR? I think this PR is coupling too much changes.

advancedxy

LGTM, pending CI passes.

Let's wait it until tonight incase others have more comments.

advancedxy · 2024-07-26T13:22:35Z

Merging this since there's no objections. Others are welcome to file additional comments if any and can be addressed in a follow-up PR.

advancedxy · 2024-07-26T13:29:37Z

BTW, thanks your contribution @xumanbu.

…huffleManagerClient creation (apache#1838) ### What changes were proposed in this pull request? 1. Introduce StatefulCloseable and ExpiringClosableSupplier 2. refactor ShuffleManagerClient to leverage ExpiringClosableSupplier ### Why are the changes needed? For better code quality ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing UTs and new UTs. (cherry picked from commit 457c865)

[apache#1608] refactor: Reuse ShuffleManageClient in ShuffleReader an…

fe994f6

…d ShuffleWriter

xumanbu changed the title ~~[#1608] refactor: refactor: Reuse ShuffleManageClient in ShuffleReader and ShuffleWriter~~ [#1608] refactor: Reuse ShuffleManageClient in ShuffleReader and ShuffleWriter Jun 26, 2024

xumanbu added 3 commits June 27, 2024 17:05

fix spark2 & refactor RssFetchFailedIterator

2e557ea

fix findbugs

38831c9

fix code style

35e9a56

advancedxy reviewed Jun 27, 2024

View reviewed changes

xumanbu added 2 commits June 28, 2024 11:49

lazyShuffleManagerClient support

7fb22ce

fix some complie bug

8a12e0c

advancedxy reviewed Jun 28, 2024

View reviewed changes

xumanbu commented Jun 28, 2024

View reviewed changes

rickyma changed the title ~~[#1608] refactor: Reuse ShuffleManageClient in ShuffleReader and ShuffleWriter~~ [#1608] refactor: Reuse ShuffleManagerClient in ShuffleReader and ShuffleWriter Jun 28, 2024

xumanbu added 11 commits July 1, 2024 21:29

add AutoCloseWrapper

9924baa

fix test

06dfe7a

fix code style

f50e466

AutoCloseWrapper add static run method

cb495ad

change managerClientAutoCloseWrapper create

dc02a3b

ShuffleManagerGrpcClient support rpcTimeout

4f1b7ca

fix bug

785e4be

stash

2906e54

Revert "stash"

cd652a0

This reverts commit 2906e54.

fix bug

8cfb044

fix spark2

5b8657c

advancedxy reviewed Jul 4, 2024

View reviewed changes

utoCloseWrapper replace with ExpireCloseableSupplier

1499396

xumanbu force-pushed the reuse_shuffle_manage_client branch from 7023449 to 1499396 Compare July 22, 2024 09:19

advancedxy reviewed Jul 23, 2024

View reviewed changes

xumanbu force-pushed the reuse_shuffle_manage_client branch from 6aa6674 to d9a1ab9 Compare July 24, 2024 06:33

add CloseStateful interface

0bb986c

xumanbu force-pushed the reuse_shuffle_manage_client branch from d9a1ab9 to 0bb986c Compare July 24, 2024 06:48

Refine ExpiringCloseableSupplier

f1cb900

advancedxy reviewed Jul 24, 2024

View reviewed changes

xumanbu and others added 4 commits July 25, 2024 10:30

Merge pull request #1 from advancedxy/reuse_shuffle_manage_client

f949fc7

Refine ExpiringCloseableSupplier

fix spark2

dc66b08

StatefulCloseable rename to StatefulCloseable

c8a35e7

fix accessTime by findbug SE_TRANSIENT_FIELD_NOT_RESTORED

5d14f36

advancedxy reviewed Jul 25, 2024

View reviewed changes

common/src/main/java/org/apache/uniffle/common/util/ExpiringCloseableSupplier.java Outdated Show resolved Hide resolved

xumanbu added 2 commits July 25, 2024 14:26

fix default rpcTimeout 60s

88a8a6c

opt accessTime

6e525de

advancedxy previously approved these changes Jul 25, 2024

View reviewed changes

common/src/main/java/org/apache/uniffle/common/util/ExpiringCloseableSupplier.java Outdated Show resolved Hide resolved

advancedxy changed the title ~~[#1608] refactor: Reuse ShuffleManagerClient in ShuffleReader and ShuffleWriter~~ [#1608] feat: Introduce ExpiringClosableSupplier and refactor ShuffleManagerClient creation Jul 25, 2024

rickyma reviewed Jul 25, 2024

View reviewed changes

purge ShuffleManagerClientFactory fun

2d657b1

xumanbu dismissed advancedxy’s stale review via 2d657b1 July 25, 2024 15:33

fix SE_TRANSIENT_FIELD_NOT_RESTORED

02ca92c

advancedxy approved these changes Jul 26, 2024

View reviewed changes

advancedxy merged commit 457c865 into apache:master Jul 26, 2024
43 checks passed

xumanbu mentioned this pull request Aug 2, 2024

Unify ShuffleManageClient config rss.client.rpc.timeout.ms and rss.rpc.client.type.grpc.timeout #2003

Closed


		import org.apache.uniffle.common.exception.RssException;

		public class AutoCloseWrapper<T extends Closeable> implements Closeable {

[#1608] feat: Introduce ExpiringClosableSupplier and refactor ShuffleManagerClient creation #1838

[#1608] feat: Introduce ExpiringClosableSupplier and refactor ShuffleManagerClient creation #1838

Conversation

xumanbu commented Jun 26, 2024 • edited by rickyma Loading

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?

github-actions bot commented Jun 26, 2024 • edited Loading

Test Results

xumanbu commented Jun 27, 2024

codecov-commenter commented Jun 27, 2024 • edited Loading

Codecov Report

Choose a reason for hiding this comment

xumanbu Jun 27, 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

advancedxy 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

rickyma commented Jun 28, 2024 • edited Loading

xumanbu commented Jul 1, 2024

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

advancedxy 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

xumanbu commented Jul 23, 2024 • edited Loading

advancedxy commented Jul 23, 2024

Choose a reason for hiding this comment

advancedxy commented Jul 24, 2024

advancedxy 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

advancedxy left a comment

Choose a reason for hiding this comment

advancedxy commented Jul 26, 2024

advancedxy commented Jul 26, 2024

xumanbu commented Jun 26, 2024 •

edited by rickyma

Loading

github-actions bot commented Jun 26, 2024 •

edited

Loading

codecov-commenter commented Jun 27, 2024 •

edited

Loading

xumanbu Jun 27, 2024 •

edited

Loading

rickyma commented Jun 28, 2024 •

edited

Loading

xumanbu commented Jul 23, 2024 •

edited

Loading