From bef8bace9015c85f352033391cbefd0707086fab Mon Sep 17 00:00:00 2001 From: yhmo Date: Thu, 26 Dec 2024 16:29:41 +0800 Subject: [PATCH] ResourceGroup interfaces for V2 Signed-off-by: yhmo --- .../milvus/common/resourcegroup/NodeInfo.java | 31 +++ .../resourcegroup/ResourceGroupConfig.java | 37 +++- .../resourcegroup/ResourceGroupLimit.java | 21 +- .../ResourceGroupNodeFilter.java | 84 ++++++++ .../resourcegroup/ResourceGroupTransfer.java | 19 ++ .../io/milvus/v2/client/MilvusClientV2.java | 88 ++++++++- .../resourcegroup/ResourceGroupService.java | 187 ++++++++++++++++++ .../request/CreateResourceGroupReq.java | 12 ++ .../request/DescribeResourceGroupReq.java | 10 + .../request/DropResourceGroupReq.java | 10 + .../request/ListResourceGroupsReq.java | 9 + .../request/TransferReplicaReq.java | 14 ++ .../request/UpdateResourceGroupsReq.java | 16 ++ .../response/DescribeResourceGroupResp.java | 26 +++ .../response/ListResourceGroupsResp.java | 15 ++ .../v2/client/MilvusClientV2DockerTest.java | 80 ++++++++ 16 files changed, 649 insertions(+), 10 deletions(-) create mode 100644 src/main/java/io/milvus/common/resourcegroup/NodeInfo.java create mode 100644 src/main/java/io/milvus/common/resourcegroup/ResourceGroupNodeFilter.java create mode 100644 src/main/java/io/milvus/v2/service/resourcegroup/ResourceGroupService.java create mode 100644 src/main/java/io/milvus/v2/service/resourcegroup/request/CreateResourceGroupReq.java create mode 100644 src/main/java/io/milvus/v2/service/resourcegroup/request/DescribeResourceGroupReq.java create mode 100644 src/main/java/io/milvus/v2/service/resourcegroup/request/DropResourceGroupReq.java create mode 100644 src/main/java/io/milvus/v2/service/resourcegroup/request/ListResourceGroupsReq.java create mode 100644 src/main/java/io/milvus/v2/service/resourcegroup/request/TransferReplicaReq.java create mode 100644 src/main/java/io/milvus/v2/service/resourcegroup/request/UpdateResourceGroupsReq.java create mode 100644 src/main/java/io/milvus/v2/service/resourcegroup/response/DescribeResourceGroupResp.java create mode 100644 src/main/java/io/milvus/v2/service/resourcegroup/response/ListResourceGroupsResp.java diff --git a/src/main/java/io/milvus/common/resourcegroup/NodeInfo.java b/src/main/java/io/milvus/common/resourcegroup/NodeInfo.java new file mode 100644 index 000000000..4d42a0746 --- /dev/null +++ b/src/main/java/io/milvus/common/resourcegroup/NodeInfo.java @@ -0,0 +1,31 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package io.milvus.common.resourcegroup; + +import lombok.Data; +import lombok.experimental.SuperBuilder; + +@Data +@SuperBuilder +public class NodeInfo { + private Long nodeId; + private String address; + private String hostname; +} diff --git a/src/main/java/io/milvus/common/resourcegroup/ResourceGroupConfig.java b/src/main/java/io/milvus/common/resourcegroup/ResourceGroupConfig.java index 4bd699245..735145820 100644 --- a/src/main/java/io/milvus/common/resourcegroup/ResourceGroupConfig.java +++ b/src/main/java/io/milvus/common/resourcegroup/ResourceGroupConfig.java @@ -1,3 +1,22 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + package io.milvus.common.resourcegroup; import java.util.stream.Collectors; @@ -12,10 +31,12 @@ public class ResourceGroupConfig { private final ResourceGroupLimit limits; private final List from; private final List 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 +60,7 @@ public static final class Builder { private ResourceGroupLimit limits; private List from; private List to; + private ResourceGroupNodeFilter nodeFilter; private Builder() { } @@ -87,6 +109,17 @@ public Builder withTo(@NonNull List 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 Builder + */ + + public Builder withNodeFilter(@NonNull ResourceGroupNodeFilter nodeFilter) { + this.nodeFilter = nodeFilter; + return this; + } + public ResourceGroupConfig build() { return new ResourceGroupConfig(this); } @@ -101,12 +134,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()); } diff --git a/src/main/java/io/milvus/common/resourcegroup/ResourceGroupLimit.java b/src/main/java/io/milvus/common/resourcegroup/ResourceGroupLimit.java index 2456f3794..f1ea78720 100644 --- a/src/main/java/io/milvus/common/resourcegroup/ResourceGroupLimit.java +++ b/src/main/java/io/milvus/common/resourcegroup/ResourceGroupLimit.java @@ -1,6 +1,23 @@ -package io.milvus.common.resourcegroup; +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ -import org.jetbrains.annotations.NotNull; +package io.milvus.common.resourcegroup; import lombok.Getter; import lombok.NonNull; diff --git a/src/main/java/io/milvus/common/resourcegroup/ResourceGroupNodeFilter.java b/src/main/java/io/milvus/common/resourcegroup/ResourceGroupNodeFilter.java new file mode 100644 index 000000000..9f90b31fc --- /dev/null +++ b/src/main/java/io/milvus/common/resourcegroup/ResourceGroupNodeFilter.java @@ -0,0 +1,84 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package io.milvus.common.resourcegroup; + +import java.util.HashMap; +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 nodeLabels; + + private ResourceGroupNodeFilter(Builder builder) { + this.nodeLabels = builder.nodeLabels; + } + + public static Builder newBuilder() { + return new Builder(); + } + + public static final class Builder { + private Map nodeLabels = new HashMap<>(); + private Builder() { + } + + /** + * Set the node label filter + * @param key label name + * @param value label value + * @return Builder + */ + 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 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)); + } + +} \ No newline at end of file diff --git a/src/main/java/io/milvus/common/resourcegroup/ResourceGroupTransfer.java b/src/main/java/io/milvus/common/resourcegroup/ResourceGroupTransfer.java index 2312598dc..bb135764d 100644 --- a/src/main/java/io/milvus/common/resourcegroup/ResourceGroupTransfer.java +++ b/src/main/java/io/milvus/common/resourcegroup/ResourceGroupTransfer.java @@ -1,3 +1,22 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + package io.milvus.common.resourcegroup; import lombok.Getter; diff --git a/src/main/java/io/milvus/v2/client/MilvusClientV2.java b/src/main/java/io/milvus/v2/client/MilvusClientV2.java index cd780c81a..e28311037 100644 --- a/src/main/java/io/milvus/v2/client/MilvusClientV2.java +++ b/src/main/java/io/milvus/v2/client/MilvusClientV2.java @@ -22,10 +22,7 @@ import io.grpc.ManagedChannel; import io.grpc.Status; import io.grpc.StatusRuntimeException; -import io.milvus.grpc.ClientInfo; -import io.milvus.grpc.ConnectRequest; -import io.milvus.grpc.ConnectResponse; -import io.milvus.grpc.MilvusServiceGrpc; +import io.milvus.grpc.*; import io.milvus.orm.iterator.QueryIterator; import io.milvus.orm.iterator.SearchIterator; @@ -46,6 +43,9 @@ import io.milvus.v2.service.rbac.RBACService; import io.milvus.v2.service.rbac.request.*; import io.milvus.v2.service.rbac.response.*; +import io.milvus.v2.service.resourcegroup.ResourceGroupService; +import io.milvus.v2.service.resourcegroup.request.*; +import io.milvus.v2.service.resourcegroup.response.*; import io.milvus.v2.service.utility.UtilityService; import io.milvus.v2.service.utility.request.*; import io.milvus.v2.service.utility.response.*; @@ -74,6 +74,7 @@ public class MilvusClientV2 { private final VectorService vectorService = new VectorService(); private final PartitionService partitionService = new PartitionService(); private final RBACService rbacService = new RBACService(); + private final ResourceGroupService rgroupService = new ResourceGroupService(); private final UtilityService utilityService = new UtilityService(); private ConnectConfig connectConfig; private RetryConfig retryConfig = RetryConfig.builder().build(); @@ -267,6 +268,9 @@ private T retry(Callable callable) { return null; } + ///////////////////////////////////////////////////////////////////////////////////////////// + // Database Operations + ///////////////////////////////////////////////////////////////////////////////////////////// /** * use Database * @param dbName databaseName @@ -341,7 +345,9 @@ public DescribeDatabaseResp describeDatabase(DescribeDatabaseReq request) { return retry(()-> databaseService.describeDatabase(this.getRpcStub(), request)); } - //Collection Operations + ///////////////////////////////////////////////////////////////////////////////////////////// + // Collection Operations + ///////////////////////////////////////////////////////////////////////////////////////////// /** * Creates a collection in Milvus. * @param request create collection request @@ -480,7 +486,9 @@ public Boolean getLoadState(GetLoadStateReq request) { return retry(()->collectionService.getLoadState(this.getRpcStub(), request)); } + ///////////////////////////////////////////////////////////////////////////////////////////// //Index Operations + ///////////////////////////////////////////////////////////////////////////////////////////// /** * Creates an index for a specified field in a collection in Milvus. * @@ -546,7 +554,9 @@ public List listIndexes(ListIndexesReq request) { return retry(()->indexService.listIndexes(this.getRpcStub(), request)); } + ///////////////////////////////////////////////////////////////////////////////////////////// // Vector Operations + ///////////////////////////////////////////////////////////////////////////////////////////// /** * Inserts vectors into a collection in Milvus. * @@ -633,7 +643,9 @@ public SearchIterator searchIterator(SearchIteratorReq request) { return retry(()->vectorService.searchIterator(this.getRpcStub(), request)); } + ///////////////////////////////////////////////////////////////////////////////////////////// // Partition Operations + ///////////////////////////////////////////////////////////////////////////////////////////// /** * Creates a partition in a collection in Milvus. * @@ -699,7 +711,9 @@ public void releasePartitions(ReleasePartitionsReq request) { retry(()->partitionService.releasePartitions(this.getRpcStub(), request)); } - // RBAC operations + ///////////////////////////////////////////////////////////////////////////////////////////// + // RBAC Operations + ///////////////////////////////////////////////////////////////////////////////////////////// /** * list users * @@ -836,8 +850,68 @@ public void revokePrivilegeV2(RevokePrivilegeReqV2 request) { retry(()->rbacService.revokePrivilegeV2(this.getRpcStub(), request)); } - // Utility Operations + ///////////////////////////////////////////////////////////////////////////////////////////// + // Resource group Operations + ///////////////////////////////////////////////////////////////////////////////////////////// + /** + * Create a resource group. + * + * @param request {@link CreateResourceGroupReq} + */ + public void createResourceGroup(CreateResourceGroupReq request){ + retry(()->rgroupService.createResourceGroup(this.getRpcStub(), request)); + } + + /** + * Update resource groups. + * + * @param request {@link UpdateResourceGroupsReq} + */ + public void updateResourceGroups(UpdateResourceGroupsReq request) { + retry(()->rgroupService.updateResourceGroups(this.getRpcStub(), request)); + } + + /** + * Drop a resource group. + * + * @param request {@link DropResourceGroupReq} + */ + public void dropResourceGroup(DropResourceGroupReq request) { + retry(()->rgroupService.dropResourceGroup(this.getRpcStub(), request)); + } + /** + * List resource groups. + * + * @param request {@link ListResourceGroupsReq} + * @return ListResourceGroupsResp + */ + ListResourceGroupsResp listResourceGroups(ListResourceGroupsReq request) { + return retry(()->rgroupService.listResourceGroups(this.getRpcStub(), request)); + } + + /** + * Describe a resource group. + * + * @param request {@link DescribeResourceGroupReq} + * @return DescribeResourceGroupResp + */ + DescribeResourceGroupResp describeResourceGroup(DescribeResourceGroupReq request) { + return retry(()->rgroupService.describeResourceGroup(this.getRpcStub(), request)); + } + + /** + * Transfer a replica from source resource group to target resource_group. + * + * @param request {@link TransferReplicaReq} + */ + public void transferReplica(TransferReplicaReq request) { + retry(()->rgroupService.transferReplica(this.getRpcStub(), request)); + } + + ///////////////////////////////////////////////////////////////////////////////////////////// + // Utility Operations + ///////////////////////////////////////////////////////////////////////////////////////////// /** * create aliases * diff --git a/src/main/java/io/milvus/v2/service/resourcegroup/ResourceGroupService.java b/src/main/java/io/milvus/v2/service/resourcegroup/ResourceGroupService.java new file mode 100644 index 000000000..5e761592d --- /dev/null +++ b/src/main/java/io/milvus/v2/service/resourcegroup/ResourceGroupService.java @@ -0,0 +1,187 @@ +package io.milvus.v2.service.resourcegroup; + +import io.milvus.grpc.*; +import io.milvus.v2.exception.ErrorCode; +import io.milvus.v2.exception.MilvusClientException; +import io.milvus.v2.service.BaseService; +import io.milvus.v2.service.resourcegroup.request.*; +import io.milvus.v2.service.resourcegroup.response.*; +import org.apache.commons.lang3.StringUtils; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +public class ResourceGroupService extends BaseService { + private static ResourceGroupConfig convertResourceGroupConfig(io.milvus.common.resourcegroup.ResourceGroupConfig config) { + if (config == null) { + throw new MilvusClientException(ErrorCode.INVALID_PARAMS, "Invalid resource group config"); + } + + ResourceGroupConfig.Builder builder = ResourceGroupConfig.newBuilder(); + builder.setRequests(ResourceGroupLimit.newBuilder() + .setNodeNum(config.getRequests().getNodeNum())) + .build(); + builder.setLimits(ResourceGroupLimit.newBuilder() + .setNodeNum(config.getLimits().getNodeNum())) + .build(); + + for (io.milvus.common.resourcegroup.ResourceGroupTransfer groupFrom : config.getFrom()) { + builder.addTransferFrom(ResourceGroupTransfer.newBuilder() + .setResourceGroup(groupFrom.getResourceGroupName())) + .build(); + } + + for (io.milvus.common.resourcegroup.ResourceGroupTransfer groupTo : config.getTo()) { + builder.addTransferTo(ResourceGroupTransfer.newBuilder() + .setResourceGroup(groupTo.getResourceGroupName())) + .build(); + } + + if (config.getNodeFilter() != null) { + builder.setNodeFilter(config.getNodeFilter().toGRPC()); + } + + return builder.build(); + } + + private static io.milvus.common.resourcegroup.ResourceGroupConfig convertResourceGroupConfig(ResourceGroupConfig config) { + List fromList = new ArrayList<>(); + config.getTransferFromList().forEach((groupFrom)->{ + fromList.add(new io.milvus.common.resourcegroup.ResourceGroupTransfer(groupFrom.getResourceGroup())); + }); + + List toList = new ArrayList<>(); + config.getTransferToList().forEach((groupTo)->{ + toList.add(new io.milvus.common.resourcegroup.ResourceGroupTransfer(groupTo.getResourceGroup())); + }); + + return io.milvus.common.resourcegroup.ResourceGroupConfig.newBuilder() + .withRequests(new io.milvus.common.resourcegroup.ResourceGroupLimit(config.getRequests().getNodeNum())) + .withLimits(new io.milvus.common.resourcegroup.ResourceGroupLimit(config.getLimits().getNodeNum())) + .withFrom(fromList) + .withTo(toList) + .withNodeFilter(new io.milvus.common.resourcegroup.ResourceGroupNodeFilter(config.getNodeFilter())) + .build(); + } + + public Void createResourceGroup(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, + CreateResourceGroupReq request) { + String title = String.format("CreateResourceGroupReq groupName:%s", request.getGroupName()); + + ResourceGroupConfig rpcConfig = convertResourceGroupConfig(request.getConfig()); + CreateResourceGroupRequest rpcRequest = CreateResourceGroupRequest.newBuilder() + .setResourceGroup(request.getGroupName()) + .setConfig(rpcConfig) + .build(); + + Status status = blockingStub.createResourceGroup(rpcRequest); + rpcUtils.handleResponse(title, status); + return null; + } + + public Void updateResourceGroups(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, + UpdateResourceGroupsReq request) { + String title = "UpdateResourceGroupsReq"; + + Map resourceGroups = request.getResourceGroups(); + if (resourceGroups.isEmpty()) { + throw new MilvusClientException(ErrorCode.INVALID_PARAMS, "Resource group configurations cannot be empty"); + } + + UpdateResourceGroupsRequest.Builder requestBuilder = UpdateResourceGroupsRequest.newBuilder(); + resourceGroups.forEach((groupName, config) -> { + ResourceGroupConfig rpcConfig = convertResourceGroupConfig(config); + requestBuilder.putResourceGroups(groupName, rpcConfig); + }); + + Status status = blockingStub.updateResourceGroups(requestBuilder.build()); + rpcUtils.handleResponse(title, status); + return null; + } + + public Void dropResourceGroup(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, + DropResourceGroupReq request) { + String title = String.format("DropResourceGroupReq groupName:%s", request.getGroupName()); + + DropResourceGroupRequest rpcRequest = DropResourceGroupRequest.newBuilder() + .setResourceGroup(request.getGroupName()) + .build(); + + Status status = blockingStub.dropResourceGroup(rpcRequest); + rpcUtils.handleResponse(title, status); + return null; + } + + public ListResourceGroupsResp listResourceGroups(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, + ListResourceGroupsReq request) { + String title = "ListResourceGroupsReq"; + ListResourceGroupsResponse response = blockingStub.listResourceGroups(ListResourceGroupsRequest.newBuilder().build()); + rpcUtils.handleResponse(title, response.getStatus()); + return ListResourceGroupsResp.builder() + .groupNames(response.getResourceGroupsList()) + .build(); + } + + public DescribeResourceGroupResp describeResourceGroup(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, + DescribeResourceGroupReq request) { + String title = String.format("DescribeResourceGroupReq groupName:%s", request.getGroupName()); + + DescribeResourceGroupRequest rpcRequest = DescribeResourceGroupRequest.newBuilder() + .setResourceGroup(request.getGroupName()) + .build(); + + DescribeResourceGroupResponse response = blockingStub.describeResourceGroup(rpcRequest); + rpcUtils.handleResponse(title, response.getStatus()); + + ResourceGroup rgroup = response.getResourceGroup(); + List nodes = new ArrayList<>(); + rgroup.getNodesList().forEach((node)->{ + nodes.add(io.milvus.common.resourcegroup.NodeInfo.builder() + .nodeId(node.getNodeId()) + .address(node.getAddress()) + .hostname(node.getHostname()) + .build()); + }); + return DescribeResourceGroupResp.builder() + .groupName(rgroup.getName()) + .capacity(rgroup.getCapacity()) + .numberOfAvailableNode(rgroup.getNumAvailableNode()) + .numberOfLoadedReplica(rgroup.getNumLoadedReplicaMap()) + .numberOfOutgoingNode(rgroup.getNumOutgoingNodeMap()) + .numberOfIncomingNode(rgroup.getNumIncomingNodeMap()) + .config(convertResourceGroupConfig(rgroup.getConfig())) + .nodes(nodes) + .build(); + } + + public Void transferReplica(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, + TransferReplicaReq request) { + if (StringUtils.isEmpty(request.getSourceGroupName())) { + throw new MilvusClientException(ErrorCode.INVALID_PARAMS, "Invalid source group name"); + } + if (StringUtils.isEmpty(request.getTargetGroupName())) { + throw new MilvusClientException(ErrorCode.INVALID_PARAMS, "Invalid target group name"); + } + if (StringUtils.isEmpty(request.getCollectionName())) { + throw new MilvusClientException(ErrorCode.INVALID_PARAMS, "Invalid collection name"); + } + + String title = String.format("TransferReplicaReq sourceGroupName:%s targetGroupName:%s collectionName:%s", + request.getSourceGroupName(), request.getTargetGroupName(), request.getCollectionName()); + + TransferReplicaRequest.Builder requestBuilder = TransferReplicaRequest.newBuilder() + .setSourceResourceGroup(request.getSourceGroupName()) + .setTargetResourceGroup(request.getTargetGroupName()) + .setCollectionName(request.getCollectionName()) + .setNumReplica(request.getNumberOfReplicas()); + + if (StringUtils.isNotEmpty(request.getDatabaseName())) { + requestBuilder.setDbName(request.getDatabaseName()); + } + + Status status = blockingStub.transferReplica(requestBuilder.build()); + rpcUtils.handleResponse(title, status); + return null; + } +} diff --git a/src/main/java/io/milvus/v2/service/resourcegroup/request/CreateResourceGroupReq.java b/src/main/java/io/milvus/v2/service/resourcegroup/request/CreateResourceGroupReq.java new file mode 100644 index 000000000..458bec16a --- /dev/null +++ b/src/main/java/io/milvus/v2/service/resourcegroup/request/CreateResourceGroupReq.java @@ -0,0 +1,12 @@ +package io.milvus.v2.service.resourcegroup.request; + +import io.milvus.common.resourcegroup.ResourceGroupConfig; +import lombok.Data; +import lombok.experimental.SuperBuilder; + +@Data +@SuperBuilder +public class CreateResourceGroupReq { + private String groupName; + private ResourceGroupConfig config; +} diff --git a/src/main/java/io/milvus/v2/service/resourcegroup/request/DescribeResourceGroupReq.java b/src/main/java/io/milvus/v2/service/resourcegroup/request/DescribeResourceGroupReq.java new file mode 100644 index 000000000..c7c5712ba --- /dev/null +++ b/src/main/java/io/milvus/v2/service/resourcegroup/request/DescribeResourceGroupReq.java @@ -0,0 +1,10 @@ +package io.milvus.v2.service.resourcegroup.request; + +import lombok.Data; +import lombok.experimental.SuperBuilder; + +@Data +@SuperBuilder +public class DescribeResourceGroupReq { + private String groupName; +} diff --git a/src/main/java/io/milvus/v2/service/resourcegroup/request/DropResourceGroupReq.java b/src/main/java/io/milvus/v2/service/resourcegroup/request/DropResourceGroupReq.java new file mode 100644 index 000000000..f7f973805 --- /dev/null +++ b/src/main/java/io/milvus/v2/service/resourcegroup/request/DropResourceGroupReq.java @@ -0,0 +1,10 @@ +package io.milvus.v2.service.resourcegroup.request; + +import lombok.Data; +import lombok.experimental.SuperBuilder; + +@Data +@SuperBuilder +public class DropResourceGroupReq { + private String groupName; +} diff --git a/src/main/java/io/milvus/v2/service/resourcegroup/request/ListResourceGroupsReq.java b/src/main/java/io/milvus/v2/service/resourcegroup/request/ListResourceGroupsReq.java new file mode 100644 index 000000000..77c119303 --- /dev/null +++ b/src/main/java/io/milvus/v2/service/resourcegroup/request/ListResourceGroupsReq.java @@ -0,0 +1,9 @@ +package io.milvus.v2.service.resourcegroup.request; + +import lombok.Data; +import lombok.experimental.SuperBuilder; + +@Data +@SuperBuilder +public class ListResourceGroupsReq { +} diff --git a/src/main/java/io/milvus/v2/service/resourcegroup/request/TransferReplicaReq.java b/src/main/java/io/milvus/v2/service/resourcegroup/request/TransferReplicaReq.java new file mode 100644 index 000000000..65d27890d --- /dev/null +++ b/src/main/java/io/milvus/v2/service/resourcegroup/request/TransferReplicaReq.java @@ -0,0 +1,14 @@ +package io.milvus.v2.service.resourcegroup.request; + +import lombok.Data; +import lombok.experimental.SuperBuilder; + +@Data +@SuperBuilder +public class TransferReplicaReq { + private String sourceGroupName; + private String targetGroupName; + private String collectionName; + private String databaseName; + private long numberOfReplicas; +} diff --git a/src/main/java/io/milvus/v2/service/resourcegroup/request/UpdateResourceGroupsReq.java b/src/main/java/io/milvus/v2/service/resourcegroup/request/UpdateResourceGroupsReq.java new file mode 100644 index 000000000..4e040faa8 --- /dev/null +++ b/src/main/java/io/milvus/v2/service/resourcegroup/request/UpdateResourceGroupsReq.java @@ -0,0 +1,16 @@ +package io.milvus.v2.service.resourcegroup.request; + +import io.milvus.common.resourcegroup.ResourceGroupConfig; +import lombok.Builder; +import lombok.Data; +import lombok.experimental.SuperBuilder; + +import java.util.HashMap; +import java.util.Map; + +@Data +@SuperBuilder +public class UpdateResourceGroupsReq { + @Builder.Default + private Map resourceGroups = new HashMap<>(); +} diff --git a/src/main/java/io/milvus/v2/service/resourcegroup/response/DescribeResourceGroupResp.java b/src/main/java/io/milvus/v2/service/resourcegroup/response/DescribeResourceGroupResp.java new file mode 100644 index 000000000..86f23d8ce --- /dev/null +++ b/src/main/java/io/milvus/v2/service/resourcegroup/response/DescribeResourceGroupResp.java @@ -0,0 +1,26 @@ +package io.milvus.v2.service.resourcegroup.response; + +import io.milvus.common.resourcegroup.NodeInfo; +import io.milvus.common.resourcegroup.ResourceGroupConfig; +import lombok.Builder; +import lombok.Data; +import lombok.experimental.SuperBuilder; + +import java.util.*; + +@Data +@SuperBuilder +public class DescribeResourceGroupResp { + private String groupName; + private Integer capacity; + private Integer numberOfAvailableNode; + @Builder.Default + private Map numberOfLoadedReplica = new HashMap<>(); + @Builder.Default + private Map numberOfOutgoingNode = new HashMap<>(); + @Builder.Default + private Map numberOfIncomingNode = new HashMap<>(); + private ResourceGroupConfig config; + @Builder.Default + private List nodes = new ArrayList<>(); +} diff --git a/src/main/java/io/milvus/v2/service/resourcegroup/response/ListResourceGroupsResp.java b/src/main/java/io/milvus/v2/service/resourcegroup/response/ListResourceGroupsResp.java new file mode 100644 index 000000000..8c57ec300 --- /dev/null +++ b/src/main/java/io/milvus/v2/service/resourcegroup/response/ListResourceGroupsResp.java @@ -0,0 +1,15 @@ +package io.milvus.v2.service.resourcegroup.response; + +import lombok.Builder; +import lombok.Data; +import lombok.experimental.SuperBuilder; + +import java.util.ArrayList; +import java.util.List; + +@Data +@SuperBuilder +public class ListResourceGroupsResp { + @Builder.Default + private List groupNames = new ArrayList<>(); +} diff --git a/src/test/java/io/milvus/v2/client/MilvusClientV2DockerTest.java b/src/test/java/io/milvus/v2/client/MilvusClientV2DockerTest.java index e605af744..d44d042f0 100644 --- a/src/test/java/io/milvus/v2/client/MilvusClientV2DockerTest.java +++ b/src/test/java/io/milvus/v2/client/MilvusClientV2DockerTest.java @@ -24,6 +24,7 @@ import com.google.gson.reflect.TypeToken; import io.milvus.TestUtils; +import io.milvus.common.resourcegroup.*; import io.milvus.common.utils.Float16Utils; import io.milvus.common.utils.JsonUtils; import io.milvus.orm.iterator.QueryIterator; @@ -46,12 +47,15 @@ import io.milvus.v2.service.rbac.PrivilegeGroup; import io.milvus.v2.service.rbac.request.*; import io.milvus.v2.service.rbac.response.*; +import io.milvus.v2.service.resourcegroup.request.*; +import io.milvus.v2.service.resourcegroup.response.*; import io.milvus.v2.service.utility.request.*; import io.milvus.v2.service.utility.response.*; import io.milvus.v2.service.vector.request.*; import io.milvus.v2.service.vector.request.data.*; import io.milvus.v2.service.vector.request.ranker.*; import io.milvus.v2.service.vector.response.*; +import org.apache.commons.lang3.StringUtils; import org.apache.commons.text.RandomStringGenerator; import org.junit.jupiter.api.Assertions; @@ -1862,4 +1866,80 @@ void testRBAC() { Assertions.assertEquals(1, groupsPlivileges.get("dummy").size()); Assertions.assertEquals("CreateCollection", groupsPlivileges.get("dummy").get(0)); } + + @Test + void testResourceGroup() { + String groupA = "group_A"; + String groupDefault = "__default_resource_group"; + client.createResourceGroup(CreateResourceGroupReq.builder() + .groupName(groupA) + .config(ResourceGroupConfig.newBuilder() + .withRequests(new ResourceGroupLimit(3)) + .withLimits(new ResourceGroupLimit(4)) + .withFrom(Collections.singletonList(new ResourceGroupTransfer(groupDefault))) + .withTo(Collections.singletonList(new ResourceGroupTransfer(groupDefault))) + .build()) + .build()); + + ListResourceGroupsResp listResp = client.listResourceGroups(ListResourceGroupsReq.builder().build()); + List groupNames = listResp.getGroupNames(); + Assertions.assertEquals(2, groupNames.size()); + Assertions.assertTrue(groupNames.contains(groupA)); + Assertions.assertTrue(groupNames.contains(groupDefault)); + + // A + DescribeResourceGroupResp descResp = client.describeResourceGroup(DescribeResourceGroupReq.builder() + .groupName(groupA) + .build()); + Assertions.assertEquals(groupA, descResp.getGroupName()); + Assertions.assertEquals(3, descResp.getCapacity()); + Assertions.assertEquals(1, descResp.getNumberOfAvailableNode()); + + ResourceGroupConfig config = descResp.getConfig(); + Assertions.assertEquals(3, config.getRequests().getNodeNum()); + Assertions.assertEquals(4, config.getLimits().getNodeNum()); + + Assertions.assertEquals(1, config.getFrom().size()); + Assertions.assertEquals(groupDefault, config.getFrom().get(0).getResourceGroupName()); + Assertions.assertEquals(1, config.getTo().size()); + Assertions.assertEquals(groupDefault, config.getTo().get(0).getResourceGroupName()); + + List nodes = descResp.getNodes(); + Assertions.assertEquals(1, nodes.size()); + Assertions.assertTrue(nodes.get(0).getNodeId() > 0L); + Assertions.assertTrue(StringUtils.isNotEmpty(nodes.get(0).getAddress())); + Assertions.assertTrue(StringUtils.isNotEmpty(nodes.get(0).getHostname())); + + // update + Map resourceGroups = new HashMap<>(); + resourceGroups.put(groupA, ResourceGroupConfig.newBuilder() + .withRequests(new ResourceGroupLimit(0)) + .withLimits(new ResourceGroupLimit(0)) + .build()); + client.updateResourceGroups(UpdateResourceGroupsReq.builder() + .resourceGroups(resourceGroups) + .build()); + + descResp = client.describeResourceGroup(DescribeResourceGroupReq.builder() + .groupName(groupA) + .build()); + + config = descResp.getConfig(); + Assertions.assertEquals(0, config.getRequests().getNodeNum()); + Assertions.assertEquals(0, config.getLimits().getNodeNum()); + Assertions.assertTrue(config.getFrom().isEmpty()); + Assertions.assertTrue(config.getTo().isEmpty()); + + // drop + client.dropResourceGroup(DropResourceGroupReq.builder() + .groupName(groupA) + .build()); + + // transfer + String collectionName = generator.generate(10); + client.createCollection(CreateCollectionReq.builder() + .collectionName(collectionName) + .dimension(DIMENSION) + .build()); + } }