integration-test/spark-common/src/test/java/org/apache/uniffle/test/RssShuffleManagerTest.java - incubator-uniffle - Git at Google

 /*
  * 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.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;

 import scala.Option;

 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.RssSparkConfig;
 import org.apache.spark.shuffle.RssSparkShuffleUtils;
 import org.apache.spark.shuffle.ShuffleHandleInfo;
 import org.apache.spark.sql.SparkSession;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
 import org.junit.jupiter.params.provider.ValueSource;
 import org.roaringbitmap.longlong.Roaring64NavigableMap;

 import org.apache.uniffle.client.api.ShuffleWriteClient;
 import org.apache.uniffle.client.factory.ShuffleClientFactory;
 import org.apache.uniffle.client.util.RssClientConfig;
 import org.apache.uniffle.common.ClientType;
 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.rpc.ServerType;
 import org.apache.uniffle.common.util.BlockId;
 import org.apache.uniffle.common.util.BlockIdLayout;
 import org.apache.uniffle.coordinator.CoordinatorConf;
 import org.apache.uniffle.shuffle.manager.RssShuffleManagerBase;
 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 beforeAll() throws Exception {
     shutdownServers();
   }

   @AfterEach
   public void after() throws Exception {
     shutdownServers();
   }

   public static Map<String, String> startServers(BlockIdLayout dynamicConfLayout) throws Exception {
     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());
     // configure block id layout (if set)
     if (dynamicConfLayout != null) {
       dynamicConf.put(
           RssClientConf.BLOCKID_SEQUENCE_NO_BITS.key(),
           String.valueOf(dynamicConfLayout.sequenceNoBits));
       dynamicConf.put(
           RssClientConf.BLOCKID_PARTITION_ID_BITS.key(),
           String.valueOf(dynamicConfLayout.partitionIdBits));
       dynamicConf.put(
           RssClientConf.BLOCKID_TASK_ATTEMPT_ID_BITS.key(),
           String.valueOf(dynamicConfLayout.taskAttemptIdBits));
     }
     CoordinatorConf coordinatorConf = getCoordinatorConf();
     addDynamicConf(coordinatorConf, dynamicConf);
     createCoordinatorServer(coordinatorConf);
     createShuffleServer(getShuffleServerConf(ServerType.GRPC));
     startServers();
     return dynamicConf;
   }

   @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();
   }

   private static final BlockIdLayout DEFAULT = BlockIdLayout.from(21, 20, 22);
   private static final BlockIdLayout CUSTOM1 = BlockIdLayout.from(20, 21, 22);
   private static final BlockIdLayout CUSTOM2 = BlockIdLayout.from(22, 18, 23);

   public static Stream<Arguments> testBlockIdLayouts() {
     return Stream.of(Arguments.of(DEFAULT), Arguments.of(CUSTOM1), Arguments.of(CUSTOM2));
   }

   @ParameterizedTest
   @ValueSource(booleans = {false, true})
   public void testRssShuffleManager(boolean enableDynamicClientConf) throws Exception {
     doTestRssShuffleManager(null, null, DEFAULT, enableDynamicClientConf);
   }

   @ParameterizedTest
   @MethodSource("testBlockIdLayouts")
   public void testRssShuffleManagerClientConf(BlockIdLayout layout) throws Exception {
     doTestRssShuffleManager(layout, null, layout, true);
   }

   @ParameterizedTest
   @MethodSource("testBlockIdLayouts")
   public void testRssShuffleManagerDynamicClientConf(BlockIdLayout layout) throws Exception {
     doTestRssShuffleManager(null, layout, layout, true);
   }

   @ParameterizedTest
   @ValueSource(booleans = {false, true})
   public void testRssShuffleManagerClientConfOverride(boolean enableDynamicClientConf)
       throws Exception {
     doTestRssShuffleManager(CUSTOM1, CUSTOM2, CUSTOM1, enableDynamicClientConf);
   }

   private void doTestRssShuffleManager(
       BlockIdLayout clientConfLayout,
       BlockIdLayout dynamicConfLayout,
       BlockIdLayout expectedLayout,
       boolean enableDynamicCLientConf)
       throws Exception {
     Map<String, String> dynamicConf = startServers(dynamicConfLayout);

     SparkConf conf = createSparkConf();
     updateSparkConfWithRss(conf);
     // enable stage recompute
     conf.set("spark." + RssClientConfig.RSS_RESUBMIT_STAGE, "true");
     // enable dynamic client conf
     conf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED, enableDynamicCLientConf);
     // configure storage type
     conf.set("spark." + RssClientConfig.RSS_STORAGE_TYPE, StorageType.MEMORY_LOCALFILE.name());
     // restarting the coordinator may cause RssException: There isn't enough shuffle servers
     // retry quickly (default is 65s interval)
     conf.set("spark." + RssClientConfig.RSS_CLIENT_ASSIGNMENT_RETRY_INTERVAL, "1000");
     conf.set("spark." + RssClientConfig.RSS_CLIENT_ASSIGNMENT_RETRY_TIMES, "10");

     // configure client conf block id layout (if set)
     if (clientConfLayout != null) {
       conf.set(
           "spark." + RssClientConf.BLOCKID_SEQUENCE_NO_BITS.key(),
           String.valueOf(clientConfLayout.sequenceNoBits));
       conf.set(
           "spark." + RssClientConf.BLOCKID_PARTITION_ID_BITS.key(),
           String.valueOf(clientConfLayout.partitionIdBits));
       conf.set(
           "spark." + RssClientConf.BLOCKID_TASK_ATTEMPT_ID_BITS.key(),
           String.valueOf(clientConfLayout.taskAttemptIdBits));
     }

     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);

       // configure expected block id layout
       conf.set(
           "spark." + RssClientConf.BLOCKID_SEQUENCE_NO_BITS.key(),
           String.valueOf(expectedLayout.sequenceNoBits));
       conf.set(
           "spark." + RssClientConf.BLOCKID_PARTITION_ID_BITS.key(),
           String.valueOf(expectedLayout.partitionIdBits));
       conf.set(
           "spark." + RssClientConf.BLOCKID_TASK_ATTEMPT_ID_BITS.key(),
           String.valueOf(expectedLayout.taskAttemptIdBits));

       // get written block ids (we know there is one shuffle where two task attempts wrote two
       // partitions)
       RssConf rssConf = RssSparkConfig.toRssConf(conf);
       if (clientConfLayout == null && dynamicConfLayout != null) {
         RssSparkShuffleUtils.applyDynamicClientConf(conf, dynamicConf);
       }
       RssShuffleManagerBase shuffleManager =
           (RssShuffleManagerBase) SparkEnv.get().shuffleManager();
       shuffleManager.configureBlockIdLayout(conf, rssConf);
       ShuffleWriteClient shuffleWriteClient =
           ShuffleClientFactory.newWriteBuilder()
               .clientType(ClientType.GRPC.name())
               .retryMax(3)
               .retryIntervalMax(2000)
               .heartBeatThreadNum(4)
               .replica(1)
               .replicaWrite(1)
               .replicaRead(1)
               .replicaSkipEnabled(true)
               .dataTransferPoolSize(1)
               .dataCommitPoolSize(1)
               .unregisterThreadPoolSize(10)
               .unregisterRequestTimeSec(10)
               .rssConf(rssConf)
               .build();
       ShuffleHandleInfo handle = shuffleManager.getShuffleHandleInfoByShuffleId(0);
       Set<ShuffleServerInfo> servers =
           handle.getPartitionToServers().values().stream()
               .flatMap(Collection::stream)
               .collect(Collectors.toSet());

       for (int partitionId : new int[] {0, 1}) {
         Roaring64NavigableMap blockIdLongs =
             shuffleWriteClient.getShuffleResult(
                 ClientType.GRPC.name(), servers, shuffleManager.getAppId(), 0, partitionId);
         List<BlockId> blockIds =
             blockIdLongs.stream()
                 .sorted()
                 .mapToObj(expectedLayout::asBlockId)
                 .collect(Collectors.toList());
         assertEquals(2, blockIds.size());
         long taskAttemptId0 = shuffleManager.getTaskAttemptIdForBlockId(0, 0);
         long taskAttemptId1 = shuffleManager.getTaskAttemptIdForBlockId(1, 0);
         assertEquals(expectedLayout.asBlockId(0, partitionId, taskAttemptId0), blockIds.get(0));
         assertEquals(expectedLayout.asBlockId(0, partitionId, taskAttemptId1), blockIds.get(1));
       }

       shuffleManager.unregisterAllMapOutput(0);
       MapOutputTrackerMaster master = (MapOutputTrackerMaster) SparkEnv.get().mapOutputTracker();
       assertTrue(master.containsShuffle(0));
     } finally {
       if (sc != null) {
         sc.stop();
       }
     }
   }
 }
	/*
	* 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.Collection;
	import java.util.HashMap;
	import java.util.List;
	import java.util.Map;
	import java.util.Set;
	import java.util.stream.Collectors;
	import java.util.stream.Stream;

	import scala.Option;

	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.RssSparkConfig;
	import org.apache.spark.shuffle.RssSparkShuffleUtils;
	import org.apache.spark.shuffle.ShuffleHandleInfo;
	import org.apache.spark.sql.SparkSession;
	import org.junit.jupiter.api.AfterEach;
	import org.junit.jupiter.api.BeforeAll;
	import org.junit.jupiter.params.ParameterizedTest;
	import org.junit.jupiter.params.provider.Arguments;
	import org.junit.jupiter.params.provider.MethodSource;
	import org.junit.jupiter.params.provider.ValueSource;
	import org.roaringbitmap.longlong.Roaring64NavigableMap;

	import org.apache.uniffle.client.api.ShuffleWriteClient;
	import org.apache.uniffle.client.factory.ShuffleClientFactory;
	import org.apache.uniffle.client.util.RssClientConfig;
	import org.apache.uniffle.common.ClientType;
	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.rpc.ServerType;
	import org.apache.uniffle.common.util.BlockId;
	import org.apache.uniffle.common.util.BlockIdLayout;
	import org.apache.uniffle.coordinator.CoordinatorConf;
	import org.apache.uniffle.shuffle.manager.RssShuffleManagerBase;
	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 beforeAll() throws Exception {
	shutdownServers();
	}

	@AfterEach
	public void after() throws Exception {
	shutdownServers();
	}

	public static Map<String, String> startServers(BlockIdLayout dynamicConfLayout) throws Exception {
	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());
	// configure block id layout (if set)
	if (dynamicConfLayout != null) {
	dynamicConf.put(
	RssClientConf.BLOCKID_SEQUENCE_NO_BITS.key(),
	String.valueOf(dynamicConfLayout.sequenceNoBits));
	dynamicConf.put(
	RssClientConf.BLOCKID_PARTITION_ID_BITS.key(),
	String.valueOf(dynamicConfLayout.partitionIdBits));
	dynamicConf.put(
	RssClientConf.BLOCKID_TASK_ATTEMPT_ID_BITS.key(),
	String.valueOf(dynamicConfLayout.taskAttemptIdBits));
	}
	CoordinatorConf coordinatorConf = getCoordinatorConf();
	addDynamicConf(coordinatorConf, dynamicConf);
	createCoordinatorServer(coordinatorConf);
	createShuffleServer(getShuffleServerConf(ServerType.GRPC));
	startServers();
	return dynamicConf;
	}

	@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();
	}

	private static final BlockIdLayout DEFAULT = BlockIdLayout.from(21, 20, 22);
	private static final BlockIdLayout CUSTOM1 = BlockIdLayout.from(20, 21, 22);
	private static final BlockIdLayout CUSTOM2 = BlockIdLayout.from(22, 18, 23);

	public static Stream<Arguments> testBlockIdLayouts() {
	return Stream.of(Arguments.of(DEFAULT), Arguments.of(CUSTOM1), Arguments.of(CUSTOM2));
	}

	@ParameterizedTest
	@ValueSource(booleans = {false, true})
	public void testRssShuffleManager(boolean enableDynamicClientConf) throws Exception {
	doTestRssShuffleManager(null, null, DEFAULT, enableDynamicClientConf);
	}

	@ParameterizedTest
	@MethodSource("testBlockIdLayouts")
	public void testRssShuffleManagerClientConf(BlockIdLayout layout) throws Exception {
	doTestRssShuffleManager(layout, null, layout, true);
	}

	@ParameterizedTest
	@MethodSource("testBlockIdLayouts")
	public void testRssShuffleManagerDynamicClientConf(BlockIdLayout layout) throws Exception {
	doTestRssShuffleManager(null, layout, layout, true);
	}

	@ParameterizedTest
	@ValueSource(booleans = {false, true})
	public void testRssShuffleManagerClientConfOverride(boolean enableDynamicClientConf)
	throws Exception {
	doTestRssShuffleManager(CUSTOM1, CUSTOM2, CUSTOM1, enableDynamicClientConf);
	}

	private void doTestRssShuffleManager(
	BlockIdLayout clientConfLayout,
	BlockIdLayout dynamicConfLayout,
	BlockIdLayout expectedLayout,
	boolean enableDynamicCLientConf)
	throws Exception {
	Map<String, String> dynamicConf = startServers(dynamicConfLayout);

	SparkConf conf = createSparkConf();
	updateSparkConfWithRss(conf);
	// enable stage recompute
	conf.set("spark." + RssClientConfig.RSS_RESUBMIT_STAGE, "true");
	// enable dynamic client conf
	conf.set(RssSparkConfig.RSS_DYNAMIC_CLIENT_CONF_ENABLED, enableDynamicCLientConf);
	// configure storage type
	conf.set("spark." + RssClientConfig.RSS_STORAGE_TYPE, StorageType.MEMORY_LOCALFILE.name());
	// restarting the coordinator may cause RssException: There isn't enough shuffle servers
	// retry quickly (default is 65s interval)
	conf.set("spark." + RssClientConfig.RSS_CLIENT_ASSIGNMENT_RETRY_INTERVAL, "1000");
	conf.set("spark." + RssClientConfig.RSS_CLIENT_ASSIGNMENT_RETRY_TIMES, "10");

	// configure client conf block id layout (if set)
	if (clientConfLayout != null) {
	conf.set(
	"spark." + RssClientConf.BLOCKID_SEQUENCE_NO_BITS.key(),
	String.valueOf(clientConfLayout.sequenceNoBits));
	conf.set(
	"spark." + RssClientConf.BLOCKID_PARTITION_ID_BITS.key(),
	String.valueOf(clientConfLayout.partitionIdBits));
	conf.set(
	"spark." + RssClientConf.BLOCKID_TASK_ATTEMPT_ID_BITS.key(),
	String.valueOf(clientConfLayout.taskAttemptIdBits));
	}

	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);

	// configure expected block id layout
	conf.set(
	"spark." + RssClientConf.BLOCKID_SEQUENCE_NO_BITS.key(),
	String.valueOf(expectedLayout.sequenceNoBits));
	conf.set(
	"spark." + RssClientConf.BLOCKID_PARTITION_ID_BITS.key(),
	String.valueOf(expectedLayout.partitionIdBits));
	conf.set(
	"spark." + RssClientConf.BLOCKID_TASK_ATTEMPT_ID_BITS.key(),
	String.valueOf(expectedLayout.taskAttemptIdBits));

	// get written block ids (we know there is one shuffle where two task attempts wrote two
	// partitions)
	RssConf rssConf = RssSparkConfig.toRssConf(conf);
	if (clientConfLayout == null && dynamicConfLayout != null) {
	RssSparkShuffleUtils.applyDynamicClientConf(conf, dynamicConf);
	}
	RssShuffleManagerBase shuffleManager =
	(RssShuffleManagerBase) SparkEnv.get().shuffleManager();
	shuffleManager.configureBlockIdLayout(conf, rssConf);
	ShuffleWriteClient shuffleWriteClient =
	ShuffleClientFactory.newWriteBuilder()
	.clientType(ClientType.GRPC.name())
	.retryMax(3)
	.retryIntervalMax(2000)
	.heartBeatThreadNum(4)
	.replica(1)
	.replicaWrite(1)
	.replicaRead(1)
	.replicaSkipEnabled(true)
	.dataTransferPoolSize(1)
	.dataCommitPoolSize(1)
	.unregisterThreadPoolSize(10)
	.unregisterRequestTimeSec(10)
	.rssConf(rssConf)
	.build();
	ShuffleHandleInfo handle = shuffleManager.getShuffleHandleInfoByShuffleId(0);
	Set<ShuffleServerInfo> servers =
	handle.getPartitionToServers().values().stream()
	.flatMap(Collection::stream)
	.collect(Collectors.toSet());

	for (int partitionId : new int[] {0, 1}) {
	Roaring64NavigableMap blockIdLongs =
	shuffleWriteClient.getShuffleResult(
	ClientType.GRPC.name(), servers, shuffleManager.getAppId(), 0, partitionId);
	List<BlockId> blockIds =
	blockIdLongs.stream()
	.sorted()
	.mapToObj(expectedLayout::asBlockId)
	.collect(Collectors.toList());
	assertEquals(2, blockIds.size());
	long taskAttemptId0 = shuffleManager.getTaskAttemptIdForBlockId(0, 0);
	long taskAttemptId1 = shuffleManager.getTaskAttemptIdForBlockId(1, 0);
	assertEquals(expectedLayout.asBlockId(0, partitionId, taskAttemptId0), blockIds.get(0));
	assertEquals(expectedLayout.asBlockId(0, partitionId, taskAttemptId1), blockIds.get(1));
	}

	shuffleManager.unregisterAllMapOutput(0);
	MapOutputTrackerMaster master = (MapOutputTrackerMaster) SparkEnv.get().mapOutputTracker();
	assertTrue(master.containsShuffle(0));
	} finally {
	if (sc != null) {
	sc.stop();
	}
	}
	}
	}