[CELEBORN-1626] group mapTask

client-spark/spark-2/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala

@@ -21,6 +21,8 @@ import java.io.IOException
 import java.util.concurrent.{ThreadPoolExecutor, TimeUnit}
 import java.util.concurrent.atomic.AtomicReference
 
+import scala.collection.mutable.ArrayBuffer
+
 import org.apache.spark.{Aggregator, InterruptibleIterator, TaskContext}
 import org.apache.spark.internal.Logging
 import org.apache.spark.shuffle.{FetchFailedException, ShuffleReader}
@@ -88,8 +90,29 @@ class CelebornShuffleReader[K, C](
       }
     }
 
+    var partitionIdList = new ArrayBuffer[Int]()
+    if (!conf.groupMapTaskEnabled) {
+      partitionIdList = ArrayBuffer[Int]() ++ (startPartition until endPartition)
+    } else {
+      val numPartitions = dep.partitioner.numPartitions
+      val numMappers = handle.numMaps
+      val partitionGroupCnt =
+        if (conf.groupMapTaskEnabled)
+          math.ceil(numMappers.toDouble / conf.groupMapTaskGroupSize).toInt
+        else 1
+      val groupNumPartitions = numPartitions * partitionGroupCnt
+      (startPartition until endPartition).foreach { originalPartitionId =>
+        (0 until partitionGroupCnt).foreach { groupCnt =>
+          val tmpPartitionId = {
+            originalPartitionId + groupCnt * (groupNumPartitions / partitionGroupCnt)
+          }
+          partitionIdList += tmpPartitionId
+        }
+      }
+    }
+
     val streams = JavaUtils.newConcurrentHashMap[Integer, CelebornInputStream]()
