This is an automated email from the ASF dual-hosted git repository.
ethanfeng pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new bfb39632 [CELEBORN-235] Implement flink plugin. (#1244)
bfb39632 is described below
commit bfb39632d95ad0c6d160f70b0a3ee302d34af8f7
Author: Ethan Feng <[email protected]>
AuthorDate: Fri Feb 17 19:31:12 2023 +0800
[CELEBORN-235] Implement flink plugin. (#1244)
---
.../celeborn/plugin/flink/buffer/Buffer.java | 49 ++
.../plugin/flink/buffer/BufferRecycler.java | 26 +
.../plugin/flink/buffer/CreditListener.java | 47 ++
.../plugin/flink/buffer/TransferBufferPool.java | 218 +++++++
.../plugin/flink/network/ReadClientHandler.java | 8 +
.../flink/readclient/FlinkShuffleClientImpl.java | 121 ++++
.../plugin/flink/readclient/RssBufferStream.java | 159 +++++
.../plugin/flink/RemoteBufferStreamReader.java | 129 ++++
.../plugin/flink/RemoteShuffleDescriptor.java | 6 +-
.../plugin/flink/RemoteShuffleEnvironment.java | 19 +-
.../plugin/flink/RemoteShuffleInputGate.java | 705 +++++++++++++++++++++
.../flink/RemoteShuffleInputGateFactory.java | 123 ++++
.../plugin/flink/RemoteShuffleServiceFactory.java | 8 +-
.../apache/celeborn/client/ShuffleClientImpl.java | 125 ++--
14 files changed, 1680 insertions(+), 63 deletions(-)
diff --git
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/buffer/Buffer.java
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/buffer/Buffer.java
new file mode 100644
index 00000000..3a4844f3
--- /dev/null
+++
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/buffer/Buffer.java
@@ -0,0 +1,49 @@
+/*
+ * 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.plugin.flink.buffer;
+
+import java.nio.ByteBuffer;
+
+import org.apache.flink.shaded.netty4.io.netty.buffer.UnpooledByteBufAllocator;
+import org.apache.flink.shaded.netty4.io.netty.buffer.UnpooledDirectByteBuf;
+
+/**
+ * We use {@link UnpooledDirectByteBuf} directly to reduce one copy from
{@link Buffer} to Netty's
+ * {@link org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf} when
transmitting data over the
+ * network.
+ */
+public class Buffer extends UnpooledDirectByteBuf {
+
+ private final ByteBuffer buffer;
+
+ private final BufferRecycler recycler;
+
+ public Buffer(ByteBuffer buffer, BufferRecycler recycler, int readableBytes)
{
+ super(UnpooledByteBufAllocator.DEFAULT, buffer, buffer.capacity());
+
+ this.buffer = buffer;
+ this.recycler = recycler;
+ writerIndex(readableBytes);
+ }
+
+ @Override
+ protected void deallocate() {
+ buffer.clear();
+ recycler.recycle(buffer);
+ }
+}
diff --git
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/buffer/BufferRecycler.java
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/buffer/BufferRecycler.java
new file mode 100644
index 00000000..47a51f97
--- /dev/null
+++
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/buffer/BufferRecycler.java
@@ -0,0 +1,26 @@
+/*
+ * 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.plugin.flink.buffer;
+
+import java.nio.ByteBuffer;
+
+/** Recycler which can recycle {@link ByteBuffer}s for reuse. */
+public interface BufferRecycler {
+
+ /** Recycles the target {@link ByteBuffer}. */
+ void recycle(ByteBuffer buffer);
+}
diff --git
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/buffer/CreditListener.java
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/buffer/CreditListener.java
new file mode 100644
index 00000000..47b9b847
--- /dev/null
+++
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/buffer/CreditListener.java
@@ -0,0 +1,47 @@
+/*
+ * 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.plugin.flink.buffer;
+
+/** Listener to be notified when there is any available credit. */
+public abstract class CreditListener {
+
+ private int numCreditsNeeded;
+
+ private boolean isRegistered;
+
+ void increaseNumCreditsNeeded(int numCredits) {
+ numCreditsNeeded += numCredits;
+ }
+
+ void decreaseNumCreditsNeeded(int numCredits) {
+ numCreditsNeeded -= numCredits;
+ }
+
+ int getNumCreditsNeeded() {
+ return numCreditsNeeded;
+ }
+
+ boolean isRegistered() {
+ return isRegistered;
+ }
+
+ void setRegistered(boolean registered) {
+ isRegistered = registered;
+ }
+
+ public abstract void notifyAvailableCredits(int numCredits);
+}
diff --git
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/buffer/TransferBufferPool.java
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/buffer/TransferBufferPool.java
new file mode 100644
index 00000000..2e07fd16
--- /dev/null
+++
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/buffer/TransferBufferPool.java
@@ -0,0 +1,218 @@
+/*
+ * 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.plugin.flink.buffer;
+
+import static org.apache.celeborn.plugin.flink.utils.Utils.checkArgument;
+import static org.apache.celeborn.plugin.flink.utils.Utils.checkState;
+
+import java.nio.ByteBuffer;
+import java.util.*;
+
+import javax.annotation.concurrent.GuardedBy;
+
+import org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** A buffer pool which will dispatch buffers to all {@link CreditListener}s.
*/
+public class TransferBufferPool implements BufferRecycler {
+ private static Logger logger =
LoggerFactory.getLogger(TransferBufferPool.class);
+
+ private static final int MIN_CREDITS_TO_NOTIFY = 2;
+
+ private final Object lock = new Object();
+
+ private final Queue<ByteBuf> buffers = new ArrayDeque<>();
+
+ @GuardedBy("lock")
+ private final Queue<CreditListener> listeners = new ArrayDeque<>();
+
+ @GuardedBy("lock")
+ private int numAvailableBuffers;
+
+ @GuardedBy("lock")
+ private boolean isDestroyed;
+
+ public TransferBufferPool(Collection<ByteBuf> initialBuffers) {
+ synchronized (lock) {
+ buffers.addAll(initialBuffers);
+ numAvailableBuffers += initialBuffers.size();
+ }
+ }
+
+ /** Requests a data transmitting unit. */
+ public ByteBuf requestBuffer() {
+ synchronized (lock) {
+ checkState(!isDestroyed, "Buffer pool has been destroyed.");
+
+ return buffers.poll();
+ }
+ }
+
+ /** Adds an available buffer to this buffer pool. */
+ public void addBuffers(List<? extends ByteBuf> byteBufs) {
+ List<CreditAssignment> creditAssignments;
+ synchronized (lock) {
+ if (isDestroyed) {
+ byteBufs.forEach(ByteBuf::release);
+ return;
+ }
+
+ buffers.addAll(byteBufs);
+ numAvailableBuffers += byteBufs.size();
+ creditAssignments = dispatchReservedCredits();
+ }
+ for (CreditAssignment creditAssignment : creditAssignments) {
+
creditAssignment.getCreditListener().notifyAvailableCredits(creditAssignment.getNumCredits());
+ }
+ }
+
+ /** Tries to reserve buffers for the target {@link CreditListener}. */
+ public void reserveBuffers(CreditListener creditListener, int
numRequiredBuffers) {
+ int numCredits;
+ CreditListener listener = null;
+ synchronized (lock) {
+ if (isDestroyed) {
+ throw new IllegalStateException("Buffer pool has been destroyed.");
+ }
+
+ if (numRequiredBuffers > numAvailableBuffers) {
+ creditListener.increaseNumCreditsNeeded(numRequiredBuffers -
numAvailableBuffers);
+ }
+
+ if (!creditListener.isRegistered() &&
creditListener.getNumCreditsNeeded() > 0) {
+ listeners.add(creditListener);
+ creditListener.setRegistered(true);
+ }
+
+ numCredits = Math.min(numAvailableBuffers, numRequiredBuffers);
+ if (numCredits > 0) {
+ numAvailableBuffers -= numCredits;
+ listener = creditListener;
+ }
+
+ logger.warn("reserveBuffers,numCredits: {}, required: {}", numCredits,
numRequiredBuffers);
+ }
+ if (listener != null) {
+ listener.notifyAvailableCredits(numCredits);
+ }
+ }
+
+ /** Returns the number of available buffers. */
+ public int numBuffers() {
+ synchronized (lock) {
+ return buffers.size();
+ }
+ }
+
+ /** Destroys buffer pool. */
+ public void destroy() {
+ synchronized (lock) {
+ isDestroyed = true;
+ listeners.clear();
+ buffers.forEach(ByteBuf::release);
+ buffers.clear();
+ }
+ }
+
+ /** Returns true if this buffer pool has been destroyed. */
+ public boolean isDestroyed() {
+ synchronized (lock) {
+ return isDestroyed;
+ }
+ }
+
+ @Override
+ public void recycle(ByteBuffer buffer) {
+ List<CreditAssignment> creditAssignments;
+ synchronized (lock) {
+ // unmanaged memory no need to recycle, currently it is used only by
tests
+ if (isDestroyed) {
+ return;
+ }
+
+ buffers.add(new Buffer(buffer, this, 0));
+ ++numAvailableBuffers;
+ creditAssignments = dispatchReservedCredits();
+ }
+ for (CreditAssignment creditAssignment : creditAssignments) {
+
creditAssignment.getCreditListener().notifyAvailableCredits(creditAssignment.getNumCredits());
+ }
+ }
+
+ private int assignCredits(CreditListener creditListener) {
+ assert Thread.holdsLock(lock);
+
+ if (creditListener == null) {
+ return 0;
+ }
+
+ int numCredits = Math.min(creditListener.getNumCreditsNeeded(),
numAvailableBuffers);
+ if (numCredits > 0) {
+ creditListener.decreaseNumCreditsNeeded(numCredits);
+ numAvailableBuffers -= numCredits;
+ }
+
+ if (creditListener.getNumCreditsNeeded() > 0) {
+ listeners.add(creditListener);
+ } else {
+ creditListener.setRegistered(false);
+ }
+ return numCredits;
+ }
+
+ private List<CreditAssignment> dispatchReservedCredits() {
+ assert Thread.holdsLock(lock);
+
+ if (numAvailableBuffers < MIN_CREDITS_TO_NOTIFY || listeners.size() <= 0) {
+ return Collections.emptyList();
+ }
+
+ List<CreditAssignment> creditAssignments = new ArrayList<>();
+ while (numAvailableBuffers > 0 && listeners.size() > 0) {
+ CreditListener creditListener = listeners.poll();
+ int numCredits = assignCredits(creditListener);
+ if (numCredits > 0) {
+ creditAssignments.add(new CreditAssignment(numCredits,
creditListener));
+ }
+ }
+ return creditAssignments;
+ }
+
+ private static class CreditAssignment {
+
+ private final int numCredits;
+ private final CreditListener creditListener;
+
+ CreditAssignment(int numCredits, CreditListener creditListener) {
+ checkArgument(numCredits > 0, "Must be positive.");
+ checkArgument(creditListener != null, "Must be not null.");
+
+ this.numCredits = numCredits;
+ this.creditListener = creditListener;
+ }
+
+ public int getNumCredits() {
+ return numCredits;
+ }
+
+ public CreditListener getCreditListener() {
+ return creditListener;
+ }
+ }
+}
diff --git
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/ReadClientHandler.java
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/ReadClientHandler.java
index fc0dfacc..0da9f67d 100644
---
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/ReadClientHandler.java
+++
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/ReadClientHandler.java
@@ -101,4 +101,12 @@ public class ReadClientHandler extends BaseMessageHandler {
public void exceptionCaught(Throwable cause, TransportClient client) {
logger.warn("exception caught {}", client.getSocketAddress(), cause);
}
+
+ public void close() {
+ streamHandlers.clear();
+ for (TransportClient value : streamClients.values()) {
+ value.close();
+ }
+ streamClients.clear();
+ }
}
diff --git
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java
new file mode 100644
index 00000000..d11b0dcb
--- /dev/null
+++
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java
@@ -0,0 +1,121 @@
+/*
+ * 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.plugin.flink.readclient;
+
+import java.io.IOException;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.client.ShuffleClientImpl;
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.identity.UserIdentifier;
+import org.apache.celeborn.common.network.TransportContext;
+import org.apache.celeborn.common.network.util.TransportConf;
+import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.protocol.TransportModuleConstants;
+import org.apache.celeborn.common.util.Utils;
+import org.apache.celeborn.plugin.flink.network.FlinkTransportClientFactory;
+import org.apache.celeborn.plugin.flink.network.ReadClientHandler;
+
+public class FlinkShuffleClientImpl extends ShuffleClientImpl {
+ public static final Logger logger =
LoggerFactory.getLogger(FlinkShuffleClientImpl.class);
+ private static volatile FlinkShuffleClientImpl _instance;
+ private static volatile boolean initialized = false;
+ private FlinkTransportClientFactory flinkTransportClientFactory;
+ private ReadClientHandler readClientHandler = new ReadClientHandler();
+
+ public static FlinkShuffleClientImpl get(
+ String driverHost, int port, CelebornConf conf, UserIdentifier
userIdentifier) {
+ if (null == _instance || !initialized) {
+ synchronized (FlinkShuffleClientImpl.class) {
+ if (null == _instance) {
+ _instance = new FlinkShuffleClientImpl(driverHost, port, conf,
userIdentifier);
+ _instance.setupMetaServiceRef(driverHost, port);
+ initialized = true;
+ } else if (!initialized) {
+ _instance.shutdown();
+ _instance = new FlinkShuffleClientImpl(driverHost, port, conf,
userIdentifier);
+ _instance.setupMetaServiceRef(driverHost, port);
+ initialized = true;
+ }
+ }
+ }
+ return _instance;
+ }
+
+ @Override
+ public void shutdown() {
+ super.shutdown();
+ if (flinkTransportClientFactory != null) {
+ flinkTransportClientFactory.close();
+ }
+ if (readClientHandler != null) {
+ readClientHandler.close();
+ }
+ }
+
+ public FlinkShuffleClientImpl(
+ String driverHost, int port, CelebornConf conf, UserIdentifier
userIdentifier) {
+ super(conf, userIdentifier);
+ String module = TransportModuleConstants.DATA_MODULE;
+ TransportConf dataTransportConf =
+ Utils.fromCelebornConf(conf, module, conf.getInt("celeborn." + module
+ ".io.threads", 8));
+ TransportContext context =
+ new TransportContext(
+ dataTransportConf, readClientHandler,
conf.clientCloseIdleConnections());
+ this.flinkTransportClientFactory = new
FlinkTransportClientFactory(context);
+ this.setupMetaServiceRef(driverHost, port);
+ }
+
+ public RssBufferStream readBufferedPartition(
+ String applicationId,
+ int shuffleId,
+ int partitionId,
+ int subPartitionIndexStart,
+ int subPartitionIndexEnd)
+ throws IOException {
+ String shuffleKey = Utils.makeShuffleKey(applicationId, shuffleId);
+ ReduceFileGroups fileGroups = loadFileGroup(applicationId, shuffleKey,
shuffleId, partitionId);
+ if (fileGroups.partitionGroups.size() == 0
+ || !fileGroups.partitionGroups.containsKey(partitionId)) {
+ logger.warn("Shuffle data is empty for shuffle {} partitionId {}.",
shuffleId, partitionId);
+ return RssBufferStream.empty();
+ } else {
+ return RssBufferStream.create(
+ conf,
+ flinkTransportClientFactory,
+ shuffleKey,
+ fileGroups.partitionGroups.get(partitionId).toArray(new
PartitionLocation[0]),
+ subPartitionIndexStart,
+ subPartitionIndexEnd);
+ }
+ }
+
+ @Override
+ protected ReduceFileGroups updateFileGroup(
+ String applicationId, String shuffleKey, int shuffleId) {
+ ReduceFileGroups newGroup = loadFileGroupInternal(applicationId,
shuffleKey, shuffleId);
+ reduceFileGroupsMap.put(shuffleId, newGroup);
+ return newGroup;
+ }
+
+ public ReadClientHandler getReadClientHandler() {
+ return readClientHandler;
+ }
+}
diff --git
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/RssBufferStream.java
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/RssBufferStream.java
new file mode 100644
index 00000000..0be757c2
--- /dev/null
+++
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/RssBufferStream.java
@@ -0,0 +1,159 @@
+/*
+ * 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.plugin.flink.readclient;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
+
+import org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf;
+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.network.client.TransportClient;
+import org.apache.celeborn.common.network.protocol.*;
+import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.plugin.flink.network.FlinkTransportClientFactory;
+
+public class RssBufferStream {
+
+ private static Logger logger =
LoggerFactory.getLogger(RssBufferStream.class);
+ private CelebornConf conf;
+ private FlinkTransportClientFactory clientFactory;
+ private String shuffleKey;
+ private PartitionLocation[] locations;
+ private int subIndexStart;
+ private int subIndexEnd;
+ private TransportClient client;
+ private int currentLocationIndex = 0;
+ private long streamId = 0;
+
+ public RssBufferStream() {}
+
+ public RssBufferStream(
+ CelebornConf conf,
+ FlinkTransportClientFactory dataClientFactory,
+ String shuffleKey,
+ PartitionLocation[] locations,
+ int subIndexStart,
+ int subIndexEnd) {
+ this.conf = conf;
+ this.clientFactory = dataClientFactory;
+ this.shuffleKey = shuffleKey;
+ this.locations = locations;
+ this.subIndexStart = subIndexStart;
+ this.subIndexEnd = subIndexEnd;
+ }
+
+ public void open(
+ Supplier<ByteBuf> supplier,
+ int initialCredit,
+ FlinkShuffleClientImpl mapShuffleClient,
+ Consumer<RequestMessage> messageConsumer)
+ throws IOException, InterruptedException {
+ this.client =
+ clientFactory.createClient(
+ locations[currentLocationIndex].getHost(),
+ locations[currentLocationIndex].getFetchPort(),
+ -1,
+ supplier);
+ OpenStreamWithCredit openBufferStream =
+ new OpenStreamWithCredit(
+ shuffleKey,
+ locations[currentLocationIndex].getFileName(),
+ subIndexStart,
+ subIndexEnd,
+ initialCredit);
+ long timeoutMs = conf.fetchTimeoutMs();
+ CountDownLatch latch = new CountDownLatch(1);
+ client.sendRpc(
+ openBufferStream.toByteBuffer(),
+ new RpcResponseCallback() {
+
+ @Override
+ public void onSuccess(ByteBuffer response) {
+ StreamHandle streamHandle = (StreamHandle)
Message.decode(response);
+ RssBufferStream.this.streamId = streamHandle.streamId;
+ mapShuffleClient
+ .getReadClientHandler()
+ .registerHandler(streamId, messageConsumer, client);
+ latch.countDown();
+ }
+
+ @Override
+ public void onFailure(Throwable e) {
+ throw new RuntimeException("OpenStream failed.", e);
+ }
+ });
+ if (!latch.await(timeoutMs, TimeUnit.MILLISECONDS)) {
+ throw new IOException("Send openStream RPC failed to " +
client.getSocketAddress());
+ }
+ logger.debug("rssbufferstream streamid:{}", streamId);
+ }
+
+ public void addCredit(ReadAddCredit addCredit) {
+ this.client.sendRpc(
+ addCredit.toByteBuffer(),
+ new RpcResponseCallback() {
+ @Override
+ public void onSuccess(ByteBuffer response) {
+ // ignore response
+ }
+
+ @Override
+ public void onFailure(Throwable e) {
+ logger.warn(
+ "Send Add Credit {} failed to {}", addCredit,
client.getSocketAddress().toString());
+ }
+ });
+ }
+
+ public static RssBufferStream empty() {
+ return emptyRssBufferStream;
+ }
+
+ public long getStreamId() {
+ return streamId;
+ }
+
+ public static RssBufferStream create(
+ CelebornConf conf,
+ FlinkTransportClientFactory dataClientFactory,
+ String shuffleKey,
+ PartitionLocation[] locations,
+ int subIndexStart,
+ int subIndexEnd) {
+ if (locations == null || locations.length == 0) {
+ return empty();
+ } else {
+ return new RssBufferStream(
+ conf, dataClientFactory, shuffleKey, locations, subIndexStart,
subIndexEnd);
+ }
+ }
+
+ private static final RssBufferStream emptyRssBufferStream = new
RssBufferStream();
+
+ public TransportClient getClient() {
+ return client;
+ }
+}
diff --git
a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteBufferStreamReader.java
b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteBufferStreamReader.java
new file mode 100644
index 00000000..80f96942
--- /dev/null
+++
b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteBufferStreamReader.java
@@ -0,0 +1,129 @@
+/*
+ * 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.plugin.flink;
+
+import java.io.IOException;
+import java.util.function.Consumer;
+
+import org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.common.network.protocol.BacklogAnnouncement;
+import org.apache.celeborn.common.network.protocol.ReadAddCredit;
+import org.apache.celeborn.common.network.protocol.RequestMessage;
+import org.apache.celeborn.plugin.flink.buffer.CreditListener;
+import org.apache.celeborn.plugin.flink.buffer.TransferBufferPool;
+import org.apache.celeborn.plugin.flink.protocol.ReadData;
+import org.apache.celeborn.plugin.flink.readclient.FlinkShuffleClientImpl;
+import org.apache.celeborn.plugin.flink.readclient.RssBufferStream;
+
+public class RemoteBufferStreamReader extends CreditListener {
+ private static Logger logger =
LoggerFactory.getLogger(RemoteBufferStreamReader.class);
+ private TransferBufferPool bufferPool;
+ private FlinkShuffleClientImpl client;
+ private String applicationId;
+ private int shuffleId;
+ private int partitionId;
+ private int subPartitionIndexStart;
+ private int subPartitionIndexEnd;
+ private boolean isOpen;
+ private Consumer<ByteBuf> dataListener;
+ private Consumer<Throwable> failureListener;
+ private RssBufferStream bufferStream;
+ private boolean closed = false;
+ private Consumer<RequestMessage> messageConsumer;
+
+ public RemoteBufferStreamReader(
+ FlinkShuffleClientImpl client,
+ ShuffleResourceDescriptor shuffleDescriptor,
+ String applicationId,
+ int startSubIdx,
+ int endSubIdx,
+ TransferBufferPool bufferPool,
+ Consumer<ByteBuf> dataListener,
+ Consumer<Throwable> failureListener) {
+ this.client = client;
+ this.applicationId = applicationId;
+ this.shuffleId = shuffleDescriptor.getShuffleId();
+ this.partitionId = shuffleDescriptor.getPartitionId();
+ this.bufferPool = bufferPool;
+ this.subPartitionIndexStart = startSubIdx;
+ this.subPartitionIndexEnd = endSubIdx;
+ this.dataListener = dataListener;
+ this.failureListener = failureListener;
+ this.messageConsumer =
+ requestMessage -> {
+ if (requestMessage instanceof ReadData) {
+ dataReceived((ReadData) requestMessage);
+ } else if (requestMessage instanceof BacklogAnnouncement) {
+ backlogReceived(((BacklogAnnouncement)
requestMessage).getBacklog());
+ }
+ };
+ }
+
+ public void open(int initialCredit) throws IOException {
+ try {
+ this.bufferStream =
+ client.readBufferedPartition(
+ applicationId, shuffleId, partitionId, subPartitionIndexStart,
subPartitionIndexEnd);
+ bufferStream.open(
+ RemoteBufferStreamReader.this::requestBuffer, initialCredit, client,
messageConsumer);
+ } catch (InterruptedException e) {
+ throw new RuntimeException("Failed to openStream.", e);
+ }
+ isOpen = true;
+ }
+
+ public void close() {
+ isOpen = false;
+
client.getReadClientHandler().removeHandler(this.bufferStream.getStreamId());
+ }
+
+ public boolean isOpened() {
+ return isOpen;
+ }
+
+ public void notifyAvailableCredits(int numCredits) {
+ ReadAddCredit addCredit = new
ReadAddCredit(this.bufferStream.getStreamId(), numCredits);
+ bufferStream.addCredit(addCredit);
+ }
+
+ public ByteBuf requestBuffer() {
+ return bufferPool.requestBuffer();
+ }
+
+ public void backlogReceived(int backlog) {
+ bufferPool.reserveBuffers(this, backlog);
+ }
+
+ public void dataReceived(ReadData readData) {
+ logger.debug(
+ "Rss buffer stream reader get streamid {} received readable bytes {}.",
+ readData.getStreamId(),
+ readData.getFlinkBuffer().readableBytes());
+ if (closed) {
+ readData.body().release();
+ return;
+ }
+ int backLogInData = readData.getBacklog();
+ if (backLogInData > 0) {
+ backlogReceived(backLogInData);
+ }
+ dataListener.accept(readData.getFlinkBuffer());
+ }
+}
diff --git
a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleDescriptor.java
b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleDescriptor.java
index 12db3495..30e050f2 100644
---
a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleDescriptor.java
+++
b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleDescriptor.java
@@ -30,10 +30,10 @@ public class RemoteShuffleDescriptor implements
ShuffleDescriptor {
private final JobID jobID;
- private final ShuffleResource shuffleResource;
+ private final RemoteShuffleResource shuffleResource;
public RemoteShuffleDescriptor(
- JobID jobID, ResultPartitionID resultPartitionID, ShuffleResource
shuffleResource) {
+ JobID jobID, ResultPartitionID resultPartitionID, RemoteShuffleResource
shuffleResource) {
this.jobID = jobID;
this.resultPartitionID = resultPartitionID;
this.shuffleResource = shuffleResource;
@@ -48,7 +48,7 @@ public class RemoteShuffleDescriptor implements
ShuffleDescriptor {
return jobID;
}
- public ShuffleResource getShuffleResource() {
+ public RemoteShuffleResource getShuffleResource() {
return shuffleResource;
}
diff --git
a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleEnvironment.java
b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleEnvironment.java
index 0599f433..7ad50963 100644
---
a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleEnvironment.java
+++
b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleEnvironment.java
@@ -77,6 +77,8 @@ public class RemoteShuffleEnvironment
private final CelebornConf conf;
+ private final RemoteShuffleInputGateFactory inputGateFactory;
+
/**
* @param networkBufferPool Network buffer pool for shuffle read and shuffle
write.
* @param resultPartitionManager A trivial {@link ResultPartitionManager}.
@@ -87,13 +89,13 @@ public class RemoteShuffleEnvironment
NetworkBufferPool networkBufferPool,
ResultPartitionManager resultPartitionManager,
RemoteShuffleResultPartitionFactory resultPartitionFactory,
- // RemoteShuffleInputGateFactory inputGateFactory,
+ RemoteShuffleInputGateFactory inputGateFactory,
CelebornConf conf) {
this.networkBufferPool = networkBufferPool;
this.resultPartitionManager = resultPartitionManager;
this.resultPartitionFactory = resultPartitionFactory;
- // this.inputGateFactory = inputGateFactory;
+ this.inputGateFactory = inputGateFactory;
this.conf = conf;
this.isClosed = false;
}
@@ -123,7 +125,18 @@ public class RemoteShuffleEnvironment
ShuffleIOOwnerContext ownerContext,
PartitionProducerStateProvider producerStateProvider,
List<InputGateDeploymentDescriptor> inputGateDescriptors) {
- return null;
+ synchronized (lock) {
+ checkState(!isClosed, "The RemoteShuffleEnvironment has already been
shut down.");
+
+ IndexedInputGate[] inputGates = new
IndexedInputGate[inputGateDescriptors.size()];
+ for (int gateIndex = 0; gateIndex < inputGates.length; gateIndex++) {
+ InputGateDeploymentDescriptor igdd =
inputGateDescriptors.get(gateIndex);
+ RemoteShuffleInputGate inputGate =
+ inputGateFactory.create(ownerContext.getOwnerName(), gateIndex,
igdd);
+ inputGates[gateIndex] = inputGate;
+ }
+ return Arrays.asList(inputGates);
+ }
}
@Override
diff --git
a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java
b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java
new file mode 100644
index 00000000..07472f40
--- /dev/null
+++
b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java
@@ -0,0 +1,705 @@
+/*
+ * 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.plugin.flink;
+
+import static org.apache.celeborn.plugin.flink.utils.Utils.checkState;
+
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.util.*;
+import java.util.concurrent.CompletableFuture;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import org.apache.commons.lang3.tuple.Pair;
+import org.apache.flink.core.memory.MemorySegment;
+import org.apache.flink.core.memory.MemorySegmentProvider;
+import org.apache.flink.metrics.SimpleCounter;
+import org.apache.flink.runtime.checkpoint.CheckpointOptions;
+import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter;
+import org.apache.flink.runtime.checkpoint.channel.InputChannelInfo;
+import org.apache.flink.runtime.checkpoint.channel.ResultSubpartitionInfo;
+import org.apache.flink.runtime.deployment.InputGateDeploymentDescriptor;
+import org.apache.flink.runtime.event.AbstractEvent;
+import org.apache.flink.runtime.event.TaskEvent;
+import org.apache.flink.runtime.io.network.ConnectionID;
+import org.apache.flink.runtime.io.network.LocalConnectionManager;
+import org.apache.flink.runtime.io.network.api.CheckpointBarrier;
+import org.apache.flink.runtime.io.network.api.EndOfData;
+import org.apache.flink.runtime.io.network.api.EndOfPartitionEvent;
+import org.apache.flink.runtime.io.network.api.serialization.EventSerializer;
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferDecompressor;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import
org.apache.flink.runtime.io.network.partition.PartitionNotFoundException;
+import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
+import org.apache.flink.runtime.io.network.partition.ResultPartitionType;
+import org.apache.flink.runtime.io.network.partition.consumer.*;
+import org.apache.flink.runtime.jobgraph.IntermediateDataSetID;
+import org.apache.flink.runtime.shuffle.ShuffleDescriptor;
+import org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf;
+import org.apache.flink.util.CloseableIterator;
+import org.apache.flink.util.ExceptionUtils;
+import org.apache.flink.util.FlinkRuntimeException;
+import org.apache.flink.util.function.SupplierWithException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.identity.UserIdentifier;
+import org.apache.celeborn.plugin.flink.buffer.BufferPacker;
+import org.apache.celeborn.plugin.flink.buffer.TransferBufferPool;
+import org.apache.celeborn.plugin.flink.readclient.FlinkShuffleClientImpl;
+import org.apache.celeborn.plugin.flink.utils.BufferUtils;
+
+/** A {@link IndexedInputGate} which ingest data from remote shuffle workers.
*/
+public class RemoteShuffleInputGate extends IndexedInputGate {
+
+ private static final Logger LOG =
LoggerFactory.getLogger(RemoteShuffleInputGate.class);
+
+ /** Lock to protect {@link #receivedBuffers} and {@link #cause} and {@link
#closed}. */
+ private final Object lock = new Object();
+
+ /** Name of the corresponding computing task. */
+ private final String taskName;
+
+ /** Index of the gate of the corresponding computing task. */
+ private final int gateIndex;
+
+ /** Deployment descriptor for a single input gate instance. */
+ private final InputGateDeploymentDescriptor gateDescriptor;
+
+ /** Buffer pool provider. */
+ private final SupplierWithException<BufferPool, IOException>
bufferPoolFactory;
+
+ /** Flink buffer pools to allocate network memory. */
+ private BufferPool bufferPool;
+
+ /** Buffer pool used by the transfer layer. */
+ private final TransferBufferPool transferBufferPool =
+ new TransferBufferPool(Collections.emptySet());
+
+ private final List<RemoteBufferStreamReader> bufferReaders = new
ArrayList<>();
+ private final List<InputChannelInfo> channelsInfo;
+ /** Map from channel index to shuffle client index. */
+ private final int[] clientIndexMap;
+
+ /** Map from shuffle client index to channel index. */
+ private final int[] channelIndexMap;
+
+ /** The number of subpartitions that has not consumed per channel. */
+ private final int[] numSubPartitionsHasNotConsumed;
+
+ /** The overall number of subpartitions that has not been consumed. */
+ private long numUnconsumedSubpartitions;
+
+ /** Received buffers from remote shuffle worker. It's consumed by upper
computing task. */
+ private final Queue<Pair<Buffer, InputChannelInfo>> receivedBuffers = new
LinkedList<>();
+
+ /** {@link Throwable} when reading failure. */
+ private Throwable cause;
+
+ /** Whether this remote input gate has been closed or not. */
+ private boolean closed;
+
+ /** Whether we have opened all initial channels or not. */
+ private boolean initialChannelsOpened;
+
+ /** Number of pending {@link EndOfData} events to be received. */
+ private long pendingEndOfDataEvents;
+ /** Max concurrent reader count */
+ private int numConcurrentReading = Integer.MAX_VALUE;
+ /** Keep compatibility with streaming mode. */
+ private boolean shouldDrainOnEndOfData = true;
+
+ /** Data decompressor. */
+ private final BufferDecompressor bufferDecompressor;
+
+ private FlinkShuffleClientImpl shuffleClient;
+
+ public RemoteShuffleInputGate(
+ CelebornConf celebornConf,
+ String taskName,
+ int gateIndex,
+ InputGateDeploymentDescriptor gateDescriptor,
+ SupplierWithException<BufferPool, IOException> bufferPoolFactory,
+ BufferDecompressor bufferDecompressor) {
+
+ this.taskName = taskName;
+ this.gateIndex = gateIndex;
+ this.gateDescriptor = gateDescriptor;
+ this.bufferPoolFactory = bufferPoolFactory;
+
+ int numChannels = gateDescriptor.getShuffleDescriptors().length;
+ this.clientIndexMap = new int[numChannels];
+ this.channelIndexMap = new int[numChannels];
+ this.numSubPartitionsHasNotConsumed = new int[numChannels];
+ this.bufferDecompressor = bufferDecompressor;
+
+ RemoteShuffleDescriptor remoteShuffleDescriptor =
+ (RemoteShuffleDescriptor) gateDescriptor.getShuffleDescriptors()[0];
+ this.shuffleClient =
+ FlinkShuffleClientImpl.get(
+
remoteShuffleDescriptor.getShuffleResource().getRssMetaServiceHost(),
+
remoteShuffleDescriptor.getShuffleResource().getRssMetaServicePort(),
+ celebornConf,
+ new UserIdentifier("default", "default"));
+
+ this.numUnconsumedSubpartitions = initShuffleReadClients();
+ this.pendingEndOfDataEvents = numUnconsumedSubpartitions;
+ this.channelsInfo = createChannelInfos();
+ }
+
+ private long initShuffleReadClients() {
+ int startSubIdx = gateDescriptor.getConsumedSubpartitionIndex();
+ int endSubIdx = gateDescriptor.getConsumedSubpartitionIndex();
+ int numSubpartitionsPerChannel = endSubIdx - startSubIdx + 1;
+ long numUnconsumedSubpartitions = 0;
+
+ // left element is index
+ List<Pair<Integer, ShuffleDescriptor>> descriptors =
+ IntStream.range(0, gateDescriptor.getShuffleDescriptors().length)
+ .mapToObj(i -> Pair.of(i,
gateDescriptor.getShuffleDescriptors()[i]))
+ .collect(Collectors.toList());
+
+ int clientIndex = 0;
+ for (Pair<Integer, ShuffleDescriptor> descriptor : descriptors) {
+ RemoteShuffleDescriptor remoteDescriptor = (RemoteShuffleDescriptor)
descriptor.getRight();
+ ShuffleResourceDescriptor shuffleDescriptor =
+
remoteDescriptor.getShuffleResource().getMapPartitionShuffleDescriptor();
+
+ LOG.debug("create shuffle reader for descriptor {}", shuffleDescriptor);
+ String applicationId = remoteDescriptor.getJobID().toString();
+
+ RemoteBufferStreamReader reader =
+ new RemoteBufferStreamReader(
+ shuffleClient,
+ shuffleDescriptor,
+ applicationId,
+ startSubIdx,
+ endSubIdx,
+ transferBufferPool,
+ getDataListener(descriptor.getLeft()),
+ getFailureListener(remoteDescriptor.getResultPartitionID()));
+
+ bufferReaders.add(reader);
+ numSubPartitionsHasNotConsumed[descriptor.getLeft()] =
numSubpartitionsPerChannel;
+ numUnconsumedSubpartitions += numSubpartitionsPerChannel;
+ clientIndexMap[descriptor.getLeft()] = clientIndex;
+ channelIndexMap[clientIndex] = descriptor.getLeft();
+ ++clientIndex;
+ }
+ return numUnconsumedSubpartitions;
+ }
+
+ /** Setup gate and build network connections. */
+ @Override
+ public void setup() throws IOException {
+ long startTime = System.nanoTime();
+
+ bufferPool = bufferPoolFactory.get();
+ BufferUtils.reserveNumRequiredBuffers(bufferPool, 16);
+
+ tryRequestBuffers();
+ // Complete availability future though handshake not fired yet, thus to
allow fetcher to
+ // 'pollNext' and fire handshake to remote. This mechanism is to avoid
bookkeeping remote
+ // reading resource before task start processing data from input gate.
+ availabilityHelper.getUnavailableToResetAvailable().complete(null);
+ LOG.info("Set up read gate by {} ms.", (System.nanoTime() - startTime) /
1000_000);
+ }
+
+ /** Index of the gate of the corresponding computing task. */
+ @Override
+ public int getGateIndex() {
+ return gateIndex;
+ }
+
+ /** Get number of input channels. A channel is a data flow from one shuffle
worker. */
+ @Override
+ public int getNumberOfInputChannels() {
+ return bufferReaders.size();
+ }
+
+ /** Whether reading is finished -- all channels are finished and cached
buffers are drained. */
+ @Override
+ public boolean isFinished() {
+ synchronized (lock) {
+ return allReadersEOF() && receivedBuffers.isEmpty();
+ }
+ }
+
+ @Override
+ public Optional<BufferOrEvent> getNext() {
+ throw new UnsupportedOperationException("Not implemented (DataSet API is
not supported).");
+ }
+
+ /** Poll a received {@link BufferOrEvent}. */
+ @Override
+ public Optional<BufferOrEvent> pollNext() throws IOException {
+ if (!initialChannelsOpened) {
+ tryOpenSomeChannels();
+ initialChannelsOpened = true;
+ // DO NOT return, method of 'getReceived' will manipulate
'availabilityHelper'.
+ }
+
+ Pair<Buffer, InputChannelInfo> pair = getReceived();
+ Optional<BufferOrEvent> bufferOrEvent = Optional.empty();
+ LOG.debug("pollNext called with pair null {}", pair == null);
+ while (pair != null) {
+ Buffer buffer = pair.getLeft();
+ InputChannelInfo channelInfo = pair.getRight();
+ LOG.debug("get buffer {} on channel {}", buffer, channelInfo);
+ if (buffer.isBuffer()) {
+ bufferOrEvent = transformBuffer(buffer, channelInfo);
+ } else {
+ bufferOrEvent = transformEvent(buffer, channelInfo);
+ LOG.info("recevied event: " +
bufferOrEvent.get().getEvent().getClass().getName());
+ }
+
+ if (bufferOrEvent.isPresent()) {
+ break;
+ }
+ pair = getReceived();
+ }
+
+ tryRequestBuffers();
+ return bufferOrEvent;
+ }
+
+ private Buffer decompressBufferIfNeeded(Buffer buffer) throws IOException {
+ if (buffer.isCompressed()) {
+ try {
+ checkState(bufferDecompressor != null, "Buffer decompressor not set.");
+ return bufferDecompressor.decompressToIntermediateBuffer(buffer);
+ } catch (Throwable t) {
+ throw new IOException("Decompress failure", t);
+ } finally {
+ buffer.recycleBuffer();
+ }
+ }
+ return buffer;
+ }
+
+ /** Close all reading channels inside this {@link RemoteShuffleInputGate}. */
+ @Override
+ public void close() throws Exception {
+ List<Buffer> buffersToRecycle;
+ Throwable closeException = null;
+ // Do not check closed flag, thus to allow calling this method from both
task thread and
+ // cancel thread.
+ for (RemoteBufferStreamReader shuffleReadClient : bufferReaders) {
+ try {
+ shuffleReadClient.close();
+ } catch (Throwable throwable) {
+ closeException = closeException == null ? throwable : closeException;
+ LOG.error("Failed to close shuffle read client.", throwable);
+ }
+ }
+ synchronized (lock) {
+ buffersToRecycle =
receivedBuffers.stream().map(Pair::getLeft).collect(Collectors.toList());
+ receivedBuffers.clear();
+ closed = true;
+ }
+
+ try {
+ buffersToRecycle.forEach(Buffer::recycleBuffer);
+ } catch (Throwable throwable) {
+ closeException = closeException == null ? throwable : closeException;
+ LOG.error("Failed to recycle buffers.", throwable);
+ }
+
+ try {
+ transferBufferPool.destroy();
+ } catch (Throwable throwable) {
+ closeException = closeException == null ? throwable : closeException;
+ LOG.error("Failed to close transfer buffer pool.", throwable);
+ }
+
+ try {
+ if (bufferPool != null) {
+ bufferPool.lazyDestroy();
+ }
+ } catch (Throwable throwable) {
+ closeException = closeException == null ? throwable : closeException;
+ LOG.error("Failed to close local buffer pool.", throwable);
+ }
+
+ if (closeException != null) {
+ ExceptionUtils.rethrowException(closeException);
+ }
+ }
+
+ /** Get {@link InputChannelInfo}s of this {@link RemoteShuffleInputGate}. */
+ @Override
+ public List<InputChannelInfo> getChannelInfos() {
+ return channelsInfo;
+ }
+
+ /** Each one corresponds to a reading channel. */
+ public List<RemoteBufferStreamReader> getBufferReaders() {
+ return bufferReaders;
+ }
+
+ private List<InputChannelInfo> createChannelInfos() {
+ return IntStream.range(0, gateDescriptor.getShuffleDescriptors().length)
+ .mapToObj(i -> new InputChannelInfo(gateIndex, i))
+ .collect(Collectors.toList());
+ }
+
+ /** Try to open more readers to {@link #numConcurrentReading}. */
+ private void tryOpenSomeChannels() throws IOException {
+ List<RemoteBufferStreamReader> clientsToOpen = new ArrayList<>();
+
+ synchronized (lock) {
+ if (closed) {
+ throw new IOException("Input gate already closed.");
+ }
+
+ LOG.debug("Try open some partition readers.");
+ int numOnGoing = 0;
+ for (int i = 0; i < bufferReaders.size(); i++) {
+ RemoteBufferStreamReader bufferStreamReader = bufferReaders.get(i);
+ LOG.debug(
+ "Trying reader: {}, isOpened={},
numSubPartitionsHasNotConsumed={}.",
+ bufferStreamReader,
+ bufferStreamReader.isOpened(),
+ numSubPartitionsHasNotConsumed[channelIndexMap[i]]);
+ if (numOnGoing >= numConcurrentReading) {
+ break;
+ }
+
+ if (bufferStreamReader.isOpened()
+ && numSubPartitionsHasNotConsumed[channelIndexMap[i]] > 0) {
+ numOnGoing++;
+ continue;
+ }
+
+ if (!bufferStreamReader.isOpened()) {
+ clientsToOpen.add(bufferStreamReader);
+ numOnGoing++;
+ }
+ }
+ }
+
+ for (RemoteBufferStreamReader reader : clientsToOpen) {
+ reader.open(0);
+ }
+ }
+
+ private void tryRequestBuffers() {
+ checkState(bufferPool != null, "Not initialized yet.");
+
+ Buffer buffer;
+ List<ByteBuf> buffers = new ArrayList<>();
+ while ((buffer = bufferPool.requestBuffer()) != null) {
+ buffers.add(buffer.asByteBuf());
+ }
+
+ if (!buffers.isEmpty()) {
+ transferBufferPool.addBuffers(buffers);
+ }
+ }
+
+ private void onBuffer(Buffer buffer, int channelIdx) {
+ synchronized (lock) {
+ if (closed || cause != null) {
+ buffer.recycleBuffer();
+ throw new IllegalStateException("Input gate already closed or
failed.");
+ }
+
+ boolean needRecycle = true;
+ try {
+ boolean wasEmpty = receivedBuffers.isEmpty();
+ InputChannelInfo channelInfo = channelsInfo.get(channelIdx);
+ checkState(channelInfo.getInputChannelIdx() == channelIdx, "Illegal
channel index.");
+ LOG.debug("ReceivedBuffers is adding buffer {} on {}", buffer,
channelInfo);
+ receivedBuffers.add(Pair.of(buffer, channelInfo));
+ needRecycle = false;
+ if (wasEmpty) {
+ availabilityHelper.getUnavailableToResetAvailable().complete(null);
+ }
+ } catch (Throwable throwable) {
+ if (needRecycle) {
+ buffer.recycleBuffer();
+ }
+ throw throwable;
+ }
+ }
+ }
+
+ private Consumer<ByteBuf> getDataListener(int channelIdx) {
+ return byteBuf -> {
+ Queue<Buffer> unpackedBuffers = null;
+ try {
+ unpackedBuffers = BufferPacker.unpack(byteBuf);
+ while (!unpackedBuffers.isEmpty()) {
+ onBuffer(unpackedBuffers.poll(), channelIdx);
+ }
+ } catch (Throwable throwable) {
+ synchronized (lock) {
+ cause = cause == null ? throwable : cause;
+ availabilityHelper.getUnavailableToResetAvailable().complete(null);
+ }
+
+ if (unpackedBuffers != null) {
+ unpackedBuffers.forEach(Buffer::recycleBuffer);
+ }
+ LOG.error("Failed to process the received buffer.", throwable);
+ }
+ };
+ }
+
+ private Consumer<Throwable> getFailureListener(ResultPartitionID rpID) {
+ return throwable -> {
+ synchronized (lock) {
+ if (cause != null) {
+ return;
+ }
+ Class<?> clazz = PartitionNotFoundException.class;
+ if (throwable.getMessage() != null &&
throwable.getMessage().contains(clazz.getName())) {
+ cause = new PartitionNotFoundException(rpID);
+ } else {
+ cause = throwable;
+ }
+ availabilityHelper.getUnavailableToResetAvailable().complete(null);
+ }
+ };
+ }
+
+ private Pair<Buffer, InputChannelInfo> getReceived() throws IOException {
+ synchronized (lock) {
+ healthCheck();
+ if (!receivedBuffers.isEmpty()) {
+ return receivedBuffers.poll();
+ } else {
+ if (!allReadersEOF()) {
+ availabilityHelper.resetUnavailable();
+ }
+ return null;
+ }
+ }
+ }
+
+ private void healthCheck() throws IOException {
+ if (closed) {
+ throw new IOException("Input gate already closed.");
+ }
+ if (cause != null) {
+ if (cause instanceof IOException) {
+ throw (IOException) cause;
+ } else {
+ throw new IOException(cause);
+ }
+ }
+ }
+
+ private boolean allReadersEOF() {
+ return numUnconsumedSubpartitions <= 0;
+ }
+
+ private Optional<BufferOrEvent> transformBuffer(Buffer buf, InputChannelInfo
info)
+ throws IOException {
+ return Optional.of(
+ new BufferOrEvent(decompressBufferIfNeeded(buf), info, !isFinished(),
false));
+ }
+
+ private Optional<BufferOrEvent> transformEvent(Buffer buffer,
InputChannelInfo channelInfo)
+ throws IOException {
+ final AbstractEvent event;
+ try {
+ event = EventSerializer.fromBuffer(buffer, getClass().getClassLoader());
+ } catch (Throwable t) {
+ throw new IOException("Deserialize failure.", t);
+ } finally {
+ buffer.recycleBuffer();
+ }
+ if (event.getClass() == EndOfPartitionEvent.class) {
+ checkState(
+ numSubPartitionsHasNotConsumed[channelInfo.getInputChannelIdx()] > 0,
+ "BUG -- EndOfPartitionEvent received repeatedly.");
+ numSubPartitionsHasNotConsumed[channelInfo.getInputChannelIdx()]--;
+ numUnconsumedSubpartitions--;
+ // not the real end.
+ if (numSubPartitionsHasNotConsumed[channelInfo.getInputChannelIdx()] !=
0) {
+ return Optional.empty();
+ } else {
+ // the real end.
+
bufferReaders.get(clientIndexMap[channelInfo.getInputChannelIdx()]).close();
+ // tryOpenSomeChannels();
+ if (allReadersEOF()) {
+ availabilityHelper.getUnavailableToResetAvailable().complete(null);
+ }
+ }
+ } else if (event.getClass() == EndOfData.class) {
+ checkState(!hasReceivedEndOfData(), "ERROR ");
+ --pendingEndOfDataEvents;
+ }
+
+ return Optional.of(
+ new BufferOrEvent(
+ event,
+ buffer.getDataType().hasPriority(),
+ channelInfo,
+ !isFinished(),
+ buffer.getSize(),
+ false));
+ }
+
+ @Override
+ public void requestPartitions() {
+ // do-nothing
+ }
+
+ @Override
+ public void checkpointStarted(CheckpointBarrier barrier) {
+ // do-nothing.
+ }
+
+ @Override
+ public void checkpointStopped(long cancelledCheckpointId) {
+ // do-nothing.
+ }
+
+ @Override
+ public int getBuffersInUseCount() {
+ return 0;
+ }
+
+ @Override
+ public void announceBufferSize(int i) {}
+
+ @Override
+ public List<InputChannelInfo> getUnfinishedChannels() {
+ return Collections.emptyList();
+ }
+
+ @Override
+ public boolean hasReceivedEndOfData() {
+ return pendingEndOfDataEvents <= 0;
+ }
+
+ @Override
+ public void finishReadRecoveredState() {
+ // do-nothing.
+ }
+
+ @Override
+ public InputChannel getChannel(int channelIndex) {
+ return new FakedRemoteInputChannel(channelIndex);
+ }
+
+ @Override
+ public void sendTaskEvent(TaskEvent event) {
+ throw new FlinkRuntimeException("Method should not be called.");
+ }
+
+ @Override
+ public void resumeConsumption(InputChannelInfo channelInfo) {
+ throw new FlinkRuntimeException("Method should not be called.");
+ }
+
+ @Override
+ public void acknowledgeAllRecordsProcessed(InputChannelInfo
inputChannelInfo) {}
+
+ @Override
+ public CompletableFuture<Void> getStateConsumedFuture() {
+ return CompletableFuture.completedFuture(null);
+ }
+
+ @Override
+ public String toString() {
+ return String.format(
+ "ReadGate [owning task: %s, gate index: %d, descriptor: %s]",
+ taskName, gateIndex, gateDescriptor.toString());
+ }
+
+ /** Accommodation for the incompleteness of Flink pluggable shuffle service.
*/
+ private class FakedRemoteInputChannel extends RemoteInputChannel {
+ FakedRemoteInputChannel(int channelIndex) {
+ super(
+ new SingleInputGate(
+ "",
+ gateIndex,
+ new IntermediateDataSetID(),
+ ResultPartitionType.BLOCKING,
+ 0,
+ 1,
+ (a, b, c) -> {},
+ () -> null,
+ null,
+ new FakedMemorySegmentProvider(),
+ 0),
+ channelIndex,
+ new ResultPartitionID(),
+ new ConnectionID(new InetSocketAddress("", 0), 0),
+ new LocalConnectionManager(),
+ 0,
+ 0,
+ 0,
+ new SimpleCounter(),
+ new SimpleCounter(),
+ new FakedChannelStateWriter());
+ }
+ }
+
+ /** Accommodation for the incompleteness of Flink pluggable shuffle service.
*/
+ private static class FakedMemorySegmentProvider implements
MemorySegmentProvider {
+
+ @Override
+ public Collection<MemorySegment> requestUnpooledMemorySegments(int i)
throws IOException {
+ return null;
+ }
+
+ @Override
+ public void recycleUnpooledMemorySegments(Collection<MemorySegment>
collection)
+ throws IOException {}
+ }
+
+ /** Accommodation for the incompleteness of Flink pluggable shuffle service.
*/
+ private static class FakedChannelStateWriter implements ChannelStateWriter {
+
+ @Override
+ public void start(long cpId, CheckpointOptions checkpointOptions) {}
+
+ @Override
+ public void addInputData(
+ long cpId, InputChannelInfo info, int startSeqNum,
CloseableIterator<Buffer> data) {}
+
+ @Override
+ public void addOutputData(
+ long cpId, ResultSubpartitionInfo info, int startSeqNum, Buffer...
data) {}
+
+ @Override
+ public void finishInput(long checkpointId) {}
+
+ @Override
+ public void finishOutput(long checkpointId) {}
+
+ @Override
+ public void abort(long checkpointId, Throwable cause, boolean cleanup) {}
+
+ @Override
+ public ChannelStateWriteResult getAndRemoveWriteResult(long checkpointId) {
+ return null;
+ }
+
+ @Override
+ public void close() {}
+ }
+}
diff --git
a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java
b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java
new file mode 100644
index 00000000..76f53549
--- /dev/null
+++
b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java
@@ -0,0 +1,123 @@
+/*
+ * 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.plugin.flink;
+
+import java.io.IOException;
+
+import org.apache.flink.runtime.deployment.InputGateDeploymentDescriptor;
+import org.apache.flink.runtime.io.network.buffer.BufferDecompressor;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.BufferPoolFactory;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.util.function.SupplierWithException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.plugin.flink.utils.Utils;
+
+/** Factory class to create {@link RemoteShuffleInputGate}. */
+public class RemoteShuffleInputGateFactory {
+
+ public static final int MIN_BUFFERS_PER_GATE = 16;
+
+ private static final Logger LOG =
LoggerFactory.getLogger(RemoteShuffleInputGateFactory.class);
+
+ /** Number of max concurrent reading channels. */
+ private final int numConcurrentReading;
+
+ /** Codec used for compression / decompression. */
+ private static final String compressionCodec = "LZ4";
+
+ /** Network buffer size. */
+ private final int networkBufferSize;
+
+ /**
+ * Network buffer pool used for shuffle read buffers. {@link BufferPool}s
will be created from it
+ * and each of them will be used by a channel exclusively.
+ */
+ private final NetworkBufferPool networkBufferPool;
+
+ /** Sum of buffers. */
+ private final int numBuffersPerGate;
+
+ public static final String MEMORY_PER_INPUT_GATE = "32m";
+ public static final String MIN_MEMORY_PER_GATE = "8m";
+ public static final int NUM_CONCURRENT_READINGS = Integer.MAX_VALUE;
+ private CelebornConf celebornConf;
+
+ public RemoteShuffleInputGateFactory(
+ CelebornConf conf, NetworkBufferPool networkBufferPool, int
networkBufferSize) {
+ long configuredMemorySize =
+
org.apache.celeborn.common.util.Utils.byteStringAsBytes(MEMORY_PER_INPUT_GATE);
+ if (configuredMemorySize < MIN_BUFFERS_PER_GATE) {
+ throw new IllegalArgumentException(
+ String.format(
+ "Insufficient network memory per input gate, please increase %s
to at " + "least %s.",
+ MEMORY_PER_INPUT_GATE, MIN_MEMORY_PER_GATE));
+ }
+
+ this.numBuffersPerGate = Utils.checkedDownCast(configuredMemorySize /
networkBufferSize);
+ if (numBuffersPerGate < MIN_BUFFERS_PER_GATE) {
+ throw new IllegalArgumentException(
+ String.format(
+ "Insufficient network memory per input gate, please increase %s
to at "
+ + "least %d bytes.",
+ MEMORY_PER_INPUT_GATE, networkBufferSize *
MIN_BUFFERS_PER_GATE));
+ }
+
+ this.networkBufferSize = networkBufferSize;
+ this.numConcurrentReading = NUM_CONCURRENT_READINGS;
+ this.networkBufferPool = networkBufferPool;
+ this.celebornConf = conf;
+ }
+
+ /** Create {@link RemoteShuffleInputGate} from {@link
InputGateDeploymentDescriptor}. */
+ public RemoteShuffleInputGate create(
+ String owningTaskName, int gateIndex, InputGateDeploymentDescriptor
igdd) {
+ LOG.info(
+ "Create input gate -- number of buffers per input gate={}, "
+ + "number of concurrent readings={}.",
+ numBuffersPerGate,
+ numConcurrentReading);
+
+ SupplierWithException<BufferPool, IOException> bufferPoolFactory =
+ createBufferPoolFactory(networkBufferPool, numBuffersPerGate);
+ BufferDecompressor bufferDecompressor =
+ new BufferDecompressor(networkBufferSize, compressionCodec);
+
+ return createInputGate(owningTaskName, gateIndex, igdd, bufferPoolFactory,
bufferDecompressor);
+ }
+
+ // For testing.
+ RemoteShuffleInputGate createInputGate(
+ String owningTaskName,
+ int gateIndex,
+ InputGateDeploymentDescriptor igdd,
+ SupplierWithException<BufferPool, IOException> bufferPoolFactory,
+ BufferDecompressor bufferDecompressor) {
+ return new RemoteShuffleInputGate(
+ this.celebornConf, owningTaskName, gateIndex, igdd, bufferPoolFactory,
bufferDecompressor);
+ }
+
+ private SupplierWithException<BufferPool, IOException>
createBufferPoolFactory(
+ BufferPoolFactory bufferPoolFactory, int numBuffers) {
+ return () -> bufferPoolFactory.createBufferPool(numBuffers, numBuffers);
+ }
+}
diff --git
a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
index 7a93511b..d8ba441c 100644
---
a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
+++
b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
@@ -74,9 +74,15 @@ public class RemoteShuffleServiceFactory
RemoteShuffleResultPartitionFactory resultPartitionFactory =
new RemoteShuffleResultPartitionFactory(
celebornConf, resultPartitionManager, networkBufferPool,
bufferSize);
+ RemoteShuffleInputGateFactory inputGateFactory =
+ new RemoteShuffleInputGateFactory(celebornConf, networkBufferPool,
bufferSize);
return new RemoteShuffleEnvironment(
- networkBufferPool, resultPartitionManager, resultPartitionFactory,
celebornConf);
+ networkBufferPool,
+ resultPartitionManager,
+ resultPartitionFactory,
+ inputGateFactory,
+ celebornConf);
}
private static int calculateNumberOfNetworkBuffers(MemorySize memorySize,
int bufferSize) {
diff --git
a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
index 8cdaa486..c44f3492 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
@@ -75,7 +75,7 @@ public class ShuffleClientImpl extends ShuffleClient {
private static final Random RND = new Random();
- private final CelebornConf conf;
+ protected final CelebornConf conf;
private final UserIdentifier userIdentifier;
@@ -116,9 +116,9 @@ public class ShuffleClientImpl extends ShuffleClient {
}
};
- private static class ReduceFileGroups {
- final Map<Integer, Set<PartitionLocation>> partitionGroups;
- final int[] mapAttempts;
+ protected static class ReduceFileGroups {
+ public final Map<Integer, Set<PartitionLocation>> partitionGroups;
+ public final int[] mapAttempts;
ReduceFileGroups(Map<Integer, Set<PartitionLocation>> partitionGroups,
int[] mapAttempts) {
this.partitionGroups = partitionGroups;
@@ -127,7 +127,7 @@ public class ShuffleClientImpl extends ShuffleClient {
}
// key: shuffleId
- private final Map<Integer, ReduceFileGroups> reduceFileGroupsMap = new
ConcurrentHashMap<>();
+ protected final Map<Integer, ReduceFileGroups> reduceFileGroupsMap = new
ConcurrentHashMap<>();
private ConcurrentHashMap<String, TransportClient> currentClient = new
ConcurrentHashMap<>();
@@ -1244,6 +1244,68 @@ public class ShuffleClientImpl extends ShuffleClient {
applicationId, shuffleId, partitionId, attemptNumber, 0,
Integer.MAX_VALUE);
}
+ protected ReduceFileGroups loadFileGroupInternal(
+ String applicationId, String shuffleKey, int shuffleId) {
+ {
+ long getReducerFileGroupStartTime = System.nanoTime();
+ try {
+ if (driverRssMetaService == null) {
+ logger.warn("Driver endpoint is null!");
+ return null;
+ }
+
+ GetReducerFileGroup getReducerFileGroup = new
GetReducerFileGroup(applicationId, shuffleId);
+
+ GetReducerFileGroupResponse response =
+ driverRssMetaService.askSync(
+ getReducerFileGroup,
+ conf.getReducerFileGroupRpcAskTimeout(),
+ ClassTag$.MODULE$.apply(GetReducerFileGroupResponse.class));
+
+ if (response.status() == StatusCode.SUCCESS) {
+ logger.info(
+ "Shuffle {} request reducer file group success using time:{} ms,
result partition ids: {}",
+ shuffleId,
+ (System.nanoTime() - getReducerFileGroupStartTime) / 1000_000,
+ response.fileGroup().keySet());
+ return new ReduceFileGroups(response.fileGroup(),
response.attempts());
+ } else if (response.status() == StatusCode.STAGE_END_TIME_OUT) {
+ logger.warn(
+ "Request {} return {} for {}",
+ getReducerFileGroup,
+ StatusCode.STAGE_END_TIME_OUT.toString(),
+ shuffleKey);
+ } else if (response.status() == StatusCode.SHUFFLE_DATA_LOST) {
+ logger.warn(
+ "Request {} return {} for {}",
+ getReducerFileGroup,
+ StatusCode.SHUFFLE_DATA_LOST.toString(),
+ shuffleKey);
+ }
+ } catch (Exception e) {
+ logger.error("Exception raised while call GetReducerFileGroup for " +
shuffleKey + ".", e);
+ }
+ return null;
+ }
+ }
+
+ protected ReduceFileGroups updateFileGroup(
+ String applicationId, String shuffleKey, int shuffleId) {
+ return reduceFileGroupsMap.computeIfAbsent(
+ shuffleId, (id) -> loadFileGroupInternal(applicationId, shuffleKey,
shuffleId));
+ }
+
+ protected ReduceFileGroups loadFileGroup(
+ String applicationId, String shuffleKey, int shuffleId, int partitionId)
throws IOException {
+ ReduceFileGroups reduceFileGroups = updateFileGroup(applicationId,
shuffleKey, shuffleId);
+ if (reduceFileGroups == null) {
+ String msg = "Shuffle data lost for shuffle " + shuffleId + " reduce " +
partitionId + "!";
+ logger.error(msg);
+ throw new IOException(msg);
+ }
+ return reduceFileGroups;
+ }
+
@Override
public RssInputStream readPartition(
String applicationId,
@@ -1254,58 +1316,9 @@ public class ShuffleClientImpl extends ShuffleClient {
int endMapIndex)
throws IOException {
String shuffleKey = Utils.makeShuffleKey(applicationId, shuffleId);
- ReduceFileGroups fileGroups =
- reduceFileGroupsMap.computeIfAbsent(
- shuffleId,
- (id) -> {
- long getReducerFileGroupStartTime = System.nanoTime();
- try {
- if (driverRssMetaService == null) {
- logger.warn("Driver endpoint is null!");
- return null;
- }
-
- GetReducerFileGroup getReducerFileGroup =
- new GetReducerFileGroup(applicationId, shuffleId);
-
- GetReducerFileGroupResponse response =
- driverRssMetaService.askSync(
- getReducerFileGroup,
- conf.getReducerFileGroupRpcAskTimeout(),
-
ClassTag$.MODULE$.apply(GetReducerFileGroupResponse.class));
-
- if (response.status() == StatusCode.SUCCESS) {
- logger.info(
- "Shuffle {} request reducer file group success using
time:{} ms, result partition ids: {}",
- shuffleId,
- (System.nanoTime() - getReducerFileGroupStartTime) /
1000_000,
- response.fileGroup().keySet());
- return new ReduceFileGroups(response.fileGroup(),
response.attempts());
- } else if (response.status() == StatusCode.STAGE_END_TIME_OUT)
{
- logger.warn(
- "Request {} return {} for {}",
- getReducerFileGroup,
- StatusCode.STAGE_END_TIME_OUT.toString(),
- shuffleKey);
- } else if (response.status() == StatusCode.SHUFFLE_DATA_LOST) {
- logger.warn(
- "Request {} return {} for {}",
- getReducerFileGroup,
- StatusCode.SHUFFLE_DATA_LOST.toString(),
- shuffleKey);
- }
- } catch (Exception e) {
- logger.error(
- "Exception raised while call GetReducerFileGroup for " +
shuffleKey + ".", e);
- }
- return null;
- });
+ ReduceFileGroups fileGroups = loadFileGroup(applicationId, shuffleKey,
shuffleId, partitionId);
- if (fileGroups == null) {
- String msg = "Shuffle data lost for shuffle " + shuffleId + " reduce " +
partitionId + "!";
- logger.error(msg);
- throw new IOException(msg);
- } else if (fileGroups.partitionGroups.size() == 0
+ if (fileGroups.partitionGroups.size() == 0
|| !fileGroups.partitionGroups.containsKey(partitionId)) {
logger.warn("Shuffle data is empty for shuffle {} partitionId {}.",
shuffleId, partitionId);
return RssInputStream.empty();
