You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@celeborn.apache.org by et...@apache.org on 2023/02/28 08:38:26 UTC
[incubator-celeborn] 01/01: [CELEBORN-346][FLINK] Support read multiple locations.
This is an automated email from the ASF dual-hosted git repository.
ethanfeng pushed a commit to branch CELEBORN-347
in repository https://gitbox.apache.org/repos/asf/incubator-celeborn.git
commit ead2c22cff22d57e7a199a3ba74ef2e71c15d60b
Author: Ethan Feng <et...@apache.org>
AuthorDate: Tue Feb 28 16:38:05 2023 +0800
[CELEBORN-346][FLINK] Support read multiple locations.
---
.../plugin/flink/network/MessageDecoderExt.java | 4 ++
.../plugin/flink/network/ReadClientHandler.java | 8 ++-
.../flink/readclient/FlinkShuffleClientImpl.java | 5 +-
.../plugin/flink/readclient/RssBufferStream.java | 68 +++++++++++++++++-----
.../plugin/flink/RemoteBufferStreamReader.java | 39 ++++++++-----
.../plugin/flink/RemoteShuffleInputGate.java | 12 +---
.../common/network/protocol/EndOfStream.java | 52 +++++++++++++++++
.../celeborn/common/network/protocol/Message.java | 8 ++-
.../common/network/server/BufferStreamManager.java | 1 +
.../common/network/server/DataPartitionReader.java | 39 +++++++------
10 files changed, 173 insertions(+), 63 deletions(-)
diff --git a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/MessageDecoderExt.java b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/MessageDecoderExt.java
index fbee2e4fd..38cecf232 100644
--- a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/MessageDecoderExt.java
+++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/MessageDecoderExt.java
@@ -81,6 +81,10 @@ public class MessageDecoderExt {
backlog = in.readInt();
return new BacklogAnnouncement(streamId, backlog);
+ case END_OF_STREAM:
+ streamId = in.readLong();
+ return new EndOfStream(streamId);
+
default:
throw new IllegalArgumentException("Unexpected message type: " + type);
}
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 8f768c26b..be82d4574 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
@@ -30,10 +30,10 @@ import org.apache.celeborn.common.network.server.BaseMessageHandler;
import org.apache.celeborn.plugin.flink.protocol.ReadData;
public class ReadClientHandler extends BaseMessageHandler {
- private static Logger logger = LoggerFactory.getLogger(ReadClientHandler.class);
- private ConcurrentHashMap<Long, Consumer<RequestMessage>> streamHandlers =
+ private static final Logger logger = LoggerFactory.getLogger(ReadClientHandler.class);
+ private final ConcurrentHashMap<Long, Consumer<RequestMessage>> streamHandlers =
new ConcurrentHashMap<>();
- private ConcurrentHashMap<Long, TransportClient> streamClients = new ConcurrentHashMap<>();
+ private final ConcurrentHashMap<Long, TransportClient> streamClients = new ConcurrentHashMap<>();
public void registerHandler(
long streamId, Consumer<RequestMessage> handle, TransportClient client) {
@@ -75,6 +75,8 @@ public class ReadClientHandler extends BaseMessageHandler {
logger.warn("Unexpected streamId received: {}", streamId);
}
break;
+ case END_OF_STREAM:
+ streamHandlers.get(streamId).accept(msg);
case ONE_WAY_MESSAGE:
// ignore it.
break;
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
index 3de2bb1f7..83736afff 100644
--- 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
@@ -37,8 +37,8 @@ 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();
+ private final FlinkTransportClientFactory flinkTransportClientFactory;
+ private final ReadClientHandler readClientHandler = new ReadClientHandler();
public static FlinkShuffleClientImpl get(
String driverHost, int port, CelebornConf conf, UserIdentifier userIdentifier) {
@@ -98,7 +98,6 @@ public class FlinkShuffleClientImpl extends ShuffleClientImpl {
return RssBufferStream.empty();
} else {
return RssBufferStream.create(
- conf,
flinkTransportClientFactory,
shuffleKey,
fileGroups.partitionGroups.get(partitionId).toArray(new PartitionLocation[0]),
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
index 16cc03419..b8514a5b8 100644
--- 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
@@ -26,7 +26,6 @@ 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.*;
@@ -35,27 +34,29 @@ import org.apache.celeborn.plugin.flink.network.FlinkTransportClientFactory;
public class RssBufferStream {
- private static Logger logger = LoggerFactory.getLogger(RssBufferStream.class);
- private CelebornConf conf;
+ private static final Logger logger = LoggerFactory.getLogger(RssBufferStream.class);
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;
+ private volatile int currentLocationIndex = 0;
+ private volatile long streamId = 0;
+ private Supplier<ByteBuf> supplier;
+ private int initialCredit;
+ private FlinkShuffleClientImpl mapShuffleClient;
+ private Consumer<RequestMessage> messageConsumer;
+ private Consumer<Throwable> failureListener;
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;
@@ -63,12 +64,20 @@ public class RssBufferStream {
this.subIndexEnd = subIndexEnd;
}
- public void open(
+ public void initByReader(
Supplier<ByteBuf> supplier,
int initialCredit,
FlinkShuffleClientImpl mapShuffleClient,
- Consumer<RequestMessage> messageConsumer)
- throws IOException, InterruptedException {
+ Consumer<RequestMessage> messageConsumer,
+ Consumer<Throwable> failureListener) {
+ this.supplier = supplier;
+ this.initialCredit = initialCredit;
+ this.mapShuffleClient = mapShuffleClient;
+ this.messageConsumer = messageConsumer;
+ this.failureListener = failureListener;
+ }
+
+ public void open() throws IOException, InterruptedException {
this.client =
clientFactory.createClient(
locations[currentLocationIndex].getHost(),
@@ -105,9 +114,8 @@ public class RssBufferStream {
.writeAndFlush(addCredit)
.addListener(
future -> {
- if (future.isSuccess()) {
+ if (!future.isSuccess()) {
// Send ReadAddCredit do not expect response.
- } else {
logger.warn(
"Send ReadAddCredit to {} failed, detail {}",
this.client.getSocketAddress().toString(),
@@ -125,7 +133,6 @@ public class RssBufferStream {
}
public static RssBufferStream create(
- CelebornConf conf,
FlinkTransportClientFactory dataClientFactory,
String shuffleKey,
PartitionLocation[] locations,
@@ -135,7 +142,7 @@ public class RssBufferStream {
return empty();
} else {
return new RssBufferStream(
- conf, dataClientFactory, shuffleKey, locations, subIndexStart, subIndexEnd);
+ dataClientFactory, shuffleKey, locations, subIndexStart, subIndexEnd);
}
}
@@ -146,6 +153,37 @@ public class RssBufferStream {
}
public void close() {
- clientFactory.unregisterSupplier(this.getStreamId());
+ freeSupplier(streamId);
+ }
+
+ private void freeSupplier(long streamId) {
+ clientFactory.unregisterSupplier(streamId);
+ }
+
+ public void endStream(long streamId) {
+ if (this.streamId != streamId) {
+ logger.error("Unexpected end of stream with {}", streamId);
+ return;
+ }
+
+ if (currentLocationIndex >= locations.length) {
+ return;
+ }
+
+ currentLocationIndex++;
+ freeSupplier(streamId);
+ // release netty thread
+ Thread thread =
+ new Thread(
+ () -> {
+ try {
+ this.open();
+ } catch (Exception e) {
+ this.failureListener.accept(
+ new IOException(
+ "Read location " + locations[currentLocationIndex] + " failed.", e));
+ }
+ });
+ thread.start();
}
}
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
index 12cc0f324..c2e04793f 100644
--- 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
@@ -24,6 +24,7 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.network.protocol.BacklogAnnouncement;
+import org.apache.celeborn.common.network.protocol.EndOfStream;
import org.apache.celeborn.common.network.protocol.ReadAddCredit;
import org.apache.celeborn.common.network.protocol.RequestMessage;
import org.apache.celeborn.plugin.flink.buffer.CreditListener;
@@ -33,20 +34,20 @@ 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 static final Logger logger = LoggerFactory.getLogger(RemoteBufferStreamReader.class);
+ private final TransferBufferPool bufferPool;
+ private final FlinkShuffleClientImpl client;
+ private final String applicationId;
+ private final int shuffleId;
+ private final int partitionId;
+ private final int subPartitionIndexStart;
+ private final int subPartitionIndexEnd;
private boolean isOpen;
- private Consumer<ByteBuf> dataListener;
- private Consumer<Throwable> failureListener;
+ private final Consumer<ByteBuf> dataListener;
private RssBufferStream bufferStream;
private volatile boolean closed = false;
- private Consumer<RequestMessage> messageConsumer;
+ private final Consumer<RequestMessage> messageConsumer;
+ private final Consumer<Throwable> failureListener;
public RemoteBufferStreamReader(
FlinkShuffleClientImpl client,
@@ -72,6 +73,8 @@ public class RemoteBufferStreamReader extends CreditListener {
dataReceived((ReadData) requestMessage);
} else if (requestMessage instanceof BacklogAnnouncement) {
backlogReceived(((BacklogAnnouncement) requestMessage).getBacklog());
+ } else if (requestMessage instanceof EndOfStream) {
+ endOfStreamReceived(((EndOfStream) requestMessage));
}
};
}
@@ -81,8 +84,13 @@ public class RemoteBufferStreamReader extends CreditListener {
this.bufferStream =
client.readBufferedPartition(
applicationId, shuffleId, partitionId, subPartitionIndexStart, subPartitionIndexEnd);
- bufferStream.open(
- RemoteBufferStreamReader.this::requestBuffer, initialCredit, client, messageConsumer);
+ bufferStream.initByReader(
+ RemoteBufferStreamReader.this::requestBuffer,
+ initialCredit,
+ client,
+ messageConsumer,
+ failureListener);
+ bufferStream.open();
} catch (InterruptedException e) {
throw new RuntimeException("Failed to openStream.", e);
}
@@ -136,4 +144,9 @@ public class RemoteBufferStreamReader extends CreditListener {
}
dataListener.accept(readData.getFlinkBuffer());
}
+
+ public void endOfStreamReceived(EndOfStream endOfStream) {
+ bufferStream.endStream(endOfStream.getStreamId());
+ client.getReadClientHandler().removeHandler(endOfStream.getStreamId());
+ }
}
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
index 1192125cf..0eecc5a54 100644
--- 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
@@ -124,14 +124,11 @@ public class RemoteShuffleInputGate extends IndexedInputGate {
/** 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;
-
+ private final int numConcurrentReading = Integer.MAX_VALUE;
/** Data decompressor. */
private final BufferDecompressor bufferDecompressor;
- private FlinkShuffleClientImpl shuffleClient;
+ private final FlinkShuffleClientImpl shuffleClient;
public RemoteShuffleInputGate(
CelebornConf celebornConf,
@@ -351,11 +348,6 @@ public class RemoteShuffleInputGate extends IndexedInputGate {
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))
diff --git a/common/src/main/java/org/apache/celeborn/common/network/protocol/EndOfStream.java b/common/src/main/java/org/apache/celeborn/common/network/protocol/EndOfStream.java
new file mode 100644
index 000000000..3bebbd3fb
--- /dev/null
+++ b/common/src/main/java/org/apache/celeborn/common/network/protocol/EndOfStream.java
@@ -0,0 +1,52 @@
+/*
+ * 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.common.network.protocol;
+
+import io.netty.buffer.ByteBuf;
+
+public class EndOfStream extends RequestMessage {
+ long streamId;
+
+ public EndOfStream(long streamId) {
+ this.streamId = streamId;
+ }
+
+ @Override
+ public int encodedLength() {
+ return 8;
+ }
+
+ @Override
+ public void encode(ByteBuf buf) {
+ buf.writeLong(streamId);
+ }
+
+ @Override
+ public Type type() {
+ return Type.END_OF_STREAM;
+ }
+
+ public long getStreamId() {
+ return streamId;
+ }
+
+ public static EndOfStream decode(ByteBuf in) {
+ long streamId = in.readLong();
+ return new EndOfStream(streamId);
+ }
+}
diff --git a/common/src/main/java/org/apache/celeborn/common/network/protocol/Message.java b/common/src/main/java/org/apache/celeborn/common/network/protocol/Message.java
index d2bbd55a6..f1562044f 100644
--- a/common/src/main/java/org/apache/celeborn/common/network/protocol/Message.java
+++ b/common/src/main/java/org/apache/celeborn/common/network/protocol/Message.java
@@ -92,7 +92,8 @@ public abstract class Message implements Encodable {
READ_ADD_CREDIT(16),
READ_DATA(17),
OPEN_STREAM_WITH_CREDIT(18),
- BACKLOG_ANNOUNCEMENT(19);
+ BACKLOG_ANNOUNCEMENT(19),
+ END_OF_STREAM(20);
private final byte id;
Type(int id) {
@@ -158,6 +159,8 @@ public abstract class Message implements Encodable {
return OPEN_STREAM_WITH_CREDIT;
case 19:
return BACKLOG_ANNOUNCEMENT;
+ case 20:
+ return END_OF_STREAM;
case -1:
throw new IllegalArgumentException("User type messages cannot be decoded.");
default:
@@ -223,6 +226,9 @@ public abstract class Message implements Encodable {
case BACKLOG_ANNOUNCEMENT:
return BacklogAnnouncement.decode(in);
+ case END_OF_STREAM:
+ return EndOfStream.decode(in);
+
default:
throw new IllegalArgumentException("Unexpected message type: " + msgType);
}
diff --git a/common/src/main/java/org/apache/celeborn/common/network/server/BufferStreamManager.java b/common/src/main/java/org/apache/celeborn/common/network/server/BufferStreamManager.java
index 9e367776d..f9acfb770 100644
--- a/common/src/main/java/org/apache/celeborn/common/network/server/BufferStreamManager.java
+++ b/common/src/main/java/org/apache/celeborn/common/network/server/BufferStreamManager.java
@@ -384,6 +384,7 @@ public class BufferStreamManager {
removeStream(streamId);
streams.remove(streamId);
servingStreams.remove(streamId);
+ dataPartitionReader.notifyEndOfStream(streamId);
return true;
}
diff --git a/common/src/main/java/org/apache/celeborn/common/network/server/DataPartitionReader.java b/common/src/main/java/org/apache/celeborn/common/network/server/DataPartitionReader.java
index 612712147..06da8657c 100644
--- a/common/src/main/java/org/apache/celeborn/common/network/server/DataPartitionReader.java
+++ b/common/src/main/java/org/apache/celeborn/common/network/server/DataPartitionReader.java
@@ -34,6 +34,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.meta.FileInfo;
import org.apache.celeborn.common.network.protocol.BacklogAnnouncement;
+import org.apache.celeborn.common.network.protocol.EndOfStream;
import org.apache.celeborn.common.network.protocol.ReadData;
import org.apache.celeborn.common.network.server.memory.Recycler;
import org.apache.celeborn.common.network.server.memory.WrappedDataBuffer;
@@ -51,12 +52,12 @@ public class DataPartitionReader implements Comparable<DataPartitionReader> {
private int currentDataRegion = -1;
private long dataConsumingOffset;
private volatile long currentPartitionRemainingBytes;
- private FileInfo fileInfo;
- private int INDEX_ENTRY_SIZE = 16;
- private long streamId;
+ private final FileInfo fileInfo;
+ private final int INDEX_ENTRY_SIZE = 16;
+ private final long streamId;
protected final Object lock = new Object();
- private AtomicInteger credits = new AtomicInteger();
+ private final AtomicInteger credits = new AtomicInteger();
@GuardedBy("lock")
protected final Queue<WrappedDataBuffer> buffersRead = new ArrayDeque<>();
@@ -80,11 +81,11 @@ public class DataPartitionReader implements Comparable<DataPartitionReader> {
private FileChannel dataFileChannel;
private FileChannel indexFileChannel;
- private Channel associatedChannel;
+ private final Channel associatedChannel;
- private Runnable recycleStream;
+ private final Runnable recycleStream;
- private AtomicInteger numInFlightRequests = new AtomicInteger(0);
+ private final AtomicInteger numInFlightRequests = new AtomicInteger(0);
public DataPartitionReader(
int startPartitionIndex,
@@ -121,11 +122,7 @@ public class DataPartitionReader implements Comparable<DataPartitionReader> {
public boolean sendWithCredit(int credit) {
int oldCredit = credits.getAndAdd(credit);
- if (oldCredit == 0) {
- return true;
- }
-
- return false;
+ return oldCredit == 0;
}
public synchronized boolean readAndSend(Queue<ByteBuf> bufferQueue, Recycler bufferRecycler)
@@ -438,11 +435,17 @@ public class DataPartitionReader implements Comparable<DataPartitionReader> {
@Override
public String toString() {
- final StringBuilder sb = new StringBuilder("DataPartitionReader{");
- sb.append("startPartitionIndex=").append(startPartitionIndex);
- sb.append(", endPartitionIndex=").append(endPartitionIndex);
- sb.append(", streamId=").append(streamId);
- sb.append('}');
- return sb.toString();
+ return "DataPartitionReader{"
+ + "startPartitionIndex="
+ + startPartitionIndex
+ + ", endPartitionIndex="
+ + endPartitionIndex
+ + ", streamId="
+ + streamId
+ + '}';
+ }
+
+ public void notifyEndOfStream(long streamId) {
+ this.associatedChannel.writeAndFlush(new EndOfStream(streamId));
}
}