[CELEBORN-2174] Remove partitionSplitEnabled from ReserveSlots and FileInfo

Author: SteNicholas

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

@@ -1320,7 +1320,6 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
             rangeReadFilter,
             userIdentifier,
             conf.pushDataTimeoutMs,
-            partitionSplitEnabled = true,
             isSegmentGranularityVisible = isSegmentGranularityVisible))
         futures.add((future, workerInfo))
       }(ec)

common/src/main/java/org/apache/celeborn/common/meta/DiskFileInfo.java

@@ -19,7 +19,7 @@
 
 import java.io.File;
 import java.util.ArrayList;
-import java.util.Arrays;
+import java.util.Collections;
 
 import com.google.common.annotations.VisibleForTesting;
 import org.apache.hadoop.fs.FileSystem;
@@ -42,24 +42,22 @@ public class DiskFileInfo extends FileInfo {
 
   public DiskFileInfo(
       UserIdentifier userIdentifier,
-      boolean partitionSplitEnabled,
       FileMeta fileMeta,
       String filePath,
       StorageInfo.Type storageType) {
-    super(userIdentifier, partitionSplitEnabled, fileMeta);
+    super(userIdentifier, fileMeta);
     this.filePath = filePath;
     this.storageType = storageType;
   }
 
   // only called when restore from pb or in UT
   public DiskFileInfo(
       UserIdentifier userIdentifier,
-      boolean partitionSplitEnabled,
       FileMeta fileMeta,
       String filePath,
       StorageInfo.Type storageType,
       long bytesFlushed) {
-    super(userIdentifier, partitionSplitEnabled, fileMeta);
+    super(userIdentifier, fileMeta);
     this.filePath = filePath;
     if (storageType != null) {
       this.storageType = storageType;
@@ -73,14 +71,13 @@ public DiskFileInfo(
   public DiskFileInfo(File file, UserIdentifier userIdentifier, CelebornConf conf) {
     this(
         userIdentifier,
-        true,
-        new ReduceFileMeta(new ArrayList<>(Arrays.asList(0L)), conf.shuffleChunkSize()),
+        new ReduceFileMeta(new ArrayList<>(Collections.singletonList(0L)), conf.shuffleChunkSize()),
         file.getAbsolutePath(),
         StorageInfo.Type.HDD);
   }
 
   public DiskFileInfo(UserIdentifier userIdentifier, FileMeta fileMeta, String filePath) {
-    super(userIdentifier, true, fileMeta);
+    super(userIdentifier, fileMeta);
     this.filePath = filePath;
     this.storageType = StorageInfo.Type.HDD;
   }

common/src/main/java/org/apache/celeborn/common/meta/FileInfo.java

@@ -24,17 +24,13 @@
 
 public abstract class FileInfo {
   private final UserIdentifier userIdentifier;
-  // whether to split is decided by client side.
-  // now it's just used for mappartition to compatible with old client which can't support split
-  private boolean partitionSplitEnabled;
   protected FileMeta fileMeta;
   protected final Set<Long> streams = ConcurrentHashMap.newKeySet();
   protected volatile long bytesFlushed;
   private boolean isReduceFileMeta;
 
-  public FileInfo(UserIdentifier userIdentifier, boolean partitionSplitEnabled, FileMeta fileMeta) {
+  public FileInfo(UserIdentifier userIdentifier, FileMeta fileMeta) {
     this.userIdentifier = userIdentifier;
-    this.partitionSplitEnabled = partitionSplitEnabled;
     this.fileMeta = fileMeta;
     this.isReduceFileMeta = fileMeta instanceof ReduceFileMeta;
   }
@@ -67,14 +63,6 @@ public UserIdentifier getUserIdentifier() {
     return userIdentifier;
   }
 
-  public boolean isPartitionSplitEnabled() {
-    return partitionSplitEnabled;
-  }
-
-  public void setPartitionSplitEnabled(boolean partitionSplitEnabled) {
-    this.partitionSplitEnabled = partitionSplitEnabled;
-  }
-
   public boolean addStream(long streamId) {
     if (!isReduceFileMeta) {
       throw new IllegalStateException("In addStream, filemeta cannot be MapFileMeta");

common/src/main/java/org/apache/celeborn/common/meta/MemoryFileInfo.java

@@ -32,18 +32,13 @@ public class MemoryFileInfo extends FileInfo {
   private CompositeByteBuf sortedBuffer;
   private Map<Integer, List<ShuffleBlockInfo>> sortedIndexes;
 
-  public MemoryFileInfo(
-      UserIdentifier userIdentifier, boolean partitionSplitEnabled, FileMeta fileMeta) {
-    super(userIdentifier, partitionSplitEnabled, fileMeta);
+  public MemoryFileInfo(UserIdentifier userIdentifier, FileMeta fileMeta) {
+    super(userIdentifier, fileMeta);
   }
 
   // This constructor is only used in partition sorter for temp new memory file
-  public MemoryFileInfo(
-      UserIdentifier userIdentifier,
-      boolean partitionSplitEnabled,
-      FileMeta fileMeta,
-      CompositeByteBuf buffer) {
-    super(userIdentifier, partitionSplitEnabled, fileMeta);
+  public MemoryFileInfo(UserIdentifier userIdentifier, FileMeta fileMeta, CompositeByteBuf buffer) {
+    super(userIdentifier, fileMeta);
     this.buffer = buffer;
   }
 

common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala

@@ -481,7 +481,6 @@ object ControlMessages extends Logging {
       rangeReadFilter: Boolean,
       userIdentifier: UserIdentifier,
       pushDataTimeout: Long,
-      partitionSplitEnabled: Boolean = false,
       isSegmentGranularityVisible: Boolean = false)
     extends WorkerMessage
 
@@ -960,7 +959,6 @@ object ControlMessages extends Logging {
           rangeReadFilter,
           userIdentifier,
           pushDataTimeout,
-          partitionSplitEnabled,
           isSegmentGranularityVisible) =>
       val payload = PbReserveSlots.newBuilder()
         .setApplicationId(applicationId)
@@ -973,7 +971,6 @@ object ControlMessages extends Logging {
         .setRangeReadFilter(rangeReadFilter)
         .setUserIdentifier(PbSerDeUtils.toPbUserIdentifier(userIdentifier))
         .setPushDataTimeout(pushDataTimeout)
-        .setPartitionSplitEnabled(partitionSplitEnabled)
         .setIsSegmentGranularityVisible(isSegmentGranularityVisible)
         .build().toByteArray
       new TransportMessage(MessageType.RESERVE_SLOTS, payload)
@@ -1438,7 +1435,6 @@ object ControlMessages extends Logging {
           pbReserveSlots.getRangeReadFilter,
           userIdentifier,
           pbReserveSlots.getPushDataTimeout,
-          pbReserveSlots.getPartitionSplitEnabled,
           pbReserveSlots.getIsSegmentGranularityVisible)
 
       case RESERVE_SLOTS_RESPONSE_VALUE =>

common/src/main/scala/org/apache/celeborn/common/util/PbSerDeUtils.scala

@@ -128,7 +128,6 @@ object PbSerDeUtils {
     }
     new DiskFileInfo(
       userIdentifier,
-      pbFileInfo.getPartitionSplitEnabled,
       meta,
       pbFileInfo.getFilePath,
       storageType,
@@ -153,21 +152,20 @@ object PbSerDeUtils {
       .setFilePath(fileInfo.getFilePath)
       .setUserIdentifier(toPbUserIdentifier(fileInfo.getUserIdentifier))
       .setBytesFlushed(fileInfo.getFileLength)
-      .setPartitionSplitEnabled(fileInfo.isPartitionSplitEnabled)
       .setStorageType(fileInfo.getStorageType.getValue)
-    if (fileInfo.getFileMeta.isInstanceOf[MapFileMeta]) {
-      val mapFileMeta = fileInfo.getFileMeta.asInstanceOf[MapFileMeta]
-      builder.setPartitionType(PartitionType.MAP.getValue)
-      builder.setBufferSize(mapFileMeta.getBufferSize)
-      builder.setNumSubpartitions(mapFileMeta.getNumSubpartitions)
-      builder.setIsSegmentGranularityVisible(mapFileMeta.isSegmentGranularityVisible)
-      builder.putAllPartitionWritingSegment(mapFileMeta.getSubPartitionWritingSegmentId)
-      builder.addAllSegmentIndex(
-        mapFileMeta.getSubPartitionSegmentIndexes.asScala.map(toPbSegmentIndex).asJava)
-    } else {
-      val reduceFileMeta = fileInfo.getFileMeta.asInstanceOf[ReduceFileMeta]
-      builder.setPartitionType(PartitionType.REDUCE.getValue)
-      builder.addAllChunkOffsets(reduceFileMeta.getChunkOffsets)
+    fileInfo.getFileMeta match {
+      case mapFileMeta: MapFileMeta =>
+        builder.setPartitionType(PartitionType.MAP.getValue)
+        builder.setBufferSize(mapFileMeta.getBufferSize)
+        builder.setNumSubpartitions(mapFileMeta.getNumSubpartitions)
+        builder.setIsSegmentGranularityVisible(mapFileMeta.isSegmentGranularityVisible)
+        builder.putAllPartitionWritingSegment(mapFileMeta.getSubPartitionWritingSegmentId)
+        builder.addAllSegmentIndex(
+          mapFileMeta.getSubPartitionSegmentIndexes.asScala.map(toPbSegmentIndex).asJava)
+      case _ =>
+        val reduceFileMeta = fileInfo.getFileMeta.asInstanceOf[ReduceFileMeta]
+        builder.setPartitionType(PartitionType.REDUCE.getValue)
+        builder.addAllChunkOffsets(reduceFileMeta.getChunkOffsets)
     }
     builder.build
   }

common/src/test/scala/org/apache/celeborn/common/util/PbSerDeUtilsTest.scala

@@ -77,57 +77,49 @@ class PbSerDeUtilsTest extends CelebornFunSuite {
 
   val fileInfo1 = new DiskFileInfo(
     userIdentifier1,
-    true,
     new ReduceFileMeta(chunkOffsets1, 123),
     file1.getAbsolutePath,
     StorageInfo.Type.HDD,
     3000L)
   val fileInfo2 = new DiskFileInfo(
     userIdentifier2,
-    true,
     new ReduceFileMeta(chunkOffsets2, 123),
     file2.getAbsolutePath,
     StorageInfo.Type.SSD,
     6000L)
   val fileInfo3 = new DiskFileInfo(
     userIdentifier3,
-    true,
     new ReduceFileMeta(chunkOffsets3, 123),
     file3,
     StorageInfo.Type.HDFS,
     6000L)
   val fileInfo4 = new DiskFileInfo(
     userIdentifier3,
-    true,
     new ReduceFileMeta(chunkOffsets3, 123),
     file4,
     StorageInfo.Type.OSS,
     6000L)
   val fileInfo5 = new DiskFileInfo(
     userIdentifier3,
-    true,
     new ReduceFileMeta(chunkOffsets3, 123),
     file5,
     StorageInfo.Type.S3,
     6000L)
   val fileInfo6 = new DiskFileInfo(
     userIdentifier3,
-    true,
     new ReduceFileMeta(chunkOffsets3, 123),
     file6,
     StorageInfo.Type.S3,
     6000L)
 
   val mapFileInfo1 = new DiskFileInfo(
     userIdentifier1,
-    true,
     new MapFileMeta(1024, 10),
     file1.getAbsolutePath,
     StorageInfo.Type.HDD,
     6000L)
   val mapFileInfo2 = new DiskFileInfo(
     userIdentifier2,
-    true,
     new MapFileMeta(1024, 10),
     file2.getAbsolutePath,
     StorageInfo.Type.SSD,
@@ -376,7 +368,6 @@ class PbSerDeUtilsTest extends CelebornFunSuite {
       .setFilePath(diskFileInfo.getFilePath)
       .setUserIdentifier(toPbUserIdentifier(diskFileInfo.getUserIdentifier))
       .setBytesFlushed(diskFileInfo.getFileLength)
-      .setPartitionSplitEnabled(diskFileInfo.isPartitionSplitEnabled)
     val reduceFileMeta = diskFileInfo.getFileMeta.asInstanceOf[ReduceFileMeta]
     builder.setPartitionType(PartitionType.REDUCE.getValue)
     builder.addAllChunkOffsets(reduceFileMeta.getChunkOffsets)

worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/MapPartitionDataReader.java

@@ -455,7 +455,7 @@ public void release() {
         logger.debug("release reader for stream {}", streamId);
         // old client can't support BufferStreamEnd, so for new client it tells client that this
         // stream is finished.
-        if (fileInfo.isPartitionSplitEnabled() && !errorNotified) {
+        if (!errorNotified) {
           associatedChannel.writeAndFlush(
               new RpcRequest(
                   TransportClient.requestId(),

worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionDataWriterContext.java

@@ -33,7 +33,6 @@ public class PartitionDataWriterContext {
   private final String appId;
   private final int shuffleId;
   private final UserIdentifier userIdentifier;
-  private final boolean partitionSplitEnabled;
   private final String shuffleKey;
   private final PartitionType partitionType;
   private final boolean isSegmentGranularityVisible;
@@ -51,7 +50,6 @@ public PartitionDataWriterContext(
       int shuffleId,
       UserIdentifier userIdentifier,
       PartitionType partitionType,
-      boolean partitionSplitEnabled,
       boolean isSegmentGranularityVisible) {
     this.splitThreshold = splitThreshold;
     this.partitionSplitMode = partitionSplitMode;
@@ -60,7 +58,6 @@ public PartitionDataWriterContext(
     this.appId = appId;
     this.shuffleId = shuffleId;
     this.userIdentifier = userIdentifier;
-    this.partitionSplitEnabled = partitionSplitEnabled;
     this.partitionType = partitionType;
     this.shuffleKey = Utils.makeShuffleKey(appId, shuffleId);
     this.isSegmentGranularityVisible = isSegmentGranularityVisible;
@@ -94,10 +91,6 @@ public UserIdentifier getUserIdentifier() {
     return userIdentifier;
   }
 
-  public boolean isPartitionSplitEnabled() {
-    return partitionSplitEnabled;
-  }
-
   public String getShuffleKey() {
     return shuffleKey;
   }
@@ -152,8 +145,6 @@ public String toString() {
         + shuffleId
         + ", userIdentifier="
         + userIdentifier
-        + ", partitionSplitEnabled="
-        + partitionSplitEnabled
         + ", shuffleKey='"
         + shuffleKey
         + '\''

worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java

@@ -227,11 +227,7 @@ public FileInfo getSortedFileInfo(
           memoryFileInfo.getSortedBuffer(),
           targetBuffer,
           shuffleChunkSize);
-      return new MemoryFileInfo(
-          memoryFileInfo.getUserIdentifier(),
-          memoryFileInfo.isPartitionSplitEnabled(),
-          reduceFileMeta,
-          targetBuffer);
+      return new MemoryFileInfo(memoryFileInfo.getUserIdentifier(), reduceFileMeta, targetBuffer);
     } else {
       DiskFileInfo diskFileInfo = ((DiskFileInfo) fileInfo);
       String fileId = shuffleKey + "-" + fileName;