This is an automated email from the ASF dual-hosted git repository.
zhouky 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 505ba804c [CELEBORN-752] Support read local shuffle file for spark
505ba804c is described below
commit 505ba804c7d7890922e09aaacd2b155c70328a35
Author: mingji <[email protected]>
AuthorDate: Wed Aug 30 18:52:18 2023 +0800
[CELEBORN-752] Support read local shuffle file for spark
### What changes were proposed in this pull request?
For spark clusters, support read local shuffle file if Celeborn is
co-deployed with yarn node managers. This PR help to reduce the number of
active connections.
### Why are the changes needed?
Ditto.
### Does this PR introduce _any_ user-facing change?
NO.
### How was this patch tested?
GA and cluster. The performance is identical whether you enable local
reader, but the active connection number may vary according to your connections
per peer.
<img width="951" alt="截屏2023-08-16 20 20 14"
src="https://github.com/apache/incubator-celeborn/assets/4150993/9106e731-28fc-4e78-9c05-ae6a269d249a";>
The active connection number changed from 3745 to 2894. This PR will help
to improve cluster stability.
Closes #1812 from FMX/CELEBORN-752.
Authored-by: mingji <[email protected]>
Signed-off-by: zky.zhoukeyong <[email protected]>
---
.../org/apache/celeborn/client/ShuffleClient.java | 22 +-
.../celeborn/client/read/CelebornInputStream.java | 38 +++-
.../celeborn/client/read/DfsPartitionReader.java | 1 +
.../celeborn/client/read/LocalPartitionReader.java | 232 +++++++++++++++++++++
.../client/read/WorkerPartitionReader.java | 2 +
.../celeborn/common/util}/FileChannelUtils.java | 2 +-
common/src/main/proto/TransportMessages.proto | 2 +-
.../org/apache/celeborn/common/CelebornConf.scala | 17 ++
.../org/apache/celeborn/common/util/Utils.scala | 4 +-
docs/configuration/client.md | 2 +
.../service/deploy/worker/storage/FileWriter.java | 5 +-
.../deploy/worker/storage/MapDataPartition.java | 1 +
.../worker/storage/MapPartitionFileWriter.java | 13 +-
.../worker/storage/ReducePartitionFileWriter.java | 24 ++-
.../service/deploy/worker/FetchHandler.scala | 42 +++-
.../cluster/ClusterReadWriteTestWithLZ4.scala | 4 +
.../cluster/ClusterReadWriteTestWithZSTD.scala | 4 +
.../service/deploy/cluster/ReadWriteTestBase.scala | 3 +-
18 files changed, 374 insertions(+), 44 deletions(-)
diff --git a/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
b/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
index 30f7af8fb..99025e037 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
@@ -19,6 +19,7 @@ package org.apache.celeborn.client;
import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.LongAdder;
import org.apache.hadoop.fs.FileSystem;
import org.slf4j.Logger;
@@ -37,10 +38,12 @@ import org.apache.celeborn.common.write.PushState;
* implementation
*/
public abstract class ShuffleClient {
+ private static Logger logger = LoggerFactory.getLogger(ShuffleClient.class);
private static volatile ShuffleClient _instance;
private static volatile boolean initialized = false;
private static volatile FileSystem hdfsFs;
- private static Logger logger = LoggerFactory.getLogger(ShuffleClient.class);
+ private static LongAdder totalReadCounter = new LongAdder();
+ private static LongAdder localShuffleReadCounter = new LongAdder();
// for testing
public static void reset() {
@@ -95,6 +98,23 @@ public abstract class ShuffleClient {
return hdfsFs;
}
+ public static void incrementLocalReadCounter() {
+ localShuffleReadCounter.increment();
+ totalReadCounter.increment();
+ }
+
+ public static void incrementTotalReadCounter() {
+ totalReadCounter.increment();
+ }
+
+ public static String getReadCounters() {
+ long totalReadCount = totalReadCounter.longValue();
+ long localReadCount = localShuffleReadCounter.longValue();
+ return String.format(
+ "Current client read %d(local)/%d(total) partitions, local ratio %.2f",
+ localReadCount, totalReadCount, (localReadCount * 1.0d /
totalReadCount) * 100);
+ }
+
public abstract void setupLifecycleManagerRef(String host, int port);
public abstract void setupLifecycleManagerRef(RpcEndpointRef endpointRef);
diff --git
a/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
b/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
index 29b4abcee..8645cc4c1 100644
---
a/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
+++
b/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
@@ -31,6 +31,7 @@ import org.roaringbitmap.RoaringBitmap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.client.compress.Decompressor;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.exception.CelebornIOException;
@@ -129,6 +130,8 @@ public abstract class CelebornInputStream extends
InputStream {
private final byte[] sizeBuf = new byte[BATCH_HEADER_SIZE];
private LongAdder skipCount = new LongAdder();
private final boolean rangeReadFilter;
+ private final boolean enabledReadLocalShuffle;
+ private final String localHostAddress;
private boolean pushReplicateEnabled;
private boolean fetchExcludeWorkerOnFailureEnabled;
@@ -156,6 +159,8 @@ public abstract class CelebornInputStream extends
InputStream {
this.startMapIndex = startMapIndex;
this.endMapIndex = endMapIndex;
this.rangeReadFilter = conf.shuffleRangeReadFilterEnabled();
+ this.enabledReadLocalShuffle = conf.enableReadLocalShuffleFile();
+ this.localHostAddress = Utils.localHostName(conf);
this.pushReplicateEnabled = conf.clientPushReplicateEnabled();
this.fetchExcludeWorkerOnFailureEnabled =
conf.clientFetchExcludeWorkerOnFailureEnabled();
this.shuffleCompressionEnabled =
@@ -389,20 +394,30 @@ public abstract class CelebornInputStream extends
InputStream {
logger.debug("Read peer {} for attempt {}.", location, attemptNumber);
}
- logger.debug("create reader for location {}", location);
+ logger.debug("Create reader for location {}", location);
StorageInfo storageInfo = location.getStorageInfo();
if (storageInfo.getType() == StorageInfo.Type.HDD
|| storageInfo.getType() == StorageInfo.Type.SSD) {
- return new WorkerPartitionReader(
- conf,
- shuffleKey,
- location,
- clientFactory,
- startMapIndex,
- endMapIndex,
- fetchChunkRetryCnt,
- fetchChunkMaxRetry);
+ logger.debug(
+ "Read local shuffle file enabled {} , {}, {}",
+ enabledReadLocalShuffle,
+ location.getWorker().host(),
+ localHostAddress);
+ if (enabledReadLocalShuffle &&
location.getWorker().host().equals(localHostAddress)) {
+ return new LocalPartitionReader(
+ conf, shuffleKey, location, clientFactory, startMapIndex,
endMapIndex);
+ } else {
+ return new WorkerPartitionReader(
+ conf,
+ shuffleKey,
+ location,
+ clientFactory,
+ startMapIndex,
+ endMapIndex,
+ fetchChunkRetryCnt,
+ fetchChunkMaxRetry);
+ }
}
if (storageInfo.getType() == StorageInfo.Type.HDFS) {
return new DfsPartitionReader(
@@ -484,6 +499,9 @@ public abstract class CelebornInputStream extends
InputStream {
currentReader.close();
currentReader = null;
}
+ if (enabledReadLocalShuffle) {
+ logger.info(ShuffleClient.getReadCounters());
+ }
}
private boolean moveToNextChunk() throws IOException {
diff --git
a/client/src/main/java/org/apache/celeborn/client/read/DfsPartitionReader.java
b/client/src/main/java/org/apache/celeborn/client/read/DfsPartitionReader.java
index d42219773..68f6308b3 100644
---
a/client/src/main/java/org/apache/celeborn/client/read/DfsPartitionReader.java
+++
b/client/src/main/java/org/apache/celeborn/client/read/DfsPartitionReader.java
@@ -170,6 +170,7 @@ public class DfsPartitionReader implements PartitionReader {
});
fetchThread.start();
logger.debug("Start dfs read on location {}", location);
+ ShuffleClient.incrementTotalReadCounter();
}
}
diff --git
a/client/src/main/java/org/apache/celeborn/client/read/LocalPartitionReader.java
b/client/src/main/java/org/apache/celeborn/client/read/LocalPartitionReader.java
new file mode 100644
index 000000000..1168c8d2c
--- /dev/null
+++
b/client/src/main/java/org/apache/celeborn/client/read/LocalPartitionReader.java
@@ -0,0 +1,232 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.celeborn.client.read;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.FileChannel;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.*;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+import io.netty.util.ReferenceCounted;
+import org.apache.commons.lang3.StringUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.client.ShuffleClient;
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.exception.CelebornIOException;
+import org.apache.celeborn.common.network.client.TransportClient;
+import org.apache.celeborn.common.network.client.TransportClientFactory;
+import org.apache.celeborn.common.network.protocol.TransportMessage;
+import org.apache.celeborn.common.protocol.MessageType;
+import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.protocol.PbOpenStream;
+import org.apache.celeborn.common.protocol.PbStreamHandler;
+import org.apache.celeborn.common.util.FileChannelUtils;
+import org.apache.celeborn.common.util.ThreadUtils;
+
+public class LocalPartitionReader implements PartitionReader {
+
+ private static final Logger logger =
LoggerFactory.getLogger(LocalPartitionReader.class);
+ private static volatile ThreadPoolExecutor readLocalShufflePool;
+ private final LinkedBlockingQueue<ByteBuf> results;
+ private final AtomicReference<IOException> exception = new
AtomicReference<>();
+ private final int fetchMaxReqsInFlight;
+ private final PartitionLocation location;
+ private volatile boolean closed = false;
+ private final int numChunks;
+ private int returnedChunks = 0;
+ private int chunkIndex = 0;
+ private final FileChannel shuffleChannel;
+ private List<Long> chunkOffsets;
+ private AtomicBoolean pendingFetchTask = new AtomicBoolean(false);
+
+ public LocalPartitionReader(
+ CelebornConf conf,
+ String shuffleKey,
+ PartitionLocation location,
+ TransportClientFactory clientFactory,
+ int startMapIndex,
+ int endMapIndex)
+ throws IOException {
+ if (readLocalShufflePool == null) {
+ synchronized (LocalPartitionReader.class) {
+ if (readLocalShufflePool == null) {
+ readLocalShufflePool =
+ ThreadUtils.newDaemonCachedThreadPool(
+ "local-shuffle-reader-thread",
conf.readLocalShuffleThreads(), 60);
+ }
+ }
+ }
+ fetchMaxReqsInFlight = conf.clientFetchMaxReqsInFlight();
+ results = new LinkedBlockingQueue<>();
+ this.location = location;
+ PbStreamHandler streamHandle;
+ long fetchTimeoutMs = conf.clientFetchTimeoutMs();
+ try {
+ TransportClient client =
+ clientFactory.createClient(location.getHost(),
location.getFetchPort(), 0);
+ TransportMessage openStreamMsg =
+ new TransportMessage(
+ MessageType.OPEN_STREAM,
+ PbOpenStream.newBuilder()
+ .setShuffleKey(shuffleKey)
+ .setFileName(location.getFileName())
+ .setStartIndex(startMapIndex)
+ .setEndIndex(endMapIndex)
+ .setReadLocalShuffle(true)
+ .build()
+ .toByteArray());
+ ByteBuffer response = client.sendRpcSync(openStreamMsg.toByteBuffer(),
fetchTimeoutMs);
+ streamHandle =
TransportMessage.fromByteBuffer(response).getParsedPayload();
+ } catch (IOException | InterruptedException e) {
+ throw new IOException(
+ "Read shuffle file from local file failed, partition location: "
+ + location
+ + " filePath: "
+ + location.getStorageInfo().getFilePath(),
+ e);
+ }
+
+ chunkOffsets = new ArrayList<>(streamHandle.getChunkOffsetsList());
+ numChunks = streamHandle.getNumChunks();
+ shuffleChannel =
FileChannelUtils.openReadableFileChannel(streamHandle.getFullPath());
+ if (endMapIndex != Integer.MAX_VALUE) {
+ shuffleChannel.position(chunkOffsets.get(0));
+ }
+
+ logger.debug(
+ "Local partition reader {} offsets:{}",
+ location.getStorageInfo().getFilePath(),
+ StringUtils.join(chunkOffsets, ","));
+
+ ShuffleClient.incrementLocalReadCounter();
+ }
+
+ private void doFetchChunks(int chunkIndex, int toFetch) {
+ try {
+ for (int i = 0; i < toFetch; i++) {
+ long offset = chunkOffsets.get(chunkIndex + i);
+ long length = chunkOffsets.get(chunkIndex + i + 1) - offset;
+ logger.debug("Read {} offset {} length {}", chunkIndex, offset,
length);
+ // A chunk must be smaller than INT.MAX_VALUE
+ ByteBuffer buffer = ByteBuffer.allocate((int) length);
+ while (buffer.hasRemaining()) {
+ if (-1 == shuffleChannel.read(buffer)) {
+ throw new CelebornIOException(
+ "Read local file " + location.getStorageInfo().getFilePath() +
" failed");
+ }
+ }
+ buffer.flip();
+ // Avoid resource leak
+ synchronized (this) {
+ if (!closed) {
+ results.put(Unpooled.wrappedBuffer(buffer));
+ logger.debug("Add index {} to results", chunkIndex + i);
+ }
+ }
+ }
+ } catch (InterruptedException e) {
+ // cancel a task for speculative, ignore this exception
+ logger.warn("Read thread is interrupted.", e);
+ } catch (Exception ioe) {
+ logger.error("Read thread encountered error.", ioe);
+ if (ioe instanceof CelebornIOException) {
+ exception.set((IOException) ioe);
+ } else {
+ exception.set(new CelebornIOException("Read thread encountered error",
ioe));
+ }
+ }
+ pendingFetchTask.compareAndSet(true, false);
+ }
+
+ private void fetchChunks() {
+ int inFlight = chunkIndex - returnedChunks;
+ if (inFlight < fetchMaxReqsInFlight) {
+ int toFetch = Math.min(fetchMaxReqsInFlight - inFlight + 1, numChunks -
chunkIndex);
+ if (pendingFetchTask.compareAndSet(false, true)) {
+ logger.debug(
+ "Trigger local reader fetch chunk with {} and fetch {} chunks",
chunkIndex, toFetch);
+ int currentIndex = chunkIndex;
+ readLocalShufflePool.submit(() -> doFetchChunks(currentIndex,
toFetch));
+ chunkIndex += toFetch;
+ }
+ }
+ }
+
+ @Override
+ public boolean hasNext() {
+ logger.debug("Check has next current index: {} chunks {}", returnedChunks,
numChunks);
+ return returnedChunks < numChunks;
+ }
+
+ @Override
+ public ByteBuf next() throws IOException, InterruptedException {
+ checkException();
+ if (chunkIndex < numChunks) {
+ fetchChunks();
+ }
+ ByteBuf chunk = null;
+ try {
+ while (chunk == null) {
+ checkException();
+ chunk = results.poll(100, TimeUnit.MILLISECONDS);
+ logger.debug("Poll result with result size: {}", results.size());
+ }
+ } catch (InterruptedException e) {
+ logger.error("PartitionReader thread interrupted while fetching data.");
+ throw e;
+ }
+ returnedChunks++;
+ return chunk;
+ }
+
+ private void checkException() throws IOException {
+ IOException e = exception.get();
+ if (e != null) {
+ throw e;
+ }
+ }
+
+ @Override
+ public void close() {
+ synchronized (this) {
+ closed = true;
+ if (!results.isEmpty()) {
+ results.forEach(ReferenceCounted::release);
+ }
+ results.clear();
+ }
+ try {
+ shuffleChannel.close();
+ } catch (IOException e) {
+ logger.warn("Close local shuffle file failed.", e);
+ }
+ }
+
+ @Override
+ public PartitionLocation getLocation() {
+ return location;
+ }
+}
diff --git
a/client/src/main/java/org/apache/celeborn/client/read/WorkerPartitionReader.java
b/client/src/main/java/org/apache/celeborn/client/read/WorkerPartitionReader.java
index c02c7dd4f..aa73297de 100644
---
a/client/src/main/java/org/apache/celeborn/client/read/WorkerPartitionReader.java
+++
b/client/src/main/java/org/apache/celeborn/client/read/WorkerPartitionReader.java
@@ -28,6 +28,7 @@ import io.netty.util.ReferenceCounted;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.exception.CelebornIOException;
import org.apache.celeborn.common.network.buffer.ManagedBuffer;
@@ -125,6 +126,7 @@ public class WorkerPartitionReader implements
PartitionReader {
this.fetchChunkRetryCnt = fetchChunkRetryCnt;
this.fetchChunkMaxRetry = fetchChunkMaxRetry;
testFetch = conf.testFetchFailure();
+ ShuffleClient.incrementTotalReadCounter();
}
public boolean hasNext() {
diff --git
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/FileChannelUtils.java
b/common/src/main/java/org/apache/celeborn/common/util/FileChannelUtils.java
similarity index 95%
rename from
worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/FileChannelUtils.java
rename to
common/src/main/java/org/apache/celeborn/common/util/FileChannelUtils.java
index 3f08c78d9..ea89a6782 100644
---
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/FileChannelUtils.java
+++ b/common/src/main/java/org/apache/celeborn/common/util/FileChannelUtils.java
@@ -15,7 +15,7 @@
* limitations under the License.
*/
-package org.apache.celeborn.service.deploy.worker.storage;
+package org.apache.celeborn.common.util;
import java.io.IOException;
import java.nio.channels.FileChannel;
diff --git a/common/src/main/proto/TransportMessages.proto
b/common/src/main/proto/TransportMessages.proto
index b303e6482..4c6ea7cc5 100644
--- a/common/src/main/proto/TransportMessages.proto
+++ b/common/src/main/proto/TransportMessages.proto
@@ -483,7 +483,7 @@ message PbOpenStream {
int32 startIndex = 3;
int32 endIndex = 4;
int32 initialCredit = 5;
- bool localRead = 6;
+ bool readLocalShuffle = 6;
}
message PbStreamHandler {
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 916d0806b..46e1c88a0 100644
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@@ -823,6 +823,8 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable
with Logging with Se
get(CLIENT_BATCH_HANDLE_RELEASE_PARTITION_THREADS)
def batchHandleReleasePartitionRequestInterval: Long =
get(CLIENT_BATCH_HANDLED_RELEASE_PARTITION_INTERVAL)
+ def enableReadLocalShuffleFile: Boolean = get(READ_LOCAL_SHUFFLE_FILE)
+ def readLocalShuffleThreads: Int = get(READ_LOCAL_SHUFFLE_THREADS)
// //////////////////////////////////////////////////////
// Worker //
@@ -3799,4 +3801,19 @@ object CelebornConf extends Logging {
.transform(_.toUpperCase(Locale.ROOT))
.createWithDefault("HDD,SSD")
+ val READ_LOCAL_SHUFFLE_FILE: ConfigEntry[Boolean] =
+ buildConf("celeborn.client.readLocalShuffleFile.enabled")
+ .categories("client")
+ .version("0.3.1")
+ .doc("Enable read local shuffle file for clusters that co-deployed with
yarn node manager.")
+ .booleanConf
+ .createWithDefault(false)
+
+ val READ_LOCAL_SHUFFLE_THREADS: ConfigEntry[Int] =
+ buildConf("celeborn.client.readLocalShuffleFile.threads")
+ .categories("client")
+ .version("0.3.1")
+ .doc("Threads count for read local shuffle file.")
+ .intConf
+ .createWithDefault(4)
}
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 4b5cfa05a..1b5717bbf 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
@@ -25,7 +25,7 @@ import java.nio.ByteBuffer
import java.nio.channels.FileChannel
import java.nio.charset.StandardCharsets
import java.util
-import java.util.{Locale, Properties, Random, UUID}
+import java.util.{ArrayList, List, Locale, Properties, Random, UUID}
import java.util.concurrent.{Callable, ThreadPoolExecutor, TimeoutException,
TimeUnit}
import scala.annotation.tailrec
@@ -38,7 +38,7 @@ import scala.util.control.{ControlThrowable, NonFatal}
import com.google.protobuf.{ByteString, GeneratedMessageV3}
import io.netty.channel.unix.Errors.NativeIoException
import org.apache.commons.lang3.SystemUtils
-import org.apache.hadoop.fs.Path
+import org.apache.hadoop.fs.{FSDataInputStream, Path}
import org.roaringbitmap.RoaringBitmap
import org.apache.celeborn.common.CelebornConf
diff --git a/docs/configuration/client.md b/docs/configuration/client.md
index 028650c7f..9c288848c 100644
--- a/docs/configuration/client.md
+++ b/docs/configuration/client.md
@@ -62,6 +62,8 @@ license: |
| celeborn.client.push.takeTaskMaxWaitAttempts | 1 | Max wait times if no task
available to push to worker. | 0.3.0 |
| celeborn.client.push.takeTaskWaitInterval | 50ms | Wait interval if no task
available to push to worker. | 0.3.0 |
| celeborn.client.push.timeout | 120s | Timeout for a task to push data rpc
message. This value should better be more than twice of
`celeborn.<module>.push.timeoutCheck.interval` | 0.3.0 |
+| celeborn.client.readLocalShuffleFile.enabled | false | Enable read local
shuffle file for clusters that co-deployed with yarn node manager. | 0.3.1 |
+| celeborn.client.readLocalShuffleFile.threads | 4 | Threads count for read
local shuffle file. | 0.3.1 |
| celeborn.client.registerShuffle.maxRetries | 3 | Max retry times for client
to register shuffle. | 0.3.0 |
| celeborn.client.registerShuffle.retryWait | 3s | Wait time before next retry
if register shuffle failed. | 0.3.0 |
| celeborn.client.requestCommitFiles.maxRetries | 4 | Max retry times for
requestCommitFiles RPC. | 0.3.0 |
diff --git
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/FileWriter.java
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/FileWriter.java
index 7054ad108..e1a8ef8a8 100644
---
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/FileWriter.java
+++
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/FileWriter.java
@@ -41,6 +41,7 @@ import org.apache.celeborn.common.protocol.PartitionSplitMode;
import org.apache.celeborn.common.protocol.PartitionType;
import org.apache.celeborn.common.protocol.StorageInfo;
import org.apache.celeborn.common.unsafe.Platform;
+import org.apache.celeborn.common.util.FileChannelUtils;
import org.apache.celeborn.service.deploy.worker.WorkerSource;
import
org.apache.celeborn.service.deploy.worker.congestcontrol.CongestionController;
import org.apache.celeborn.service.deploy.worker.memory.MemoryManager;
@@ -282,9 +283,7 @@ public abstract class FileWriter implements DeviceObserver {
if (channel != null) {
channel.close();
}
- if (fileInfo.isHdfs()) {
- streamClose.run();
- }
+ streamClose.run();
} catch (IOException e) {
logger.warn("close file writer {} failed", this, e);
}
diff --git
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/MapDataPartition.java
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/MapDataPartition.java
index e02e2a362..c423aacb0 100644
---
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/MapDataPartition.java
+++
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/MapDataPartition.java
@@ -37,6 +37,7 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.meta.FileInfo;
+import org.apache.celeborn.common.util.FileChannelUtils;
import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.service.deploy.worker.memory.BufferQueue;
import org.apache.celeborn.service.deploy.worker.memory.BufferRecycler;
diff --git
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/MapPartitionFileWriter.java
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/MapPartitionFileWriter.java
index 848b5f32b..98ba4a13b 100644
---
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/MapPartitionFileWriter.java
+++
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/MapPartitionFileWriter.java
@@ -34,6 +34,7 @@ import org.apache.celeborn.common.meta.FileInfo;
import org.apache.celeborn.common.metrics.source.AbstractSource;
import org.apache.celeborn.common.protocol.PartitionSplitMode;
import org.apache.celeborn.common.protocol.PartitionType;
+import org.apache.celeborn.common.util.FileChannelUtils;
import org.apache.celeborn.common.util.Utils;
/*
@@ -130,11 +131,13 @@ public final class MapPartitionFileWriter extends
FileWriter {
flushIndex();
},
() -> {
- if
(StorageManager.hadoopFs().exists(fileInfo.getHdfsPeerWriterSuccessPath())) {
- StorageManager.hadoopFs().delete(fileInfo.getHdfsPath(), false);
- deleted = true;
- } else {
-
StorageManager.hadoopFs().create(fileInfo.getHdfsWriterSuccessPath()).close();
+ if (fileInfo.isHdfs()) {
+ if
(StorageManager.hadoopFs().exists(fileInfo.getHdfsPeerWriterSuccessPath())) {
+ StorageManager.hadoopFs().delete(fileInfo.getHdfsPath(), false);
+ deleted = true;
+ } else {
+
StorageManager.hadoopFs().create(fileInfo.getHdfsWriterSuccessPath()).close();
+ }
}
},
() -> {
diff --git
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/ReducePartitionFileWriter.java
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/ReducePartitionFileWriter.java
index 06db73b0a..d9b89849b 100644
---
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/ReducePartitionFileWriter.java
+++
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/ReducePartitionFileWriter.java
@@ -95,18 +95,20 @@ public final class ReducePartitionFileWriter extends
FileWriter {
}
},
() -> {
- if
(StorageManager.hadoopFs().exists(fileInfo.getHdfsPeerWriterSuccessPath())) {
- StorageManager.hadoopFs().delete(fileInfo.getHdfsPath(), false);
- deleted = true;
- } else {
-
StorageManager.hadoopFs().create(fileInfo.getHdfsWriterSuccessPath()).close();
- FSDataOutputStream indexOutputStream =
- StorageManager.hadoopFs().create(fileInfo.getHdfsIndexPath());
- indexOutputStream.writeInt(fileInfo.getChunkOffsets().size());
- for (Long offset : fileInfo.getChunkOffsets()) {
- indexOutputStream.writeLong(offset);
+ if (fileInfo.isHdfs()) {
+ if
(StorageManager.hadoopFs().exists(fileInfo.getHdfsPeerWriterSuccessPath())) {
+ StorageManager.hadoopFs().delete(fileInfo.getHdfsPath(), false);
+ deleted = true;
+ } else {
+
StorageManager.hadoopFs().create(fileInfo.getHdfsWriterSuccessPath()).close();
+ FSDataOutputStream indexOutputStream =
+
StorageManager.hadoopFs().create(fileInfo.getHdfsIndexPath());
+ indexOutputStream.writeInt(fileInfo.getChunkOffsets().size());
+ for (Long offset : fileInfo.getChunkOffsets()) {
+ indexOutputStream.writeLong(offset);
+ }
+ indexOutputStream.close();
}
- indexOutputStream.close();
}
},
() -> {});
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 624c44655..019c8871e 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
@@ -17,12 +17,14 @@
package org.apache.celeborn.service.deploy.worker
-import java.{lang, util}
import java.io.{FileNotFoundException, IOException}
import java.nio.charset.StandardCharsets
+import java.util
import java.util.concurrent.atomic.AtomicBoolean
import java.util.function.Consumer
+import scala.collection.JavaConverters.asScalaBufferConverter
+
import com.google.common.base.Throwables
import io.netty.util.concurrent.{Future, GenericFutureListener}
@@ -100,13 +102,14 @@ class FetchHandler(val conf: CelebornConf, val
transportConf: TransportConf)
try {
val pbMsg = TransportMessage.fromByteBuffer(r.body().nioByteBuffer())
val pbOpenStream = pbMsg.getParsedPayload[PbOpenStream]
- val (shuffleKey, fileName, startIndex, endIndex, initialCredit) =
+ val (shuffleKey, fileName, startIndex, endIndex, initialCredit,
readLocalShuffle) =
(
pbOpenStream.getShuffleKey,
pbOpenStream.getFileName,
pbOpenStream.getStartIndex,
pbOpenStream.getEndIndex,
- pbOpenStream.getInitialCredit)
+ pbOpenStream.getInitialCredit,
+ pbOpenStream.getReadLocalShuffle)
streamShuffleKey = shuffleKey
streamFileName = fileName
workerSource.startTimer(WorkerSource.OPEN_STREAM_TIME,
streamShuffleKey)
@@ -118,7 +121,8 @@ class FetchHandler(val conf: CelebornConf, val
transportConf: TransportConf)
endIndex,
initialCredit,
r,
- false)
+ false,
+ readLocalShuffle)
} catch {
case _: Exception =>
// process legacy OpenStream RPCs
@@ -182,7 +186,8 @@ class FetchHandler(val conf: CelebornConf, val
transportConf: TransportConf)
endIndex: Int,
initialCredit: Int,
request: RpcRequest,
- isLegacy: Boolean): Unit = {
+ isLegacy: Boolean,
+ readLocalShuffle: Boolean = false): Unit = {
try {
var fileInfo = getRawFileInfo(shuffleKey, fileName)
fileInfo.getPartitionType match {
@@ -198,7 +203,16 @@ class FetchHandler(val conf: CelebornConf, val
transportConf: TransportConf)
logDebug(s"Received chunk fetch request $shuffleKey $fileName
$startIndex " +
s"$endIndex get file info $fileInfo from client channel " +
s"${NettyUtils.getRemoteAddress(client.getChannel)}")
- if (fileInfo.isHdfs) {
+ if (readLocalShuffle) {
+ replyStreamHandler(
+ client,
+ request.requestId,
+ -1,
+ fileInfo.numChunks(),
+ isLegacy,
+ fileInfo.getChunkOffsets,
+ fileInfo.getFilePath)
+ } else if (fileInfo.isHdfs) {
replyStreamHandler(client, request.requestId, 0, 0, isLegacy)
} else {
val buffers = new FileManagedBuffers(fileInfo, transportConf)
@@ -245,18 +259,28 @@ class FetchHandler(val conf: CelebornConf, val
transportConf: TransportConf)
requestId: Long,
streamId: Long,
numChunks: Int,
- isLegacy: Boolean): Unit = {
+ isLegacy: Boolean,
+ offsets: util.List[java.lang.Long] = null,
+ filepath: String = ""): Unit = {
if (isLegacy) {
client.getChannel.writeAndFlush(new RpcResponse(
requestId,
new NioManagedBuffer(new StreamHandle(streamId,
numChunks).toByteBuffer)))
} else {
+ val pbStreamHandlerBuilder =
PbStreamHandler.newBuilder.setStreamId(streamId).setNumChunks(
+ numChunks)
+ if (offsets != null) {
+ pbStreamHandlerBuilder.addAllChunkOffsets(offsets)
+ }
+ if (filepath.nonEmpty) {
+ pbStreamHandlerBuilder.setFullPath(filepath)
+ }
+ val pbStreamHandler = pbStreamHandlerBuilder.build()
client.getChannel.writeAndFlush(new RpcResponse(
requestId,
new NioManagedBuffer(new TransportMessage(
MessageType.STREAM_HANDLER,
- PbStreamHandler.newBuilder.setStreamId(streamId).setNumChunks(
- numChunks).build.toByteArray).toByteBuffer)))
+ pbStreamHandler.toByteArray).toByteBuffer)))
}
}
diff --git
a/worker/src/test/scala/org/apache/celeborn/service/deploy/cluster/ClusterReadWriteTestWithLZ4.scala
b/worker/src/test/scala/org/apache/celeborn/service/deploy/cluster/ClusterReadWriteTestWithLZ4.scala
index cfc14b1d2..bc142e1cc 100644
---
a/worker/src/test/scala/org/apache/celeborn/service/deploy/cluster/ClusterReadWriteTestWithLZ4.scala
+++
b/worker/src/test/scala/org/apache/celeborn/service/deploy/cluster/ClusterReadWriteTestWithLZ4.scala
@@ -25,4 +25,8 @@ class ClusterReadWriteTestWithLZ4 extends ReadWriteTestBase {
testReadWriteByCode(CompressionCodec.LZ4)
}
+ test(s"test MiniCluster With LZ4 and local read for spark") {
+ testReadWriteByCode(CompressionCodec.LZ4, true)
+ }
+
}
diff --git
a/worker/src/test/scala/org/apache/celeborn/service/deploy/cluster/ClusterReadWriteTestWithZSTD.scala
b/worker/src/test/scala/org/apache/celeborn/service/deploy/cluster/ClusterReadWriteTestWithZSTD.scala
index 8d207cb25..e2fe77928 100644
---
a/worker/src/test/scala/org/apache/celeborn/service/deploy/cluster/ClusterReadWriteTestWithZSTD.scala
+++
b/worker/src/test/scala/org/apache/celeborn/service/deploy/cluster/ClusterReadWriteTestWithZSTD.scala
@@ -25,4 +25,8 @@ class ClusterReadWriteTestWithZSTD extends ReadWriteTestBase {
testReadWriteByCode(CompressionCodec.ZSTD)
}
+ test(s"test MiniCluster With ZSTD and local read for spark") {
+ testReadWriteByCode(CompressionCodec.ZSTD, true)
+ }
+
}
diff --git
a/worker/src/test/scala/org/apache/celeborn/service/deploy/cluster/ReadWriteTestBase.scala
b/worker/src/test/scala/org/apache/celeborn/service/deploy/cluster/ReadWriteTestBase.scala
index 433f89ea1..e5e3c50c7 100644
---
a/worker/src/test/scala/org/apache/celeborn/service/deploy/cluster/ReadWriteTestBase.scala
+++
b/worker/src/test/scala/org/apache/celeborn/service/deploy/cluster/ReadWriteTestBase.scala
@@ -51,7 +51,7 @@ trait ReadWriteTestBase extends AnyFunSuite
shutdownMiniCluster()
}
- def testReadWriteByCode(codec: CompressionCodec): Unit = {
+ def testReadWriteByCode(codec: CompressionCodec, readLocalShuffle: Boolean =
false): Unit = {
val APP = "app-1"
val clientConf = new CelebornConf()
@@ -59,6 +59,7 @@ trait ReadWriteTestBase extends AnyFunSuite
.set(CelebornConf.SHUFFLE_COMPRESSION_CODEC.key, codec.name)
.set(CelebornConf.CLIENT_PUSH_REPLICATE_ENABLED.key, "true")
.set(CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE.key, "256K")
+ .set(CelebornConf.READ_LOCAL_SHUFFLE_FILE, readLocalShuffle)
.set("celeborn.data.io.numConnectionsPerPeer", "1")
val lifecycleManager = new LifecycleManager(APP, clientConf)
val shuffleClient = new ShuffleClientImpl(APP, clientConf,
UserIdentifier("mock", "mock"))
