enhance: add resource group declarative api (#313)

* init resource group api update Signed-off-by: ryjiang <jiangruiyi@gmail.com> * finish Signed-off-by: ryjiang <jiangruiyi@gmail.com> * remove console Signed-off-by: ryjiang <jiangruiyi@gmail.com> * recover Signed-off-by: ryjiang <jiangruiyi@gmail.com> --------- Signed-off-by: ryjiang <jiangruiyi@gmail.com>
milvus-io · Jun 4, 2024 · b1a2cd0 · b1a2cd0
1 parent 9478768
commit b1a2cd0
Show file tree

Hide file tree

Showing 3 changed files with 85 additions and 13 deletions.
diff --git a/milvus/grpc/Resource.ts b/milvus/grpc/Resource.ts
@@ -8,6 +8,7 @@ import {
   ListResourceGroupsResponse,
   DescribeResourceGroupsReq,
   DescribeResourceGroupResponse,
+  UpdateRresourceGroupReq,
   TransferNodeReq,
   TransferReplicaReq,
   promisify,
@@ -19,6 +20,7 @@ export class Resource extends Partition {
    *
    * @param {Object} data - The data for the resource group.
    * @param {string} data.resource_group - The name of the resource group.
+   * @param {ResourceGroupConfig} data.config - The configuration for the resource group.
    * @param {number} [data.timeout] - An optional duration of time in milliseconds to allow for the RPC. If it is set to undefined, the client keeps waiting until the server responds or an error occurs. Default is undefined.
    *
    * @returns {Promise<ResStatus>} A promise that resolves to the response status.
@@ -104,6 +106,26 @@ export class Resource extends Partition {
     return promise;
   }
 
+  /**
+   * Updates a resource group.
+   * @param {Object} data - The data for the resource group.
+   * @param {{ [key: string]: ResourceGroupConfig }} data.resource_groups - A map from resource group name to the configuration for the resource group.
+   * @param {number} [data.timeout] - An optional duration of time in milliseconds to allow for the RPC. If it is set to undefined, the client keeps waiting until the server responds or an error occurs. Default is undefined.
+   * @returns {Promise<ResStatus>} A promise that resolves to the response status.
+   */
+  /* istanbul ignore next */
+  async updateResourceGroups(
+    data: UpdateRresourceGroupReq
+  ): Promise<ResStatus> {
+    const promise = await promisify(
+      this.channelPool,
+      'UpdateResourceGroups',
+      data,
+      data.timeout || this.timeout
+    );
+    return promise;
+  }
+
   /**
    * Drops a resource group.
    *

diff --git a/milvus/types/Resource.ts b/milvus/types/Resource.ts
@@ -1,10 +1,34 @@
 import { GrpcTimeOut, resStatusResponse } from './Common';
 
-export interface CreateResourceGroupReq extends GrpcTimeOut {
+type ResourceGroupConfig = {
+  requests?: { node_num: number }; // requests node num in resource group, if node num is less than requests.nodeNum, it will be transfer from other resource group.
+  limits?: { node_num: number }; // limited node num in resource group, if node num is more than limits.nodeNum, it will be transfer to other resource group.
+  transfer_from?: { resource_group: string }[]; // missing node should be transfer from given resource group at high priority in repeated list.
+  transfer_to?: { resource_group: string }[]; // redundant node should be transfer to given resource group at high priority in repeated list.
+};
+
+type ResourceGroup = {
+  name: string;
+  capacity: number;
+  num_available_node: number;
+  num_loaded_replica: { [key: string]: number };
+  num_outgoing_node: { [key: string]: number };
+  num_incoming_node: { [key: string]: number };
+  config?: ResourceGroupConfig;
+};
+
+interface BaseResourceGroupReq extends GrpcTimeOut {
   resource_group: string;
 }
-export interface DropResourceGroupsReq extends CreateResourceGroupReq {}
-export interface DescribeResourceGroupsReq extends CreateResourceGroupReq {}
+
+export interface CreateResourceGroupReq extends BaseResourceGroupReq {
+  config?: ResourceGroupConfig;
+}
+export interface DescribeResourceGroupsReq extends BaseResourceGroupReq {}
+export interface UpdateRresourceGroupReq extends GrpcTimeOut {
+  resource_groups: { [key: string]: ResourceGroupConfig };
+}
+export interface DropResourceGroupsReq extends BaseResourceGroupReq {}
 
 export interface TransferNodeReq extends GrpcTimeOut {
   source_resource_group: string; // required, source resource group
@@ -23,15 +47,6 @@ export interface ListResourceGroupsResponse extends resStatusResponse {
   resource_groups: string[]; // resource groups
 }
 
-type ResourceGroup = {
-  name: string;
-  capacity: number;
-  num_available_node: number;
-  num_loaded_replica: { [key: string]: number };
-  num_outgoing_node: { [key: string]: number };
-  num_incoming_node: { [key: string]: number };
-};
-
 export interface DescribeResourceGroupResponse extends resStatusResponse {
   resource_group: ResourceGroup;
 }
diff --git a/test/grpc/Resource.spec.ts b/test/grpc/Resource.spec.ts
@@ -19,11 +19,11 @@ const dbParam = {
 let runRgTransferTest = false;
 
 describe(`Resource API`, () => {
+  let queryNodeNum = 0;
   beforeAll(async () => {
     await milvusClient.createDatabase(dbParam);
     await milvusClient.use(dbParam);
 
-    // [TODO]: move getMetric to client
     const metrics = await milvusClient.getMetric({
       request: {
         metric_type: 'system_info',
@@ -34,6 +34,8 @@ describe(`Resource API`, () => {
       (node: any) => node.infos.type === 'querynode'
     );
 
+    queryNodeNum = queryNodes.length;
+
     runRgTransferTest = queryNodes.length > 1;
 
     // create collection
@@ -66,6 +68,12 @@ describe(`Resource API`, () => {
 
     const res2 = await milvusClient.createResourceGroup({
       resource_group: resource_group2,
+      // config: {
+      //   requests: { node_num: 1 },
+      //   limits: { node_num: 10 },
+      //   transfer_from: [{ resource_group: DEFAULT_RESOURCE_GROUP }],
+      //   transfer_to: [{ resource_group: DEFAULT_RESOURCE_GROUP }],
+      // },
     });
 
     const res3 = await milvusClient.createResourceGroup({
@@ -89,13 +97,39 @@ describe(`Resource API`, () => {
     const res = await milvusClient.describeResourceGroup({
       resource_group: DEFAULT_RESOURCE_GROUP,
     });
+
+    // console.dir(res, { depth: null });
     expect(res.status.error_code).toEqual(ErrorCode.SUCCESS);
     expect(res.resource_group).toHaveProperty('name');
     expect(res.resource_group).toHaveProperty('num_available_node');
     expect(res.resource_group).toHaveProperty('capacity');
+
     expect(typeof res.resource_group.name).toBe('string');
+    expect(res.resource_group.name).toBe(DEFAULT_RESOURCE_GROUP);
     expect(typeof res.resource_group.num_available_node).toBe('number');
+    expect(res.resource_group.num_available_node).toBe(queryNodeNum);
     expect(typeof res.resource_group.capacity).toBe('number');
+
+    await milvusClient.describeResourceGroup({
+      resource_group: resource_group2,
+    });
+  });
+
+  it(`Update rg should be successful`, async () => {
+    if (runRgTransferTest) {
+      const res = await milvusClient.updateResourceGroups({
+        resource_groups: {
+          [resource_group]: {
+            requests: { node_num: 1 },
+            limits: { node_num: 2 },
+            transfer_from: [{ resource_group: DEFAULT_RESOURCE_GROUP }],
+            transfer_to: [{ resource_group: DEFAULT_RESOURCE_GROUP }],
+          },
+        },
+      });
+
+      expect(res.error_code).toEqual(ErrorCode.SUCCESS);
+    }
   });
 
   it(`Transfer node to another rg should be successful`, async () => {
@@ -145,6 +179,7 @@ describe(`Resource API`, () => {
 
   it(`Drop all resource groups should be successful`, async () => {
     const res = await milvusClient.dropAllResourceGroups();
+
     res.forEach(r => {
       expect(r.error_code).toEqual(ErrorCode.SUCCESS);
     });