-    (startPartition until endPartition).map(partitionId => {
+    partitionIdList.foreach(partitionId => {
       streamCreatorPool.submit(new Runnable {
         override def run(): Unit = {
           if (exceptionRef.get() == null) {
@@ -115,7 +138,7 @@ class CelebornShuffleReader[K, C](
       })
     })
 
-    val recordIter = (startPartition until endPartition).iterator.map(partitionId => {
+    val recordIter = partitionIdList.iterator.map(partitionId => {
       if (handle.numMaps > 0) {
         val startFetchWait = System.nanoTime()
         var inputStream: CelebornInputStream = streams.get(partitionId)

client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala

@@ -23,6 +23,7 @@ import java.util.concurrent.{ConcurrentHashMap, ThreadPoolExecutor, TimeUnit}
 import java.util.concurrent.atomic.AtomicReference
 
 import scala.collection.JavaConverters._
+import scala.collection.mutable.ArrayBuffer
 
 import org.apache.spark.{Aggregator, InterruptibleIterator, ShuffleDependency, TaskContext}
 import org.apache.spark.celeborn.ExceptionMakerHelper
@@ -107,7 +108,7 @@ class CelebornShuffleReader[K, C](
     val shuffleKey = Utils.makeShuffleKey(handle.appUniqueId, shuffleId)
     var fileGroups: ReduceFileGroups = null
     try {
-      // startPartition is irrelevant
+      // startPartition is irrelevant, for error log print
       fileGroups = shuffleClient.updateFileGroup(shuffleId, startPartition)
     } catch {
       case ce @ (_: CelebornIOException | _: PartitionUnRetryAbleException) =>
@@ -121,8 +122,25 @@ class CelebornShuffleReader[K, C](
       (TransportClient, util.ArrayList[PartitionLocation], PbOpenStreamList.Builder)]()
 
     var partCnt = 0
+    var groupPartitionIdList = new ArrayBuffer[Int]()
+    if (!conf.groupMapTaskEnabled) {
+      groupPartitionIdList = ArrayBuffer[Int]() ++ (startPartition until endPartition)
+    } else {
+      val numPartitions = dep.partitioner.numPartitions
+      val numMappers = handle.numMappers
+      val partitionGroupCnt = math.ceil(numMappers.toDouble / conf.groupMapTaskGroupSize).toInt
+      val groupNumPartitions = numPartitions * partitionGroupCnt
+      (startPartition until endPartition).foreach { originalPartitionId =>
+        (0 until partitionGroupCnt).foreach { groupCnt =>
+          val tmpPartitionId = {
+            originalPartitionId + groupCnt * (groupNumPartitions / partitionGroupCnt)
+          }
+          groupPartitionIdList += tmpPartitionId
+        }
+      }
+    }
 
-    (startPartition until endPartition).foreach { partitionId =>
+    groupPartitionIdList.foreach { partitionId =>
       if (fileGroups.partitionGroups.containsKey(partitionId)) {
         fileGroups.partitionGroups.get(partitionId).asScala.foreach { location =>
           partCnt += 1
@@ -236,17 +254,17 @@ class CelebornShuffleReader[K, C](
     }
 
     val inputStreamCreationWindow = conf.clientInputStreamCreationWindow
-    (startPartition until Math.min(
-      startPartition + inputStreamCreationWindow,
-      endPartition)).foreach(partitionId => {
+
+    (0 until Math.min(inputStreamCreationWindow, groupPartitionIdList.size)).foreach(listIndex => {
       streamCreatorPool.submit(new Runnable {
         override def run(): Unit = {
-          createInputStream(partitionId)
+          createInputStream(groupPartitionIdList(listIndex))
         }
       })
     })
 
-    val recordIter = (startPartition until endPartition).iterator.map(partitionId => {
+    var curIndex = 0
+    val recordIter = groupPartitionIdList.iterator.map(partitionId => {
       if (handle.numMappers > 0) {
         val startFetchWait = System.nanoTime()
         var inputStream: CelebornInputStream = streams.get(partitionId)
@@ -258,7 +276,7 @@ class CelebornShuffleReader[K, C](
               case e => throw e
             }
           }
-          logInfo("inputStream is null, sleeping...")
+          logInfo(s"partitionId ${partitionId} inputStream is null, sleeping...")
           Thread.sleep(50)
           inputStream = streams.get(partitionId)
         }
@@ -268,16 +286,19 @@ class CelebornShuffleReader[K, C](
         context.addTaskCompletionListener[Unit](_ => inputStream.close())
 
         // Advance the input creation window
-        if (partitionId + inputStreamCreationWindow < endPartition) {
+        if (curIndex + inputStreamCreationWindow < groupPartitionIdList.size) {
+          val nextPartitionId = groupPartitionIdList(curIndex + inputStreamCreationWindow)
           streamCreatorPool.submit(new Runnable {
             override def run(): Unit = {
-              createInputStream(partitionId + inputStreamCreationWindow)
+              createInputStream(nextPartitionId)
             }
           })
         }
 
+        curIndex = curIndex + 1
         (partitionId, inputStream)
       } else {
+        curIndex = curIndex + 1
         (partitionId, CelebornInputStream.empty())
       }
     }).filter {

client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java

@@ -965,11 +965,17 @@ public int pushOrMergeData(
       return 0;
     }
 
-    final PartitionLocation loc = map.get(partitionId);
+    int innerPartitionId =
+        conf.groupMapTaskEnabled()
+            ? partitionId + (mapId / conf.groupMapTaskGroupSize()) * numPartitions
+            : partitionId;
+    PartitionLocation loc = map.get(innerPartitionId);
+
     if (loc == null) {
       throw new CelebornIOException(
           String.format(
-              "Partition location for shuffle %s partition %d is NULL!", shuffleId, partitionId));
+              "Partition location for shuffle %s partition %d groupPartition %d is NULL!",
+              shuffleId, partitionId, innerPartitionId));
     }
 
     PushState pushState = getPushState(mapKey);
@@ -1017,21 +1023,22 @@ public void onSuccess(ByteBuffer response) {
                     .add(mapId);
               }
               logger.debug(
-                  "Push data to {} success for shuffle {} map {} attempt {} partition {} batch {}.",
+                  "Push data to {} success for shuffle {} map {} attempt {} partition {} groupPartition {} batch {}.",
                   loc.hostAndPushPort(),
                   shuffleId,
                   mapId,
                   attemptId,
                   partitionId,
+                  innerPartitionId,
                   nextBatchId);
             }
 
             @Override
             public void onFailure(Throwable e) {
               String errorMsg =
                   String.format(
-                      "Push data to %s failed for shuffle %d map %d attempt %d partition %d batch %d.",
-                      loc, shuffleId, mapId, attemptId, partitionId, nextBatchId);
+                      "Push data to %s failed for shuffle %d map %d attempt %d partition %d groupPartition %d batch %d.",
+                      loc, shuffleId, mapId, attemptId, partitionId, innerPartitionId, nextBatchId);
               pushState.exception.compareAndSet(null, new CelebornIOException(errorMsg, e));
             }
           };
@@ -1054,21 +1061,25 @@ public void onSuccess(ByteBuffer response) {
                 byte reason = response.get();
                 if (reason == StatusCode.SOFT_SPLIT.getValue()) {
                   logger.debug(
-                      "Push data to {} soft split required for shuffle {} map {} attempt {} partition {} batch {}.",
+                      "Push data to {} soft split required for shuffle {} map {} attempt {} partition {} groupPartition {} batch {}.",
                       latest.hostAndPushPort(),
                       shuffleId,
                       mapId,
                       attemptId,
                       partitionId,
+                      innerPartitionId,
                       nextBatchId);
                   if (!newerPartitionLocationExists(
-                      reducePartitionMap.get(shuffleId), partitionId, latest.getEpoch(), false)) {
+                      reducePartitionMap.get(shuffleId),
+                      innerPartitionId,
+                      latest.getEpoch(),
+                      false)) {
                     ReviveRequest reviveRequest =
                         new ReviveRequest(
                             shuffleId,
                             mapId,
                             attemptId,
-                            partitionId,
+                            innerPartitionId,
                             latest.getEpoch(),
                             latest,
                             StatusCode.SOFT_SPLIT);
@@ -1079,19 +1090,20 @@ public void onSuccess(ByteBuffer response) {
                   callback.onSuccess(response);
                 } else if (reason == StatusCode.HARD_SPLIT.getValue()) {
                   logger.debug(
-                      "Push data to {} hard split required for shuffle {} map {} attempt {} partition {} batch {}.",
+                      "Push data to {} hard split required for shuffle {} map {} attempt {} partition {} groupPartition {} batch {}.",
                       latest.hostAndPushPort(),
                       shuffleId,
                       mapId,
                       attemptId,
                       partitionId,
+                      innerPartitionId,
                       nextBatchId);
                   ReviveRequest reviveRequest =
                       new ReviveRequest(
                           shuffleId,
                           mapId,
                           attemptId,
-                          partitionId,
+                          innerPartitionId,
                           latest.getEpoch(),
                           latest,
                           StatusCode.HARD_SPLIT);
@@ -1114,24 +1126,26 @@ public void onSuccess(ByteBuffer response) {
                               dueTime));
                 } else if (reason == StatusCode.PUSH_DATA_SUCCESS_PRIMARY_CONGESTED.getValue()) {
                   logger.debug(
-                      "Push data to {} primary congestion required for shuffle {} map {} attempt {} partition {} batch {}.",
+                      "Push data to {} primary congestion required for shuffle {} map {} attempt {} partition {} groupPartition {} batch {}.",
                       latest.hostAndPushPort(),
                       shuffleId,
                       mapId,
                       attemptId,
                       partitionId,
+                      innerPartitionId,
                       nextBatchId);
                   pushState.onCongestControl(latest.hostAndPushPort());
                   pushState.removeBatch(nextBatchId, latest.hostAndPushPort());
                   callback.onSuccess(response);
                 } else if (reason == StatusCode.PUSH_DATA_SUCCESS_REPLICA_CONGESTED.getValue()) {
                   logger.debug(
-                      "Push data to {} replica congestion required for shuffle {} map {} attempt {} partition {} batch {}.",
+                      "Push data to {} replica congestion required for shuffle {} map {} attempt {} partition {} groupPartition {} batch {}.",
                       latest.hostAndPushPort(),
                       shuffleId,
                       mapId,
                       attemptId,
                       partitionId,
+                      innerPartitionId,
                       nextBatchId);
                   pushState.onCongestControl(latest.hostAndPushPort());
                   pushState.removeBatch(nextBatchId, latest.hostAndPushPort());
@@ -1166,12 +1180,13 @@ public void onFailure(Throwable e) {
               }
 
               logger.error(
-                  "Push data to {} failed for shuffle {} map {} attempt {} partition {} batch {}, remain revive times {}.",
+                  "Push data to {} failed for shuffle {} map {} attempt {} partition {} groupPartition {} batch {}, remain revive times {}.",
                   latest.hostAndPushPort(),
                   shuffleId,
                   mapId,
                   attemptId,
                   partitionId,
+                  innerPartitionId,
                   nextBatchId,
                   remainReviveTimes,
                   e);
@@ -1180,7 +1195,13 @@ public void onFailure(Throwable e) {
                 remainReviveTimes = remainReviveTimes - 1;
                 ReviveRequest reviveRequest =
                     new ReviveRequest(
-                        shuffleId, mapId, attemptId, partitionId, latest.getEpoch(), latest, cause);
+                        shuffleId,
+                        mapId,
+                        attemptId,
+                        innerPartitionId,
+                        latest.getEpoch(),
+                        latest,
+                        cause);
                 reviveManager.addRequest(reviveRequest);
                 long dueTime =
                     System.currentTimeMillis()
@@ -1217,7 +1238,7 @@ public void onFailure(Throwable e) {
           if (!testRetryRevive) {
             assert dataClientFactory != null;
             TransportClient client =
-                dataClientFactory.createClient(loc.getHost(), loc.getPushPort(), partitionId);
+                dataClientFactory.createClient(loc.getHost(), loc.getPushPort(), innerPartitionId);
             client.pushData(pushData, pushDataTimeout, wrappedCallback);
           } else {
             wrappedCallback.onFailure(
@@ -1228,11 +1249,12 @@ public void onFailure(Throwable e) {
         }
       } catch (Exception e) {
         logger.error(
-            "Exception raised while pushing data for shuffle {} map {} attempt {} partition {} batch {} location {}.",
+            "Exception raised while pushing data for shuffle {} map {} attempt {} partition {} groupPartition {} batch {} location {}.",
             shuffleId,
             mapId,
             attemptId,
             partitionId,
+            innerPartitionId,
             nextBatchId,
             loc,
             e);

client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala

@@ -79,6 +79,8 @@ class ChangePartitionManager(
   private val dynamicResourceEnabled = conf.clientShuffleDynamicResourceEnabled
   private val dynamicResourceUnavailableFactor = conf.clientShuffleDynamicResourceFactor
 
+  private val groupWorkerResources = conf.groupWorkerResources
+
   def start(): Unit = {
     batchHandleChangePartition = batchHandleChangePartitionSchedulerThread.map {
       // noinspection ConvertExpressionToSAM
@@ -363,11 +365,16 @@ class ChangePartitionManager(
       return
     }
 
+    val groupWorkerMap = lifecycleManager.groupedWorkers.getOrDefault(
+      shuffleId,
+      new ConcurrentHashMap[Int, util.HashSet[WorkerInfo]]())
+
     // PartitionSplit all contains oldPartition
     val newlyAllocatedLocations =
       reallocateChangePartitionRequestSlotsFromCandidates(
         changePartitions.toList,
-        candidates.asScala.toList)
+        candidates.asScala.toList,
+        groupWorkerMap)
 
     if (!lifecycleManager.reserveSlotsWithRetry(
         shuffleId,
@@ -413,14 +420,29 @@ class ChangePartitionManager(
 
   private def reallocateChangePartitionRequestSlotsFromCandidates(
       changePartitionRequests: List[ChangePartitionRequest],
-      candidates: List[WorkerInfo]): WorkerResource = {
+      candidates: List[WorkerInfo],
+      groupWorkerMap: ConcurrentHashMap[Int, util.HashSet[WorkerInfo]]): WorkerResource = {
     val slots = new WorkerResource()
     changePartitionRequests.foreach { partition =>
+      val partitionId = partition.partitionId
+      val groupWorkerList =
+        if (groupWorkerResources) {
+          Option(lifecycleManager.partitionGroupMap.get(partitionId)) match {
+            case Some(partitionGroup) =>
+              groupWorkerMap.getOrDefault(partitionGroup, new util.HashSet[WorkerInfo]()).asScala
+                .filter(lifecycleManager.workerStatusTracker.workerAvailable)
+                .toList
+            case None => List()
+          }
+        } else {
+          List()
+        }
       lifecycleManager.allocateFromCandidates(
-        partition.partitionId,
+        partitionId,
         partition.epoch,
         candidates,
-        slots)
+        slots,
+        groupWorkerList)
     }
     slots
   }