Skip to content

[CELEBORN-2371] Bound Spark batch-open client creation retries and stop them on interruption #3746

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sunchao wants to merge 2 commits into
base: main
Choose a base branch
from
Open

[CELEBORN-2371] Bound Spark batch-open client creation retries and stop them on interruption #3746

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8225f4f
[CELEBORN-2371] Bound Spark batch-open client creation retries and st…
sunchao Jun 25, 2026
311bc8c
[CELEBORN-2371] Address batch-open interruption review feedback
sunchao Jul 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 38 additions & 23 deletions ...park/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -215,21 +215,23 @@ class CelebornShuffleReader[K, C](
partCnt += 1
val hostPort = location.hostAndFetchPort
if (!workerRequestMap.containsKey(hostPort)) {
try {
val client = shuffleClient.getDataClientFactory().createClient(
location.getHost,
location.getFetchPort)
CelebornShuffleReader.tryCreateClient(

@SteNicholas SteNicholas Jul 8, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The retry-bounding fix (and the headOption change in createClientsInParallel) only covers the parallel client-creation path. The sequential makeOpenStreamList still attempts creation per partition-location: when tryCreateClient returns None, onClientCreateFailure only excludes + warns and never populates workerRequestMap, so the next location sharing the same hostPort re-enters if (!workerRequestMap.containsKey(hostPort)) (line 217) and calls createClient against the same dead endpoint again — the N × maxRetries blow-up this PR sets out to remove.

So with celeborn.client.spark.batch.openStream.parallelClientCreation.enabled=false, the stated goal ("Bound … client creation retries") isn't achieved. This isn't a regression (it matches the old behavior), but since the parallel path is now bounded it's an inconsistency worth either closing (e.g. record the failed hostPort so subsequent same-host locations short-circuit) or calling out as a known limitation.

location,
loc =>
shuffleClient.getDataClientFactory().createClient(
loc.getHost,
loc.getFetchPort),
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.")
}).foreach { client =>
val pbOpenStreamList = PbOpenStreamList.newBuilder()
pbOpenStreamList.setShuffleKey(shuffleKey)
workerRequestMap.put(
hostPort,
(client, new JArrayList[PartitionLocation], pbOpenStreamList))
} catch {
case ex: Exception =>
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.")
}
}
workerRequestMap.get(hostPort) match {
Expand Down Expand Up @@ -600,6 +602,26 @@ class CelebornShuffleReader[K, C](
object CelebornShuffleReader {
var streamCreatorPool: ThreadPoolExecutor = null

@VisibleForTesting
private[celeborn] def tryCreateClient(
location: PartitionLocation,
createClient: PartitionLocation => TransportClient,
onClientCreateFailure: Exception => Unit): Option[TransportClient] = {
if (Thread.currentThread().isInterrupted) {
throw new InterruptedException("Client creation interrupted")
}
try {
Some(createClient(location))
} catch {
case ex: InterruptedException =>

@SteNicholas SteNicholas Jul 8, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consistency: this catches only a bare InterruptedException. A wrapped one (IOException(InterruptedException)) falls through to case ex: Exception below and is reported as a worker failure via excludeFailedFetchLocation — the exact misclassification that TransportClientFactory.findInterruptedException and CelebornInputStream.isInterruption (both added in this PR) walk the cause chain to prevent. It's masked today only because the real createClient path goes through retryCreateClient, which already unwraps to a bare InterruptedException — a fragile coupling.

Note the three interrupt-detection helpers this PR introduces already disagree: findInterruptedException ignores the thread flag, isInterruption checks it first, and tryCreateClient checks neither the chain nor the flag mid-createClient. Consider extracting one shared helper (e.g. in common ExceptionUtils) and using it in all three sites so they can't drift.

Related latent gap in createClientsInParallel below: a pool task that rethrows InterruptedException surfaces at futures.foreach(_.get()) as a java.util.concurrent.ExecutionException, which the catch { case ex: InterruptedException } there won't match — so the re-interrupt is skipped and a raw ExecutionException escapes read(). It's hard to hit today (the waiting thread's own .get() usually throws InterruptedException first, and pool-thread flags are cleared between tasks), but unwrapping ExecutionException(InterruptedException) there would close it.

Thread.currentThread().interrupt()
throw ex
case ex: Exception =>
onClientCreateFailure(ex)
None
}
}

@VisibleForTesting
private[celeborn] def createClientsInParallel(
locationsByHostPort: Seq[(String, Seq[PartitionLocation])],
Expand All @@ -611,19 +633,12 @@ object CelebornShuffleReader {
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)
locations.headOption.foreach { location =>
tryCreateClient(
location,
createClient,
ex => onClientCreateFailure(hostPort, location, ex)).foreach { client =>
clientsByHostPort.put(hostPort, client)
}
}
}
Expand Down
101 changes: 69 additions & 32 deletions ...spark-3/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReaderSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ package org.apache.spark.shuffle.celeborn
import java.io.IOException
import java.nio.file.Files
import java.util.concurrent.{CountDownLatch, TimeoutException, TimeUnit}
import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger, AtomicReference}

import org.apache.spark.{Dependency, ShuffleDependency, TaskContext}
import org.apache.spark.shuffle.ShuffleReadMetricsReporter
Expand Down Expand Up @@ -164,39 +164,82 @@ class CelebornShuffleReaderSuite extends AnyFunSuite {
}
}

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])
test("attempt batch open stream client creation once per worker") {
val locations = (0 until 100).map(id => newLocation(id, "worker-0", 19098))
val streamCreatorPool = ThreadUtils.newDaemonCachedThreadPool("test-create-client", 1, 60)
var failureCount = 0
val attempts = new AtomicInteger()
val failures = new AtomicInteger()

try {
val clients = CelebornShuffleReader.createClientsInParallel(
Seq(failedLocation.hostAndFetchPort -> Seq(failedLocation, retryLocation)),
Seq(locations.head.hostAndFetchPort -> locations),
streamCreatorPool,
location => {
if (location eq failedLocation) throw new IOException("boom")
client
_ => {
attempts.incrementAndGet()
throw new IOException("boom")
},
(_, _, _) => failureCount += 1)
(_, _, _) => failures.incrementAndGet())

assert(failureCount === 1)
assert(clients(failedLocation.hostAndFetchPort) eq client)
assert(clients.isEmpty)
assert(attempts.get() === 1)
assert(failures.get() === 1)
} finally {
streamCreatorPool.shutdownNow()
}
}

test("propagate client creation interruption without reporting a worker failure") {
val location = newLocation(0, "worker-0", 19098)
val failureReported = new AtomicBoolean(false)

try {
val exception = intercept[InterruptedException] {
CelebornShuffleReader.tryCreateClient(
location,
_ => throw new InterruptedException("test"),
_ => failureReported.set(true))
}

assert(exception.getMessage === "test")
assert(Thread.currentThread().isInterrupted)
assert(!failureReported.get())
} finally {
Thread.interrupted()
}
}

test("skip client creation when the thread is already interrupted") {
val location = newLocation(0, "worker-0", 19098)
val clientCreationAttempted = new AtomicBoolean(false)
val failureReported = new AtomicBoolean(false)

try {
Thread.currentThread().interrupt()
intercept[InterruptedException] {
CelebornShuffleReader.tryCreateClient(
location,
_ => {
clientCreationAttempted.set(true)
Mockito.mock(classOf[TransportClient])
},
_ => failureReported.set(true))
}

assert(Thread.currentThread().isInterrupted)
assert(!clientCreationAttempted.get())
assert(!failureReported.get())
} finally {
Thread.interrupted()
}
}

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 client = 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]()
Expand All @@ -205,22 +248,17 @@ class CelebornShuffleReaderSuite extends AnyFunSuite {
override def run(): Unit = {
try {
CelebornShuffleReader.createClientsInParallel(
Seq(blockedLocation.hostAndFetchPort -> Seq(blockedLocation, retryLocation)),
Seq(blockedLocation.hostAndFetchPort -> Seq(blockedLocation)),
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
_ => {
clientStarted.countDown()
try {
releaseClient.await()
client
} catch {
case ex: InterruptedException =>
clientInterrupted.countDown()
throw ex
}
},
(_, _, _) => failureReported.set(true))
Expand All @@ -245,7 +283,6 @@ class CelebornShuffleReaderSuite extends AnyFunSuite {
assert(streamCreatorPool.awaitTermination(5, TimeUnit.SECONDS))
assert(callerFailure.get() == null)
assert(callerInterrupted.get())
assert(!retried.get())
assert(!failureReported.get())
} finally {
releaseClient.countDown()
Expand Down
40 changes: 40 additions & 0 deletions client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
import scala.Tuple2;

import com.github.luben.zstd.ZstdException;
import com.google.common.base.Throwables;
import com.google.common.util.concurrent.Uninterruptibles;
import io.netty.buffer.ByteBuf;
import net.jpountz.lz4.LZ4Exception;
Expand All @@ -48,6 +49,7 @@
import org.apache.celeborn.common.protocol.*;
import org.apache.celeborn.common.unsafe.Platform;
import org.apache.celeborn.common.util.ExceptionMaker;
import org.apache.celeborn.common.util.ExceptionUtils;
import org.apache.celeborn.common.util.Utils;
import org.apache.celeborn.common.write.LocationPushFailedBatches;

Expand Down Expand Up @@ -438,6 +440,23 @@ private boolean isExcluded(PartitionLocation location) {
}
}

private static boolean isInterruption(Throwable throwable) {
if (Thread.currentThread().isInterrupted()) {

@SteNicholas SteNicholas Jul 8, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

isInterruption returns true when either (a) the current thread's interrupt flag is set, or (b) any exception in throwable's cause chain is an InterruptedException. Both are broader than "this failure is a task cancellation."

Concern: a genuine, retryable fetch failure whose cause chain happens to wrap an InterruptedException (e.g. a transient Netty blocking-await interruption that got wrapped, without the task actually being cancelled) is now treated as an interruption in getNextChunk / createReaderWithRetry — the code closes the reader and throws without attempting the peer/replica, whereas the pre-PR path would have failed over. It also flips the live thread's interrupt flag via interruptedIOException. In the common cancellation case that's exactly right; the risk is the false-positive where an InterruptedException in the chain doesn't represent a real cancellation, which would strand a recoverable read on the primary and turn it into a task failure. Worth confirming that scenario can't arise from a non-cancellation source.

return true;
}
for (Throwable cause : Throwables.getCausalChain(throwable)) {
if (cause instanceof InterruptedException) {
return true;
}
}
return false;
}

private static IOException interruptedIOException(Exception exception) {
Thread.currentThread().interrupt();
return ExceptionUtils.wrapThrowableToIOException(exception);
}

private PartitionReader createReaderWithRetry(
PartitionLocation location, PbStreamHandler pbStreamHandler) throws IOException {
return createReaderWithRetry(location, pbStreamHandler, Optional.empty());
Expand All @@ -464,6 +483,9 @@ private PartitionReader createReaderWithRetry(
checkpointMetadata);
return reader;
} catch (Exception e) {
if (isInterruption(e)) {
throw interruptedIOException(e);
}
lastException = e;
shuffleClient.excludeFailedFetchLocation(location.hostAndFetchPort(), e);
fetchChunkRetryCnt++;
Expand Down Expand Up @@ -493,6 +515,9 @@ private PartitionReader createReaderWithRetry(
.toByteArray());
client.sendRpc(bufferStreamEnd.toByteBuffer());
} catch (InterruptedException | IOException | RuntimeException ex) {
if (isInterruption(ex)) {
throw interruptedIOException(ex);
}
logger.warn(
"Close {} stream {} failed",
location.hostAndFetchPort(),
Expand Down Expand Up @@ -528,6 +553,15 @@ private ByteBuf getNextChunk() throws IOException {
}
return currentReader.next();
} catch (Exception e) {
if (isInterruption(e)) {
IOException interrupted = interruptedIOException(e);
try {
currentReader.close();
} catch (RuntimeException closeException) {
interrupted.addSuppressed(closeException);
}
throw interrupted;
}
shuffleClient.excludeFailedFetchLocation(
currentReader.getLocation().hostAndFetchPort(), e);
fetchChunkRetryCnt++;
Expand Down Expand Up @@ -890,6 +924,9 @@ private boolean fillBuffer() throws IOException {
}
return hasData;
} catch (LZ4Exception | ZstdException | IOException e) {
if (isInterruption(e)) {
throw interruptedIOException(e);
}
logger.error(
"Failed to fill buffer from chunk. AppShuffleId {}, shuffleId {}, partitionId {}, location {}",
appShuffleId,
Expand Down Expand Up @@ -918,6 +955,9 @@ private boolean fillBuffer() throws IOException {
}
throw ioe;
} catch (Exception e) {
if (isInterruption(e)) {
throw interruptedIOException(e);
}
logger.error(
"Failed to fill buffer from chunk. AppShuffleId {}, shuffleId {}, partitionId {}, location {}",
appShuffleId,
Expand Down
Loading
Loading