This is an automated email from the ASF dual-hosted git repository.
SteNicholas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new 759e7b547 [CELEBORN-2331] Parallelize batch open stream client creation
759e7b547 is described below
commit 759e7b547641139576f21ddecbf1db5b42f83b01
Author: Chao Sun <[email protected]>
AuthorDate: Tue May 26 10:21:44 2026 +0800
[CELEBORN-2331] Parallelize batch open stream client creation
## Why are the changes needed?
`CelebornShuffleReader` batches stream-open requests by worker, but it
previously created the data client for each worker serially before sending
those already-parallel batch requests. When a reducer reads from multiple
workers, connection setup for a slow or unavailable worker can delay useful
work against the remaining healthy workers.
Parallelizing this setup removes the worker-by-worker wait from the normal
path. Because this changes task-side connection scheduling, the optimization
also needs an operational fallback that restores the prior behavior without
requiring a code rollback.
## What changes were proposed in this PR?
The reader now first gathers pending stream-open locations by worker
address, then creates one data client per distinct worker concurrently using
the existing stream-creator pool. Once client setup completes, it sends the
existing `BATCH_OPEN_STREAM` requests only for workers with an available
client, allowing healthy workers to proceed even if another worker fails during
setup.
The client-creation phase preserves the prior retry behavior for later
locations on the same worker when an earlier client attempt fails. It also
handles task cancellation explicitly: if the waiting Spark task is interrupted,
it restores the interrupt status and cancels unfinished setup work; worker-side
interruption is propagated rather than treated as an ordinary retryable failure.
This optimization is controlled by
`celeborn.client.spark.batch.openStream.parallelClientCreation.enabled`, which
defaults to `true`. Setting it to `false` selects the original serial
client-creation and request-building flow, giving deployments a targeted
rollback switch if parallel connection setup causes unexpected operational
behavior.
## How was this PR tested?
- Unit tests for parallel client setup, failure/retry handling,
cancellation on interruption, and the new configuration default and override.
- Configuration documentation generation validation for the new client
setting.
- Spotless formatting validation.
Closes #3692 from sunchao/dev/chao/codex/port-pr72-to-oss-main.
Authored-by: Chao Sun <[email protected]>
Signed-off-by: SteNicholas <[email protected]>
---
.../shuffle/celeborn/CelebornShuffleReader.scala | 99 +++++++++++-
.../celeborn/CelebornShuffleReaderSuite.scala | 167 ++++++++++++++++++++-
.../org/apache/celeborn/common/CelebornConf.scala | 11 ++
.../apache/celeborn/common/CelebornConfSuite.scala | 9 ++
docs/configuration/client.md | 1 +
5 files changed, 285 insertions(+), 2 deletions(-)
diff --git
a/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala
b/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala
index 55e036155..155fc0886 100644
---
a/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala
+++
b/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala
@@ -24,6 +24,7 @@ import java.util.concurrent.atomic.AtomicReference
import java.util.function.BiFunction
import scala.collection.JavaConverters._
+import scala.collection.mutable
import com.google.common.annotations.VisibleForTesting
import org.apache.commons.lang3.tuple.Pair
@@ -205,6 +206,10 @@ class CelebornShuffleReader[K, C](
val partitionIdList = List.range(startPartition, endPartition).filter(p =>
fileGroups.partitionGroups.containsKey(p))
+ val parallelClientCreationEnabled =
conf.batchOpenStreamParallelClientCreationEnabled
+ val locationsByHostPort =
+ new mutable.LinkedHashMap[String, JArrayList[PartitionLocation]]()
+
def makeOpenStreamList(locations: JSet[PartitionLocation]): Unit = {
locations.asScala.foreach { location =>
partCnt += 1
@@ -241,6 +246,16 @@ class CelebornShuffleReader[K, C](
}
}
+ def groupOpenStreamLocations(locations: JSet[PartitionLocation]): Unit = {
+ locations.asScala.foreach { location =>
+ partCnt += 1
+ val hostPort = location.hostAndFetchPort
+ locationsByHostPort
+ .getOrElseUpdate(hostPort, new JArrayList[PartitionLocation]())
+ .add(location)
+ }
+ }
+
partitionIdList.foreach { partitionId =>
if (fileGroups.partitionGroups.containsKey(partitionId)) {
// CELEBORN-2032. For the first time of open stream and
@@ -272,7 +287,43 @@ class CelebornShuffleReader[K, C](
locations = filterLocations.asJava
}
partitionId2PartitionLocations.put(partitionId, locations)
- makeOpenStreamList(locations)
+ if (parallelClientCreationEnabled) {
+ groupOpenStreamLocations(locations)
+ } else {
+ makeOpenStreamList(locations)
+ }
+ }
+ }
+
+ if (parallelClientCreationEnabled) {
+ val clientsByHostPort = CelebornShuffleReader.createClientsInParallel(
+ locationsByHostPort.map { case (hostPort, locations) =>
+ (hostPort, locations.asScala.toSeq)
+ }.toSeq,
+ streamCreatorPool,
+ location =>
+ shuffleClient.getDataClientFactory().createClient(
+ location.getHost,
+ location.getFetchPort),
+ (hostPort, location, ex) => {
+ shuffleClient.excludeFailedFetchLocation(hostPort, ex)
+ logWarning(
+ s"Failed to create client for $shuffleKey-${location.getId} from
host: ${hostPort}. " +
+ s"Shuffle reader will try its replica if exists.")
+ })
+
+ clientsByHostPort.foreach { case (hostPort, client) =>
+ val locArr = locationsByHostPort(hostPort)
+ val pbOpenStreamList = PbOpenStreamList.newBuilder()
+ pbOpenStreamList.setShuffleKey(shuffleKey)
+ locArr.asScala.foreach { location =>
+ pbOpenStreamList.addFileName(location.getFileName)
+ .addStartIndex(startMapIndex)
+ .addEndIndex(endMapIndex)
+ pbOpenStreamList.addReadLocalShuffle(
+ localFetchEnabled && location.getHost.equals(localHostAddress))
+ }
+ workerRequestMap.put(hostPort, (client, locArr, pbOpenStreamList))
}
}
@@ -548,6 +599,52 @@ class CelebornShuffleReader[K, C](
object CelebornShuffleReader {
var streamCreatorPool: ThreadPoolExecutor = null
+
+ @VisibleForTesting
+ private[celeborn] def createClientsInParallel(
+ locationsByHostPort: Seq[(String, Seq[PartitionLocation])],
+ streamCreatorPool: ThreadPoolExecutor,
+ createClient: PartitionLocation => TransportClient,
+ onClientCreateFailure: (String, PartitionLocation, Exception) => Unit)
+ : Map[String, TransportClient] = {
+ val clientsByHostPort = JavaUtils.newConcurrentHashMap[String,
TransportClient]()
+ val futures = locationsByHostPort.map { case (hostPort, locations) =>
+ streamCreatorPool.submit(new Runnable {
+ override def run(): Unit = {
+ val locationsIterator = locations.iterator
+ var clientCreated = false
+ while (!clientCreated && locationsIterator.hasNext) {
+ val location = locationsIterator.next()
+ try {
+ clientsByHostPort.put(hostPort, createClient(location))
+ clientCreated = true
+ } catch {
+ case ex: InterruptedException =>
+ Thread.currentThread().interrupt()
+ throw ex
+ case ex: Exception =>
+ onClientCreateFailure(hostPort, location, ex)
+ }
+ }
+ }
+ })
+ }
+ var waitCompleted = false
+ try {
+ futures.foreach(_.get())
+ waitCompleted = true
+ } catch {
+ case ex: InterruptedException =>
+ Thread.currentThread().interrupt()
+ throw ex
+ } finally {
+ if (!waitCompleted) {
+ futures.foreach(_.cancel(true))
+ }
+ }
+ clientsByHostPort.asScala.toMap
+ }
+
// Register the deserializer for GetReducerFileGroupResponse broadcast
ShuffleClient.registerDeserializeReducerFileGroupResponseFunction(new
BiFunction[
Integer,
diff --git
a/client-spark/spark-3/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReaderSuite.scala
b/client-spark/spark-3/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReaderSuite.scala
index 29878fd76..d2cec3abf 100644
---
a/client-spark/spark-3/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReaderSuite.scala
+++
b/client-spark/spark-3/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReaderSuite.scala
@@ -17,8 +17,10 @@
package org.apache.spark.shuffle.celeborn
+import java.io.IOException
import java.nio.file.Files
-import java.util.concurrent.TimeoutException
+import java.util.concurrent.{CountDownLatch, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
import org.apache.spark.{Dependency, ShuffleDependency, TaskContext}
import org.apache.spark.shuffle.ShuffleReadMetricsReporter
@@ -31,6 +33,9 @@ import org.apache.celeborn.client.{DummyShuffleClient,
ShuffleClient}
import org.apache.celeborn.common.CelebornConf
import org.apache.celeborn.common.exception.CelebornIOException
import org.apache.celeborn.common.identity.UserIdentifier
+import org.apache.celeborn.common.network.client.TransportClient
+import org.apache.celeborn.common.protocol.PartitionLocation
+import org.apache.celeborn.common.util.ThreadUtils
class CelebornShuffleReaderSuite extends AnyFunSuite {
@@ -92,4 +97,164 @@ class CelebornShuffleReaderSuite extends AnyFunSuite {
shuffleReader.shuffleClient.asInstanceOf[DummyShuffleClient].fetchFailureCount.get()
=== 2)
}
+
+ test("create batch open stream clients in parallel per worker") {
+ val worker0 = newLocation(0, "worker-0", 19098)
+ val worker1 = newLocation(0, "worker-1", 19098)
+ val worker0Client = Mockito.mock(classOf[TransportClient])
+ val worker1Client = Mockito.mock(classOf[TransportClient])
+ val streamCreatorPool =
ThreadUtils.newDaemonCachedThreadPool("test-create-client", 2, 60)
+ val started = new CountDownLatch(2)
+ val release = new CountDownLatch(1)
+
+ try {
+ val clientsFuture = scala.concurrent.Future {
+ CelebornShuffleReader.createClientsInParallel(
+ Seq(
+ worker0.hostAndFetchPort -> Seq(worker0),
+ worker1.hostAndFetchPort -> Seq(worker1)),
+ streamCreatorPool,
+ location => {
+ started.countDown()
+ assert(started.await(5, TimeUnit.SECONDS))
+ assert(release.await(5, TimeUnit.SECONDS))
+ if (location eq worker0) worker0Client else worker1Client
+ },
+ (_, _, ex) => fail("Unexpected client creation failure", ex))
+ }(scala.concurrent.ExecutionContext.global)
+
+ assert(started.await(5, TimeUnit.SECONDS))
+ release.countDown()
+ val clients =
+ scala.concurrent.Await.result(
+ clientsFuture,
+ scala.concurrent.duration.Duration(5, "seconds"))
+
+ assert(clients(worker0.hostAndFetchPort) eq worker0Client)
+ assert(clients(worker1.hostAndFetchPort) eq worker1Client)
+ } finally {
+ streamCreatorPool.shutdownNow()
+ }
+ }
+
+ test("skip failed batch open stream client creation while keeping healthy
workers") {
+ val failedWorker = newLocation(0, "worker-0", 19098)
+ val healthyWorker = newLocation(0, "worker-1", 19098)
+ val healthyClient = Mockito.mock(classOf[TransportClient])
+ val streamCreatorPool =
ThreadUtils.newDaemonCachedThreadPool("test-create-client", 2, 60)
+ var failedHostPort: String = null
+
+ try {
+ val clients = CelebornShuffleReader.createClientsInParallel(
+ Seq(
+ failedWorker.hostAndFetchPort -> Seq(failedWorker),
+ healthyWorker.hostAndFetchPort -> Seq(healthyWorker)),
+ streamCreatorPool,
+ location => {
+ if (location eq failedWorker) throw new IOException("boom")
+ healthyClient
+ },
+ (hostPort, _, _) => failedHostPort = hostPort)
+
+ assert(failedHostPort === failedWorker.hostAndFetchPort)
+ assert(!clients.contains(failedWorker.hostAndFetchPort))
+ assert(clients(healthyWorker.hostAndFetchPort) eq healthyClient)
+ } finally {
+ streamCreatorPool.shutdownNow()
+ }
+ }
+
+ test("retry failed batch open stream client creation for the same worker") {
+ val failedLocation = newLocation(0, "worker-0", 19098)
+ val retryLocation = newLocation(1, "worker-0", 19098)
+ val client = Mockito.mock(classOf[TransportClient])
+ val streamCreatorPool =
ThreadUtils.newDaemonCachedThreadPool("test-create-client", 1, 60)
+ var failureCount = 0
+
+ try {
+ val clients = CelebornShuffleReader.createClientsInParallel(
+ Seq(failedLocation.hostAndFetchPort -> Seq(failedLocation,
retryLocation)),
+ streamCreatorPool,
+ location => {
+ if (location eq failedLocation) throw new IOException("boom")
+ client
+ },
+ (_, _, _) => failureCount += 1)
+
+ assert(failureCount === 1)
+ assert(clients(failedLocation.hostAndFetchPort) eq client)
+ } finally {
+ streamCreatorPool.shutdownNow()
+ }
+ }
+
+ test("cancel batch open stream client creation when waiting thread is
interrupted") {
+ val blockedLocation = newLocation(0, "worker-0", 19098)
+ val retryLocation = newLocation(1, "worker-0", 19098)
+ val retryClient = Mockito.mock(classOf[TransportClient])
+ val streamCreatorPool =
ThreadUtils.newDaemonCachedThreadPool("test-create-client", 1, 60)
+ val clientStarted = new CountDownLatch(1)
+ val releaseClient = new CountDownLatch(1)
+ val clientInterrupted = new CountDownLatch(1)
+ val retried = new AtomicBoolean(false)
+ val failureReported = new AtomicBoolean(false)
+ val callerInterrupted = new AtomicBoolean(false)
+ val callerFailure = new AtomicReference[Throwable]()
+
+ val caller = new Thread(new Runnable {
+ override def run(): Unit = {
+ try {
+ CelebornShuffleReader.createClientsInParallel(
+ Seq(blockedLocation.hostAndFetchPort -> Seq(blockedLocation,
retryLocation)),
+ streamCreatorPool,
+ location => {
+ if (location eq blockedLocation) {
+ clientStarted.countDown()
+ try {
+ releaseClient.await()
+ retryClient
+ } catch {
+ case ex: InterruptedException =>
+ clientInterrupted.countDown()
+ throw ex
+ }
+ } else {
+ retried.set(true)
+ retryClient
+ }
+ },
+ (_, _, _) => failureReported.set(true))
+ callerFailure.set(new AssertionError("Expected waiting thread to be
interrupted"))
+ } catch {
+ case _: InterruptedException =>
+ callerInterrupted.set(Thread.currentThread().isInterrupted)
+ case ex: Throwable =>
+ callerFailure.set(ex)
+ }
+ }
+ })
+
+ try {
+ caller.start()
+ assert(clientStarted.await(5, TimeUnit.SECONDS))
+ caller.interrupt()
+ caller.join(TimeUnit.SECONDS.toMillis(5))
+ assert(!caller.isAlive)
+ assert(clientInterrupted.await(5, TimeUnit.SECONDS))
+ streamCreatorPool.shutdown()
+ assert(streamCreatorPool.awaitTermination(5, TimeUnit.SECONDS))
+ assert(callerFailure.get() == null)
+ assert(callerInterrupted.get())
+ assert(!retried.get())
+ assert(!failureReported.get())
+ } finally {
+ releaseClient.countDown()
+ streamCreatorPool.shutdownNow()
+ caller.interrupt()
+ caller.join(TimeUnit.SECONDS.toMillis(5))
+ }
+ }
+
+ private def newLocation(id: Int, host: String, fetchPort: Int):
PartitionLocation =
+ new PartitionLocation(id, 0, host, 0, 0, fetchPort, 0,
PartitionLocation.Mode.PRIMARY)
}
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 40d06617f..32c4eb73d 100644
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@@ -1174,6 +1174,8 @@ class CelebornConf(loadDefaults: Boolean) extends
Cloneable with Logging with Se
def enableReadLocalShuffleFile: Boolean = get(READ_LOCAL_SHUFFLE_FILE)
def readLocalShuffleThreads: Int = get(READ_LOCAL_SHUFFLE_THREADS)
def readStreamCreatorPoolThreads: Int = get(READ_STREAM_CREATOR_POOL_THREADS)
+ def batchOpenStreamParallelClientCreationEnabled: Boolean =
+ get(CLIENT_SPARK_BATCH_OPEN_STREAM_PARALLEL_CLIENT_CREATION_ENABLED)
def registerShuffleFilterExcludedWorkerEnabled: Boolean =
get(REGISTER_SHUFFLE_FILTER_EXCLUDED_WORKER_ENABLED)
@@ -6335,6 +6337,15 @@ object CelebornConf extends Logging {
.intConf
.createWithDefault(32)
+ val CLIENT_SPARK_BATCH_OPEN_STREAM_PARALLEL_CLIENT_CREATION_ENABLED:
ConfigEntry[Boolean] =
+
buildConf("celeborn.client.spark.batch.openStream.parallelClientCreation.enabled")
+ .categories("client")
+ .version("0.6.3")
+ .doc("Whether to create data clients in parallel before sending Spark
batch open-stream requests. " +
+ "When false, data clients are created serially.")
+ .booleanConf
+ .createWithDefault(true)
+
val CLIENT_CHUNK_PREFETCH_ENABLED: ConfigEntry[Boolean] =
buildConf("celeborn.client.chunk.prefetch.enabled")
.categories("client")
diff --git
a/common/src/test/scala/org/apache/celeborn/common/CelebornConfSuite.scala
b/common/src/test/scala/org/apache/celeborn/common/CelebornConfSuite.scala
index 953255cde..5c0bf32df 100644
--- a/common/src/test/scala/org/apache/celeborn/common/CelebornConfSuite.scala
+++ b/common/src/test/scala/org/apache/celeborn/common/CelebornConfSuite.scala
@@ -492,4 +492,13 @@ class CelebornConfSuite extends CelebornFunSuite {
}
}
+ test("parallel Spark batch open stream client creation can be disabled") {
+ val conf = new CelebornConf()
+
+ assert(conf.batchOpenStreamParallelClientCreationEnabled)
+
+
conf.set(CLIENT_SPARK_BATCH_OPEN_STREAM_PARALLEL_CLIENT_CREATION_ENABLED.key,
"false")
+ assert(!conf.batchOpenStreamParallelClientCreationEnabled)
+ }
+
}
diff --git a/docs/configuration/client.md b/docs/configuration/client.md
index ece9503bd..52d3d984e 100644
--- a/docs/configuration/client.md
+++ b/docs/configuration/client.md
@@ -124,6 +124,7 @@ license: |
| celeborn.client.shuffle.reviseLostShuffles.enabled | false | false | Whether
to revise lost shuffles. | 0.6.0 | |
| celeborn.client.shuffleDataLostOnUnknownWorker.enabled | false | false |
Whether to mark shuffle data lost when unknown worker is detected. | 0.6.3 | |
| celeborn.client.slot.assign.maxWorkers | 10000 | false | Max workers that
slots of one shuffle can be allocated on. Will choose the smaller positive one
from Master side and Client side, see `celeborn.master.slot.assign.maxWorkers`.
| 0.3.1 | |
+| celeborn.client.spark.batch.openStream.parallelClientCreation.enabled | true
| false | Whether to create data clients in parallel before sending Spark batch
open-stream requests. When false, data clients are created serially. | 0.6.3 |
|
| celeborn.client.spark.fetch.cleanFailedShuffle | false | false | whether to
clean those disk space occupied by shuffles which cannot be fetched | 0.6.0 |
|
| celeborn.client.spark.fetch.cleanFailedShuffleInterval | 1s | false | the
interval to clean the failed-to-fetch shuffle files, only valid when
celeborn.client.spark.fetch.cleanFailedShuffle is enabled | 0.6.0 | |
| celeborn.client.spark.push.dynamicWriteMode.enabled | false | false |
Whether to dynamically switch push write mode based on conditions.If true,
shuffle mode will be only determined by partition count | 0.5.0 | |
