You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@celeborn.apache.org by zh...@apache.org on 2024/03/07 01:54:45 UTC
(incubator-celeborn) branch main updated: [CELEBORN-1261] Add auth support to client
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 835437f0b [CELEBORN-1261] Add auth support to client
835437f0b is described below
commit 835437f0b9d89e5c3ce44f5aacad2a7d396da1a8
Author: Chandni Singh <si...@gmail.com>
AuthorDate: Thu Mar 7 09:54:28 2024 +0800
[CELEBORN-1261] Add auth support to client
### What changes were proposed in this pull request?
This enables client to push and fetch shuffle data securely to Celeborn Workers.
### Why are the changes needed?
This change is required for adding authentication. ([CELEBORN-1011](https://issues.apache.org/jira/browse/CELEBORN-1011)).
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
It is part of bigger change which will be tested end to end.
Closes #2360 from otterc/CELEBORN-1261.
Authored-by: Chandni Singh <si...@gmail.com>
Signed-off-by: zky.zhoukeyong <zk...@alibaba-inc.com>
---
.../apache/celeborn/client/ShuffleClientImpl.java | 50 +++++++++++++++++++---
.../apache/celeborn/client/LifecycleManager.scala | 16 +++++--
.../common/network/protocol/TransportMessage.java | 3 ++
common/src/main/proto/TransportMessages.proto | 4 ++
.../org/apache/celeborn/common/CelebornConf.scala | 4 ++
5 files changed, 68 insertions(+), 9 deletions(-)
diff --git a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
index 6f38741eb..3a1ac6af8 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
@@ -29,6 +29,7 @@ import scala.Tuple2;
import scala.reflect.ClassTag$;
import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.Lists;
import io.netty.buffer.CompositeByteBuf;
import io.netty.buffer.Unpooled;
import org.apache.commons.lang3.StringUtils;
@@ -46,9 +47,12 @@ import org.apache.celeborn.common.network.TransportContext;
import org.apache.celeborn.common.network.buffer.NettyManagedBuffer;
import org.apache.celeborn.common.network.client.RpcResponseCallback;
import org.apache.celeborn.common.network.client.TransportClient;
+import org.apache.celeborn.common.network.client.TransportClientBootstrap;
import org.apache.celeborn.common.network.client.TransportClientFactory;
import org.apache.celeborn.common.network.protocol.PushData;
import org.apache.celeborn.common.network.protocol.PushMergedData;
+import org.apache.celeborn.common.network.sasl.SaslClientBootstrap;
+import org.apache.celeborn.common.network.sasl.SaslCredentials;
import org.apache.celeborn.common.network.server.BaseMessageHandler;
import org.apache.celeborn.common.network.util.TransportConf;
import org.apache.celeborn.common.protocol.*;
@@ -119,6 +123,8 @@ public class ShuffleClientImpl extends ShuffleClient {
private final Map<Integer, Set<Integer>> splitting = JavaUtils.newConcurrentHashMap();
protected final String appUniqueId;
+ private final boolean authEnabled;
+ private final TransportConf dataTransportConf;
private final ThreadLocal<Compressor> compressorThreadLocal =
new ThreadLocal<Compressor>() {
@@ -178,6 +184,7 @@ public class ShuffleClientImpl extends ShuffleClient {
} else {
pushDataTimeout = conf.pushDataTimeoutMs();
}
+ authEnabled = conf.authEnabledOnClient();
// init rpc env
rpcEnv =
@@ -189,13 +196,9 @@ public class ShuffleClientImpl extends ShuffleClient {
scala.None$.empty());
String module = TransportModuleConstants.DATA_MODULE;
- TransportConf dataTransportConf =
+ dataTransportConf =
Utils.fromCelebornConf(conf, module, conf.getInt("celeborn." + module + ".io.threads", 8));
- TransportContext context =
- new TransportContext(
- dataTransportConf, new BaseMessageHandler(), conf.clientCloseIdleConnections());
- dataClientFactory = context.createClientFactory();
-
+ initDataClientFactoryIfNeeded();
int pushDataRetryThreads = conf.clientPushRetryThreads();
pushDataRetryPool =
ThreadUtils.newDaemonCachedThreadPool("celeborn-retry-sender", pushDataRetryThreads, 60);
@@ -209,6 +212,35 @@ public class ShuffleClientImpl extends ShuffleClient {
logger.info("Created ShuffleClientImpl, appUniqueId: {}", appUniqueId);
}
+ private void initDataClientFactoryIfNeeded() {
+ if (dataClientFactory != null) {
+ return;
+ }
+ TransportContext context =
+ new TransportContext(
+ dataTransportConf, new BaseMessageHandler(), conf.clientCloseIdleConnections());
+ if (!authEnabled) {
+ logger.info("Initializing data client factory for {}.", appUniqueId);
+ dataClientFactory = context.createClientFactory();
+ } else if (lifecycleManagerRef != null) {
+ PbApplicationMetaRequest pbApplicationMetaRequest =
+ PbApplicationMetaRequest.newBuilder().setAppId(appUniqueId).build();
+ PbApplicationMeta pbApplicationMeta =
+ lifecycleManagerRef.askSync(
+ pbApplicationMetaRequest,
+ conf.clientRpcRegisterShuffleRpcAskTimeout(),
+ ClassTag$.MODULE$.apply(PbApplicationMeta.class));
+ logger.info("Initializing data client factory for secured {}.", appUniqueId);
+ List<TransportClientBootstrap> bootstraps = Lists.newArrayList();
+ bootstraps.add(
+ new SaslClientBootstrap(
+ dataTransportConf,
+ appUniqueId,
+ new SaslCredentials(appUniqueId, pbApplicationMeta.getSecret())));
+ dataClientFactory = context.createClientFactory(bootstraps);
+ }
+ }
+
private boolean isPushTargetWorkerExcluded(
PartitionLocation location, RpcResponseCallback wrappedCallback) {
// If pushExcludeWorkerOnFailureEnabled = false, pushExcludedWorkers should be empty.
@@ -290,6 +322,7 @@ public class ShuffleClientImpl extends ShuffleClient {
try {
if (!isPushTargetWorkerExcluded(newLoc, pushDataRpcResponseCallback)) {
if (!testRetryRevive || remainReviveTimes < 1) {
+ assert dataClientFactory != null;
TransportClient client =
dataClientFactory.createClient(newLoc.getHost(), newLoc.getPushPort(), partitionId);
NettyManagedBuffer newBuffer = new NettyManagedBuffer(Unpooled.wrappedBuffer(body));
@@ -1116,6 +1149,7 @@ public class ShuffleClientImpl extends ShuffleClient {
try {
if (!isPushTargetWorkerExcluded(loc, wrappedCallback)) {
if (!testRetryRevive) {
+ assert dataClientFactory != null;
TransportClient client =
dataClientFactory.createClient(loc.getHost(), loc.getPushPort(), partitionId);
client.pushData(pushData, pushDataTimeout, wrappedCallback);
@@ -1494,6 +1528,7 @@ public class ShuffleClientImpl extends ShuffleClient {
try {
if (!isPushTargetWorkerExcluded(batches.get(0).loc, wrappedCallback)) {
if (!testRetryRevive || remainReviveTimes < 1) {
+ assert dataClientFactory != null;
TransportClient client = dataClientFactory.createClient(host, port);
client.pushMergedData(mergedData, pushDataTimeout, wrappedCallback);
} else {
@@ -1678,6 +1713,7 @@ public class ShuffleClientImpl extends ShuffleClient {
return CelebornInputStream.empty();
} else {
String shuffleKey = Utils.makeShuffleKey(appUniqueId, shuffleId);
+ assert dataClientFactory != null;
return CelebornInputStream.create(
conf,
dataClientFactory,
@@ -1733,11 +1769,13 @@ public class ShuffleClientImpl extends ShuffleClient {
public void setupLifecycleManagerRef(String host, int port) {
lifecycleManagerRef =
rpcEnv.setupEndpointRef(new RpcAddress(host, port), RpcNameConstants.LIFECYCLE_MANAGER_EP);
+ initDataClientFactoryIfNeeded();
}
@Override
public void setupLifecycleManagerRef(RpcEndpointRef endpointRef) {
lifecycleManagerRef = endpointRef;
+ initDataClientFactoryIfNeeded();
}
@Override
diff --git a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
index a75f68b40..92d28447d 100644
--- a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
+++ b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
@@ -42,7 +42,7 @@ import org.apache.celeborn.common.CelebornConf
import org.apache.celeborn.common.client.MasterClient
import org.apache.celeborn.common.identity.{IdentityProvider, UserIdentifier}
import org.apache.celeborn.common.internal.Logging
-import org.apache.celeborn.common.meta.{ShufflePartitionLocationInfo, WorkerInfo}
+import org.apache.celeborn.common.meta.{ApplicationMeta, ShufflePartitionLocationInfo, WorkerInfo}
import org.apache.celeborn.common.network.sasl.registration.RegistrationInfo
import org.apache.celeborn.common.protocol._
import org.apache.celeborn.common.protocol.RpcNameConstants.WORKER_EP
@@ -112,8 +112,8 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
.build().asInstanceOf[Cache[Int, ByteBuffer]]
private val mockDestroyFailure = conf.testMockDestroySlotsFailure
- private val authEnabled = conf.authEnabled
-
+ private val authEnabled = conf.authEnabledOnClient
+ private var applicationMeta: ApplicationMeta = _
@VisibleForTesting
def workerSnapshots(shuffleId: Int): util.Map[WorkerInfo, ShufflePartitionLocationInfo] =
shuffleAllocatedWorkers.get(shuffleId)
@@ -173,6 +173,7 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
if (authEnabled) {
logInfo(s"Authentication is enabled; setting up master and worker RPC environments")
val appSecret = createSecret()
+ applicationMeta = ApplicationMeta(appUniqueId, appSecret)
val registrationInfo = new RegistrationInfo()
masterRpcEnvInUse =
RpcEnv.create(
@@ -412,6 +413,15 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
val shuffleId = pb.getShuffleId
logDebug(s"Received ReportShuffleFetchFailure request, appShuffleId $appShuffleId shuffleId $shuffleId")
handleReportShuffleFetchFailure(context, appShuffleId, shuffleId)
+
+ case pb: PbApplicationMetaRequest =>
+ logDebug(s"Received request for meta info ${pb.getAppId}.")
+ if (applicationMeta == null) {
+ context.sendFailure(
+ new IllegalArgumentException("Application meta is not initialized for this app."))
+ } else {
+ context.reply(PbSerDeUtils.toPbApplicationMeta(applicationMeta))
+ }
}
def setupEndpoints(
diff --git a/common/src/main/java/org/apache/celeborn/common/network/protocol/TransportMessage.java b/common/src/main/java/org/apache/celeborn/common/network/protocol/TransportMessage.java
index 67c3f329d..392edce85 100644
--- a/common/src/main/java/org/apache/celeborn/common/network/protocol/TransportMessage.java
+++ b/common/src/main/java/org/apache/celeborn/common/network/protocol/TransportMessage.java
@@ -30,6 +30,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.exception.CelebornIOException;
import org.apache.celeborn.common.protocol.MessageType;
import org.apache.celeborn.common.protocol.PbApplicationMeta;
+import org.apache.celeborn.common.protocol.PbApplicationMetaRequest;
import org.apache.celeborn.common.protocol.PbAuthenticationInitiationRequest;
import org.apache.celeborn.common.protocol.PbAuthenticationInitiationResponse;
import org.apache.celeborn.common.protocol.PbBacklogAnnouncement;
@@ -120,6 +121,8 @@ public class TransportMessage implements Serializable {
return (T) PbRegisterApplicationResponse.parseFrom(payload);
case APPLICATION_META_VALUE:
return (T) PbApplicationMeta.parseFrom(payload);
+ case APPLICATION_META_REQUEST_VALUE:
+ return (T) PbApplicationMetaRequest.parseFrom(payload);
default:
logger.error("Unexpected type {}", type);
}
diff --git a/common/src/main/proto/TransportMessages.proto b/common/src/main/proto/TransportMessages.proto
index c16077e76..e949aa9d4 100644
--- a/common/src/main/proto/TransportMessages.proto
+++ b/common/src/main/proto/TransportMessages.proto
@@ -99,6 +99,7 @@ enum MessageType {
WORKER_EVENT_REQUEST = 76;
WORKER_EVENT_RESPONSE = 77;
APPLICATION_META = 78;
+ APPLICATION_META_REQUEST = 79;
}
enum StreamType {
@@ -732,3 +733,6 @@ message PbApplicationMeta {
string secret = 2;
}
+message PbApplicationMetaRequest {
+ string appId = 1;
+}
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 d402173fd..adf4fc84b 100644
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@@ -1144,6 +1144,10 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable with Logging with Se
def masterSendApplicationMetaThreads: Int = get(MASTER_SEND_APPLICATION_META_THREADS)
+ def authEnabledOnClient: Boolean = {
+ get(AUTH_ENABLED)
+ }
+
// //////////////////////////////////////////////////////
// Internal Port //
// //////////////////////////////////////////////////////