[GitHub] [incubator-celeborn] waitinfuture commented on a diff in pull request #1102: [CELEBORN-55][FEATURE] Split maxReqsInFlight limitation into level of target worker

[GitBox <gi...@apache.org>]

waitinfuture commented on code in PR #1102:
URL: https://github.com/apache/incubator-celeborn/pull/1102#discussion_r1069294838


##########
client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java:
##########
@@ -0,0 +1,124 @@
+/*
+ * 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.celeborn.client.write;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.LinkedBlockingQueue;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.client.ShuffleClient;
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.util.Utils;
+
+/*
+ * Queue for push data,
+ * it can take one special PushTask queue by partitionId,
+ * and it can take one random PushTask queue.
+ * workingQueuePerPartition: for PushTask queue per every partition.
+ *
+ * */
+public class DataPushQueue {
+  private static final Logger logger = LoggerFactory.getLogger(DataPushQueue.class);
+
+  private final List<LinkedBlockingQueue<PushTask>> workingQueuePerPartition;

Review Comment:
   We should not assume that partitionId is sequential starting from 0 with step 1. Better to use a map



##########
client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java:
##########
@@ -0,0 +1,124 @@
+/*
+ * 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.celeborn.client.write;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.LinkedBlockingQueue;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.client.ShuffleClient;
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.util.Utils;
+
+/*
+ * Queue for push data,
+ * it can take one special PushTask queue by partitionId,
+ * and it can take one random PushTask queue.
+ * workingQueuePerPartition: for PushTask queue per every partition.
+ *
+ * */
+public class DataPushQueue {
+  private static final Logger logger = LoggerFactory.getLogger(DataPushQueue.class);
+
+  private final List<LinkedBlockingQueue<PushTask>> workingQueuePerPartition;
+  private final PushState pushState;
+  private final DataPusher dataPusher;
+  private final int maxInFlight;
+
+  private final String appId;
+  private final int shuffleId;
+  private final int numMappers;
+  private final int numPartitions;
+  private final ShuffleClient client;
+  private int partitionIdIdx = 0;
+
+  public DataPushQueue(
+      CelebornConf conf,
+      DataPusher dataPusher,
+      ShuffleClient client,
+      String appId,
+      int shuffleId,
+      int mapId,
+      int attemptId,
+      int numMappers,
+      int numPartitions) {
+    this.appId = appId;
+    this.shuffleId = shuffleId;
+    this.numMappers = numMappers;
+    this.numPartitions = numPartitions;
+    this.client = client;
+    this.dataPusher = dataPusher;
+    final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
+    this.pushState = client.getOrRegisterPushState(mapKey);
+    this.maxInFlight = conf.pushMaxReqsInFlight();
+    final int capacity = conf.pushQueueCapacity();
+    workingQueuePerPartition = new ArrayList<>(numPartitions);
+    for (int i = 0; i < numPartitions; i++) {
+      workingQueuePerPartition.add(new LinkedBlockingQueue<>(capacity));
+    }
+  }
+
+  public LinkedBlockingQueue<PushTask> takeAnyWorkingQueue() throws IOException {
+    while (!dataPusher.terminatedOrHasException()) {
+      int partitionId = nextPartitionId();
+      LinkedBlockingQueue<PushTask> pushTasks;
+      try {
+        pushTasks = workingQueuePerPartition.get(partitionId);
+      } catch (IndexOutOfBoundsException ex) {
+        if (dataPusher.terminatedOrHasException()) {
+          return null;
+        }
+        throw ex;
+      }
+      if (Objects.isNull(pushTasks) && dataPusher.terminatedOrHasException()) {

Review Comment:
   We don't need Objects.isNull(pushTasks) since the DataPusher thread will stop before we clear DataPushQueue



##########
client/src/main/java/org/apache/celeborn/client/write/InFlightRequestTracker.java:
##########
@@ -0,0 +1,209 @@
+/*
+ * 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.celeborn.client.write;
+
+import java.io.IOException;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import io.netty.channel.ChannelFuture;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.network.client.RpcResponseCallback;
+import org.apache.celeborn.common.protocol.message.StatusCode;
+
+/*
+ * This class is for track in flight request and limit request.
+ * */
+public class InFlightRequestTracker {
+  private static final Logger logger = LoggerFactory.getLogger(InFlightRequestTracker.class);
+
+  private final long timeoutMs;
+  private final long pushTimeout;
+  private final long delta;
+  private final PushState pushState;
+
+  private final AtomicInteger batchId = new AtomicInteger();
+  private final ConcurrentHashMap<String, ConcurrentHashMap<Integer, BatchInfo>>
+      batchIdPerAddressPair = new ConcurrentHashMap<>();

Review Comment:
   better rename to inflightBatchesPerAddress



##########
client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java:
##########
@@ -192,7 +192,7 @@ private void submitRetryPushData(
     } else if (mapperEnded(shuffleId, mapId, attemptId)) {
       logger.debug(
           "Retrying push data, but the mapper(map {} attempt {}) has ended.", mapId, attemptId);
-      pushState.removeBatch(batchId);
+      pushState.removeBatch(batchId, loc.hostAndPushPort());

Review Comment:
   We can make hostAndPushPort a transient member and initiates in Constructor to speedup this method.



##########
client/src/main/java/org/apache/celeborn/client/write/InFlightRequestTracker.java:
##########
@@ -0,0 +1,209 @@
+/*
+ * 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.celeborn.client.write;
+
+import java.io.IOException;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import io.netty.channel.ChannelFuture;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.network.client.RpcResponseCallback;
+import org.apache.celeborn.common.protocol.message.StatusCode;
+
+/*
+ * This class is for track in flight request and limit request.
+ * */
+public class InFlightRequestTracker {
+  private static final Logger logger = LoggerFactory.getLogger(InFlightRequestTracker.class);
+
+  private final long timeoutMs;
+  private final long pushTimeout;
+  private final long delta;
+  private final PushState pushState;
+
+  private final AtomicInteger batchId = new AtomicInteger();
+  private final ConcurrentHashMap<String, ConcurrentHashMap<Integer, BatchInfo>>
+      batchIdPerAddressPair = new ConcurrentHashMap<>();
+
+  public InFlightRequestTracker(CelebornConf conf, PushState pushState) {
+    this.timeoutMs = conf.pushLimitInFlightTimeoutMs();
+    this.pushTimeout = conf.pushDataTimeoutMs();
+    this.delta = conf.pushLimitInFlightSleepDeltaMs();
+    this.pushState = pushState;
+  }
+
+  public void addFlightBatches(int batchId, String hostAndPushPort) {

Review Comment:
   addFlightBatch



##########
client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java:
##########
@@ -0,0 +1,124 @@
+/*
+ * 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.celeborn.client.write;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.LinkedBlockingQueue;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.client.ShuffleClient;
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.util.Utils;
+
+/*
+ * Queue for push data,
+ * it can take one special PushTask queue by partitionId,
+ * and it can take one random PushTask queue.
+ * workingQueuePerPartition: for PushTask queue per every partition.
+ *
+ * */
+public class DataPushQueue {
+  private static final Logger logger = LoggerFactory.getLogger(DataPushQueue.class);
+
+  private final List<LinkedBlockingQueue<PushTask>> workingQueuePerPartition;
+  private final PushState pushState;
+  private final DataPusher dataPusher;
+  private final int maxInFlight;
+
+  private final String appId;
+  private final int shuffleId;
+  private final int numMappers;
+  private final int numPartitions;
+  private final ShuffleClient client;
+  private int partitionIdIdx = 0;
+
+  public DataPushQueue(
+      CelebornConf conf,
+      DataPusher dataPusher,
+      ShuffleClient client,
+      String appId,
+      int shuffleId,
+      int mapId,
+      int attemptId,
+      int numMappers,
+      int numPartitions) {
+    this.appId = appId;
+    this.shuffleId = shuffleId;
+    this.numMappers = numMappers;
+    this.numPartitions = numPartitions;
+    this.client = client;
+    this.dataPusher = dataPusher;
+    final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
+    this.pushState = client.getOrRegisterPushState(mapKey);
+    this.maxInFlight = conf.pushMaxReqsInFlight();
+    final int capacity = conf.pushQueueCapacity();
+    workingQueuePerPartition = new ArrayList<>(numPartitions);
+    for (int i = 0; i < numPartitions; i++) {
+      workingQueuePerPartition.add(new LinkedBlockingQueue<>(capacity));
+    }
+  }
+
+  public LinkedBlockingQueue<PushTask> takeAnyWorkingQueue() throws IOException {
+    while (!dataPusher.terminatedOrHasException()) {
+      int partitionId = nextPartitionId();
+      LinkedBlockingQueue<PushTask> pushTasks;
+      try {
+        pushTasks = workingQueuePerPartition.get(partitionId);
+      } catch (IndexOutOfBoundsException ex) {
+        if (dataPusher.terminatedOrHasException()) {
+          return null;
+        }
+        throw ex;
+      }
+      if (Objects.isNull(pushTasks) && dataPusher.terminatedOrHasException()) {
+        return null;
+      }
+      if (!pushTasks.isEmpty()) {
+        Map<Integer, PartitionLocation> partitionLocationMap =
+            client.getOrRegisterShuffle(appId, shuffleId, numMappers, numPartitions);

Review Comment:
   IMO we should not expose this method to PushDataQueue. We can add a method in client like
   ```
   PartitionLocation getPartitionLocation(partitionId);
   ```
   which can call registerShuffle if needed



##########
tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/HugeDataTest.scala:
##########
@@ -38,6 +38,7 @@ class HugeDataTest extends AnyFunSuite
 
   test("celeborn spark integration test - huge data") {
     val sparkConf = new SparkConf().setAppName("rss-demo").setMaster("local[4]")
+      .set("spark.celeborn.shuffle.register.maxRetries", "10")

Review Comment:
   just curious, what's the purpose to set this config?



##########
client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java:
##########
@@ -0,0 +1,124 @@
+/*
+ * 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.celeborn.client.write;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.LinkedBlockingQueue;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.client.ShuffleClient;
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.util.Utils;
+
+/*
+ * Queue for push data,
+ * it can take one special PushTask queue by partitionId,
+ * and it can take one random PushTask queue.
+ * workingQueuePerPartition: for PushTask queue per every partition.
+ *
+ * */
+public class DataPushQueue {
+  private static final Logger logger = LoggerFactory.getLogger(DataPushQueue.class);
+
+  private final List<LinkedBlockingQueue<PushTask>> workingQueuePerPartition;
+  private final PushState pushState;
+  private final DataPusher dataPusher;
+  private final int maxInFlight;
+
+  private final String appId;
+  private final int shuffleId;
+  private final int numMappers;
+  private final int numPartitions;
+  private final ShuffleClient client;
+  private int partitionIdIdx = 0;
+
+  public DataPushQueue(
+      CelebornConf conf,
+      DataPusher dataPusher,
+      ShuffleClient client,
+      String appId,
+      int shuffleId,
+      int mapId,
+      int attemptId,
+      int numMappers,
+      int numPartitions) {
+    this.appId = appId;
+    this.shuffleId = shuffleId;
+    this.numMappers = numMappers;
+    this.numPartitions = numPartitions;
+    this.client = client;
+    this.dataPusher = dataPusher;
+    final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
+    this.pushState = client.getOrRegisterPushState(mapKey);
+    this.maxInFlight = conf.pushMaxReqsInFlight();
+    final int capacity = conf.pushQueueCapacity();
+    workingQueuePerPartition = new ArrayList<>(numPartitions);
+    for (int i = 0; i < numPartitions; i++) {
+      workingQueuePerPartition.add(new LinkedBlockingQueue<>(capacity));
+    }
+  }
+
+  public LinkedBlockingQueue<PushTask> takeAnyWorkingQueue() throws IOException {
+    while (!dataPusher.terminatedOrHasException()) {
+      int partitionId = nextPartitionId();
+      LinkedBlockingQueue<PushTask> pushTasks;
+      try {
+        pushTasks = workingQueuePerPartition.get(partitionId);
+      } catch (IndexOutOfBoundsException ex) {
+        if (dataPusher.terminatedOrHasException()) {
+          return null;
+        }
+        throw ex;
+      }
+      if (Objects.isNull(pushTasks) && dataPusher.terminatedOrHasException()) {
+        return null;
+      }
+      if (!pushTasks.isEmpty()) {
+        Map<Integer, PartitionLocation> partitionLocationMap =
+            client.getOrRegisterShuffle(appId, shuffleId, numMappers, numPartitions);
+        PartitionLocation loc = partitionLocationMap.get(partitionId);
+        boolean reachLimit = pushState.limitMaxInFlight(loc.hostAndPushPort(), maxInFlight);

Review Comment:
   We can add a method in PushState such as
   ```
   boolean reachLimit(address)
   ```
   which indicates whether the location reaches limit



##########
client/src/main/java/org/apache/celeborn/client/ShuffleClient.java:
##########
@@ -235,4 +237,9 @@ public abstract void pushDataHandShake(
 
   public abstract PartitionLocation registerMapPartitionTask(
       String appId, int shuffleId, int numMappers, int mapId, int attemptId);
+
+  public abstract ConcurrentHashMap<Integer, PartitionLocation> getOrRegisterShuffle(

Review Comment:
   we should call this in addTask()



##########
client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java:
##########
@@ -0,0 +1,124 @@
+/*
+ * 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.celeborn.client.write;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.LinkedBlockingQueue;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.client.ShuffleClient;
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.util.Utils;
+
+/*
+ * Queue for push data,
+ * it can take one special PushTask queue by partitionId,
+ * and it can take one random PushTask queue.
+ * workingQueuePerPartition: for PushTask queue per every partition.
+ *
+ * */
+public class DataPushQueue {
+  private static final Logger logger = LoggerFactory.getLogger(DataPushQueue.class);
+
+  private final List<LinkedBlockingQueue<PushTask>> workingQueuePerPartition;
+  private final PushState pushState;
+  private final DataPusher dataPusher;
+  private final int maxInFlight;
+
+  private final String appId;
+  private final int shuffleId;
+  private final int numMappers;
+  private final int numPartitions;
+  private final ShuffleClient client;
+  private int partitionIdIdx = 0;
+
+  public DataPushQueue(
+      CelebornConf conf,
+      DataPusher dataPusher,
+      ShuffleClient client,
+      String appId,
+      int shuffleId,
+      int mapId,
+      int attemptId,
+      int numMappers,
+      int numPartitions) {
+    this.appId = appId;
+    this.shuffleId = shuffleId;
+    this.numMappers = numMappers;
+    this.numPartitions = numPartitions;
+    this.client = client;
+    this.dataPusher = dataPusher;
+    final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
+    this.pushState = client.getOrRegisterPushState(mapKey);
+    this.maxInFlight = conf.pushMaxReqsInFlight();
+    final int capacity = conf.pushQueueCapacity();

Review Comment:
   We should decrease the default value since it is multiplied by number of partitions. I think we should have a global capacity and a local capacity.



##########
client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java:
##########
@@ -0,0 +1,124 @@
+/*
+ * 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.celeborn.client.write;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.LinkedBlockingQueue;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.client.ShuffleClient;
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.util.Utils;
+
+/*
+ * Queue for push data,
+ * it can take one special PushTask queue by partitionId,
+ * and it can take one random PushTask queue.
+ * workingQueuePerPartition: for PushTask queue per every partition.
+ *
+ * */
+public class DataPushQueue {
+  private static final Logger logger = LoggerFactory.getLogger(DataPushQueue.class);
+
+  private final List<LinkedBlockingQueue<PushTask>> workingQueuePerPartition;

Review Comment:
   In fact, I think we can still use one LinkedBlockingQueue. Each time we want to push data, we traverse the linked queue from the beginning until we find a task that can be pushed, then remove it from the queue and push it. PushState should expose a method to indicate whether the partitionId can be pushed, like:
   ```
   boolean canPush(partitionId)
   ```



##########
client/src/main/java/org/apache/celeborn/client/write/InFlightRequestTracker.java:
##########
@@ -0,0 +1,209 @@
+/*
+ * 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.celeborn.client.write;
+
+import java.io.IOException;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import io.netty.channel.ChannelFuture;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.network.client.RpcResponseCallback;
+import org.apache.celeborn.common.protocol.message.StatusCode;
+
+/*
+ * This class is for track in flight request and limit request.
+ * */
+public class InFlightRequestTracker {
+  private static final Logger logger = LoggerFactory.getLogger(InFlightRequestTracker.class);
+
+  private final long timeoutMs;
+  private final long pushTimeout;
+  private final long delta;
+  private final PushState pushState;
+
+  private final AtomicInteger batchId = new AtomicInteger();
+  private final ConcurrentHashMap<String, ConcurrentHashMap<Integer, BatchInfo>>
+      batchIdPerAddressPair = new ConcurrentHashMap<>();
+
+  public InFlightRequestTracker(CelebornConf conf, PushState pushState) {
+    this.timeoutMs = conf.pushLimitInFlightTimeoutMs();
+    this.pushTimeout = conf.pushDataTimeoutMs();
+    this.delta = conf.pushLimitInFlightSleepDeltaMs();
+    this.pushState = pushState;
+  }
+
+  public void addFlightBatches(int batchId, String hostAndPushPort) {
+    ConcurrentHashMap<Integer, BatchInfo> batchIdSetPerPair =
+        batchIdPerAddressPair.computeIfAbsent(hostAndPushPort, id -> new ConcurrentHashMap<>());
+    batchIdSetPerPair.computeIfAbsent(batchId, id -> new BatchInfo());
+  }
+
+  public void removeFlightBatches(int batchId, String hostAndPushPort) {

Review Comment:
   removeFlightBatch



##########
client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java:
##########
@@ -0,0 +1,124 @@
+/*
+ * 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.celeborn.client.write;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.LinkedBlockingQueue;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.client.ShuffleClient;
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.util.Utils;
+
+/*
+ * Queue for push data,
+ * it can take one special PushTask queue by partitionId,
+ * and it can take one random PushTask queue.
+ * workingQueuePerPartition: for PushTask queue per every partition.
+ *
+ * */
+public class DataPushQueue {
+  private static final Logger logger = LoggerFactory.getLogger(DataPushQueue.class);
+
+  private final List<LinkedBlockingQueue<PushTask>> workingQueuePerPartition;
+  private final PushState pushState;
+  private final DataPusher dataPusher;
+  private final int maxInFlight;
+
+  private final String appId;
+  private final int shuffleId;
+  private final int numMappers;
+  private final int numPartitions;
+  private final ShuffleClient client;
+  private int partitionIdIdx = 0;
+
+  public DataPushQueue(
+      CelebornConf conf,
+      DataPusher dataPusher,
+      ShuffleClient client,
+      String appId,
+      int shuffleId,
+      int mapId,
+      int attemptId,
+      int numMappers,
+      int numPartitions) {
+    this.appId = appId;
+    this.shuffleId = shuffleId;
+    this.numMappers = numMappers;
+    this.numPartitions = numPartitions;
+    this.client = client;
+    this.dataPusher = dataPusher;
+    final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
+    this.pushState = client.getOrRegisterPushState(mapKey);
+    this.maxInFlight = conf.pushMaxReqsInFlight();
+    final int capacity = conf.pushQueueCapacity();
+    workingQueuePerPartition = new ArrayList<>(numPartitions);
+    for (int i = 0; i < numPartitions; i++) {
+      workingQueuePerPartition.add(new LinkedBlockingQueue<>(capacity));
+    }
+  }
+
+  public LinkedBlockingQueue<PushTask> takeAnyWorkingQueue() throws IOException {
+    while (!dataPusher.terminatedOrHasException()) {
+      int partitionId = nextPartitionId();
+      LinkedBlockingQueue<PushTask> pushTasks;
+      try {
+        pushTasks = workingQueuePerPartition.get(partitionId);
+      } catch (IndexOutOfBoundsException ex) {
+        if (dataPusher.terminatedOrHasException()) {
+          return null;
+        }
+        throw ex;
+      }
+      if (Objects.isNull(pushTasks) && dataPusher.terminatedOrHasException()) {
+        return null;
+      }
+      if (!pushTasks.isEmpty()) {
+        Map<Integer, PartitionLocation> partitionLocationMap =
+            client.getOrRegisterShuffle(appId, shuffleId, numMappers, numPartitions);
+        PartitionLocation loc = partitionLocationMap.get(partitionId);
+        boolean reachLimit = pushState.limitMaxInFlight(loc.hostAndPushPort(), maxInFlight);

Review Comment:
   limitMaxInFlight will wait until the max in flight count is below the threshold, which is not our purpose. our purpose is fast-check and skip if the current partition reaches the threshold.



##########
client/src/main/java/org/apache/celeborn/client/write/DataPusher.java:
##########
@@ -158,7 +166,7 @@ public void waitOnTermination() throws IOException {
 
     terminated = true;
     idleQueue.clear();
-    workingQueue.clear();
+    dataPushQueue.clear();

Review Comment:
   IMO we should wait the DataPusher thread to join before clear the dataPushQueue to avoid potential concurrent issues. The code may look like this:
   ```
   Thread runner = new Thread("DataPusher-" + taskId) {...};
   runner.start();
   
   
    public void waitOnTermination() throws IOException {
        ....
        terminated = true;
        runner.join();
   
        idleQueue.clear();
        dataPushQueue.clear()
    }
   ```



##########
client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java:
##########
@@ -0,0 +1,124 @@
+/*
+ * 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.celeborn.client.write;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.LinkedBlockingQueue;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.client.ShuffleClient;
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.util.Utils;
+
+/*
+ * Queue for push data,
+ * it can take one special PushTask queue by partitionId,
+ * and it can take one random PushTask queue.
+ * workingQueuePerPartition: for PushTask queue per every partition.
+ *
+ * */
+public class DataPushQueue {
+  private static final Logger logger = LoggerFactory.getLogger(DataPushQueue.class);
+
+  private final List<LinkedBlockingQueue<PushTask>> workingQueuePerPartition;

Review Comment:
   I think we should maintain a queue per address instead of per partition, since there can be tens of thousands of partitions, which will be too much overhead.



##########
client/src/main/java/org/apache/celeborn/client/ShuffleClient.java:
##########
@@ -235,4 +237,9 @@ public abstract void pushDataHandShake(
 
   public abstract PartitionLocation registerMapPartitionTask(
       String appId, int shuffleId, int numMappers, int mapId, int attemptId);
+
+  public abstract ConcurrentHashMap<Integer, PartitionLocation> getOrRegisterShuffle(
+      String applicationId, int shuffleId, int numMappers, int numPartitions);
+
+  public abstract PushState getOrRegisterPushState(String mapKey);

Review Comment:
   better to rename to getPushState



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org