enhance: Enable set node label filter in resource group config

Signed-off-by: Wei Liu <[email protected]>

src/main/java/io/milvus/common/resourcegroup/ResourceGroupConfig.java

@@ -12,10 +12,12 @@ public class ResourceGroupConfig {
     private final ResourceGroupLimit limits;
     private final List<ResourceGroupTransfer> from;
     private final List<ResourceGroupTransfer> to;
+    private final ResourceGroupNodeFilter nodeFilter;
 
     private ResourceGroupConfig(Builder builder) {
         this.requests = builder.requests;
         this.limits = builder.limits;
+        this.nodeFilter = builder.nodeFilter;
 
         if (null == builder.from) {
             this.from = new ArrayList<>();
@@ -39,6 +41,7 @@ public static final class Builder {
         private ResourceGroupLimit limits;
         private List<ResourceGroupTransfer> from;
         private List<ResourceGroupTransfer> to;
+        private ResourceGroupNodeFilter nodeFilter;
 
         private Builder() {
         }
@@ -87,6 +90,17 @@ public Builder withTo(@NonNull List<ResourceGroupTransfer> to) {
             return this;
         }
 
+        /**
+         * Set the node filter.
+         * @param nodeFilter if node filter set, resource group will prefer to accept node which match node filter.
+         * @return <code>Builder</code>
+         */
+
+        public Builder withNodeFilter(@NonNull ResourceGroupNodeFilter nodeFilter) {
+            this.nodeFilter = nodeFilter;
+            return this;
+        }
+
         public ResourceGroupConfig build() {
             return new ResourceGroupConfig(this);
         }
@@ -101,12 +115,14 @@ public ResourceGroupConfig(@NonNull io.milvus.grpc.ResourceGroupConfig grpcConfi
         this.to = grpcConfig.getTransferToList().stream()
                 .map(transfer -> new ResourceGroupTransfer(transfer))
                 .collect(Collectors.toList());
+        this.nodeFilter = new ResourceGroupNodeFilter(grpcConfig.getNodeFilter());
     }
 
     public @NonNull io.milvus.grpc.ResourceGroupConfig toGRPC() {
         io.milvus.grpc.ResourceGroupConfig.Builder builder = io.milvus.grpc.ResourceGroupConfig.newBuilder()
                 .setRequests(io.milvus.grpc.ResourceGroupLimit.newBuilder().setNodeNum(requests.getNodeNum()))
-                .setLimits(io.milvus.grpc.ResourceGroupLimit.newBuilder().setNodeNum(limits.getNodeNum()));
+                .setLimits(io.milvus.grpc.ResourceGroupLimit.newBuilder().setNodeNum(limits.getNodeNum()))
+                .setNodeFilter(nodeFilter.toGRPC());
         for (ResourceGroupTransfer transfer : from) {
             builder.addTransferFrom(transfer.toGRPC());
         }

src/main/java/io/milvus/common/resourcegroup/ResourceGroupNodeFilter.java

@@ -0,0 +1,64 @@
+package io.milvus.common.resourcegroup;
+
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import io.milvus.grpc.KeyValuePair;
+import io.milvus.param.ParamUtils;
+import lombok.Getter;
+import lombok.NonNull;
+
+@Getter
+public class ResourceGroupNodeFilter {
+    private final Map<String, String> nodeLabels;
+
+    private ResourceGroupNodeFilter(Builder builder) {
+        this.nodeLabels = builder.nodeLabels;
+    }
+
+    public static Builder newBuilder() {
+        return new Builder();
+    }
+
+    public static final class Builder {
+        private Map<String, String> nodeLabels;
+        private Builder() {
+        }
+
+        /**
+         * Set the node label filter
+         * @param key label name
+         * @param value label value
+         * @return <code>Builder</code>
+         */
+        public Builder withNodeLabel(@NonNull String key, @NonNull String value) {
+            this.nodeLabels.put(key, value);
+            return this;
+        }
+
+        public ResourceGroupNodeFilter build() {
+            return new ResourceGroupNodeFilter(this);
+        }
+    }
+
+    /**
+     * Transfer to grpc
+     * @return io.milvus.grpc.ResourceGroupNodeFilter
+     */
+    public @NonNull io.milvus.grpc.ResourceGroupNodeFilter toGRPC() {
+        List<KeyValuePair> pair = ParamUtils.AssembleKvPair(nodeLabels);
+        return io.milvus.grpc.ResourceGroupNodeFilter.newBuilder()
+                .addAllNodeLabels(pair)
+               .build();
+    }
+
+    /**
+     * Constructor from grpc
+     * @param filter grpc filter object
+     */
+    public ResourceGroupNodeFilter(io.milvus.grpc.ResourceGroupNodeFilter filter) {
+        this.nodeLabels = filter.getNodeLabelsList().stream().collect(Collectors.toMap(KeyValuePair::getKey, KeyValuePair::getValue));
+    }
+
+}