You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@uniffle.apache.org by ro...@apache.org on 2023/04/10 02:03:20 UTC
[incubator-uniffle] branch master updated: [#477][part-1] feat: support stage resubmit in spark clients (#787)
This is an automated email from the ASF dual-hosted git repository.
roryqi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-uniffle.git
The following commit(s) were added to refs/heads/master by this push:
new 97524710 [#477][part-1] feat: support stage resubmit in spark clients (#787)
97524710 is described below
commit 97524710cc3b95a1ce97072b84aa613789739bb3
Author: advancedxy <xi...@apache.org>
AuthorDate: Mon Apr 10 10:03:14 2023 +0800
[#477][part-1] feat: support stage resubmit in spark clients (#787)
### What changes were proposed in this pull request?
1. utility methods to create FetchFailedException and unregisterAllMapOutput dynamically
2. RssShuffleManagerBase to create a shareable base for both Spark 2 and Spark3
3. a wrapper to wrap shuffle data reader iterator that throws FetchFailedException if necessary
### Why are the changes needed?
This is the second part of #477
After this commit, spark with rss could re-submit stages if there's rss fetch failure.
### Does this PR introduce _any_ user-facing change?
Yes.
A new configuration would be introduced, which would be covered in another PR.
### How was this patch tested?
Added unit tests and integration tests.
---
.../org/apache/spark/shuffle/RssSparkConfig.java | 7 +
.../apache/spark/shuffle/RssSparkShuffleUtils.java | 53 ++++++++
.../apache/spark/shuffle/SparkVersionUtils.java | 52 ++++++++
.../shuffle/reader/RssFetchFailedIterator.java | 146 +++++++++++++++++++++
.../shuffle/manager/RssShuffleManagerBase.java | 143 ++++++++++++++++++++
.../manager/ShuffleManagerServerFactory.java | 13 +-
.../apache/spark/shuffle/RssShuffleManager.java | 78 ++++++++++-
.../spark/shuffle/reader/RssShuffleReader.java | 17 +++
.../spark/shuffle/RssSpark2ShuffleUtilsTest.java | 38 ++++++
.../spark/shuffle/SparkVersionUtilsTest.java | 29 ++++
.../apache/spark/shuffle/RssShuffleManager.java | 86 +++++++++++-
.../spark/shuffle/reader/RssShuffleReader.java | 16 +++
.../shuffle/DelegationRssShuffleManagerTest.java | 69 +---------
.../spark/shuffle/RssShuffleManagerTest.java | 44 ++++++-
.../spark/shuffle/RssShuffleManagerTestBase.java | 67 ++++++++++
.../spark/shuffle/RssSpark3ShuffleUtilsTest.java | 37 ++++++
.../spark/shuffle/SparkVersionUtilsTest.java | 29 ++++
.../uniffle/client/util/RssClientConfig.java | 1 +
.../uniffle/common/config/RssClientConf.java | 7 +
integration-test/common/pom.xml | 1 -
integration-test/spark-common/pom.xml | 6 +
.../apache/uniffle/test/RSSStageResubmitTest.java | 94 +++++++++++++
.../apache/uniffle/test/RssShuffleManagerTest.java | 96 ++++++++++++++
integration-test/spark3/pom.xml | 1 -
pom.xml | 7 +
.../server/MockedShuffleServerGrpcService.java | 27 ++++
26 files changed, 1091 insertions(+), 73 deletions(-)
diff --git a/client-spark/common/src/main/java/org/apache/spark/shuffle/RssSparkConfig.java b/client-spark/common/src/main/java/org/apache/spark/shuffle/RssSparkConfig.java
index 6aa90e1a..27fa4f7f 100644
--- a/client-spark/common/src/main/java/org/apache/spark/shuffle/RssSparkConfig.java
+++ b/client-spark/common/src/main/java/org/apache/spark/shuffle/RssSparkConfig.java
@@ -31,6 +31,7 @@ import org.apache.uniffle.client.util.RssClientConfig;
import org.apache.uniffle.common.config.ConfigOption;
import org.apache.uniffle.common.config.ConfigOptions;
import org.apache.uniffle.common.config.ConfigUtils;
+import org.apache.uniffle.common.config.RssClientConf;
import org.apache.uniffle.common.config.RssConf;
public class RssSparkConfig {
@@ -275,6 +276,12 @@ public class RssSparkConfig {
+ " spark.rss.estimate.server.assignment.enabled"))
.createWithDefault(RssClientConfig.RSS_ESTIMATE_TASK_CONCURRENCY_PER_SERVER_DEFAULT_VALUE);
+ public static final ConfigEntry<Integer> RSS_SHUFFLE_MANAGER_GRPC_PORT = createIntegerBuilder(
+ new ConfigBuilder(SPARK_RSS_CONFIG_PREFIX + RssClientConf.SHUFFLE_MANAGER_GRPC_PORT.key())
+ .internal()
+ .doc(RssClientConf.SHUFFLE_MANAGER_GRPC_PORT.description())
+ ).createWithDefault(-1);
+
// spark2 doesn't have this key defined
public static final String SPARK_SHUFFLE_COMPRESS_KEY = "spark.shuffle.compress";
diff --git a/client-spark/common/src/main/java/org/apache/spark/shuffle/RssSparkShuffleUtils.java b/client-spark/common/src/main/java/org/apache/spark/shuffle/RssSparkShuffleUtils.java
index dabe0e4b..7481be9b 100644
--- a/client-spark/common/src/main/java/org/apache/spark/shuffle/RssSparkShuffleUtils.java
+++ b/client-spark/common/src/main/java/org/apache/spark/shuffle/RssSparkShuffleUtils.java
@@ -18,6 +18,7 @@
package org.apache.spark.shuffle;
import java.lang.reflect.Constructor;
+import java.lang.reflect.InvocationTargetException;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
@@ -30,8 +31,10 @@ import org.apache.spark.SparkConf;
import org.apache.spark.SparkContext;
import org.apache.spark.broadcast.Broadcast;
import org.apache.spark.deploy.SparkHadoopUtil;
+import org.apache.spark.storage.BlockManagerId;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import scala.Option;
import scala.reflect.ClassTag;
import org.apache.uniffle.client.api.CoordinatorClient;
@@ -226,4 +229,54 @@ public class RssSparkShuffleUtils {
ShuffleHandleInfo handleInfo = new ShuffleHandleInfo(shuffleId, partitionToServers, storageInfo);
return sc.broadcast(handleInfo, SHUFFLE_HANDLER_INFO_CLASS_TAG);
}
+
+ private static <T> T instantiateFetchFailedException(
+ BlockManagerId dummy, int shuffleId, int mapIndex, int reduceId, Throwable cause) {
+ String className = FetchFailedException.class.getName();
+ T instance;
+ Class<?> klass;
+ try {
+ klass = Class.forName(className);
+ } catch (ClassNotFoundException e) {
+ // ever happens;
+ throw new RssException(e);
+ }
+ try {
+ instance = (T) klass
+ .getConstructor(dummy.getClass(), Integer.TYPE, Long.TYPE, Integer.TYPE, Integer.TYPE, Throwable.class)
+ .newInstance(dummy, shuffleId, (long) mapIndex, mapIndex, reduceId, cause);
+ } catch (NoSuchMethodException | IllegalAccessException
+ | IllegalArgumentException | InstantiationException
+ | InvocationTargetException e) { // anything goes wrong, fallback to the another constructor.
+ try {
+ instance = (T) klass
+ .getConstructor(dummy.getClass(), Integer.TYPE, Integer.TYPE, Integer.TYPE, Throwable.class)
+ .newInstance(dummy, shuffleId, mapIndex, reduceId, cause);
+ } catch (Exception ae) {
+ LOG.error("Fail to new instance.", ae);
+ throw new RssException(ae);
+ }
+ }
+ return instance;
+ }
+
+ public static FetchFailedException createFetchFailedException(
+ int shuffleId, int mapIndex, int reduceId, Throwable cause) {
+ final String dummyHost = "dummy_host";
+ final int dummyPort = 9999;
+ BlockManagerId dummy = BlockManagerId.apply("exec-dummy", dummyHost, dummyPort, Option.empty());
+ // if no cause
+ cause = cause == null ? new Throwable("No cause") : cause;
+ return instantiateFetchFailedException(dummy, shuffleId, mapIndex, reduceId, cause);
+ }
+
+ public static boolean isStageResubmitSupported() {
+ // Stage re-computation requires the ShuffleMapTask to throw a FetchFailedException, which would be produced by the
+ // shuffle data reader iterator. However, the shuffle reader iterator interface is defined in Scala, which doesn't
+ // have checked exceptions. This makes it hard to throw a FetchFailedException on the Java side.
+ // Fortunately, starting from Spark 2.3 (or maybe even Spark 2.2), it is possible to create a FetchFailedException
+ // and wrap it into a runtime exception. Spark will consider this exception as a FetchFailedException.
+ // Therefore, the stage re-computation feature is only enabled for Spark versions larger than or equal to 2.3.
+ return SparkVersionUtils.isSpark3() || (SparkVersionUtils.isSpark2() && SparkVersionUtils.MINOR_VERSION >= 3);
+ }
}
diff --git a/client-spark/common/src/main/java/org/apache/spark/shuffle/SparkVersionUtils.java b/client-spark/common/src/main/java/org/apache/spark/shuffle/SparkVersionUtils.java
new file mode 100644
index 00000000..d44a481b
--- /dev/null
+++ b/client-spark/common/src/main/java/org/apache/spark/shuffle/SparkVersionUtils.java
@@ -0,0 +1,52 @@
+/*
+ * 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;
+
+import org.apache.spark.package$;
+import org.apache.spark.util.VersionUtils;
+
+public class SparkVersionUtils {
+ public static final String SPARK_VERSION = package$.MODULE$.SPARK_VERSION();
+ public static final int MAJOR_VERSION;
+ public static final int MINOR_VERSION;
+
+ static {
+ int majorVersion;
+ int minorVersion;
+ try {
+ majorVersion = VersionUtils.majorVersion(SPARK_VERSION);
+ minorVersion = VersionUtils.minorVersion(SPARK_VERSION);
+ } catch (Exception e) {
+ // ignore, just in case some wild spark version is passed.
+ majorVersion = -1;
+ minorVersion = -1;
+ }
+ MAJOR_VERSION = majorVersion;
+ MINOR_VERSION = minorVersion;
+ }
+
+
+ public static boolean isSpark2() {
+ return MAJOR_VERSION == 2;
+ }
+
+ public static boolean isSpark3() {
+ return MAJOR_VERSION == 3;
+ }
+
+}
diff --git a/client-spark/common/src/main/java/org/apache/spark/shuffle/reader/RssFetchFailedIterator.java b/client-spark/common/src/main/java/org/apache/spark/shuffle/reader/RssFetchFailedIterator.java
new file mode 100644
index 00000000..bf4a4c50
--- /dev/null
+++ b/client-spark/common/src/main/java/org/apache/spark/shuffle/reader/RssFetchFailedIterator.java
@@ -0,0 +1,146 @@
+/*
+ * 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.reader;
+
+import java.io.IOException;
+import java.util.Objects;
+
+import org.apache.spark.shuffle.FetchFailedException;
+import org.apache.spark.shuffle.RssSparkShuffleUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import scala.Product2;
+import scala.collection.AbstractIterator;
+import scala.collection.Iterator;
+
+import org.apache.uniffle.client.api.ShuffleManagerClient;
+import org.apache.uniffle.client.factory.ShuffleManagerClientFactory;
+import org.apache.uniffle.client.request.RssReportShuffleFetchFailureRequest;
+import org.apache.uniffle.client.response.RssReportShuffleFetchFailureResponse;
+import org.apache.uniffle.common.ClientType;
+import org.apache.uniffle.common.exception.RssException;
+import org.apache.uniffle.common.exception.RssFetchFailedException;
+
+public class RssFetchFailedIterator<K, C> extends AbstractIterator<Product2<K, C>> {
+ private static final Logger LOG = LoggerFactory.getLogger(RssFetchFailedIterator.class);
+ final Iterator<Product2<K, C>> iter;
+ final Builder builder;
+
+ private RssFetchFailedIterator(Builder builder, Iterator<Product2<K, C>> iterator) {
+ this.builder = builder;
+ this.iter = iterator;
+ }
+
+ public static class Builder {
+ private String appId;
+ private int shuffleId;
+ private int partitionId;
+ private int stageAttemptId;
+ private String reportServerHost;
+ private int reportServerPort;
+
+ private Builder() {
+
+ }
+
+ Builder appId(String appId) {
+ this.appId = appId;
+ return this;
+ }
+
+ Builder shuffleId(int shuffleId) {
+ this.shuffleId = shuffleId;
+ return this;
+ }
+
+ Builder partitionId(int partitionId) {
+ this.partitionId = partitionId;
+ return this;
+ }
+
+ Builder stageAttemptId(int stageAttemptId) {
+ this.stageAttemptId = stageAttemptId;
+ return this;
+ }
+
+ Builder reportServerHost(String host) {
+ this.reportServerHost = host;
+ return this;
+ }
+
+ Builder port(int port) {
+ this.reportServerPort = port;
+ return this;
+ }
+
+ <K, C> RssFetchFailedIterator<K, C> build(Iterator<Product2<K, C>> iter) {
+ Objects.requireNonNull(this.appId);
+ Objects.requireNonNull(this.reportServerHost);
+ return new RssFetchFailedIterator<>(this, iter);
+ }
+ }
+
+ static Builder newBuilder() {
+ return new Builder();
+ }
+
+ private static ShuffleManagerClient createShuffleManagerClient(String host, int port) throws IOException {
+ ClientType grpc = ClientType.GRPC;
+ // host is passed from spark.driver.bindAddress, which would be set when SparkContext is constructed.
+ return ShuffleManagerClientFactory.getInstance().createShuffleManagerClient(grpc, host, port);
+ }
+
+ private RssException generateFetchFailedIfNecessary(RssFetchFailedException e) {
+ String driver = builder.reportServerHost;
+ int port = builder.reportServerPort;
+ // todo: reuse this manager client if this is a bottleneck.
+ try (ShuffleManagerClient client = createShuffleManagerClient(driver, port)) {
+ RssReportShuffleFetchFailureRequest req = new RssReportShuffleFetchFailureRequest(
+ builder.appId, builder.shuffleId, builder.stageAttemptId, builder.partitionId, e.getMessage());
+ RssReportShuffleFetchFailureResponse response = client.reportShuffleFetchFailure(req);
+ if (response.getReSubmitWholeStage()) {
+ // since we are going to roll out the whole stage, mapIndex shouldn't matter, hence -1 is provided.
+ FetchFailedException ffe =
+ RssSparkShuffleUtils.createFetchFailedException(builder.shuffleId, -1, builder.partitionId, e);
+ return new RssException(ffe);
+ }
+ } catch (IOException ioe) {
+ LOG.info("Error closing shuffle manager client with error:", ioe);
+ }
+ return e;
+ }
+
+ @Override
+ public boolean hasNext() {
+ try {
+ return this.iter.hasNext();
+ } catch (RssFetchFailedException e) {
+ throw generateFetchFailedIfNecessary(e);
+ }
+ }
+
+ @Override
+ public Product2<K, C> next() {
+ try {
+ return this.iter.next();
+ } catch (RssFetchFailedException e) {
+ throw generateFetchFailedIfNecessary(e);
+ }
+ }
+
+}
diff --git a/client-spark/common/src/main/java/org/apache/uniffle/shuffle/manager/RssShuffleManagerBase.java b/client-spark/common/src/main/java/org/apache/uniffle/shuffle/manager/RssShuffleManagerBase.java
new file mode 100644
index 00000000..5d3bd2a8
--- /dev/null
+++ b/client-spark/common/src/main/java/org/apache/uniffle/shuffle/manager/RssShuffleManagerBase.java
@@ -0,0 +1,143 @@
+/*
+ * 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.uniffle.shuffle.manager;
+
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import org.apache.spark.MapOutputTracker;
+import org.apache.spark.MapOutputTrackerMaster;
+import org.apache.spark.SparkEnv;
+import org.apache.spark.SparkException;
+import org.apache.spark.shuffle.RssSparkShuffleUtils;
+import org.apache.spark.shuffle.ShuffleManager;
+import org.apache.spark.shuffle.SparkVersionUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.common.exception.RssException;
+
+public abstract class RssShuffleManagerBase implements RssShuffleManagerInterface, ShuffleManager {
+ private static final Logger LOG = LoggerFactory.getLogger(RssShuffleManagerBase.class);
+ private AtomicBoolean isInitialized = new AtomicBoolean(false);
+ private Method unregisterAllMapOutputMethod;
+ private Method registerShuffleMethod;
+
+ @Override
+ public void unregisterAllMapOutput(int shuffleId) throws SparkException {
+ if (!RssSparkShuffleUtils.isStageResubmitSupported()) {
+ return;
+ }
+ MapOutputTrackerMaster tracker = getMapOutputTrackerMaster();
+ if (isInitialized.compareAndSet(false, true)) {
+ unregisterAllMapOutputMethod = getUnregisterAllMapOutputMethod(tracker);
+ registerShuffleMethod = getRegisterShuffleMethod(tracker);
+ }
+ if (unregisterAllMapOutputMethod != null) {
+ try {
+ unregisterAllMapOutputMethod.invoke(tracker, shuffleId);
+ } catch (InvocationTargetException | IllegalAccessException e) {
+ throw new RssException("Invoke unregisterAllMapOutput method failed", e);
+ }
+ } else {
+ int numMaps = getNumMaps(shuffleId);
+ int numReduces = getPartitionNum(shuffleId);
+ defaultUnregisterAllMapOutput(tracker, registerShuffleMethod, shuffleId, numMaps, numReduces);
+ }
+ }
+
+ private static void defaultUnregisterAllMapOutput(MapOutputTrackerMaster tracker, Method registerShuffle,
+ int shuffleId, int numMaps, int numReduces) throws SparkException {
+ if (tracker != null && registerShuffle != null) {
+ tracker.unregisterShuffle(shuffleId);
+ // re-register this shuffle id into map output tracker
+ try {
+ if (SparkVersionUtils.MAJOR_VERSION > 3
+ || (SparkVersionUtils.isSpark3() && SparkVersionUtils.MINOR_VERSION >= 2)) {
+ registerShuffle.invoke(tracker, shuffleId, numMaps, numReduces);
+ } else {
+ registerShuffle.invoke(tracker, shuffleId, numMaps);
+ }
+ } catch (InvocationTargetException | IllegalAccessException e) {
+ throw new RssException("Invoke registerShuffle method failed", e);
+ }
+ tracker.incrementEpoch();
+ } else {
+ throw new SparkException("default unregisterAllMapOutput should only be called on the driver side");
+ }
+ }
+
+ private static Method getUnregisterAllMapOutputMethod(MapOutputTrackerMaster tracker) {
+ if (tracker != null) {
+ Class<? extends MapOutputTrackerMaster> klass = tracker.getClass();
+ Method m = null;
+ try {
+ if (SparkVersionUtils.isSpark2() && SparkVersionUtils.MINOR_VERSION <= 3) {
+ // for spark version less than 2.3, there's no unregisterAllMapOutput support
+ LOG.warn("Spark version <= 2.3, fallback to default method");
+ } else if (SparkVersionUtils.isSpark2()) {
+ // this method is added in Spark 2.4+
+ m = klass.getDeclaredMethod("unregisterAllMapOutput", int.class);
+ } else if (SparkVersionUtils.isSpark3() && SparkVersionUtils.MINOR_VERSION <= 1) {
+ // spark 3.1 will have unregisterAllMapOutput method
+ m = klass.getDeclaredMethod("unregisterAllMapOutput", int.class);
+ } else if (SparkVersionUtils.isSpark3()) {
+ m = klass.getDeclaredMethod("unregisterAllMapAndMergeOutput", int.class);
+ } else {
+ LOG.warn("Unknown spark version({}), fallback to default method", SparkVersionUtils.SPARK_VERSION);
+ }
+ } catch (NoSuchMethodException e) {
+ LOG.warn("Got no such method error when get unregisterAllMapOutput method for spark version({})",
+ SparkVersionUtils.SPARK_VERSION);
+ }
+ return m;
+ } else {
+ return null;
+ }
+ }
+
+ private static Method getRegisterShuffleMethod(MapOutputTrackerMaster tracker) {
+ if (tracker != null) {
+ Class<? extends MapOutputTrackerMaster> klass = tracker.getClass();
+ Method m = null;
+ try {
+ if (SparkVersionUtils.MAJOR_VERSION > 3
+ || (SparkVersionUtils.isSpark3() && SparkVersionUtils.MINOR_VERSION >= 2)) {
+ // for spark >= 3.2, the register shuffle method is changed to signature:
+ // registerShuffle(shuffleId, numMapTasks, numReduceTasks);
+ m = klass.getDeclaredMethod("registerShuffle", int.class, int.class, int.class);
+ } else {
+ m = klass.getDeclaredMethod("registerShuffle", int.class, int.class);
+ }
+ } catch (NoSuchMethodException e) {
+ LOG.warn("Got no such method error when get registerShuffle method for spark version({})",
+ SparkVersionUtils.SPARK_VERSION);
+ }
+ return m;
+ } else {
+ return null;
+ }
+ }
+
+ private static MapOutputTrackerMaster getMapOutputTrackerMaster() {
+ MapOutputTracker tracker = Optional.ofNullable(SparkEnv.get()).map(SparkEnv::mapOutputTracker).orElse(null);
+ return tracker instanceof MapOutputTrackerMaster ? (MapOutputTrackerMaster) tracker : null;
+ }
+}
diff --git a/client-spark/common/src/main/java/org/apache/uniffle/shuffle/manager/ShuffleManagerServerFactory.java b/client-spark/common/src/main/java/org/apache/uniffle/shuffle/manager/ShuffleManagerServerFactory.java
index 5d0174e6..2e1c7470 100644
--- a/client-spark/common/src/main/java/org/apache/uniffle/shuffle/manager/ShuffleManagerServerFactory.java
+++ b/client-spark/common/src/main/java/org/apache/uniffle/shuffle/manager/ShuffleManagerServerFactory.java
@@ -33,13 +33,24 @@ public class ShuffleManagerServerFactory {
this.conf.addAll(conf);
}
+ public ShuffleManagerGrpcService getService() {
+ return new ShuffleManagerGrpcService(shuffleManager);
+ }
+
public GrpcServer getServer() {
+ return getServer(null);
+ }
+
+ public GrpcServer getServer(ShuffleManagerGrpcService service) {
ServerType type = conf.get(RssBaseConf.RPC_SERVER_TYPE);
if (type == ServerType.GRPC) {
+ if (service == null) {
+ service = new ShuffleManagerGrpcService(shuffleManager);
+ }
return GrpcServer.Builder.newBuilder()
.conf(conf)
.grpcMetrics(GRPCMetrics.getEmptyGRPCMetrics())
- .addService(new ShuffleManagerGrpcService(shuffleManager))
+ .addService(service)
.build();
} else {
throw new UnsupportedOperationException("Unsupported server type " + type);
diff --git a/client-spark/spark2/src/main/java/org/apache/spark/shuffle/RssShuffleManager.java b/client-spark/spark2/src/main/java/org/apache/spark/shuffle/RssShuffleManager.java
index 77b9ef0d..2142810c 100644
--- a/client-spark/spark2/src/main/java/org/apache/spark/shuffle/RssShuffleManager.java
+++ b/client-spark/spark2/src/main/java/org/apache/spark/shuffle/RssShuffleManager.java
@@ -28,6 +28,7 @@ import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import org.apache.hadoop.conf.Configuration;
import org.apache.spark.ShuffleDependency;
@@ -55,18 +56,26 @@ import scala.collection.Seq;
import org.apache.uniffle.client.api.ShuffleWriteClient;
import org.apache.uniffle.client.factory.ShuffleClientFactory;
import org.apache.uniffle.client.util.ClientUtils;
+import org.apache.uniffle.client.util.RssClientConfig;
import org.apache.uniffle.common.PartitionRange;
import org.apache.uniffle.common.RemoteStorageInfo;
import org.apache.uniffle.common.ShuffleAssignmentsInfo;
import org.apache.uniffle.common.ShuffleDataDistributionType;
import org.apache.uniffle.common.ShuffleServerInfo;
+import org.apache.uniffle.common.config.RssConf;
import org.apache.uniffle.common.exception.RssException;
+import org.apache.uniffle.common.rpc.GrpcServer;
import org.apache.uniffle.common.util.JavaUtils;
import org.apache.uniffle.common.util.RetryUtils;
import org.apache.uniffle.common.util.RssUtils;
import org.apache.uniffle.common.util.ThreadUtils;
+import org.apache.uniffle.shuffle.manager.RssShuffleManagerBase;
+import org.apache.uniffle.shuffle.manager.ShuffleManagerGrpcService;
+import org.apache.uniffle.shuffle.manager.ShuffleManagerServerFactory;
-public class RssShuffleManager implements ShuffleManager {
+import static org.apache.uniffle.common.config.RssBaseConf.RPC_SERVER_PORT;
+
+public class RssShuffleManager extends RssShuffleManagerBase {
private static final Logger LOG = LoggerFactory.getLogger(RssShuffleManager.class);
private final long heartbeatInterval;
@@ -91,6 +100,11 @@ public class RssShuffleManager implements ShuffleManager {
private final String uuid;
private DataPusher dataPusher;
+ private final Map<Integer, Integer> shuffleIdToPartitionNum = Maps.newConcurrentMap();
+ private final Map<Integer, Integer> shuffleIdToNumMapTasks = Maps.newConcurrentMap();
+ private GrpcServer shuffleManagerServer;
+ private ShuffleManagerGrpcService service;
+
public RssShuffleManager(SparkConf sparkConf, boolean isDriver) {
if (sparkConf.getBoolean("spark.sql.adaptive.enabled", false)) {
throw new IllegalArgumentException("Spark2 doesn't support AQE, spark.sql.adaptive.enabled should be false.");
@@ -142,6 +156,24 @@ public class RssShuffleManager implements ShuffleManager {
if (isDriver) {
heartBeatScheduledExecutorService =
ThreadUtils.getDaemonSingleThreadScheduledExecutor("rss-heartbeat");
+ RssConf rssConf = RssSparkConfig.toRssConf(sparkConf);
+ if (rssConf.getBoolean(RssClientConfig.RSS_RESUBMIT_STAGE, false)
+ && RssSparkShuffleUtils.isStageResubmitSupported()) {
+ LOG.info("stage resubmit is supported and enabled");
+ // start shuffle manager server
+ rssConf.set(RPC_SERVER_PORT, 0);
+ ShuffleManagerServerFactory factory = new ShuffleManagerServerFactory(this, rssConf);
+ service = factory.getService();
+ shuffleManagerServer = factory.getServer(service);
+ try {
+ shuffleManagerServer.start();
+ // pass this as a spark.rss.shuffle.manager.grpc.port config, so it can be propagated to executor properly.
+ sparkConf.set(RssSparkConfig.RSS_SHUFFLE_MANAGER_GRPC_PORT, shuffleManagerServer.getPort());
+ } catch (Exception e) {
+ LOG.error("Failed to start shuffle manager server", e);
+ throw new RssException(e);
+ }
+ }
}
// for non-driver executor, start a thread for sending shuffle data to shuffle server
LOG.info("RSS data pusher is starting...");
@@ -190,6 +222,8 @@ public class RssShuffleManager implements ShuffleManager {
}
if (dependency.partitioner().numPartitions() == 0) {
+ shuffleIdToPartitionNum.putIfAbsent(shuffleId, 0);
+ shuffleIdToNumMapTasks.putIfAbsent(shuffleId, dependency.rdd().partitions().length);
LOG.info("RegisterShuffle with ShuffleId[" + shuffleId + "], partitionNum is 0, "
+ "return the empty RssShuffleHandle directly");
Broadcast<ShuffleHandleInfo> hdlInfoBd = RssSparkShuffleUtils.broadcastShuffleHdlInfo(
@@ -236,6 +270,8 @@ public class RssShuffleManager implements ShuffleManager {
startHeartbeat();
+ shuffleIdToPartitionNum.putIfAbsent(shuffleId, dependency.partitioner().numPartitions());
+ shuffleIdToNumMapTasks.putIfAbsent(shuffleId, dependency.rdd().partitions().length);
Broadcast<ShuffleHandleInfo> hdlInfoBd = RssSparkShuffleUtils.broadcastShuffleHdlInfo(
RssSparkShuffleUtils.getActiveSparkContext(), shuffleId, partitionToServers, remoteStorage);
LOG.info("RegisterShuffle with ShuffleId[" + shuffleId + "], partitionNum[" + partitionToServers.size() + "]");
@@ -395,6 +431,11 @@ public class RssShuffleManager implements ShuffleManager {
try {
if (SparkEnv.get().executorId().equals("driver")) {
shuffleWriteClient.unregisterShuffle(appId, shuffleId);
+ shuffleIdToNumMapTasks.remove(shuffleId);
+ shuffleIdToPartitionNum.remove(shuffleId);
+ if (service != null) {
+ service.unregisterShuffle(shuffleId);
+ }
}
} catch (Exception e) {
LOG.warn("Errors on unregister to remote shuffle-servers", e);
@@ -508,4 +549,39 @@ public class RssShuffleManager implements ShuffleManager {
public void setDataPusher(DataPusher dataPusher) {
this.dataPusher = dataPusher;
}
+
+ /**
+ * @return the unique spark id for rss shuffle
+ */
+ @Override
+ public String getAppId() {
+ return appId;
+ }
+
+ /**
+ * @return the maximum number of fetch failures per shuffle partition before that shuffle stage should be recomputed
+ */
+ @Override
+ public int getMaxFetchFailures() {
+ final String TASK_MAX_FAILURE = "spark.task.maxFailures";
+ return Math.max(1, sparkConf.getInt(TASK_MAX_FAILURE, 4) - 1);
+ }
+
+ /**
+ * @param shuffleId the shuffleId to query
+ * @return the num of partitions(a.k.a reduce tasks) for shuffle with shuffle id.
+ */
+ @Override
+ public int getPartitionNum(int shuffleId) {
+ return shuffleIdToPartitionNum.getOrDefault(shuffleId, 0);
+ }
+
+ /**
+ * @param shuffleId the shuffle id to query
+ * @return the num of map tasks for current shuffle with shuffle id.
+ */
+ @Override
+ public int getNumMaps(int shuffleId) {
+ return shuffleIdToNumMapTasks.getOrDefault(shuffleId, 0);
+ }
}
diff --git a/client-spark/spark2/src/main/java/org/apache/spark/shuffle/reader/RssShuffleReader.java b/client-spark/spark2/src/main/java/org/apache/spark/shuffle/reader/RssShuffleReader.java
index 2a8fd81d..5e6d2240 100644
--- a/client-spark/spark2/src/main/java/org/apache/spark/shuffle/reader/RssShuffleReader.java
+++ b/client-spark/spark2/src/main/java/org/apache/spark/shuffle/reader/RssShuffleReader.java
@@ -45,7 +45,9 @@ import scala.runtime.BoxedUnit;
import org.apache.uniffle.client.api.ShuffleReadClient;
import org.apache.uniffle.client.factory.ShuffleClientFactory;
import org.apache.uniffle.client.request.CreateShuffleReadClientRequest;
+import org.apache.uniffle.client.util.RssClientConfig;
import org.apache.uniffle.common.ShuffleServerInfo;
+import org.apache.uniffle.common.config.RssClientConf;
import org.apache.uniffle.common.config.RssConf;
public class RssShuffleReader<K, C> implements ShuffleReader<K, C> {
@@ -186,6 +188,21 @@ public class RssShuffleReader<K, C> implements ShuffleReader<K, C> {
if (!(resultIter instanceof InterruptibleIterator)) {
resultIter = new InterruptibleIterator<>(context, resultIter);
}
+
+ // stage re-compute and shuffle manager server port are both set
+ if (rssConf.getBoolean(RssClientConfig.RSS_RESUBMIT_STAGE, false)
+ && rssConf.getInteger(RssClientConf.SHUFFLE_MANAGER_GRPC_PORT, 0) > 0) {
+ String driver = rssConf.getString("driver.host", "");
+ int port = rssConf.get(RssClientConf.SHUFFLE_MANAGER_GRPC_PORT);
+ resultIter = RssFetchFailedIterator.newBuilder()
+ .appId(appId)
+ .shuffleId(shuffleId)
+ .partitionId(startPartition)
+ .stageAttemptId(context.stageAttemptNumber())
+ .reportServerHost(driver)
+ .port(port)
+ .build(resultIter);
+ }
return resultIter;
}
diff --git a/client-spark/spark2/src/test/java/org/apache/spark/shuffle/RssSpark2ShuffleUtilsTest.java b/client-spark/spark2/src/test/java/org/apache/spark/shuffle/RssSpark2ShuffleUtilsTest.java
new file mode 100644
index 00000000..c9c89453
--- /dev/null
+++ b/client-spark/spark2/src/test/java/org/apache/spark/shuffle/RssSpark2ShuffleUtilsTest.java
@@ -0,0 +1,38 @@
+/*
+ * 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;
+
+import org.junit.jupiter.api.Test;
+
+import org.apache.uniffle.common.exception.RssFetchFailedException;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class RssSpark2ShuffleUtilsTest {
+ @Test
+ public void testCreateFetchFailedException() {
+ FetchFailedException ffe = RssSparkShuffleUtils.createFetchFailedException(0, -1, 10, null);
+ ffe = RssSparkShuffleUtils.createFetchFailedException(0, -1, 100, new RssFetchFailedException("xx"));
+ }
+
+ @Test
+ public void testIsStageResubmitSupported() {
+ // by default spark 2.3/2.4 is used, which supports stage resubmit.
+ assertTrue(RssSparkShuffleUtils.isStageResubmitSupported());
+ }
+}
diff --git a/client-spark/spark2/src/test/java/org/apache/spark/shuffle/SparkVersionUtilsTest.java b/client-spark/spark2/src/test/java/org/apache/spark/shuffle/SparkVersionUtilsTest.java
new file mode 100644
index 00000000..776b6dec
--- /dev/null
+++ b/client-spark/spark2/src/test/java/org/apache/spark/shuffle/SparkVersionUtilsTest.java
@@ -0,0 +1,29 @@
+/*
+ * 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;
+
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class SparkVersionUtilsTest {
+ @Test
+ public void testSparkVersion() {
+ assertTrue(SparkVersionUtils.isSpark2());
+ }
+}
diff --git a/client-spark/spark3/src/main/java/org/apache/spark/shuffle/RssShuffleManager.java b/client-spark/spark3/src/main/java/org/apache/spark/shuffle/RssShuffleManager.java
index e875132d..dfe8813f 100644
--- a/client-spark/spark3/src/main/java/org/apache/spark/shuffle/RssShuffleManager.java
+++ b/client-spark/spark3/src/main/java/org/apache/spark/shuffle/RssShuffleManager.java
@@ -31,6 +31,7 @@ import java.util.function.Function;
import java.util.stream.Collectors;
import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import org.apache.hadoop.conf.Configuration;
@@ -62,6 +63,7 @@ import scala.collection.Seq;
import org.apache.uniffle.client.api.ShuffleWriteClient;
import org.apache.uniffle.client.factory.ShuffleClientFactory;
import org.apache.uniffle.client.util.ClientUtils;
+import org.apache.uniffle.client.util.RssClientConfig;
import org.apache.uniffle.common.PartitionRange;
import org.apache.uniffle.common.RemoteStorageInfo;
import org.apache.uniffle.common.ShuffleAssignmentsInfo;
@@ -70,12 +72,18 @@ import org.apache.uniffle.common.ShuffleServerInfo;
import org.apache.uniffle.common.config.RssClientConf;
import org.apache.uniffle.common.config.RssConf;
import org.apache.uniffle.common.exception.RssException;
+import org.apache.uniffle.common.rpc.GrpcServer;
import org.apache.uniffle.common.util.JavaUtils;
import org.apache.uniffle.common.util.RetryUtils;
import org.apache.uniffle.common.util.RssUtils;
import org.apache.uniffle.common.util.ThreadUtils;
+import org.apache.uniffle.shuffle.manager.RssShuffleManagerBase;
+import org.apache.uniffle.shuffle.manager.ShuffleManagerGrpcService;
+import org.apache.uniffle.shuffle.manager.ShuffleManagerServerFactory;
-public class RssShuffleManager implements ShuffleManager {
+import static org.apache.uniffle.common.config.RssBaseConf.RPC_SERVER_PORT;
+
+public class RssShuffleManager extends RssShuffleManagerBase {
private static final Logger LOG = LoggerFactory.getLogger(RssShuffleManager.class);
private final String clientType;
@@ -101,6 +109,11 @@ public class RssShuffleManager implements ShuffleManager {
private Set<String> failedTaskIds = Sets.newConcurrentHashSet();
private DataPusher dataPusher;
+ private final Map<Integer, Integer> shuffleIdToPartitionNum = Maps.newConcurrentMap();
+ private final Map<Integer, Integer> shuffleIdToNumMapTasks = Maps.newConcurrentMap();
+ private ShuffleManagerGrpcService service;
+ private GrpcServer shuffleManagerServer;
+
public RssShuffleManager(SparkConf conf, boolean isDriver) {
this.sparkConf = conf;
boolean supportsRelocation = Optional.ofNullable(SparkEnv.get())
@@ -161,6 +174,24 @@ public class RssShuffleManager implements ShuffleManager {
if (isDriver) {
heartBeatScheduledExecutorService =
ThreadUtils.getDaemonSingleThreadScheduledExecutor("rss-heartbeat");
+ RssConf rssConf = RssSparkConfig.toRssConf(sparkConf);
+ if (rssConf.getBoolean(RssClientConfig.RSS_RESUBMIT_STAGE, false)
+ && RssSparkShuffleUtils.isStageResubmitSupported()) {
+ LOG.info("stage resubmit is supported and enabled");
+ // start shuffle manager server
+ rssConf.set(RPC_SERVER_PORT, 0);
+ ShuffleManagerServerFactory factory = new ShuffleManagerServerFactory(this, rssConf);
+ service = factory.getService();
+ shuffleManagerServer = factory.getServer(service);
+ try {
+ shuffleManagerServer.start();
+ // pass this as a spark.rss.shuffle.manager.grpc.port config, so it can be propagated to executor properly.
+ sparkConf.set(RssSparkConfig.RSS_SHUFFLE_MANAGER_GRPC_PORT, shuffleManagerServer.getPort());
+ } catch (Exception e) {
+ LOG.error("Failed to start shuffle manager server", e);
+ throw new RssException(e);
+ }
+ }
}
LOG.info("Rss data pusher is starting...");
int poolSize = sparkConf.get(RssSparkConfig.RSS_CLIENT_SEND_THREAD_POOL_SIZE);
@@ -270,6 +301,8 @@ public class RssShuffleManager implements ShuffleManager {
LOG.info("Generate application id used in rss: " + id.get());
if (dependency.partitioner().numPartitions() == 0) {
+ shuffleIdToPartitionNum.putIfAbsent(shuffleId, 0);
+ shuffleIdToNumMapTasks.putIfAbsent(shuffleId, dependency.rdd().partitions().length);
LOG.info("RegisterShuffle with ShuffleId[" + shuffleId + "], partitionNum is 0, "
+ "return the empty RssShuffleHandle directly");
Broadcast<ShuffleHandleInfo> hdlInfoBd = RssSparkShuffleUtils.broadcastShuffleHdlInfo(
@@ -317,6 +350,8 @@ public class RssShuffleManager implements ShuffleManager {
}
startHeartbeat();
+ shuffleIdToPartitionNum.putIfAbsent(shuffleId, dependency.partitioner().numPartitions());
+ shuffleIdToNumMapTasks.putIfAbsent(shuffleId, dependency.rdd().partitions().length);
Broadcast<ShuffleHandleInfo> hdlInfoBd = RssSparkShuffleUtils.broadcastShuffleHdlInfo(
RssSparkShuffleUtils.getActiveSparkContext(), shuffleId, partitionToServers, remoteStorage);
LOG.info("RegisterShuffle with ShuffleId[" + shuffleId + "], partitionNum[" + partitionToServers.size()
@@ -594,6 +629,11 @@ public class RssShuffleManager implements ShuffleManager {
try {
if (SparkEnv.get().executorId().equals("driver")) {
shuffleWriteClient.unregisterShuffle(id.get(), shuffleId);
+ shuffleIdToPartitionNum.remove(shuffleId);
+ shuffleIdToNumMapTasks.remove(shuffleId);
+ if (service != null) {
+ service.unregisterShuffle(shuffleId);
+ }
}
} catch (Exception e) {
LOG.warn("Errors on unregister to remote shuffle-servers", e);
@@ -621,6 +661,15 @@ public class RssShuffleManager implements ShuffleManager {
LOG.warn("Errors on closing data pusher", e);
}
}
+
+ if (shuffleManagerServer != null) {
+ try {
+ shuffleManagerServer.stop();
+ } catch (InterruptedException e) {
+ // ignore
+ LOG.info("shuffle manager server is interrupted during stop");
+ }
+ }
}
public void clearTaskMeta(String taskId) {
@@ -703,6 +752,37 @@ public class RssShuffleManager implements ShuffleManager {
return result;
}
+ /**
+ * @return the unique spark id for rss shuffle
+ */
+ @Override
+ public String getAppId() {
+ return id.get();
+ }
+
+ /**
+ * @return the maximum number of fetch failures per shuffle partition before that shuffle stage should be recomputed
+ */
+ @Override
+ public int getMaxFetchFailures() {
+ final String TASK_MAX_FAILURE = "spark.task.maxFailures";
+ return Math.max(1, sparkConf.getInt(TASK_MAX_FAILURE, 4) - 1);
+ }
+
+ @Override
+ public int getPartitionNum(int shuffleId) {
+ return shuffleIdToPartitionNum.getOrDefault(shuffleId, 0);
+ }
+
+ /**
+ * @param shuffleId the shuffle id to query
+ * @return the num of map tasks for current shuffle with shuffle id.
+ */
+ @Override
+ public int getNumMaps(int shuffleId) {
+ return shuffleIdToNumMapTasks.getOrDefault(shuffleId, 0);
+ }
+
static class ReadMetrics extends ShuffleReadMetrics {
private ShuffleReadMetricsReporter reporter;
@@ -754,10 +834,6 @@ public class RssShuffleManager implements ShuffleManager {
this.id = new AtomicReference<>(appId);
}
- public String getId() {
- return id.get();
- }
-
public boolean markFailedTask(String taskId) {
LOG.info("Mark the task: {} failed.", taskId);
failedTaskIds.add(taskId);
diff --git a/client-spark/spark3/src/main/java/org/apache/spark/shuffle/reader/RssShuffleReader.java b/client-spark/spark3/src/main/java/org/apache/spark/shuffle/reader/RssShuffleReader.java
index 4734b1aa..f2b56937 100644
--- a/client-spark/spark3/src/main/java/org/apache/spark/shuffle/reader/RssShuffleReader.java
+++ b/client-spark/spark3/src/main/java/org/apache/spark/shuffle/reader/RssShuffleReader.java
@@ -49,8 +49,10 @@ import scala.runtime.BoxedUnit;
import org.apache.uniffle.client.api.ShuffleReadClient;
import org.apache.uniffle.client.factory.ShuffleClientFactory;
import org.apache.uniffle.client.request.CreateShuffleReadClientRequest;
+import org.apache.uniffle.client.util.RssClientConfig;
import org.apache.uniffle.common.ShuffleDataDistributionType;
import org.apache.uniffle.common.ShuffleServerInfo;
+import org.apache.uniffle.common.config.RssClientConf;
import org.apache.uniffle.common.config.RssConf;
public class RssShuffleReader<K, C> implements ShuffleReader<K, C> {
@@ -173,6 +175,20 @@ public class RssShuffleReader<K, C> implements ShuffleReader<K, C> {
if (!(resultIter instanceof InterruptibleIterator)) {
resultIter = new InterruptibleIterator<>(context, resultIter);
}
+ // resubmit stage and shuffle manager server port are both set
+ if (rssConf.getBoolean(RssClientConfig.RSS_RESUBMIT_STAGE, false)
+ && rssConf.getInteger(RssClientConf.SHUFFLE_MANAGER_GRPC_PORT, 0) > 0) {
+ String driver = rssConf.getString("driver.host", "");
+ int port = rssConf.get(RssClientConf.SHUFFLE_MANAGER_GRPC_PORT);
+ resultIter = RssFetchFailedIterator.newBuilder()
+ .appId(appId)
+ .shuffleId(shuffleId)
+ .partitionId(startPartition)
+ .stageAttemptId(context.stageAttemptNumber())
+ .reportServerHost(driver)
+ .port(port)
+ .build(resultIter);
+ }
return resultIter;
}
diff --git a/client-spark/spark3/src/test/java/org/apache/spark/shuffle/DelegationRssShuffleManagerTest.java b/client-spark/spark3/src/test/java/org/apache/spark/shuffle/DelegationRssShuffleManagerTest.java
index 9e50db55..5fdc29ce 100644
--- a/client-spark/spark3/src/test/java/org/apache/spark/shuffle/DelegationRssShuffleManagerTest.java
+++ b/client-spark/spark3/src/test/java/org/apache/spark/shuffle/DelegationRssShuffleManagerTest.java
@@ -17,20 +17,12 @@
package org.apache.spark.shuffle;
-import java.util.List;
import java.util.NoSuchElementException;
-import com.google.common.collect.Lists;
import org.apache.spark.SparkConf;
import org.apache.spark.shuffle.sort.SortShuffleManager;
-import org.junit.jupiter.api.AfterAll;
-import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
-import org.mockito.MockedStatic;
-import org.mockito.Mockito;
-import org.apache.uniffle.client.api.CoordinatorClient;
-import org.apache.uniffle.client.response.RssAccessClusterResponse;
import org.apache.uniffle.storage.util.StorageType;
import static org.apache.uniffle.common.rpc.StatusCode.ACCESS_DENIED;
@@ -38,32 +30,12 @@ import static org.apache.uniffle.common.rpc.StatusCode.SUCCESS;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertTrue;
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.mockStatic;
-import static org.mockito.Mockito.when;
-public class DelegationRssShuffleManagerTest {
- private static MockedStatic<RssSparkShuffleUtils> mockedStaticRssShuffleUtils;
-
- @BeforeAll
- public static void setUp() {
- mockedStaticRssShuffleUtils = mockStatic(RssSparkShuffleUtils.class, Mockito.CALLS_REAL_METHODS);
- }
-
- @AfterAll
- public static void tearDown() {
- mockedStaticRssShuffleUtils.close();
- }
+public class DelegationRssShuffleManagerTest extends RssShuffleManagerTestBase {
@Test
public void testCreateInDriverDenied() throws Exception {
- CoordinatorClient mockCoordinatorClient = mock(CoordinatorClient.class);
- when(mockCoordinatorClient.accessCluster(any())).thenReturn(new RssAccessClusterResponse(ACCESS_DENIED, ""));
- List<CoordinatorClient> coordinatorClients = Lists.newArrayList();
- coordinatorClients.add(mockCoordinatorClient);
- mockedStaticRssShuffleUtils.when(() ->
- RssSparkShuffleUtils.createCoordinatorClients(any())).thenReturn(coordinatorClients);
+ setupMockedRssShuffleUtils(ACCESS_DENIED);
SparkConf conf = new SparkConf();
conf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED.key(), "false");
assertCreateSortShuffleManager(conf);
@@ -71,13 +43,7 @@ public class DelegationRssShuffleManagerTest {
@Test
public void testCreateInDriver() throws Exception {
- CoordinatorClient mockCoordinatorClient = mock(CoordinatorClient.class);
- when(mockCoordinatorClient.accessCluster(any())).thenReturn(
- new RssAccessClusterResponse(SUCCESS, ""));
- List<CoordinatorClient> coordinatorClients = Lists.newArrayList();
- coordinatorClients.add(mockCoordinatorClient);
- mockedStaticRssShuffleUtils.when(() ->
- RssSparkShuffleUtils.createCoordinatorClients(any())).thenReturn(coordinatorClients);
+ setupMockedRssShuffleUtils(SUCCESS);
SparkConf conf = new SparkConf();
assertCreateSortShuffleManager(conf);
@@ -91,8 +57,6 @@ public class DelegationRssShuffleManagerTest {
conf = new SparkConf();
conf.set(RssSparkConfig.RSS_COORDINATOR_QUORUM.key(), "m1:8001,m2:8002");
- when(mockCoordinatorClient.accessCluster(any())).thenReturn(
- new RssAccessClusterResponse(SUCCESS, ""));
assertCreateSortShuffleManager(conf);
}
@@ -108,12 +72,7 @@ public class DelegationRssShuffleManagerTest {
@Test
public void testCreateFallback() throws Exception {
- CoordinatorClient mockCoordinatorClient = mock(CoordinatorClient.class);
- when(mockCoordinatorClient.accessCluster(any())).thenReturn(new RssAccessClusterResponse(SUCCESS, ""));
- List<CoordinatorClient> coordinatorClients = Lists.newArrayList();
- coordinatorClients.add(mockCoordinatorClient);
- mockedStaticRssShuffleUtils.when(() ->
- RssSparkShuffleUtils.createCoordinatorClients(any())).thenReturn(coordinatorClients);
+ setupMockedRssShuffleUtils(SUCCESS);
SparkConf conf = new SparkConf();
conf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED.key(), "false");
@@ -137,15 +96,7 @@ public class DelegationRssShuffleManagerTest {
@Test
public void testTryAccessCluster() throws Exception {
- CoordinatorClient mockDeniedCoordinatorClient = mock(CoordinatorClient.class);
- when(mockDeniedCoordinatorClient.accessCluster(any()))
- .thenReturn(new RssAccessClusterResponse(ACCESS_DENIED, ""))
- .thenReturn(new RssAccessClusterResponse(ACCESS_DENIED, ""))
- .thenReturn(new RssAccessClusterResponse(SUCCESS, ""));
- List<CoordinatorClient> coordinatorClients = Lists.newArrayList();
- coordinatorClients.add(mockDeniedCoordinatorClient);
- mockedStaticRssShuffleUtils.when(() ->
- RssSparkShuffleUtils.createCoordinatorClients(any())).thenReturn(coordinatorClients);
+ setupMockedRssShuffleUtils(SUCCESS);
SparkConf conf = new SparkConf();
conf.set(RssSparkConfig.RSS_CLIENT_ACCESS_RETRY_INTERVAL_MS, 3000L);
conf.set(RssSparkConfig.RSS_CLIENT_ACCESS_RETRY_TIMES, 3);
@@ -156,15 +107,7 @@ public class DelegationRssShuffleManagerTest {
conf.set(RssSparkConfig.RSS_TEST_MODE_ENABLE, true);
assertCreateRssShuffleManager(conf);
- CoordinatorClient mockCoordinatorClient = mock(CoordinatorClient.class);
- when(mockCoordinatorClient.accessCluster(any()))
- .thenReturn(new RssAccessClusterResponse(ACCESS_DENIED, ""))
- .thenReturn(new RssAccessClusterResponse(ACCESS_DENIED, ""))
- .thenReturn(new RssAccessClusterResponse(ACCESS_DENIED, ""));
- List<CoordinatorClient> secondCoordinatorClients = Lists.newArrayList();
- secondCoordinatorClients.add(mockCoordinatorClient);
- mockedStaticRssShuffleUtils.when(() ->
- RssSparkShuffleUtils.createCoordinatorClients(any())).thenReturn(secondCoordinatorClients);
+ setupMockedRssShuffleUtils(ACCESS_DENIED);
SparkConf secondConf = new SparkConf();
secondConf.set(RssSparkConfig.RSS_CLIENT_ACCESS_RETRY_INTERVAL_MS, 3000L);
secondConf.set(RssSparkConfig.RSS_CLIENT_ACCESS_RETRY_TIMES, 3);
diff --git a/client-spark/spark3/src/test/java/org/apache/spark/shuffle/RssShuffleManagerTest.java b/client-spark/spark3/src/test/java/org/apache/spark/shuffle/RssShuffleManagerTest.java
index 00bc591c..6970842d 100644
--- a/client-spark/spark3/src/test/java/org/apache/spark/shuffle/RssShuffleManagerTest.java
+++ b/client-spark/spark3/src/test/java/org/apache/spark/shuffle/RssShuffleManagerTest.java
@@ -21,12 +21,18 @@ import org.apache.spark.SparkConf;
import org.apache.spark.sql.internal.SQLConf;
import org.junit.jupiter.api.Test;
+import org.apache.uniffle.client.util.RssClientConfig;
import org.apache.uniffle.common.ShuffleDataDistributionType;
import org.apache.uniffle.common.config.RssClientConf;
+import org.apache.uniffle.common.rpc.StatusCode;
+import org.apache.uniffle.storage.util.StorageType;
+import static org.apache.spark.shuffle.RssSparkConfig.RSS_SHUFFLE_MANAGER_GRPC_PORT;
import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
-public class RssShuffleManagerTest {
+public class RssShuffleManagerTest extends RssShuffleManagerTestBase {
private static final String SPARK_ADAPTIVE_EXECUTION_ENABLED_KEY = "spark.sql.adaptive.enabled";
@Test
@@ -83,4 +89,40 @@ public class RssShuffleManagerTest {
);
}
}
+
+ @Test
+ public void testCreateShuffleManagerServer() {
+ setupMockedRssShuffleUtils(StatusCode.SUCCESS);
+
+ SparkConf conf = new SparkConf();
+ conf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED.key(), "false");
+ conf.set(RssSparkConfig.RSS_COORDINATOR_QUORUM.key(), "m1:8001,m2:8002");
+ conf.set("spark.rss.storage.type", StorageType.LOCALFILE.name());
+ conf.set(RssSparkConfig.RSS_TEST_MODE_ENABLE, true);
+ // enable stage recompute
+ conf.set("spark." + RssClientConfig.RSS_RESUBMIT_STAGE, "true");
+
+ RssShuffleManager shuffleManager = new RssShuffleManager(conf, true);
+
+ assertTrue(conf.get(RSS_SHUFFLE_MANAGER_GRPC_PORT) > 0);
+ }
+
+ @Test
+ public void testRssShuffleManagerInterface() throws Exception {
+ setupMockedRssShuffleUtils(StatusCode.SUCCESS);
+
+ SparkConf conf = new SparkConf();
+ conf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED.key(), "false");
+ conf.set(RssSparkConfig.RSS_COORDINATOR_QUORUM.key(), "m1:8001,m2:8002");
+ conf.set("spark.rss.storage.type", StorageType.LOCALFILE.name());
+ conf.set(RssSparkConfig.RSS_TEST_MODE_ENABLE, true);
+
+ conf.set("spark.task.maxFailures", "3");
+ RssShuffleManager shuffleManager = new RssShuffleManager(conf, true);
+ assertEquals(shuffleManager.getMaxFetchFailures(), 2);
+ // by default, the appId is null
+ assertNull(shuffleManager.getAppId());
+
+ }
+
}
diff --git a/client-spark/spark3/src/test/java/org/apache/spark/shuffle/RssShuffleManagerTestBase.java b/client-spark/spark3/src/test/java/org/apache/spark/shuffle/RssShuffleManagerTestBase.java
new file mode 100644
index 00000000..f7fe4816
--- /dev/null
+++ b/client-spark/spark3/src/test/java/org/apache/spark/shuffle/RssShuffleManagerTestBase.java
@@ -0,0 +1,67 @@
+/*
+ * 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;
+
+import java.util.List;
+
+import com.google.common.collect.Lists;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+
+import org.apache.uniffle.client.api.CoordinatorClient;
+import org.apache.uniffle.client.response.RssAccessClusterResponse;
+import org.apache.uniffle.common.rpc.StatusCode;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.mockStatic;
+import static org.mockito.Mockito.when;
+
+public class RssShuffleManagerTestBase {
+ protected static MockedStatic<RssSparkShuffleUtils> mockedStaticRssShuffleUtils;
+
+ @BeforeAll
+ public static void setUp() {
+ mockedStaticRssShuffleUtils = mockStatic(RssSparkShuffleUtils.class, Mockito.CALLS_REAL_METHODS);
+ }
+
+ @AfterAll
+ public static void tearDown() {
+ mockedStaticRssShuffleUtils.close();
+ }
+
+ protected CoordinatorClient createCoordinatorClient(StatusCode status) {
+ CoordinatorClient mockedCoordinatorClient = mock(CoordinatorClient.class);
+ when(mockedCoordinatorClient.accessCluster(any())).thenReturn(
+ new RssAccessClusterResponse(status, ""));
+ return mockedCoordinatorClient;
+ }
+
+ void setupMockedRssShuffleUtils(StatusCode status) {
+ CoordinatorClient mockCoordinatorClient = createCoordinatorClient(status);
+ List<CoordinatorClient> coordinatorClients = Lists.newArrayList();
+ coordinatorClients.add(mockCoordinatorClient);
+ mockedStaticRssShuffleUtils
+ .when(() -> RssSparkShuffleUtils.createCoordinatorClients(any()))
+ .thenReturn(coordinatorClients);
+
+ }
+
+}
diff --git a/client-spark/spark3/src/test/java/org/apache/spark/shuffle/RssSpark3ShuffleUtilsTest.java b/client-spark/spark3/src/test/java/org/apache/spark/shuffle/RssSpark3ShuffleUtilsTest.java
new file mode 100644
index 00000000..91dd600e
--- /dev/null
+++ b/client-spark/spark3/src/test/java/org/apache/spark/shuffle/RssSpark3ShuffleUtilsTest.java
@@ -0,0 +1,37 @@
+/*
+ * 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;
+
+import org.junit.jupiter.api.Test;
+
+import org.apache.uniffle.common.exception.RssFetchFailedException;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class RssSpark3ShuffleUtilsTest {
+ @Test
+ public void testCreateFetchFailedException() {
+ FetchFailedException ffe = RssSparkShuffleUtils.createFetchFailedException(0, -1, 10, null);
+ ffe = RssSparkShuffleUtils.createFetchFailedException(0, -1, 100, new RssFetchFailedException("xx"));
+ }
+
+ @Test
+ public void testIsStageResubmitSupported() {
+ assertTrue(RssSparkShuffleUtils.isStageResubmitSupported());
+ }
+}
diff --git a/client-spark/spark3/src/test/java/org/apache/spark/shuffle/SparkVersionUtilsTest.java b/client-spark/spark3/src/test/java/org/apache/spark/shuffle/SparkVersionUtilsTest.java
new file mode 100644
index 00000000..dbc6bb3d
--- /dev/null
+++ b/client-spark/spark3/src/test/java/org/apache/spark/shuffle/SparkVersionUtilsTest.java
@@ -0,0 +1,29 @@
+/*
+ * 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;
+
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class SparkVersionUtilsTest {
+ @Test
+ public void testSparkVersion() {
+ assertTrue(SparkVersionUtils.isSpark3());
+ }
+}
diff --git a/client/src/main/java/org/apache/uniffle/client/util/RssClientConfig.java b/client/src/main/java/org/apache/uniffle/client/util/RssClientConfig.java
index 5ab3a6fb..8de505b8 100644
--- a/client/src/main/java/org/apache/uniffle/client/util/RssClientConfig.java
+++ b/client/src/main/java/org/apache/uniffle/client/util/RssClientConfig.java
@@ -86,4 +86,5 @@ public class RssClientConfig {
public static final String RSS_ESTIMATE_TASK_CONCURRENCY_PER_SERVER = "rss.estimate.task.concurrency.per.server";
public static final int RSS_ESTIMATE_TASK_CONCURRENCY_PER_SERVER_DEFAULT_VALUE = 80;
+ public static final String RSS_RESUBMIT_STAGE = "rss.resubmit.stage";
}
diff --git a/common/src/main/java/org/apache/uniffle/common/config/RssClientConf.java b/common/src/main/java/org/apache/uniffle/common/config/RssClientConf.java
index 10f56bd2..33d20ddf 100644
--- a/common/src/main/java/org/apache/uniffle/common/config/RssClientConf.java
+++ b/common/src/main/java/org/apache/uniffle/common/config/RssClientConf.java
@@ -95,4 +95,11 @@ public class RssClientConf {
.defaultValue(0)
.withDescription("Send buffer size (SO_SNDBUF).");
+
+ // this is reversed for internal settings, and should never set by user.
+ public static final ConfigOption<Integer> SHUFFLE_MANAGER_GRPC_PORT = ConfigOptions
+ .key("rss.shuffle.manager.grpc.port")
+ .intType()
+ .noDefaultValue()
+ .withDescription("internal configuration to indicate which port is actually bind for shuffle manager service.");
}
diff --git a/integration-test/common/pom.xml b/integration-test/common/pom.xml
index 37b7eb6e..abdb0de0 100644
--- a/integration-test/common/pom.xml
+++ b/integration-test/common/pom.xml
@@ -54,7 +54,6 @@
<groupId>org.apache.uniffle</groupId>
<artifactId>shuffle-server</artifactId>
<type>test-jar</type>
- <version>${project.version}</version>
<scope>test</scope>
</dependency>
<dependency>
diff --git a/integration-test/spark-common/pom.xml b/integration-test/spark-common/pom.xml
index ee8bb0ad..77b9496a 100644
--- a/integration-test/spark-common/pom.xml
+++ b/integration-test/spark-common/pom.xml
@@ -56,6 +56,12 @@
<artifactId>shuffle-server</artifactId>
<scope>test</scope>
</dependency>
+ <dependency>
+ <groupId>org.apache.uniffle</groupId>
+ <artifactId>shuffle-server</artifactId>
+ <type>test-jar</type>
+ <scope>test</scope>
+ </dependency>
<dependency>
<groupId>org.apache.uniffle</groupId>
<artifactId>coordinator</artifactId>
diff --git a/integration-test/spark-common/src/test/java/org/apache/uniffle/test/RSSStageResubmitTest.java b/integration-test/spark-common/src/test/java/org/apache/uniffle/test/RSSStageResubmitTest.java
new file mode 100644
index 00000000..2d83962a
--- /dev/null
+++ b/integration-test/spark-common/src/test/java/org/apache/uniffle/test/RSSStageResubmitTest.java
@@ -0,0 +1,94 @@
+/*
+ * 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.uniffle.test;
+
+import java.util.List;
+import java.util.Map;
+
+import com.google.common.collect.Maps;
+import org.apache.spark.SparkConf;
+import org.apache.spark.shuffle.RssSparkConfig;
+import org.apache.spark.sql.Row;
+import org.apache.spark.sql.SparkSession;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+import org.apache.uniffle.client.util.RssClientConfig;
+import org.apache.uniffle.coordinator.CoordinatorConf;
+import org.apache.uniffle.server.MockedGrpcServer;
+import org.apache.uniffle.server.ShuffleServer;
+import org.apache.uniffle.server.ShuffleServerConf;
+import org.apache.uniffle.storage.util.StorageType;
+
+public class RSSStageResubmitTest extends SparkIntegrationTestBase {
+
+ private static int maxTaskFailures = 3;
+
+ @BeforeAll
+ public static void setupServers() throws Exception {
+ CoordinatorConf coordinatorConf = getCoordinatorConf();
+ Map<String, String> dynamicConf = Maps.newHashMap();
+ dynamicConf.put(CoordinatorConf.COORDINATOR_REMOTE_STORAGE_PATH.key(), HDFS_URI + "rss/test");
+ dynamicConf.put(RssSparkConfig.RSS_STORAGE_TYPE.key(), StorageType.MEMORY_LOCALFILE.name());
+ addDynamicConf(coordinatorConf, dynamicConf);
+ createCoordinatorServer(coordinatorConf);
+ ShuffleServerConf shuffleServerConf = getShuffleServerConf();
+ createMockedShuffleServer(shuffleServerConf);
+ enableFirstReadRequest(2 * maxTaskFailures);
+ startServers();
+ }
+
+ private static void enableFirstReadRequest(int failCount) {
+ for (ShuffleServer server: shuffleServers) {
+ ((MockedGrpcServer) server.getServer()).getService().enableFirstNReadRequestToFail(failCount);
+ }
+ }
+
+ @Override
+ public Map runTest(SparkSession spark, String fileName) throws Exception {
+ List<Row> rows = spark.range(0, 1000, 1, 4)
+ .repartition(2).groupBy("id").count()
+ .collectAsList();
+ Map<String, Long> result = Maps.newHashMap();
+ for (Row row: rows) {
+ result.put(row.get(0).toString(), row.getLong(1));
+ }
+ return result;
+ }
+
+ @Override
+ protected SparkConf createSparkConf() {
+ return new SparkConf()
+ .setAppName(this.getClass().getSimpleName())
+ .setMaster(String.format("local[4,%d]", maxTaskFailures));
+ }
+
+ @Override
+ public void updateSparkConfCustomer(SparkConf sparkConf) {
+ sparkConf.set(RssSparkConfig.SPARK_RSS_CONFIG_PREFIX + RssClientConfig.RSS_RESUBMIT_STAGE, "true");
+ sparkConf.set("spark.task.maxFailures", String.valueOf(maxTaskFailures));
+ }
+
+
+ @Test
+ public void testRSSStageResubmit() throws Exception {
+ run();
+ }
+
+}
+
diff --git a/integration-test/spark-common/src/test/java/org/apache/uniffle/test/RssShuffleManagerTest.java b/integration-test/spark-common/src/test/java/org/apache/uniffle/test/RssShuffleManagerTest.java
new file mode 100644
index 00000000..1a59e694
--- /dev/null
+++ b/integration-test/spark-common/src/test/java/org/apache/uniffle/test/RssShuffleManagerTest.java
@@ -0,0 +1,96 @@
+/*
+ * 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.uniffle.test;
+
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Map;
+
+import com.google.common.collect.Maps;
+import org.apache.spark.MapOutputTrackerMaster;
+import org.apache.spark.SparkConf;
+import org.apache.spark.SparkEnv;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.shuffle.RssShuffleManager;
+import org.apache.spark.shuffle.RssSparkConfig;
+import org.apache.spark.sql.SparkSession;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import scala.Option;
+
+import org.apache.uniffle.client.util.RssClientConfig;
+import org.apache.uniffle.coordinator.CoordinatorConf;
+import org.apache.uniffle.server.ShuffleServerConf;
+import org.apache.uniffle.storage.util.StorageType;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class RssShuffleManagerTest extends SparkIntegrationTestBase {
+
+ @BeforeAll
+ public static void setupServers() throws Exception {
+ shutdownServers();
+ CoordinatorConf coordinatorConf = getCoordinatorConf();
+ Map<String, String> dynamicConf = Maps.newHashMap();
+ dynamicConf.put(CoordinatorConf.COORDINATOR_REMOTE_STORAGE_PATH.key(), HDFS_URI + "rss/test");
+ dynamicConf.put(RssSparkConfig.RSS_STORAGE_TYPE.key(), StorageType.MEMORY_LOCALFILE_HDFS.name());
+ addDynamicConf(coordinatorConf, dynamicConf);
+ createCoordinatorServer(coordinatorConf);
+ ShuffleServerConf shuffleServerConf = getShuffleServerConf();
+ createShuffleServer(shuffleServerConf);
+ startServers();
+ }
+
+ @Override
+ Map runTest(SparkSession spark, String fileName) throws Exception {
+ // we don't need to run any spark application here, just return an empty map.
+ return new HashMap();
+ }
+
+ @Test
+ public void testRssShuffleManager() throws Exception {
+ SparkConf conf = createSparkConf();
+ updateSparkConfWithRss(conf);
+ // enable stage recompute
+ conf.set("spark." + RssClientConfig.RSS_RESUBMIT_STAGE, "true");
+ // disable remote storage fetch
+ conf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED, false);
+ conf.set("spark." + RssClientConfig.RSS_STORAGE_TYPE, StorageType.MEMORY_LOCALFILE.name());
+
+ JavaSparkContext sc = null;
+ try {
+ Option<SparkSession> spark = SparkSession.getActiveSession();
+ if (spark.nonEmpty()) {
+ spark.get().stop();
+ }
+ sc = new JavaSparkContext(SparkSession.builder().config(conf).getOrCreate().sparkContext());
+ // create a rdd that triggers shuffle registration
+ long count = sc.parallelize(Arrays.asList(1, 2, 3, 4, 5), 2).groupBy(x -> x).count();
+ assertEquals(5, count);
+ RssShuffleManager shuffleManager = (RssShuffleManager) SparkEnv.get().shuffleManager();
+ shuffleManager.unregisterAllMapOutput(0);
+ MapOutputTrackerMaster master = (MapOutputTrackerMaster) SparkEnv.get().mapOutputTracker();
+ assertTrue(master.containsShuffle(0));
+ } finally {
+ if (sc != null) {
+ sc.stop();
+ }
+ }
+ }
+}
diff --git a/integration-test/spark3/pom.xml b/integration-test/spark3/pom.xml
index f9c05d59..c72f6d86 100644
--- a/integration-test/spark3/pom.xml
+++ b/integration-test/spark3/pom.xml
@@ -77,7 +77,6 @@
<groupId>org.apache.uniffle</groupId>
<artifactId>shuffle-server</artifactId>
<type>test-jar</type>
- <version>${project.version}</version>
<scope>test</scope>
</dependency>
<dependency>
diff --git a/pom.xml b/pom.xml
index 64f3a76f..503d7d6b 100644
--- a/pom.xml
+++ b/pom.xml
@@ -634,6 +634,13 @@
<artifactId>snappy-java</artifactId>
<version>${snappy-java.version}</version>
</dependency>
+ <dependency>
+ <groupId>org.apache.uniffle</groupId>
+ <artifactId>shuffle-server</artifactId>
+ <version>${project.version}</version>
+ <type>test-jar</type>
+ <scope>test</scope>
+ </dependency>
</dependencies>
</dependencyManagement>
diff --git a/server/src/test/java/org/apache/uniffle/server/MockedShuffleServerGrpcService.java b/server/src/test/java/org/apache/uniffle/server/MockedShuffleServerGrpcService.java
index 95b814bd..54816e1b 100644
--- a/server/src/test/java/org/apache/uniffle/server/MockedShuffleServerGrpcService.java
+++ b/server/src/test/java/org/apache/uniffle/server/MockedShuffleServerGrpcService.java
@@ -42,6 +42,9 @@ public class MockedShuffleServerGrpcService extends ShuffleServerGrpcService {
private boolean recordGetShuffleResult = false;
+ private long numOfFailedReadRequest = 0;
+ private AtomicInteger failedReadRequest = new AtomicInteger(0);
+
public void enableMockedTimeout(long timeout) {
mockedTimeout = timeout;
}
@@ -54,6 +57,15 @@ public class MockedShuffleServerGrpcService extends ShuffleServerGrpcService {
mockedTimeout = -1;
}
+ public void enableFirstNReadRequestToFail(int n) {
+ numOfFailedReadRequest = n;
+ }
+
+ public void resetFirstNReadRequestToFail() {
+ numOfFailedReadRequest = 0;
+ failedReadRequest.set(0);
+ }
+
public MockedShuffleServerGrpcService(ShuffleServer shuffleServer) {
super(shuffleServer);
}
@@ -109,4 +121,19 @@ public class MockedShuffleServerGrpcService extends ShuffleServerGrpcService {
public Map<String, Map<Integer, AtomicInteger>> getShuffleIdToPartitionRequest() {
return appToPartitionRequest;
}
+
+ @Override
+ public void getMemoryShuffleData(RssProtos.GetMemoryShuffleDataRequest request,
+ StreamObserver<RssProtos.GetMemoryShuffleDataResponse> responseObserver) {
+ if (numOfFailedReadRequest > 0) {
+ int currentFailedReadRequest = failedReadRequest.getAndIncrement();
+ if (currentFailedReadRequest < numOfFailedReadRequest) {
+ LOG.info("This request is failed as mocked failure, current/firstN: {}/{}",
+ currentFailedReadRequest, numOfFailedReadRequest);
+ throw new RuntimeException("This request is failed as mocked failure");
+ }
+ }
+ super.getMemoryShuffleData(request, responseObserver);
+ }
+
}