[CELEBORN-2349] Support worker registered tags during startup

Author: s0nskar

common/src/main/proto/TransportMessages.proto

@@ -189,6 +189,7 @@ message PbWorkerInfo {
   int32 internalPort = 8;
   string networkLocation = 9;
   int64 nextInterruptionNotice = 10; // Unix timestamp when disruption is expected to be initiated
+  repeated string tags = 11;
 }
 
 message PbFileGroup {
@@ -207,6 +208,7 @@ message PbRegisterWorker {
   map<string, PbResourceConsumption> userResourceConsumption = 8;
   int32 internalPort = 10;
   string networkLocation = 11;
+  repeated string tags = 12;
 }
 
 message PbMetaRegisterWorkerRequest {
@@ -219,6 +221,7 @@ message PbMetaRegisterWorkerRequest {
   map<string, PbResourceConsumption> userResourceConsumption = 7;
   int32 internalPort = 8;
   string networkLocation = 9;
+  repeated string tags = 12;
 }
 
 message PbHeartbeatFromWorker {

common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala

@@ -1527,6 +1527,8 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable with Logging with Se
   def clientInputStreamCreationWindow = get(CLIENT_INPUTSTREAM_CREATION_WINDOW)
 
   def tagsEnabled: Boolean = get(TAGS_ENABLED)
+  def tagsWorkerRegistrationEnabled: Boolean = get(TAGS_WORKER_REGISTRATION_ENABLED)
+  def workerTags: Seq[String] = get(WORKER_TAGS)
   def tagsExpr: String = get(TAGS_EXPR)
   def preferClientTagsExpr: Boolean = get(PREFER_CLIENT_TAGS_EXPR)
 
@@ -6901,6 +6903,24 @@ object CelebornConf extends Logging {
       .booleanConf
       .createWithDefault(true)
 
+  val TAGS_WORKER_REGISTRATION_ENABLED: ConfigEntry[Boolean] =
+    buildConf("celeborn.tags.worker.registration.enabled")
+      .categories("master")
+      .version("0.7.0")
+      .doc("When true, the master honors tags advertised by workers at registration " +
+        "(merged with the config-store tags). When false, worker-supplied tags are ignored.")
+      .booleanConf
+      .createWithDefault(true)
+
+  val WORKER_TAGS: ConfigEntry[Seq[String]] =
+    buildConf("celeborn.worker.tags")
+      .categories("worker")
+      .version("0.7.0")
+      .doc("Comma-separated tags this worker supplies to the master at registration.")
+      .stringConf
+      .toSequence
+      .createWithDefault(Seq.empty)
+
   val TAGS_EXPR: ConfigEntry[String] =
     buildConf("celeborn.tags.tagsExpr")
       .categories("master", "client")

common/src/main/scala/org/apache/celeborn/common/meta/WorkerInfo.scala

@@ -47,7 +47,8 @@ class WorkerInfo(
   var nextInterruptionNotice = Long.MaxValue
   var lastHeartbeat: Long = 0
   var workerStatus = WorkerStatus.normalWorkerStatus()
-  var isHighWorkLoad: Boolean = false;
+  var isHighWorkLoad: Boolean = false
+  var tags: util.Set[String] = new util.HashSet[String]()
   val diskInfos = {
     if (_diskInfos != null) JavaUtils.newConcurrentHashMap[String, DiskInfo](_diskInfos)
     else null

common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala

@@ -93,6 +93,7 @@ object ControlMessages extends Logging {
         networkLocation: String,
         disks: Map[String, DiskInfo],
         userResourceConsumption: Map[UserIdentifier, ResourceConsumption],
+        tags: Set[String],
         requestId: String): PbRegisterWorker = {
       val pbDisks = disks.values.map(PbSerDeUtils.toPbDiskInfo).asJava
       val pbUserResourceConsumption =
@@ -107,6 +108,7 @@ object ControlMessages extends Logging {
         .setNetworkLocation(networkLocation)
         .addAllDisks(pbDisks)
         .putAllUserResourceConsumption(pbUserResourceConsumption)
+        .addAllTags(tags.asJava)
         .setRequestId(requestId)
         .build()
     }

common/src/main/scala/org/apache/celeborn/common/util/PbSerDeUtils.scala

@@ -296,6 +296,7 @@ object PbSerDeUtils {
     } else {
       workerInfo.nextInterruptionNotice = pbWorkerInfo.getNextInterruptionNotice
     }
+    workerInfo.tags = new util.HashSet[String](pbWorkerInfo.getTagsList)
     workerInfo
   }
 
@@ -310,6 +311,7 @@ object PbSerDeUtils {
       .setPushPort(workerInfo.pushPort)
       .setReplicatePort(workerInfo.replicatePort)
       .setInternalPort(workerInfo.internalPort)
+      .addAllTags(workerInfo.tags)
     if (masterPersistWorkerNetworkLocation) {
       builder.setNetworkLocation(workerInfo.networkLocation)
     }

master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/AbstractMetaManager.java

@@ -357,7 +357,8 @@ public void updateRegisterWorkerMeta(
       int replicatePort,
       int internalPort,
       String networkLocation,
-      Map<String, DiskInfo> disks) {
+      Map<String, DiskInfo> disks,
+      Set<String> tags) {
     WorkerInfo workerInfo =
         new WorkerInfo(
             host,
@@ -369,6 +370,7 @@ public void updateRegisterWorkerMeta(
             disks,
             new HashMap<>());
     workerInfo.lastHeartbeat_$eq(System.currentTimeMillis());
+    workerInfo.tags_$eq(new HashSet<>(tags));
     if (networkLocation != null
         && !networkLocation.isEmpty()
         && !NetworkTopology.DEFAULT_RACK.equals(networkLocation)) {

master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/IMetadataHandler.java

@@ -17,8 +17,10 @@
 
 package org.apache.celeborn.service.deploy.master.clustermeta;
 
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import org.apache.celeborn.common.identity.UserIdentifier;
 import org.apache.celeborn.common.meta.ApplicationMeta;
@@ -77,6 +79,31 @@ void handleWorkerHeartbeat(
       WorkerStatus workerStatus,
       String requestId);
 
+  default void handleRegisterWorker(
+      String host,
+      int rpcPort,
+      int pushPort,
+      int fetchPort,
+      int replicatePort,
+      int internalPort,
+      String networkLocation,
+      Map<String, DiskInfo> disks,
+      Map<UserIdentifier, ResourceConsumption> userResourceConsumption,
+      String requestId) {
+    handleRegisterWorker(
+        host,
+        rpcPort,
+        pushPort,
+        fetchPort,
+        replicatePort,
+        internalPort,
+        networkLocation,
+        disks,
+        userResourceConsumption,
+        Collections.emptySet(),
+        requestId);
+  }
+
   void handleRegisterWorker(
       String host,
       int rpcPort,
@@ -87,6 +114,7 @@ void handleRegisterWorker(
       String networkLocation,
       Map<String, DiskInfo> disks,
       Map<UserIdentifier, ResourceConsumption> userResourceConsumption,
+      Set<String> tags,
       String requestId);
 
   void handleReportWorkerUnavailable(List<WorkerInfo> failedNodes, String requestId);

master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/SingleMasterMetaManager.java

@@ -19,6 +19,7 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -162,9 +163,18 @@ public void handleRegisterWorker(
       String networkLocation,
       Map<String, DiskInfo> disks,
       Map<UserIdentifier, ResourceConsumption> userResourceConsumption,
+      Set<String> tags,
       String requestId) {
     updateRegisterWorkerMeta(
-        host, rpcPort, pushPort, fetchPort, replicatePort, internalPort, networkLocation, disks);
+        host,
+        rpcPort,
+        pushPort,
+        fetchPort,
+        replicatePort,
+        internalPort,
+        networkLocation,
+        disks,
+        tags);
     updateWorkerResourceConsumptions(
         host, rpcPort, pushPort, fetchPort, replicatePort, userResourceConsumption);
   }

master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/HAMasterMetaManager.java

@@ -19,6 +19,7 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.stream.Collectors;
 
 import org.slf4j.Logger;
@@ -349,6 +350,7 @@ public void handleRegisterWorker(
       String networkLocation,
       Map<String, DiskInfo> disks,
       Map<UserIdentifier, ResourceConsumption> userResourceConsumption,
+      Set<String> tags,
       String requestId) {
     try {
       ratisServer.submitRequest(
@@ -365,6 +367,7 @@ public void handleRegisterWorker(
                       .setInternalPort(internalPort)
                       .setNetworkLocation(networkLocation)
                       .putAllDisks(MetaUtil.toPbDiskInfos(disks))
+                      .addAllTags(tags)
                       .build())
               .build());
       updateWorkerResourceConsumptions(

master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/MetaHandler.java

@@ -258,6 +258,7 @@ public org.apache.celeborn.common.protocol.PbMetaRequestResponse handleWriteRequ
           int internalPort = request.getRegisterWorkerRequest().getInternalPort();
           Map<String, PbDiskInfo> pbDiskInfo = request.getRegisterWorkerRequest().getDisksMap();
           diskInfos = MetaUtil.fromPbDiskInfoMap(pbDiskInfo);
+          Set<String> tags = new HashSet<>(request.getRegisterWorkerRequest().getTagsList());
           LOG.debug(
               "Handle worker register for {} {} {} {} {} {} {}",
               host,
@@ -275,7 +276,8 @@ public org.apache.celeborn.common.protocol.PbMetaRequestResponse handleWriteRequ
               replicatePort,
               internalPort,
               networkLocation,
-              diskInfos);
+              diskInfos,
+              tags);
           break;
 
         case ReportWorkerUnavailable: