You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@celeborn.apache.org by ni...@apache.org on 2026/05/25 05:59:47 UTC
(celeborn) branch main updated: [CELEBORN-2330] Fix HA master bootstrap redirect handling
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 1c807d99a [CELEBORN-2330] Fix HA master bootstrap redirect handling
1c807d99a is described below
commit 1c807d99a16e59fca0022b1524ec158bfa869514
Author: Chao Sun <ch...@openai.com>
AuthorDate: Mon May 25 13:59:36 2026 +0800
[CELEBORN-2330] Fix HA master bootstrap redirect handling
## Why are the changes needed?
In HA mode, clients can hit a failover window where the master they contact is no longer the leader and returns a `MasterNotLeaderException` with a suggested leader address. The client-facing symptom is an RPC failure surfaced as:
`CelebornException: Exception thrown in awaitResult`
The redirect signal is still present underneath that wrapper, but the existing bootstrap and retry logic does not consistently preserve and follow it. In particular:
- bootstrap-time redirects can be treated like generic connection failures instead of explicit leader hints
- suggested leaders can themselves redirect again or fail setup
- after such failures, the client may not continue cleanly to the remaining configured masters
The consequence is that a client can fail to establish a master connection during HA leader transitions even when a reachable leader or another configured master is available. That turns a recoverable redirect/failover event into an avoidable client-visible failure and makes rolling upgrades or leader changes noisier than necessary.
## What changes were proposed in this PR?
This port brings the HA redirect handling fix from `openai/celeborn#70` onto upstream `main`.
When a master tells the client which leader to use, the client now keeps that redirect information even if it is wrapped inside another exception, and it actively follows the suggested leader during bootstrap and failover. If that suggested leader points to another leader, the client can continue along that redirect chain instead of giving up too early.
If the redirect path is no longer useful, for example because the suggested leader cannot be reached, no leader is currently presented, or redirects start looping, the client falls back to the remaining configured masters and keeps searching for a usable endpoint. The retry scan also advances correctly after endpoint setup fails, so one bad redirect does not prevent the client from trying the next viable master.
The PR also adds focused HA tests that cover the bootstrap redirect cases, chained redirects, redirect cycles, and fallback to configured masters.
## How was this PR tested?
- `build/mvn test -pl common -am -Dtest=MasterClientSuiteJ -DwildcardSuites=org.apache.celeborn.common.client.__NoSuchSuite__`
Closes #3691 from sunchao/dev/chao/codex/port-pr70-to-oss-main.
Authored-by: Chao Sun <ch...@openai.com>
Signed-off-by: SteNicholas <pr...@163.com>
---
.../celeborn/common/client/MasterClient.java | 78 ++++++-
.../celeborn/common/client/MasterClientSuiteJ.java | 259 ++++++++++++++++++++-
2 files changed, 326 insertions(+), 11 deletions(-)
diff --git a/common/src/main/java/org/apache/celeborn/common/client/MasterClient.java b/common/src/main/java/org/apache/celeborn/common/client/MasterClient.java
index 626c112a5..e57460c49 100644
--- a/common/src/main/java/org/apache/celeborn/common/client/MasterClient.java
+++ b/common/src/main/java/org/apache/celeborn/common/client/MasterClient.java
@@ -169,16 +169,18 @@ public class MasterClient {
}
private boolean shouldRetry(@Nullable RpcEndpointRef oldRef, Throwable e) {
- // It will always throw celeborn exception , so we need to get the cause
- // 'CelebornException: Exception thrown in awaitResult'
- if (e.getCause() instanceof MasterNotLeaderException) {
- MasterNotLeaderException exception = (MasterNotLeaderException) e.getCause();
+ // A redirect can arrive wrapped by awaitResult or by endpoint setup during bootstrap.
+ // Search the full cause chain so both paths can retry against the suggested leader.
+ MasterNotLeaderException exception = findMasterNotLeaderException(e);
+ if (exception != null) {
String leaderAddr =
isWorker
? exception.getSuggestedInternalLeaderAddress()
: exception.getSuggestedLeaderAddress();
if (!leaderAddr.equals(MasterNotLeaderException.LEADER_NOT_PRESENTED)) {
- setRpcEndpointRef(leaderAddr);
+ if (!setRpcEndpointRef(leaderAddr)) {
+ resetRpcEndpointRef(oldRef);
+ }
} else {
LOG.warn("Master leader is not present currently, please check masters' status!");
resetRpcEndpointRef(oldRef);
@@ -191,13 +193,55 @@ public class MasterClient {
return false;
}
- private void setRpcEndpointRef(String masterEndpoint) {
+ @Nullable
+ private MasterNotLeaderException findMasterNotLeaderException(Throwable throwable) {
+ Throwable current = throwable;
+ while (current != null) {
+ if (current instanceof MasterNotLeaderException) {
+ return (MasterNotLeaderException) current;
+ }
+ current = current.getCause();
+ }
+ return null;
+ }
+
+ private boolean setRpcEndpointRef(String masterEndpoint) {
// This method should never care newer or old value, we just set the suggested master endpoint.
// If an error occurs when setting the suggested Master, it means that the Master may be down.
// At this time, we just set `rpcEndpointRef` to null. Then next time, we will re-select the
// Master and get the correct leader.
- rpcEndpointRef.set(setupEndpointRef(masterEndpoint));
- LOG.info("Fail over to master {}.", masterEndpoint);
+ String nextMasterEndpoint = masterEndpoint;
+ Set<String> triedMasterEndpoints = new HashSet<>();
+ while (triedMasterEndpoints.add(nextMasterEndpoint)) {
+ try {
+ RpcEndpointRef endpointRef = setupEndpointRef(nextMasterEndpoint);
+ if (endpointRef != null) {
+ rpcEndpointRef.set(endpointRef);
+ LOG.info("Fail over to master {}.", nextMasterEndpoint);
+ return true;
+ }
+ break;
+ } catch (RuntimeException e) {
+ MasterNotLeaderException exception = findMasterNotLeaderException(e);
+ if (exception == null) {
+ break;
+ }
+
+ String leaderAddr =
+ isWorker
+ ? exception.getSuggestedInternalLeaderAddress()
+ : exception.getSuggestedLeaderAddress();
+ if (MasterNotLeaderException.LEADER_NOT_PRESENTED.equals(leaderAddr)) {
+ break;
+ }
+ nextMasterEndpoint = leaderAddr;
+ }
+ }
+ rpcEndpointRef.set(null);
+ LOG.info(
+ "Fail over to master {} failed during endpoint setup; will retry with another master.",
+ masterEndpoint);
+ return false;
}
private void resetRpcEndpointRef(@Nullable RpcEndpointRef oldRef) {
@@ -225,6 +269,8 @@ public class MasterClient {
* @param currentIndex current attempt master address index.
* @throws IllegalStateException If after several attempts, the non-empty RpcEndpointRef still
* cannot be obtained.
+ * @throws RuntimeException If endpoint setup receives a leader redirect that should be retried by
+ * the outer send loop.
* @return non-empty RpcEndpointRef.
*/
private RpcEndpointRef getOrSetupRpcEndpointRef(AtomicInteger currentIndex) {
@@ -241,7 +287,13 @@ public class MasterClient {
if (endpointRef == null) {
int index = currentIndex.get();
do {
- RpcEndpointRef tempEndpointRef = setupEndpointRef(activeMasterEndpoints.get(index));
+ RpcEndpointRef tempEndpointRef;
+ try {
+ tempEndpointRef = setupEndpointRef(activeMasterEndpoints.get(index));
+ } catch (RuntimeException e) {
+ currentIndex.set((index + 1) % activeMasterEndpoints.size());
+ throw e;
+ }
if (rpcEndpointRef.compareAndSet(null, tempEndpointRef)) {
index = (index + 1) % activeMasterEndpoints.size();
}
@@ -269,6 +321,14 @@ public class MasterClient {
rpcEnv.setupEndpointRef(
RpcAddress.fromHostAndPort(endpoint), masterEndpointResolver.masterEndpointName());
} catch (Exception e) {
+ MasterNotLeaderException exception = findMasterNotLeaderException(e);
+ if (exception != null
+ && !MasterNotLeaderException.LEADER_NOT_PRESENTED.equals(
+ isWorker
+ ? exception.getSuggestedInternalLeaderAddress()
+ : exception.getSuggestedLeaderAddress())) {
+ throw new RuntimeException(e);
+ }
// Catch all exceptions. Because we don't care whether this exception is IOException or
// TimeoutException or other exceptions, so we just try to connect to host:port, if fail,
// we try next address.
diff --git a/common/src/test/java/org/apache/celeborn/common/client/MasterClientSuiteJ.java b/common/src/test/java/org/apache/celeborn/common/client/MasterClientSuiteJ.java
index bb74d99b2..26ad11c43 100644
--- a/common/src/test/java/org/apache/celeborn/common/client/MasterClientSuiteJ.java
+++ b/common/src/test/java/org/apache/celeborn/common/client/MasterClientSuiteJ.java
@@ -245,6 +245,257 @@ public class MasterClientSuiteJ {
checkOneMasterAskFailedInHA(new RpcTimeoutException("test", new TimeoutException("test")));
}
+ @Test
+ public void testBootstrapMasterNotLeaderRedirectsToSuggestedLeaderInHA() {
+ final CelebornConf conf =
+ conf()
+ .set(CelebornConf.HA_ENABLED().key(), "true")
+ .set(CelebornConf.MASTER_ENDPOINTS().key(), "host1:9097,host3:9097")
+ .set(CelebornConf.MASTER_CLIENT_MAX_RETRIES().key(), "5");
+
+ final RpcEndpointRef master2 = Mockito.mock(RpcEndpointRef.class);
+ Mockito.doReturn(Future$.MODULE$.successful(mockResponse))
+ .when(master2)
+ .ask(Mockito.any(), Mockito.any(), Mockito.any());
+
+ Mockito.doAnswer(
+ invocation -> {
+ RpcAddress address = invocation.getArgument(0, RpcAddress.class);
+ switch (address.host()) {
+ case "host1":
+ throw new RuntimeException(
+ new MasterNotLeaderException("host1:9097", "host2:9097", null));
+ case "host2":
+ return master2;
+ case "host3":
+ throw new AssertionError("Should follow the suggested leader before host3.");
+ default:
+ fail("Should use master host1/host2:" + masterPort + ", but use " + address);
+ }
+ return null;
+ })
+ .when(rpcEnv)
+ .setupEndpointRef(Mockito.any(RpcAddress.class), Mockito.anyString());
+
+ MasterClient client = new MasterClient(rpcEnv, conf, false);
+ HeartbeatFromWorker message = Mockito.mock(HeartbeatFromWorker.class);
+
+ HeartbeatFromWorkerResponse response = null;
+ try {
+ response = client.askSync(message, HeartbeatFromWorkerResponse.class);
+ } catch (Throwable t) {
+ LOG.error("It should redirect to the suggested leader during bootstrap.", t);
+ fail("It should redirect to the suggested leader during bootstrap.");
+ }
+
+ assertEquals(mockResponse, response);
+ }
+
+ @Test
+ public void testBootstrapMasterNotLeaderWithoutLeaderContinuesToOtherConfiguredMastersInHA() {
+ final CelebornConf conf = prepareForCelebornConfWithHA();
+
+ final RpcEndpointRef master2 = Mockito.mock(RpcEndpointRef.class);
+ Mockito.doReturn(Future$.MODULE$.successful(mockResponse))
+ .when(master2)
+ .ask(Mockito.any(), Mockito.any(), Mockito.any());
+
+ Mockito.doAnswer(
+ invocation -> {
+ RpcAddress address = invocation.getArgument(0, RpcAddress.class);
+ switch (address.host()) {
+ case "host1":
+ throw new RuntimeException(
+ new MasterNotLeaderException(
+ "host1:9097", MasterNotLeaderException.LEADER_NOT_PRESENTED, null));
+ case "host2":
+ return master2;
+ default:
+ fail("Should use master host1/host2:" + masterPort + ", but use " + address);
+ }
+ return null;
+ })
+ .when(rpcEnv)
+ .setupEndpointRef(Mockito.any(RpcAddress.class), Mockito.anyString());
+
+ MasterClient client = new MasterClient(rpcEnv, conf, false);
+ HeartbeatFromWorker message = Mockito.mock(HeartbeatFromWorker.class);
+
+ HeartbeatFromWorkerResponse response = null;
+ try {
+ response = client.askSync(message, HeartbeatFromWorkerResponse.class);
+ } catch (Throwable t) {
+ LOG.error("It should continue to another configured master during bootstrap.", t);
+ fail("It should continue to another configured master during bootstrap.");
+ }
+
+ assertEquals(mockResponse, response);
+ }
+
+ @Test
+ public void testSuggestedLeaderSetupFailureRetriesAnotherMasterInHA() {
+ final CelebornConf conf =
+ conf()
+ .set(CelebornConf.HA_ENABLED().key(), "true")
+ .set(CelebornConf.MASTER_ENDPOINTS().key(), "host1:9097,host4:9097")
+ .set(CelebornConf.MASTER_CLIENT_MAX_RETRIES().key(), "5");
+
+ final RpcEndpointRef master1 = Mockito.mock(RpcEndpointRef.class);
+ final RpcEndpointRef master3 = Mockito.mock(RpcEndpointRef.class);
+
+ Mockito.doReturn(
+ Future$.MODULE$.failed(new MasterNotLeaderException("host1:9097", "host2:9097", null)))
+ .when(master1)
+ .ask(Mockito.any(), Mockito.any(), Mockito.any());
+ Mockito.doReturn(Future$.MODULE$.successful(mockResponse))
+ .when(master3)
+ .ask(Mockito.any(), Mockito.any(), Mockito.any());
+
+ Mockito.doAnswer(
+ invocation -> {
+ RpcAddress address = invocation.getArgument(0, RpcAddress.class);
+ switch (address.host()) {
+ case "host1":
+ return master1;
+ case "host2":
+ throw new RuntimeException(
+ new MasterNotLeaderException("host2:9097", "host3:9097", null));
+ case "host3":
+ return master3;
+ case "host4":
+ throw new AssertionError("Should follow chained suggested leaders before host4.");
+ default:
+ fail(
+ "Should use master host1/host2/host3:" + masterPort + ", but use " + address);
+ }
+ return null;
+ })
+ .when(rpcEnv)
+ .setupEndpointRef(Mockito.any(RpcAddress.class), Mockito.anyString());
+
+ MasterClient client = new MasterClient(rpcEnv, conf, false);
+ HeartbeatFromWorker message = Mockito.mock(HeartbeatFromWorker.class);
+
+ HeartbeatFromWorkerResponse response = null;
+ try {
+ response = client.askSync(message, HeartbeatFromWorkerResponse.class);
+ } catch (Throwable t) {
+ LOG.error("It should retry with another master when suggested leader setup races.", t);
+ fail("It should retry with another master when suggested leader setup races.");
+ }
+
+ assertEquals(mockResponse, response);
+ }
+
+ @Test
+ public void testSuggestedLeaderRedirectCycleRetriesConfiguredMasterInHA() {
+ final CelebornConf conf =
+ conf()
+ .set(CelebornConf.HA_ENABLED().key(), "true")
+ .set(CelebornConf.MASTER_ENDPOINTS().key(), "host1:9097,host4:9097")
+ .set(CelebornConf.MASTER_CLIENT_MAX_RETRIES().key(), "5");
+
+ final RpcEndpointRef master1 = Mockito.mock(RpcEndpointRef.class);
+ final RpcEndpointRef master4 = Mockito.mock(RpcEndpointRef.class);
+
+ Mockito.doReturn(
+ Future$.MODULE$.failed(new MasterNotLeaderException("host1:9097", "host2:9097", null)))
+ .when(master1)
+ .ask(Mockito.any(), Mockito.any(), Mockito.any());
+ Mockito.doReturn(Future$.MODULE$.successful(mockResponse))
+ .when(master4)
+ .ask(Mockito.any(), Mockito.any(), Mockito.any());
+
+ Mockito.doAnswer(
+ invocation -> {
+ RpcAddress address = invocation.getArgument(0, RpcAddress.class);
+ switch (address.host()) {
+ case "host1":
+ return master1;
+ case "host2":
+ throw new RuntimeException(
+ new MasterNotLeaderException("host2:9097", "host3:9097", null));
+ case "host3":
+ throw new RuntimeException(
+ new MasterNotLeaderException("host3:9097", "host2:9097", null));
+ case "host4":
+ return master4;
+ default:
+ fail(
+ "Should use master host1/host2/host3/host4:"
+ + masterPort
+ + ", but use "
+ + address);
+ }
+ return null;
+ })
+ .when(rpcEnv)
+ .setupEndpointRef(Mockito.any(RpcAddress.class), Mockito.anyString());
+
+ MasterClient client = new MasterClient(rpcEnv, conf, false);
+ HeartbeatFromWorker message = Mockito.mock(HeartbeatFromWorker.class);
+
+ HeartbeatFromWorkerResponse response = null;
+ try {
+ response = client.askSync(message, HeartbeatFromWorkerResponse.class);
+ } catch (Throwable t) {
+ LOG.error("It should retry another configured master after a redirect cycle.", t);
+ fail("It should retry another configured master after a redirect cycle.");
+ }
+
+ assertEquals(mockResponse, response);
+ Mockito.verify(rpcEnv, Mockito.times(1))
+ .setupEndpointRef(
+ Mockito.eq(RpcAddress.fromHostAndPort("host2:9097")), Mockito.anyString());
+ Mockito.verify(rpcEnv, Mockito.times(1))
+ .setupEndpointRef(
+ Mockito.eq(RpcAddress.fromHostAndPort("host3:9097")), Mockito.anyString());
+ }
+
+ @Test
+ public void testBootstrapSuggestedLeaderConnectionFailureRetriesConfiguredMasterInHA() {
+ final CelebornConf conf = prepareForCelebornConfWithHA();
+
+ final RpcEndpointRef master3 = Mockito.mock(RpcEndpointRef.class);
+
+ Mockito.doReturn(Future$.MODULE$.successful(mockResponse))
+ .when(master3)
+ .ask(Mockito.any(), Mockito.any(), Mockito.any());
+
+ Mockito.doAnswer(
+ invocation -> {
+ RpcAddress address = invocation.getArgument(0, RpcAddress.class);
+ switch (address.host()) {
+ case "host1":
+ throw new RuntimeException(
+ new MasterNotLeaderException("host1:9097", "host2:9097", null));
+ case "host2":
+ throw new IOException("test");
+ case "host3":
+ return master3;
+ default:
+ fail(
+ "Should use master host1/host2/host3:" + masterPort + ", but use " + address);
+ }
+ return null;
+ })
+ .when(rpcEnv)
+ .setupEndpointRef(Mockito.any(RpcAddress.class), Mockito.anyString());
+
+ MasterClient client = new MasterClient(rpcEnv, conf, false);
+ HeartbeatFromWorker message = Mockito.mock(HeartbeatFromWorker.class);
+
+ HeartbeatFromWorkerResponse response = null;
+ try {
+ response = client.askSync(message, HeartbeatFromWorkerResponse.class);
+ } catch (Throwable t) {
+ LOG.error("It should retry another configured master after bootstrap redirect fails.", t);
+ fail("It should retry another configured master after bootstrap redirect fails.");
+ }
+
+ assertEquals(mockResponse, response);
+ }
+
private void checkOneMasterDownInHA(Exception causedByException) {
final CelebornConf conf = prepareForCelebornConfWithHA();
@@ -440,15 +691,19 @@ public class MasterClientSuiteJ {
}
private CelebornConf prepareForCelebornConfWithoutHA() {
- return conf.clone()
+ return conf()
.set(CelebornConf.HA_ENABLED().key(), "false")
.set(CelebornConf.MASTER_ENDPOINTS().key(), masterHost + ":" + masterPort);
}
private CelebornConf prepareForCelebornConfWithHA() {
- return conf.clone()
+ return conf()
.set(CelebornConf.HA_ENABLED().key(), "true")
.set(CelebornConf.MASTER_ENDPOINTS().key(), "host1:9097,host2:9097,host3:9097")
.set(CelebornConf.MASTER_CLIENT_MAX_RETRIES().key(), "5");
}
+
+ private CelebornConf conf() {
+ return conf.clone();
+ }
}