Allow to specify custom tags to decide the assignment of servers

client-mr/src/main/java/org/apache/hadoop/mapreduce/RssMRConfig.java

@@ -124,6 +124,9 @@ public class RssMRConfig {
   public static String RSS_ACCESS_TIMEOUT_MS = MR_RSS_CONFIG_PREFIX + RssClientConfig.RSS_ACCESS_TIMEOUT_MS;
   public static int RSS_ACCESS_TIMEOUT_MS_DEFAULT_VALUE = RssClientConfig.RSS_ACCESS_TIMEOUT_MS_DEFAULT_VALUE;
 
+  public static final String RSS_CLIENT_ASSIGNMENT_TAGS =
+          MR_RSS_CONFIG_PREFIX + RssClientConfig.RSS_CLIENT_ASSIGNMENT_TAGS;
+
   public static String RSS_CONF_FILE = "rss_conf.xml";
 
   public static Set<String> RSS_MANDATORY_CLUSTER_CONF = Sets.newHashSet(

client-mr/src/main/java/org/apache/hadoop/mapreduce/v2/app/RssMRAppMaster.java

@@ -19,8 +19,11 @@
 
 import java.io.IOException;
 import java.lang.reflect.Field;
+import java.util.Arrays;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ThreadFactory;
@@ -115,11 +118,19 @@ public static void main(String[] args) {
       LOG.info("Registering coordinators {}", coordinators);
       client.registerCoordinators(coordinators);
 
+      // Get the configured server assignment tags and it will also add default shuffle version tag.
+      Set<String> assignmentTags = new HashSet<>();
+      String rawTags = conf.get(RssMRConfig.RSS_CLIENT_ASSIGNMENT_TAGS, "");
+      if (StringUtils.isNotEmpty(rawTags)) {
+        rawTags = rawTags.trim();
+        assignmentTags.addAll(Arrays.asList(rawTags.split(",")));
+      }
+      assignmentTags.add(Constants.SHUFFLE_SERVER_VERSION);
+
       ApplicationAttemptId applicationAttemptId = RssMRUtils.getApplicationAttemptId();
       String appId = applicationAttemptId.toString();
       ShuffleAssignmentsInfo response = client.getShuffleAssignments(
-          appId, 0, numReduceTasks,
-          1, Sets.newHashSet(Constants.SHUFFLE_SERVER_VERSION));
+          appId, 0, numReduceTasks, 1, Sets.newHashSet(assignmentTags));
 
       Map<ShuffleServerInfo, List<PartitionRange>> serverToPartitionRanges = response.getServerToPartitionRanges();
       final ScheduledExecutorService scheduledExecutorService = Executors.newSingleThreadScheduledExecutor(

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

@@ -133,6 +133,8 @@ public class RssSparkConfig {
       SPARK_RSS_CONFIG_PREFIX + RssClientConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED;
   public static final boolean RSS_DYNAMIC_CLIENT_CONF_ENABLED_DEFAULT_VALUE =
       RssClientConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED_DEFAULT_VALUE;
+  public static final String RSS_CLIENT_ASSIGNMENT_TAGS =
+      SPARK_RSS_CONFIG_PREFIX + RssClientConfig.RSS_CLIENT_ASSIGNMENT_TAGS;
 
   public static final Set<String> RSS_MANDATORY_CLUSTER_CONF =
       Sets.newHashSet(RSS_STORAGE_TYPE, RSS_REMOTE_STORAGE_PATH);

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

@@ -18,9 +18,13 @@
 package org.apache.spark.shuffle;
 
 import java.lang.reflect.Constructor;
+import java.util.Arrays;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
+import org.apache.commons.lang3.StringUtils;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.spark.SparkConf;
 import org.apache.spark.deploy.SparkHadoopUtil;
@@ -30,6 +34,7 @@
 import org.apache.uniffle.client.api.CoordinatorClient;
 import org.apache.uniffle.client.factory.CoordinatorClientFactory;
 import org.apache.uniffle.common.RemoteStorageInfo;
+import org.apache.uniffle.common.util.Constants;
 
 public class RssSparkShuffleUtils {
 
@@ -123,4 +128,15 @@ public static Configuration getRemoteStorageHadoopConf(
     }
     return readerHadoopConf;
   }
+
+  public static Set<String> getAssignmentTags(SparkConf sparkConf) {
+    Set<String> assignmentTags = new HashSet<>();
+    String rawTags = sparkConf.get(RssSparkConfig.RSS_CLIENT_ASSIGNMENT_TAGS, "");
+    if (StringUtils.isNotEmpty(rawTags)) {
+      rawTags = rawTags.trim();
+      assignmentTags.addAll(Arrays.asList(rawTags.split(",")));
+    }
+    assignmentTags.add(Constants.SHUFFLE_SERVER_VERSION);
+    return assignmentTags;
+  }
 }

client-spark/common/src/test/java/org/apache/spark/shuffle/RssSparkShuffleUtilsTest.java

@@ -17,11 +17,15 @@
 
 package org.apache.spark.shuffle;
 
+import java.util.Iterator;
 import java.util.Map;
+import java.util.Set;
 
 import com.google.common.collect.Maps;
+
 import org.apache.hadoop.conf.Configuration;
 import org.apache.spark.SparkConf;
+import org.apache.uniffle.common.util.Constants;
 import org.junit.jupiter.api.Test;
 
 import org.apache.uniffle.client.util.RssClientConfig;
@@ -32,6 +36,30 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class RssSparkShuffleUtilsTest {
+
+  @Test
+  public void testAssignmentTags() {
+    SparkConf conf = new SparkConf();
+
+    /**
+     * Case1: dont set the tag implicitly and will return the {@code Constants.SHUFFLE_SERVER_VERSION}
+      */
+    Set<String> tags = RssSparkShuffleUtils.getAssignmentTags(conf);
+    assertEquals(Constants.SHUFFLE_SERVER_VERSION, tags.iterator().next());
+
+    /**
+     * Case2: set the multiple tags implicitly and will return the {@code Constants.SHUFFLE_SERVER_VERSION}
+     * and configured tags.
+     */
+    conf.set(RssSparkConfig.RSS_CLIENT_ASSIGNMENT_TAGS, " a,b");
+    tags = RssSparkShuffleUtils.getAssignmentTags(conf);
+    assertEquals(3, tags.size());
+    Iterator<String> iterator = tags.iterator();
+    assertEquals("a", iterator.next());
+    assertEquals("b", iterator.next());
+    assertEquals(Constants.SHUFFLE_SERVER_VERSION, iterator.next());
+  }
+
   @Test
   public void odfsConfigurationTest() {
     SparkConf conf = new SparkConf();

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

@@ -18,9 +18,9 @@
 package org.apache.spark.shuffle;
 
 import java.util.List;
+import java.util.Set;
 
 import com.google.common.collect.Lists;
-import com.google.common.collect.Sets;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.spark.ShuffleDependency;
 import org.apache.spark.SparkConf;
@@ -100,9 +100,11 @@ private boolean tryAccessCluster() {
 
     for (CoordinatorClient coordinatorClient : coordinatorClients) {
       try {
+        Set<String> assignmentTags = RssSparkShuffleUtils.getAssignmentTags(sparkConf);
+
         RssAccessClusterResponse response =
             coordinatorClient.accessCluster(new RssAccessClusterRequest(
-                accessId, Sets.newHashSet(Constants.SHUFFLE_SERVER_VERSION), accessTimeoutMs));
+                accessId, assignmentTags, accessTimeoutMs));
         if (response.getStatusCode() == ResponseStatusCode.SUCCESS) {
           LOG.warn("Success to access cluster {} using {}", coordinatorClient.getDesc(), accessId);
           return true;

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

@@ -61,7 +61,6 @@
 import org.apache.uniffle.common.ShuffleAssignmentsInfo;
 import org.apache.uniffle.common.ShuffleBlockInfo;
 import org.apache.uniffle.common.ShuffleServerInfo;
-import org.apache.uniffle.common.util.Constants;
 import org.apache.uniffle.common.util.RssUtils;
 
 public class RssShuffleManager implements ShuffleManager {
@@ -230,9 +229,11 @@ public <K, V, C> ShuffleHandle registerShuffle(int shuffleId, int numMaps, Shuff
         RssSparkConfig.RSS_PARTITION_NUM_PER_RANGE_DEFAULT_VALUE);
 
     // get all register info according to coordinator's response
+    Set<String> assignmentTags = RssSparkShuffleUtils.getAssignmentTags(sparkConf);
+
     ShuffleAssignmentsInfo response = shuffleWriteClient.getShuffleAssignments(
         appId, shuffleId, dependency.partitioner().numPartitions(),
-        partitionNumPerRange, Sets.newHashSet(Constants.SHUFFLE_SERVER_VERSION));
+        partitionNumPerRange, assignmentTags);
     Map<Integer, List<ShuffleServerInfo>> partitionToServers = response.getPartitionToServers();
 
     startHeartbeat();

client-spark/spark3/src/main/java/org/apache/spark/shuffle/DelegationRssShuffleManager.java

@@ -18,9 +18,9 @@
 package org.apache.spark.shuffle;
 
 import java.util.List;
+import java.util.Set;
 
 import com.google.common.collect.Lists;
-import com.google.common.collect.Sets;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.spark.ShuffleDependency;
 import org.apache.spark.SparkConf;
@@ -100,9 +100,11 @@ private boolean tryAccessCluster() {
 
     for (CoordinatorClient coordinatorClient : coordinatorClients) {
       try {
+        Set<String> assignmentTags = RssSparkShuffleUtils.getAssignmentTags(sparkConf);
+
         RssAccessClusterResponse response =
             coordinatorClient.accessCluster(new RssAccessClusterRequest(
-                accessId, Sets.newHashSet(Constants.SHUFFLE_SERVER_VERSION), accessTimeoutMs));
+                accessId, assignmentTags, accessTimeoutMs));
         if (response.getStatusCode() == ResponseStatusCode.SUCCESS) {
           LOG.warn("Success to access cluster {} using {}", coordinatorClient.getDesc(), accessId);
           return true;

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

@@ -65,7 +65,6 @@
 import org.apache.uniffle.common.ShuffleAssignmentsInfo;
 import org.apache.uniffle.common.ShuffleBlockInfo;
 import org.apache.uniffle.common.ShuffleServerInfo;
-import org.apache.uniffle.common.util.Constants;
 import org.apache.uniffle.common.util.RssUtils;
 
 public class RssShuffleManager implements ShuffleManager {
@@ -276,12 +275,14 @@ public <K, V, C> ShuffleHandle registerShuffle(int shuffleId, ShuffleDependency<
     remoteStorage = ClientUtils.fetchRemoteStorage(
         id.get(), remoteStorage, dynamicConfEnabled, storageType, shuffleWriteClient);
 
+    Set<String> assignmentTags = RssSparkShuffleUtils.getAssignmentTags(sparkConf);
+
     ShuffleAssignmentsInfo response = shuffleWriteClient.getShuffleAssignments(
         id.get(),
         shuffleId,
         dependency.partitioner().numPartitions(),
         1,
-        Sets.newHashSet(Constants.SHUFFLE_SERVER_VERSION));
+        assignmentTags);
     Map<Integer, List<ShuffleServerInfo>> partitionToServers = response.getPartitionToServers();
 
     startHeartbeat();

client/src/main/java/org/apache/uniffle/client/impl/ShuffleWriteClientImpl.java

@@ -124,7 +124,7 @@ private boolean sendShuffleDataAsync(
                   appId, retryMax, retryIntervalMax, shuffleIdToBlocks);
           long s = System.currentTimeMillis();
           RssSendShuffleDataResponse response = getShuffleServerClient(ssi).sendShuffleData(request);
-          LOG.info("ShuffleWriteClientImpl sendShuffleData cost:" + (System.currentTimeMillis() - s));
+          LOG.info("ShuffleWriteClientImpl sendShuffleData cost:" + (System.currentTimeMillis() - s) + "(ms)");
 
           if (response.getStatusCode() == ResponseStatusCode.SUCCESS) {
             // mark a replica of block that has been sent

client/src/main/java/org/apache/uniffle/client/util/RssClientConfig.java

@@ -56,6 +56,8 @@ public class RssClientConfig {
   // When the size of read buffer reaches the half of JVM region (i.e., 32m),
   // it will incur humongous allocation, so we set it to 14m.
   public static String RSS_CLIENT_READ_BUFFER_SIZE_DEFAULT_VALUE = "14m";
+  // The tags specified by rss client to determine server assignment.
+  public static String RSS_CLIENT_ASSIGNMENT_TAGS = "rss.client.assignment.tags";
 
   public static String RSS_ACCESS_TIMEOUT_MS = "rss.access.timeout.ms";
   public static int RSS_ACCESS_TIMEOUT_MS_DEFAULT_VALUE = 10000;