Skip to content

Commit

Permalink
enhance: add updates resource group for java sdk
Browse files Browse the repository at this point in the history
- Add UpdateResourceGroups and modify AddResourceGroup api.

Signed-off-by: chyezh <[email protected]>
  • Loading branch information
chyezh committed Feb 20, 2024
1 parent ea92b22 commit e8f98a6
Show file tree
Hide file tree
Showing 19 changed files with 873 additions and 5 deletions.
2 changes: 1 addition & 1 deletion .gitmodules
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
[submodule "src/main/milvus-proto"]
path = src/main/milvus-proto
url = https://github.com/milvus-io/milvus-proto.git
url = https://github.com/chyezh/milvus-proto.git
65 changes: 65 additions & 0 deletions examples/main/java/io/milvus/ResourceGroupExample.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
package io.milvus;

import com.google.gson.Gson;

import io.milvus.client.MilvusServiceClient;
import io.milvus.resourcegroup.ResourceGroupManagement;
import io.milvus.param.ConnectParam;

public class ResourceGroupExample {
private static final ResourceGroupManagement manager;
private static final String rgName1 = "rg1";
private static final String rgName2 = "rg2";
private static Gson gson = new Gson();

static {
ConnectParam connectParam = ConnectParam.newBuilder()
.withHost("localhost")
.withPort(19530)
.withAuthorization("root", "Milvus")
.build();
manager = new ResourceGroupManagement(new MilvusServiceClient(connectParam));
}

private static void printResourceGroupInfo() {
manager.listResourceGroups().forEach((name, rg) -> {
System.out.println(name);
System.out.println(gson.toJson(rg));
});
}

public static void main(String[] args) throws Exception {
printResourceGroupInfo();
// Initialize the cluster with 1 resource group
// default_rg: 1
manager.initializeCluster(1);
printResourceGroupInfo();

// Add one more node to default rg.
manager.scaleResourceGroupTo(ResourceGroupManagement.DEFAULT_RG, 2);
// add new query node.
// default_rg: 2
printResourceGroupInfo();

// Add a new resource group.
manager.createResourceGroup(rgName1, 1);
// default_rg: 2, rg1: 0
// add new query node.
// default_rg: 2, rg1: 1
printResourceGroupInfo();

// Add a new resource group.
manager.createResourceGroup(rgName2, 2);
// default_rg: 2, rg1: 1, rg2: 0
// add new query node.
// default_rg: 2, rg1: 1, rg2: 1
// add new query node.
// default_rg: 2, rg1: 1, rg2: 2
printResourceGroupInfo();

// downscale default_rg to 1
manager.scaleResourceGroupTo(ResourceGroupManagement.DEFAULT_RG, 1);
// default_rg: 1, rg1: 1, rg2: 2, recycle_rg: 1
printResourceGroupInfo();
}
}
46 changes: 46 additions & 0 deletions examples/main/java/io/milvus/resourcegroup/NodeInfo.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
package io.milvus.resourcegroup;

import lombok.Getter;
import lombok.NonNull;

@Getter
public class NodeInfo {
private long nodeId;
private String address;
private String hostname;

private NodeInfo(Builder builder) {
this.nodeId = builder.nodeId;
this.address = builder.address;
this.hostname = builder.hostname;
}

public static Builder newBuilder() {
return new Builder();
}

public static class Builder {
private long nodeId;
private String address;
private String hostname;

public Builder withNodeId(long nodeId) {
this.nodeId = nodeId;
return this;
}

public Builder withAddress(@NonNull String address) {
this.address = address;
return this;
}

public Builder withHostname(@NonNull String hostname) {
this.hostname = hostname;
return this;
}

public NodeInfo build() {
return new NodeInfo(this);
}
}
}
111 changes: 111 additions & 0 deletions examples/main/java/io/milvus/resourcegroup/ResourceGroupInfo.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,111 @@
package io.milvus.resourcegroup;

import java.util.HashSet;
import java.util.Set;

import io.milvus.common.resourcegroup.ResourceGroupConfig;
import lombok.Getter;
import lombok.NonNull;

@Getter
public class ResourceGroupInfo {
private String resourceGroupName;
private ResourceGroupConfig resourceGroupConfig;
private Set<String> fullDatabases; // databases belong to this resource group completely.
private Set<String> partialDatabases; // databases belong to this resource group partially, some collection is in
// other resource group.
private Set<NodeInfo> nodes; // actual query node in this resource group.
private Integer requestsNodeNum; // max query node in this resource group.

private ResourceGroupInfo(@NonNull Builder builder) {
this.resourceGroupName = builder.resourceGroupName;
this.resourceGroupConfig = builder.resourceGroupConfig;
this.fullDatabases = builder.fullDatabases;
if (this.fullDatabases == null) {
this.fullDatabases = new HashSet<String>();
}
this.partialDatabases = builder.partialDatabases;
if (this.partialDatabases == null) {
this.partialDatabases = new HashSet<String>();
}
this.nodes = builder.nodes;
if (this.nodes == null) {
this.nodes = new HashSet<NodeInfo>();
}
this.requestsNodeNum = builder.requestsNodeNum;
}

public static Builder newBuilder() {
return new Builder();
}

public static final class Builder {
private String resourceGroupName;
private ResourceGroupConfig resourceGroupConfig;
private Set<String> fullDatabases;
private Set<String> partialDatabases;
private Set<NodeInfo> nodes; // actual query node in this resource group.
private Integer requestsNodeNum;

public Builder withResourceGroupName(@NonNull String resourceGroupName) {
this.resourceGroupName = resourceGroupName;
return this;
}

public Builder addFullDatabases(@NonNull String databaseName) {
if (this.fullDatabases == null) {
this.fullDatabases = new HashSet<String>();
}
this.fullDatabases.add(databaseName);
return this;
}

public Builder addPartialDatabases(@NonNull String databaseName) {
if (this.partialDatabases == null) {
this.partialDatabases = new HashSet<String>();
}
this.partialDatabases.add(databaseName);
return this;
}

public Builder addAvailableNode(@NonNull NodeInfo node) {
if (this.nodes == null) {
this.nodes = new HashSet<NodeInfo>();
}
this.nodes.add(node);
return this;
}

public Builder withRequestsNodeNum(@NonNull Integer requestsNodeNum) {
this.requestsNodeNum = requestsNodeNum;
return this;
}

public Builder withConfig(@NonNull ResourceGroupConfig resourceGroupConfig) {
this.resourceGroupConfig = resourceGroupConfig;
return this;
}

public ResourceGroupInfo build() {
return new ResourceGroupInfo(this);
}
}

/**
* Check if this resource group is the default resource group.
*
* @return true if this resource group is the default resource group.
*/
public boolean isDefaultResourceGroup() {
return this.resourceGroupName == ResourceGroupManagement.DEFAULT_RG;
}

/**
* Check if this resource group is the recycle resource group.
*
* @return true if this resource group is the recycle resource group.
*/
public boolean isRecycleResourceGroup() {
return this.resourceGroupName == ResourceGroupManagement.RECYCLE_RG;
}
}
Loading

0 comments on commit e8f98a6

Please sign in to comment.