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/17 09:38:41 UTC
[incubator-celeborn] branch main updated: [CELEBORN-234] Implement buffer stream. (#1221)
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 0c8bb831 [CELEBORN-234] Implement buffer stream. (#1221)
0c8bb831 is described below
commit 0c8bb831149b27b18e82048a49ed92a904be4889
Author: Ethan Feng <et...@apache.org>
AuthorDate: Fri Feb 17 17:38:36 2023 +0800
[CELEBORN-234] Implement buffer stream. (#1221)
---
.../org/apache/celeborn/common/meta/FileInfo.java | 10 +
.../celeborn/common/network/protocol/Message.java | 3 -
.../celeborn/common/network/protocol/ReadData.java | 34 +-
.../common/network/server/BufferStreamManager.java | 568 ++++++++++++++++++++-
.../server/memory/ReadBufferDispatcher.java | 12 +-
.../common/network/util/TransportConf.java | 44 +-
.../org/apache/celeborn/common/CelebornConf.scala | 29 ++
.../org/apache/celeborn/common/util/Utils.scala | 23 +-
docs/configuration/worker.md | 3 +
.../service/deploy/worker/FetchHandler.scala | 29 +-
.../deploy/worker/storage/StorageManager.scala | 2 +-
11 files changed, 687 insertions(+), 70 deletions(-)
diff --git a/common/src/main/java/org/apache/celeborn/common/meta/FileInfo.java b/common/src/main/java/org/apache/celeborn/common/meta/FileInfo.java
index dac7eb15..64106179 100644
--- a/common/src/main/java/org/apache/celeborn/common/meta/FileInfo.java
+++ b/common/src/main/java/org/apache/celeborn/common/meta/FileInfo.java
@@ -40,6 +40,8 @@ public class FileInfo {
private final UserIdentifier userIdentifier;
private final PartitionType partitionType;
private int bufferSize;
+ private String mountPoint;
+ // Map partition related param
private int numReducerPartitions;
public FileInfo(String filePath, List<Long> chunkOffsets, UserIdentifier userIdentifier) {
@@ -195,4 +197,12 @@ public class FileInfo {
public void setNumReducerPartitions(int numReducerPartitions) {
this.numReducerPartitions = numReducerPartitions;
}
+
+ public String getMountPoint() {
+ return mountPoint;
+ }
+
+ public void setMountPoint(String mountPoint) {
+ this.mountPoint = mountPoint;
+ }
}
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 b15453a1..d2bbd55a 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
@@ -217,9 +217,6 @@ public abstract class Message implements Encodable {
case READ_ADD_CREDIT:
return ReadAddCredit.decode(in);
- case READ_DATA:
- return ReadData.decode(in);
-
case OPEN_STREAM_WITH_CREDIT:
return OpenStreamWithCredit.decode(in);
diff --git a/common/src/main/java/org/apache/celeborn/common/network/protocol/ReadData.java b/common/src/main/java/org/apache/celeborn/common/network/protocol/ReadData.java
index 92e107cc..ff286a9d 100644
--- a/common/src/main/java/org/apache/celeborn/common/network/protocol/ReadData.java
+++ b/common/src/main/java/org/apache/celeborn/common/network/protocol/ReadData.java
@@ -14,29 +14,32 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
+
package org.apache.celeborn.common.network.protocol;
import java.util.Objects;
import io.netty.buffer.ByteBuf;
-import io.netty.buffer.Unpooled;
+import org.apache.celeborn.common.network.buffer.NettyManagedBuffer;
+
+// This is buffer wrapper used in celeborn worker only
+// It doesn't need decode in worker.
public class ReadData extends RequestMessage {
private long streamId;
private int backlog;
private long offset;
- private ByteBuf buf;
public ReadData(long streamId, int backlog, long offset, ByteBuf buf) {
+ super(new NettyManagedBuffer(buf));
this.streamId = streamId;
this.backlog = backlog;
this.offset = offset;
- this.buf = buf;
}
@Override
public int encodedLength() {
- return 8 + 4 + 4 + 8 + 4 + buf.readableBytes();
+ return 8 + 4 + 8;
}
@Override
@@ -44,18 +47,6 @@ public class ReadData extends RequestMessage {
buf.writeLong(streamId);
buf.writeInt(backlog);
buf.writeLong(offset);
- buf.writeInt(this.buf.readableBytes());
- buf.writeBytes(this.buf);
- }
-
- public static ReadData decode(ByteBuf buf) {
- long streamId = buf.readLong();
- int backlog = buf.readInt();
- long offset = buf.readLong();
- int tmpBufSize = buf.readInt();
- ByteBuf tmpBuf = Unpooled.buffer(tmpBufSize, tmpBufSize);
- buf.readBytes(tmpBuf);
- return new ReadData(streamId, backlog, offset, tmpBuf);
}
public long getStreamId() {
@@ -70,10 +61,6 @@ public class ReadData extends RequestMessage {
return offset;
}
- public ByteBuf getBuf() {
- return buf;
- }
-
@Override
public Type type() {
return Type.READ_DATA;
@@ -86,13 +73,12 @@ public class ReadData extends RequestMessage {
ReadData readData = (ReadData) o;
return streamId == readData.streamId
&& backlog == readData.backlog
- && offset == readData.offset
- && Objects.equals(buf, readData.buf);
+ && offset == readData.offset;
}
@Override
public int hashCode() {
- return Objects.hash(streamId, backlog, offset, buf);
+ return Objects.hash(streamId, backlog, offset);
}
@Override
@@ -104,8 +90,6 @@ public class ReadData extends RequestMessage {
+ backlog
+ ", offset="
+ offset
- + ", buf="
- + buf
+ '}';
}
}
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 bba500b8..ad5516f2 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
@@ -17,19 +17,47 @@
package org.apache.celeborn.common.network.server;
-import java.util.Map;
-import java.util.Random;
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.FileChannel;
+import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.LinkedBlockingDeque;
+import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import javax.annotation.concurrent.GuardedBy;
+
+import com.google.common.util.concurrent.ThreadFactoryBuilder;
+import io.netty.buffer.ByteBuf;
import io.netty.channel.Channel;
+import io.netty.channel.ChannelFutureListener;
+import org.apache.commons.io.IOUtils;
import org.slf4j.Logger;
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.ReadData;
+import org.apache.celeborn.common.network.server.memory.MemoryManager;
+import org.apache.celeborn.common.util.Utils;
+
public class BufferStreamManager {
private static final Logger logger = LoggerFactory.getLogger(BufferStreamManager.class);
private final AtomicLong nextStreamId;
protected final ConcurrentHashMap<Long, StreamState> streams;
+ protected final ConcurrentHashMap<Long, AtomicInteger> streamCredits;
+ protected final ConcurrentHashMap<Long, MapDataPartition> servingStreams;
+ protected final ConcurrentHashMap<FileInfo, MapDataPartition> activeMapPartitions;
+ protected final MemoryManager memoryManager = MemoryManager.instance();
+ protected final StorageFetcherPool storageFetcherPool = new StorageFetcherPool();
+ protected int minReadBuffers;
+ protected int maxReadBuffers;
+ protected int threadsPerMountPoint;
protected class StreamState {
private Channel associatedChannel;
@@ -49,24 +77,550 @@ public class BufferStreamManager {
}
}
- public BufferStreamManager() {
+ public BufferStreamManager(int minReadBuffers, int maxReadBuffers, int threadsPerMountpoint) {
nextStreamId = new AtomicLong((long) new Random().nextInt(Integer.MAX_VALUE) * 1000);
streams = new ConcurrentHashMap<>();
+ streamCredits = new ConcurrentHashMap<>();
+ servingStreams = new ConcurrentHashMap<>();
+ activeMapPartitions = new ConcurrentHashMap<>();
+ this.minReadBuffers = minReadBuffers;
+ this.maxReadBuffers = maxReadBuffers;
+ this.threadsPerMountPoint = threadsPerMountpoint;
}
- public long registerStream(Channel channel, int bufferSize) {
+ public long registerStream(
+ Channel channel, int initialCredit, int startSubIndex, int endSubIndex, FileInfo fileInfo)
+ throws IOException {
long streamId = nextStreamId.getAndIncrement();
- streams.put(streamId, new StreamState(channel, bufferSize));
+ streams.put(streamId, new StreamState(channel, fileInfo.getBufferSize()));
+
+ synchronized (activeMapPartitions) {
+ MapDataPartition mapDataPartition =
+ activeMapPartitions.computeIfAbsent(fileInfo, (f) -> new MapDataPartition(fileInfo));
+
+ mapDataPartition.addStream(streamId);
+ addCredit(initialCredit, streamId);
+ servingStreams.put(streamId, mapDataPartition);
+ mapDataPartition.setupDataPartitionReader(startSubIndex, endSubIndex, streamId);
+ }
+
+ logger.debug("Register stream streamId: {}, fileInfo: {}", streamId, fileInfo);
+
return streamId;
}
- public void addCredit(int numCredit, long streamId) {}
+ public void addCredit(int numCredit, long streamId) {
+ logger.debug("streamId: {}, add credit: {}", streamId, numCredit);
+ streamCredits.compute(
+ streamId,
+ (aLong, atomicInteger) -> {
+ if (atomicInteger == null) {
+ return new AtomicInteger(numCredit);
+ } else {
+ atomicInteger.getAndAdd(numCredit);
+ return atomicInteger;
+ }
+ });
+
+ MapDataPartition mapDataPartition = servingStreams.get(streamId);
+ if (mapDataPartition != null) {
+ DataPartitionReader streamReader = mapDataPartition.getStreamReader(streamId);
+ if (streamReader != null) {
+ storageFetcherPool
+ .getExecutorPool(streamReader.fileInfo.getMountPoint())
+ .submit(() -> streamReader.sendData());
+ }
+ }
+ }
public void connectionTerminated(Channel channel) {
for (Map.Entry<Long, StreamState> entry : streams.entrySet()) {
if (entry.getValue().getAssociatedChannel() == channel) {
- streams.remove(entry.getKey());
+ logger.info("connection closed, clean streamId: {}", entry.getKey());
+ cleanResource(entry.getKey());
+ }
+ }
+ }
+
+ public void cleanResource(long streamId) {
+ logger.debug("clean stream: {}", streamId);
+ streams.remove(streamId);
+ streamCredits.remove(streamId);
+ FileInfo fileInfo = servingStreams.remove(streamId).fileInfo;
+ MapDataPartition mapDataPartition = activeMapPartitions.get(fileInfo);
+ mapDataPartition.removeStream(streamId);
+ synchronized (activeMapPartitions) {
+ if (mapDataPartition.activeStreamIds.isEmpty()) {
+ for (ByteBuf buffer : mapDataPartition.buffers) {
+ memoryManager.recycleReadBuffer(buffer);
+ }
+ activeMapPartitions.remove(fileInfo);
+ }
+ }
+ }
+
+ // this means active data partition
+ protected class MapDataPartition {
+ private final List<Long> activeStreamIds = new ArrayList<>();
+ private final FileInfo fileInfo;
+ private final Set<DataPartitionReader> readers = new HashSet<>();
+ private final ExecutorService readExecutor;
+ private final ConcurrentHashMap<Long, DataPartitionReader> streamReaders =
+ new ConcurrentHashMap<>();
+
+ /** All available buffers can be used by the partition readers for reading. */
+ protected Queue<ByteBuf> buffers;
+
+ public MapDataPartition(FileInfo fileInfo) {
+ this.fileInfo = fileInfo;
+ readExecutor = storageFetcherPool.getExecutorPool(fileInfo.getFilePath());
+ }
+
+ public synchronized void setupDataPartitionReader(
+ int startSubIndex, int endSubIndex, long streamId) throws IOException {
+ DataPartitionReader dataPartitionReader =
+ new DataPartitionReader(startSubIndex, endSubIndex, fileInfo, streamId);
+ dataPartitionReader.open();
+ // allocate resources when the first reader is registered
+ boolean allocateResources = readers.isEmpty();
+ readers.add(dataPartitionReader);
+ streamReaders.put(streamId, dataPartitionReader);
+
+ // create initial buffers for read
+ if (allocateResources) {
+ memoryManager.requestReadBuffers(
+ minReadBuffers,
+ maxReadBuffers,
+ fileInfo.getBufferSize(),
+ (allocatedBuffers, throwable) ->
+ MapDataPartition.this.onBuffer(new LinkedBlockingDeque<>(allocatedBuffers)));
+ } else {
+ triggerRead();
+ }
+ }
+
+ public DataPartitionReader getStreamReader(long streamId) {
+ return streamReaders.get(streamId);
+ }
+
+ // Read logic is executed on another thread.
+ public void onBuffer(Queue<ByteBuf> buffers) {
+ this.buffers = buffers;
+ triggerRead();
+ }
+
+ public void recycle(ByteBuf buffer, Queue<ByteBuf> bufferQueue) {
+ buffer.clear();
+ bufferQueue.add(buffer);
+ // avoid unnecessary thread switch
+ readBuffers();
+ }
+
+ public void readBuffers() {
+ PriorityQueue<DataPartitionReader> sortedReaders = new PriorityQueue<>(readers);
+ while (buffers.size() > 0 && !sortedReaders.isEmpty()) {
+ DataPartitionReader reader = sortedReaders.poll();
+ try {
+ if (!reader.readAndSend(buffers, (buffer) -> this.recycle(buffer, buffers))) {
+ readers.remove(reader);
+ }
+ } catch (IOException e) {
+ logger.error("Read thread error occurred, {}", e);
+ throw new RuntimeException(e);
+ }
+ }
+ }
+
+ public void triggerRead() {
+ readExecutor.submit(
+ () -> {
+ // Key for IO schedule.
+ synchronized (MapDataPartition.this) {
+ readBuffers();
+ }
+ });
+ }
+
+ public synchronized void addStream(long streamId) {
+ activeStreamIds.add(streamId);
+ }
+
+ public synchronized void removeStream(long streamId) {
+ activeStreamIds.remove(streamId);
+ }
+ }
+
+ private class Buffer {
+ private ByteBuf byteBuf;
+ private Consumer<ByteBuf> byteBufferConsumer;
+
+ public Buffer(ByteBuf byteBuf, Consumer<ByteBuf> byteBufferConsumer) {
+ this.byteBuf = byteBuf;
+ this.byteBufferConsumer = byteBufferConsumer;
+ }
+ }
+
+ // this is a specific partition reader
+ protected class DataPartitionReader implements Comparable<DataPartitionReader> {
+ private final ByteBuffer indexBuffer;
+ private final ByteBuffer headerBuffer;
+ private final int startPartitionIndex;
+ private final int endPartitionIndex;
+ private int numRegions;
+ private FileChannel dataFileChannel;
+ private FileChannel indexFileChannel;
+ private int numRemainingPartitions;
+ private int currentDataRegion = -1;
+ private long dataConsumingOffset;
+ private volatile long currentPartitionRemainingBytes;
+ private boolean isClosed;
+ private FileInfo fileInfo;
+ private int INDEX_ENTRY_SIZE = 16;
+ private long streamId;
+ protected final Object lock = new Object();
+
+ @GuardedBy("lock")
+ protected final Queue<Buffer> buffersRead = new ArrayDeque<>();
+
+ /** Whether all the data has been successfully read or not. */
+ @GuardedBy("lock")
+ protected boolean isFinished;
+
+ /** Whether this partition reader has been released or not. */
+ @GuardedBy("lock")
+ protected boolean isReleased;
+
+ /** Exception causing the release of this partition reader. */
+ @GuardedBy("lock")
+ protected Throwable errorCause;
+
+ /** Whether there is any error at the consumer side or not. */
+ @GuardedBy("lock")
+ protected boolean isError;
+
+ public DataPartitionReader(
+ int startPartitionIndex, int endPartitionIndex, FileInfo fileInfo, long streamId) {
+ this.startPartitionIndex = startPartitionIndex;
+ this.endPartitionIndex = endPartitionIndex;
+
+ int indexBufferSize = 16 * (endPartitionIndex - startPartitionIndex + 1);
+ this.indexBuffer = ByteBuffer.allocateDirect(indexBufferSize);
+
+ this.headerBuffer = ByteBuffer.allocateDirect(16);
+ this.streamId = streamId;
+
+ this.fileInfo = fileInfo;
+ this.isClosed = false;
+ }
+
+ public void open() throws IOException {
+ this.dataFileChannel = new FileInputStream(fileInfo.getFile()).getChannel();
+ this.indexFileChannel = new FileInputStream(fileInfo.getIndexPath()).getChannel();
+
+ long indexFileSize = indexFileChannel.size();
+ // index is (offset,length)
+ long indexRegionSize = fileInfo.getNumReducerPartitions() * (long) INDEX_ENTRY_SIZE;
+ this.numRegions = Utils.checkedDownCast(indexFileSize / indexRegionSize);
+
+ updateConsumingOffset();
+ }
+
+ public long getIndexRegionSize() {
+ return fileInfo.getNumReducerPartitions() * (long) INDEX_ENTRY_SIZE;
+ }
+
+ private void readHeaderOrIndexBuffer(FileChannel channel, ByteBuffer buffer, int length)
+ throws IOException {
+ Utils.checkFileIntegrity(channel, length);
+ buffer.clear();
+ buffer.limit(length);
+ while (buffer.hasRemaining()) {
+ channel.read(buffer);
+ }
+ buffer.flip();
+ }
+
+ private void readBufferIntoReadBuffer(FileChannel channel, ByteBuf buf, int length)
+ throws IOException {
+ Utils.checkFileIntegrity(channel, length);
+ ByteBuffer tmpBuffer = ByteBuffer.allocate(length);
+ while (tmpBuffer.hasRemaining()) {
+ channel.read(tmpBuffer);
+ }
+ tmpBuffer.flip();
+ buf.writeBytes(tmpBuffer);
+ }
+
+ private int readBuffer(
+ String filename, FileChannel channel, ByteBuffer header, ByteBuf buffer, int headerSize)
+ throws IOException {
+ readHeaderOrIndexBuffer(channel, header, headerSize);
+ // header is combined of mapId(4),attemptId(4),nextBatchId(4) and totcalCompresszedLength(4)
+ // we need size here,so we read length directly
+ int bufferLenght = header.getInt(12);
+ if (bufferLenght <= 0 || bufferLenght > buffer.capacity()) {
+ logger.error("Incorrect buffer header, buffer length: {}.", bufferLenght);
+ throw new RuntimeException("File " + filename + " is corrupted");
+ }
+ buffer.writeBytes(header);
+ readBufferIntoReadBuffer(channel, buffer, bufferLenght);
+ return bufferLenght + headerSize;
+ }
+
+ private void updateConsumingOffset() throws IOException {
+ while (currentPartitionRemainingBytes == 0
+ && (currentDataRegion < numRegions - 1 || numRemainingPartitions > 0)) {
+ if (numRemainingPartitions <= 0) {
+ ++currentDataRegion;
+ numRemainingPartitions = endPartitionIndex - startPartitionIndex + 1;
+
+ // read the target index entry to the target index buffer
+ indexFileChannel.position(
+ currentDataRegion * getIndexRegionSize()
+ + (long) startPartitionIndex * INDEX_ENTRY_SIZE);
+ readHeaderOrIndexBuffer(indexFileChannel, indexBuffer, indexBuffer.capacity());
+ }
+
+ // get the data file offset and the data size
+ dataConsumingOffset = indexBuffer.getLong();
+ currentPartitionRemainingBytes = indexBuffer.getLong();
+ --numRemainingPartitions;
+
+ logger.debug(
+ "readBuffer updateConsumingOffset, {}, {}, {}, {}",
+ streamId,
+ dataFileChannel.size(),
+ dataConsumingOffset,
+ currentPartitionRemainingBytes);
+
+ // if these checks fail, the partition file must be corrupted
+ if (dataConsumingOffset < 0
+ || dataConsumingOffset + currentPartitionRemainingBytes > dataFileChannel.size()
+ || currentPartitionRemainingBytes < 0) {
+ throw new RuntimeException("File " + fileInfo.getFilePath() + " is corrupted");
+ }
+ }
+ }
+
+ public boolean readBuffer(ByteBuf buffer) throws IOException {
+ try {
+ dataFileChannel.position(dataConsumingOffset);
+
+ int readSize =
+ readBuffer(
+ fileInfo.getFilePath(),
+ dataFileChannel,
+ headerBuffer,
+ buffer,
+ headerBuffer.capacity());
+ currentPartitionRemainingBytes -= readSize;
+
+ logger.debug(
+ "readBuffer data: {}, {}, {}, {}, {}, {}",
+ streamId,
+ currentPartitionRemainingBytes,
+ readSize,
+ dataConsumingOffset,
+ fileInfo.getFilePath(),
+ System.identityHashCode(buffer));
+
+ // if this check fails, the partition file must be corrupted
+ if (currentPartitionRemainingBytes < 0) {
+ throw new RuntimeException("File is corrupted");
+ } else if (currentPartitionRemainingBytes == 0) {
+ logger.debug(
+ "readBuffer end, {}, {}, {}, {}",
+ streamId,
+ dataFileChannel.size(),
+ dataConsumingOffset,
+ currentPartitionRemainingBytes);
+ int prevDataRegion = currentDataRegion;
+ updateConsumingOffset();
+ return prevDataRegion == currentDataRegion && currentPartitionRemainingBytes > 0;
+ }
+
+ dataConsumingOffset = dataFileChannel.position();
+
+ logger.debug(
+ "readBuffer run: {}, {}, {}, {}",
+ streamId,
+ dataFileChannel.size(),
+ dataConsumingOffset,
+ currentPartitionRemainingBytes);
+ return true;
+ } catch (Throwable throwable) {
+ logger.debug("Failed to read partition file.", throwable);
+ isReleased = true;
+ throw throwable;
+ }
+ }
+
+ public boolean hasRemaining() {
+ return currentPartitionRemainingBytes > 0;
+ }
+
+ public synchronized boolean readAndSend(Queue<ByteBuf> bufferQueue, Consumer<ByteBuf> consumer)
+ throws IOException {
+ boolean hasReaming = hasRemaining();
+ boolean continueReading = hasReaming;
+ int numDataBuffers = 0;
+ while (continueReading) {
+
+ ByteBuf buffer = bufferQueue.poll();
+ // this is used for control bytebuf manually.
+ buffer.retain();
+ if (buffer == null) {
+ break;
+ }
+
+ try {
+ continueReading = readBuffer(buffer);
+ } catch (Throwable throwable) {
+ memoryManager.recycleReadBuffer(buffer);
+ throw throwable;
+ }
+
+ hasReaming = hasRemaining();
+ addBuffer(buffer, hasReaming, consumer);
+ ++numDataBuffers;
}
+ if (numDataBuffers > 0) {
+ notifyBacklog(numDataBuffers);
+ }
+
+ if (!hasReaming) {
+ closeReader();
+ }
+
+ return hasReaming;
+ }
+
+ protected void addBuffer(ByteBuf buffer, boolean hasRemaining, Consumer<ByteBuf> consumer) {
+ if (buffer == null) {
+ return;
+ }
+ final boolean recycleBuffer;
+ boolean notifyDataAvailable = false;
+ final Throwable throwable;
+ synchronized (lock) {
+ recycleBuffer = isReleased || isFinished || isError;
+ throwable = errorCause;
+ isFinished = !hasRemaining;
+
+ if (!recycleBuffer) {
+ notifyDataAvailable = buffersRead.isEmpty();
+ buffersRead.add(new Buffer(buffer, consumer));
+ }
+ }
+
+ if (recycleBuffer) {
+ memoryManager.recycleReadBuffer(buffer);
+ throw new RuntimeException("Partition reader has been failed or finished.", throwable);
+ }
+ if (!buffersRead.isEmpty()) {
+ sendData();
+ }
+ }
+
+ public synchronized void sendData() {
+ while (!buffersRead.isEmpty()) {
+ logger.debug("senddata streamid:{}, {}", streamId, streamCredits.get(streamId));
+ if (streamCredits.get(streamId).get() > 0) {
+ Buffer readBuf = buffersRead.poll();
+ logger.debug(
+ "send {} to stream {} ,fileInfo: {}",
+ readBuf.byteBuf.readableBytes(),
+ streamId,
+ fileInfo.getFilePath());
+ ReadData readData = new ReadData(streamId, buffersRead.size(), 0, readBuf.byteBuf);
+ streams
+ .get(streamId)
+ .associatedChannel
+ .writeAndFlush(readData)
+ .addListener(
+ (ChannelFutureListener)
+ future -> {
+ if (!future.isSuccess()) {
+ isError = true;
+ errorCause = future.cause();
+ memoryManager.recycleReadBuffer(readBuf.byteBuf);
+ throw new RuntimeException(future.cause());
+ } else {
+ // netty has release the reference
+ }
+ logger.debug(
+ "recycle {} , {} to stream {} ,fileinfo {}",
+ readBuf.byteBuf.readableBytes(),
+ System.identityHashCode(readBuf.byteBuf),
+ streamId,
+ fileInfo.getFilePath());
+ readBuf.byteBufferConsumer.accept(readBuf.byteBuf);
+ });
+ logger.debug(
+ "streamId: {}, decrement credit: {}",
+ streamId,
+ streamCredits.get(streamId).decrementAndGet());
+ } else {
+ logger.debug("streamId: {}, no credit: {}", streamId, streamCredits.get(streamId));
+ // no credit
+ break;
+ }
+ }
+
+ logger.debug(
+ "streamId: {}, remaining: {}, bufferSize: {}",
+ streamId,
+ currentPartitionRemainingBytes,
+ buffersRead.size());
+
+ if (isClosed && buffersRead.isEmpty()) {
+ cleanResource(streamId);
+ }
+ }
+
+ public void closeReader() {
+ isClosed = true;
+ IOUtils.closeQuietly(indexFileChannel);
+ IOUtils.closeQuietly(dataFileChannel);
+ logger.debug("Closed read for stream {}", this.streamId);
+ }
+
+ protected void notifyBacklog(int backlog) {
+ logger.debug("stream manager streamid {} backlog:{}", streamId, backlog);
+ StreamState streamState = streams.get(streamId);
+ if (streamState == null) {
+ throw new RuntimeException("StreamId " + streamId + " should not be null");
+ }
+ streamState.associatedChannel.writeAndFlush(new BacklogAnnouncement(streamId, backlog));
+ }
+
+ public long getPriority() {
+ return dataConsumingOffset;
+ }
+
+ @Override
+ public int compareTo(DataPartitionReader that) {
+ return Long.compare(getPriority(), that.getPriority());
+ }
+ }
+
+ class StorageFetcherPool {
+ private final HashMap<String, ExecutorService> executorPools = new HashMap<>();
+
+ public ExecutorService getExecutorPool(String mountPoint) {
+ // it's ok if the mountpoint is unknown
+ return executorPools.computeIfAbsent(
+ mountPoint,
+ k ->
+ Executors.newFixedThreadPool(
+ threadsPerMountPoint,
+ new ThreadFactoryBuilder()
+ .setNameFormat("reader-thread-%d")
+ .setUncaughtExceptionHandler(
+ (t1, t2) -> {
+ logger.warn("StorageFetcherPool thread:{}:{}", t1, t2);
+ })
+ .build()));
}
}
}
diff --git a/common/src/main/java/org/apache/celeborn/common/network/server/memory/ReadBufferDispatcher.java b/common/src/main/java/org/apache/celeborn/common/network/server/memory/ReadBufferDispatcher.java
index 0ea4d332..ca76a25f 100644
--- a/common/src/main/java/org/apache/celeborn/common/network/server/memory/ReadBufferDispatcher.java
+++ b/common/src/main/java/org/apache/celeborn/common/network/server/memory/ReadBufferDispatcher.java
@@ -47,8 +47,8 @@ public class ReadBufferDispatcher extends Thread {
}
public void recycle(ByteBuf buf) {
- buf.release();
int bufferSize = buf.capacity();
+ buf.release();
memoryManager.changeReadBufferCounter(-1 * bufferSize);
}
@@ -67,9 +67,13 @@ public class ReadBufferDispatcher extends Thread {
while (buffers.size() < request.getMin()) {
if (memoryManager.readBufferAvailable(bufferSize)) {
memoryManager.incrementDiskBuffer(bufferSize);
- buffers.add(readBufferAllocator.buffer(bufferSize, bufferSize));
+ ByteBuf buf = readBufferAllocator.buffer(bufferSize, bufferSize);
+ buf.retain();
+ buffers.add(buf);
} else {
try {
+ // If dispatcher can not allocate minimum buffers, it will wait here until necessary
+ // buffers are get.
Thread.sleep(3);
} catch (InterruptedException e) {
logger.info("Buffer dispatcher is closing");
@@ -81,7 +85,9 @@ public class ReadBufferDispatcher extends Thread {
while (memoryManager.readBufferAvailable(request.getBufferSize())
&& buffers.size() < request.getMax()) {
memoryManager.changeReadBufferCounter(bufferSize);
- buffers.add(readBufferAllocator.buffer(bufferSize, bufferSize));
+ ByteBuf buf = readBufferAllocator.buffer(bufferSize, bufferSize);
+ buf.retain();
+ buffers.add(buf);
}
request.getBufferListener().notifyBuffers(buffers, null);
} else {
diff --git a/common/src/main/java/org/apache/celeborn/common/network/util/TransportConf.java b/common/src/main/java/org/apache/celeborn/common/network/util/TransportConf.java
index 2971096d..a75372db 100644
--- a/common/src/main/java/org/apache/celeborn/common/network/util/TransportConf.java
+++ b/common/src/main/java/org/apache/celeborn/common/network/util/TransportConf.java
@@ -22,13 +22,13 @@ import org.apache.celeborn.common.CelebornConf;
/** A central location that tracks all the settings we expose to users. */
public class TransportConf {
- private final CelebornConf conf;
+ private final CelebornConf celebornConf;
private final String module;
- public TransportConf(String module, CelebornConf conf) {
+ public TransportConf(String module, CelebornConf celebornConf) {
this.module = module;
- this.conf = conf;
+ this.celebornConf = celebornConf;
}
public String getModuleName() {
@@ -37,42 +37,42 @@ public class TransportConf {
/** IO mode: nio or epoll */
public String ioMode() {
- return conf.networkIoMode(module);
+ return celebornConf.networkIoMode(module);
}
/** If true, we will prefer allocating off-heap byte buffers within Netty. */
public boolean preferDirectBufs() {
- return conf.networkIoPreferDirectBufs(module);
+ return celebornConf.networkIoPreferDirectBufs(module);
}
/** Connect timeout in milliseconds. Default 10 secs. */
public int connectTimeoutMs() {
- return conf.networkIoConnectTimeoutMs(module);
+ return celebornConf.networkIoConnectTimeoutMs(module);
}
/** Connection active timeout in milliseconds. Default 240 secs. */
public int connectionTimeoutMs() {
- return conf.networkIoConnectionTimeoutMs(module);
+ return celebornConf.networkIoConnectionTimeoutMs(module);
}
/** Number of concurrent connections between two nodes for fetching data. */
public int numConnectionsPerPeer() {
- return conf.networkIoNumConnectionsPerPeer(module);
+ return celebornConf.networkIoNumConnectionsPerPeer(module);
}
/** Requested maximum length of the queue of incoming connections. Default 0 for no backlog. */
public int backLog() {
- return conf.networkIoBacklog(module);
+ return celebornConf.networkIoBacklog(module);
}
/** Number of threads used in the server thread pool. Default to 0, which is 2x#cores. */
public int serverThreads() {
- return conf.networkIoServerThreads(module);
+ return celebornConf.networkIoServerThreads(module);
}
/** Number of threads used in the client thread pool. Default to 0, which is 2x#cores. */
public int clientThreads() {
- return conf.networkIoClientThreads(module);
+ return celebornConf.networkIoClientThreads(module);
}
/**
@@ -81,12 +81,12 @@ public class TransportConf {
* buffer size should be ~ 1.25MB
*/
public int receiveBuf() {
- return conf.networkIoReceiveBuf(module);
+ return celebornConf.networkIoReceiveBuf(module);
}
/** Send buffer size (SO_SNDBUF). */
public int sendBuf() {
- return conf.networkIoSendBuf(module);
+ return celebornConf.networkIoSendBuf(module);
}
/**
@@ -94,7 +94,7 @@ public class TransportConf {
* to 0, we will not do any retries.
*/
public int maxIORetries() {
- return conf.networkIoMaxRetries(module);
+ return celebornConf.networkIoMaxRetries(module);
}
/**
@@ -102,7 +102,7 @@ public class TransportConf {
* relevant if maxIORetries > 0.
*/
public int ioRetryWaitTimeMs() {
- return conf.networkIoRetryWaitMs(module);
+ return celebornConf.networkIoRetryWaitMs(module);
}
/**
@@ -111,7 +111,7 @@ public class TransportConf {
* memory mapping has high overhead for blocks close to or below the page size of the OS.
*/
public int memoryMapBytes() {
- return conf.networkIoMemoryMapBytes(module);
+ return celebornConf.networkIoMemoryMapBytes(module);
}
/**
@@ -119,7 +119,7 @@ public class TransportConf {
* when data is going to be transferred. This can reduce the number of open files.
*/
public boolean lazyFileDescriptor() {
- return conf.networkIoLazyFileDescriptor(module);
+ return celebornConf.networkIoLazyFileDescriptor(module);
}
/**
@@ -127,7 +127,7 @@ public class TransportConf {
* PoolByteBufAllocator will be gotten, otherwise only general memory usage will be tracked.
*/
public boolean verboseMetrics() {
- return conf.networkIoVerboseMetrics(module);
+ return celebornConf.networkIoVerboseMetrics(module);
}
/**
@@ -138,10 +138,14 @@ public class TransportConf {
* failure.
*/
public long maxChunksBeingTransferred() {
- return conf.networkIoMaxChunksBeingTransferred(module);
+ return celebornConf.networkIoMaxChunksBeingTransferred(module);
+ }
+
+ public CelebornConf getCelebornConf() {
+ return celebornConf;
}
public long pushDataTimeoutCheckIntervalMs() {
- return conf.pushTimeoutCheckInterval();
+ return celebornConf.pushTimeoutCheckInterval();
}
}
diff --git a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
index 59ae4bcd..a3febb8b 100644
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@@ -752,6 +752,11 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable with Logging with Se
get(WORKER_DIRECT_MEMORY_RATIO_FOR_SHUFFLE_STORAGE)
def memoryPerResultPartition: String = get(MEMORY_PER_RESULT_PARTITION)
+ def partitionReadBuffersMin: Int = get(WORKER_PARTITION_READ_BUFFERS_MIN)
+
+ def partitionReadBuffersMax: Int = get(WORKER_PARTITION_READ_BUFFERS_MAX)
+ def bufferStreamThreadsPerMountpoint: Int = get(WORKER_BUFFERSTREAM_THREADS_PER_MOUNTPOINT)
+
// //////////////////////////////////////////////////////
// Rate Limit controller //
// //////////////////////////////////////////////////////
@@ -2935,4 +2940,28 @@ object CelebornConf extends Logging {
.doc("The size of network buffers required per result partition. The minimum valid value is 8M. Usually, several hundreds of megabytes memory is enough for large scale batch jobs.")
.stringConf
.createWithDefault("64m")
+
+ val WORKER_PARTITION_READ_BUFFERS_MIN: ConfigEntry[Int] =
+ buildConf("celeborn.worker.partition.initial.readBuffersMin")
+ .categories("worker")
+ .version("0.3.0")
+ .doc("Min number of initial read buffers")
+ .intConf
+ .createWithDefault(8)
+
+ val WORKER_PARTITION_READ_BUFFERS_MAX: ConfigEntry[Int] =
+ buildConf("celeborn.worker.partition.initial.readBuffersMax")
+ .categories("worker")
+ .version("0.3.0")
+ .doc("Max number of initial read buffers")
+ .intConf
+ .createWithDefault(8)
+
+ val WORKER_BUFFERSTREAM_THREADS_PER_MOUNTPOINT: ConfigEntry[Int] =
+ buildConf("celeborn.worker.bufferStream.threadsPerMountpoint")
+ .categories("worker")
+ .version("0.3.0")
+ .doc("Threads count for read buffer per mount point.")
+ .intConf
+ .createWithDefault(8)
}
diff --git a/common/src/main/scala/org/apache/celeborn/common/util/Utils.scala b/common/src/main/scala/org/apache/celeborn/common/util/Utils.scala
index f0696e13..951afa8d 100644
--- a/common/src/main/scala/org/apache/celeborn/common/util/Utils.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/util/Utils.scala
@@ -22,6 +22,7 @@ import java.lang.management.ManagementFactory
import java.math.{MathContext, RoundingMode}
import java.net._
import java.nio.ByteBuffer
+import java.nio.channels.FileChannel
import java.nio.charset.StandardCharsets
import java.text.SimpleDateFormat
import java.util
@@ -44,9 +45,8 @@ import org.apache.celeborn.common.exception.CelebornException
import org.apache.celeborn.common.internal.Logging
import org.apache.celeborn.common.meta.{DiskStatus, FileInfo, WorkerInfo}
import org.apache.celeborn.common.network.protocol.TransportMessage
-import org.apache.celeborn.common.network.util.{ConfigProvider, JavaUtils, TransportConf}
+import org.apache.celeborn.common.network.util.{JavaUtils, TransportConf}
import org.apache.celeborn.common.protocol.{PartitionLocation, PartitionSplitMode, PartitionType}
-import org.apache.celeborn.common.protocol.PbWorkerResource
import org.apache.celeborn.common.protocol.message.{ControlMessages, Message, StatusCode}
import org.apache.celeborn.common.protocol.message.ControlMessages.WorkerResource
@@ -994,11 +994,28 @@ object Utils extends Logging {
}
}
+ def checkedDownCast(value: Long): Int = {
+ val downCast = value.toInt
+ if (downCast.toLong != value) {
+ throw new IllegalArgumentException("Cannot downcast long value " + value + " to integer.")
+ }
+ downCast
+ }
+
+ @throws[IOException]
+ def checkFileIntegrity(fileChannel: FileChannel, length: Int): Unit = {
+ val remainingBytes = fileChannel.size - fileChannel.position
+ if (remainingBytes < length) {
+ logError(
+ s"File remaining bytes not not enough, remaining: ${remainingBytes}, wanted: ${length}.")
+ throw new RuntimeException(s"File is corrupted ${fileChannel}")
+ }
+ }
+
def getShortFormattedFileName(fileInfo: FileInfo): String = {
val parentFile = fileInfo.getFile.getParent
parentFile.substring(
parentFile.lastIndexOf("/"),
parentFile.length) + "/" + fileInfo.getFile.getName
}
-
}
diff --git a/docs/configuration/worker.md b/docs/configuration/worker.md
index 0c7493b4..2fdafe8f 100644
--- a/docs/configuration/worker.md
+++ b/docs/configuration/worker.md
@@ -32,6 +32,7 @@ license: |
| celeborn.storage.hdfs.dir | <undefined> | HDFS dir configuration for Celeborn to access HDFS. | 0.2.0 |
| celeborn.test.pushMasterDataTimeout | false | Whether to test push master data timeout | 0.2.0 |
| celeborn.test.pushSlaveDataTimeout | false | Whether to test push slave data timeout | 0.3.0 |
+| celeborn.worker.bufferStream.threadsPerMountpoint | 8 | Threads count for read buffer per mount point. | 0.3.0 |
| celeborn.worker.closeIdleConnections | false | Whether worker will close idle connections. | 0.2.0 |
| celeborn.worker.commit.threads | 32 | Thread number of worker to commit shuffle data files asynchronously. | 0.2.0 |
| celeborn.worker.congestionControl.enabled | false | Whether to enable congestion control or not. | 0.3.0 |
@@ -75,6 +76,8 @@ license: |
| celeborn.worker.monitor.disk.notifyError.threshold | 64 | Device monitor will only notify critical error once the accumulated valid non-critical error number exceeding this threshold. | 0.3.0 |
| celeborn.worker.monitor.disk.sys.block.dir | /sys/block | The directory where linux file block information is stored. | 0.2.0 |
| celeborn.worker.noneEmptyDirExpireDuration | 1d | If a non-empty application shuffle data dir have not been operated during le duration time, will mark this application as expired. | 0.2.0 |
+| celeborn.worker.partition.initial.readBuffersMax | 8 | Max number of initial read buffers | 0.3.0 |
+| celeborn.worker.partition.initial.readBuffersMin | 8 | Min number of initial read buffers | 0.3.0 |
| celeborn.worker.partitionSorter.directMemoryRatioThreshold | 0.1 | Max ratio of partition sorter's memory for sorting, when reserved memory is higher than max partition sorter memory, partition sorter will stop sorting. | 0.2.0 |
| celeborn.worker.partitionSorter.reservedMemoryPerPartition | 1mb | Reserved memory when sorting a shuffle file off-heap. | 0.2.0 |
| celeborn.worker.partitionSorter.sort.timeout | 220s | Timeout for a shuffle file to sort. | 0.2.0 |
diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/FetchHandler.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/FetchHandler.scala
index 36f1629a..bd045be3 100644
--- a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/FetchHandler.scala
+++ b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/FetchHandler.scala
@@ -41,7 +41,10 @@ import org.apache.celeborn.service.deploy.worker.storage.{PartitionFilesSorter,
class FetchHandler(val conf: TransportConf) extends BaseMessageHandler with Logging {
var chunkStreamManager = new ChunkStreamManager()
- val bufferStreamManager = new BufferStreamManager()
+ val bufferStreamManager = new BufferStreamManager(
+ conf.getCelebornConf.partitionReadBuffersMin,
+ conf.getCelebornConf.partitionReadBuffersMax,
+ conf.getCelebornConf.bufferStreamThreadsPerMountpoint)
var workerSource: WorkerSource = _
var rpcSource: RPCSource = _
var storageManager: StorageManager = _
@@ -78,14 +81,20 @@ class FetchHandler(val conf: TransportConf) extends BaseMessageHandler with Logg
rpcSource.updateMessageMetrics(r, 0)
handleChunkFetchRequest(client, r)
case r: RpcRequest =>
- handleOpenStream(client, r)
+ handleRpcs(client, r)
case unknown: RequestMessage =>
throw new IllegalArgumentException(s"Unknown message type id: ${unknown.`type`.id}")
}
}
- def handleOpenStream(client: TransportClient, request: RpcRequest): Unit = {
+ // here are BackLogAnnouncement,OpenStream and OpenStreamWithCredit RPCs to handle
+ def handleRpcs(client: TransportClient, request: RpcRequest): Unit = {
val msg = Message.decode(request.body().nioByteBuffer())
+ if (msg.`type`() == Type.BACKLOG_ANNOUNCEMENT) {
+ rpcSource.updateMessageMetrics(msg, 0)
+ handleReadAddCredit(client, msg.asInstanceOf[ReadAddCredit])
+ return
+ }
val (shuffleKey, fileName) =
if (msg.`type`() == Type.OPEN_STREAM) {
val openStream = msg.asInstanceOf[OpenStream]
@@ -136,16 +145,20 @@ class FetchHandler(val conf: TransportConf) extends BaseMessageHandler with Logg
new NioManagedBuffer(streamHandle.toByteBuffer)))
}
case PartitionType.MAP =>
- // return stream id
+ val initialCredit = msg.asInstanceOf[OpenStreamWithCredit].initialCredit
val startIndex = msg.asInstanceOf[OpenStreamWithCredit].startIndex
val endIndex = msg.asInstanceOf[OpenStreamWithCredit].endIndex
val streamId =
- bufferStreamManager.registerStream(client.getChannel, fileInfo.getBufferSize)
- val res = ByteBuffer.allocate(8)
- res.putLong(streamId)
+ bufferStreamManager.registerStream(
+ client.getChannel,
+ initialCredit,
+ startIndex,
+ endIndex,
+ fileInfo)
+ val bufferStreamHandle = new StreamHandle(streamId, 0)
client.getChannel.writeAndFlush(new RpcResponse(
request.requestId,
- new NioManagedBuffer(res)))
+ new NioManagedBuffer(bufferStreamHandle.toByteBuffer)))
case PartitionType.MAPGROUP =>
} catch {
case e: IOException =>
diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
index 8003182d..6e57baa0 100644
--- a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
+++ b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
@@ -339,7 +339,7 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
}
}
val fileInfo = new FileInfo(file.getAbsolutePath, userIdentifier, partitionType)
-
+ fileInfo.setMountPoint(mountPoint)
val fileWriter = partitionType match {
case PartitionType.MAP => new MapPartitionFileWriter(
fileInfo,