Skip to content

Commit

Permalink
Add Resource Group API (#417)
Browse files Browse the repository at this point in the history
Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
  • Loading branch information
congqixia authored Mar 8, 2023
1 parent 69f7d14 commit 01b5dd8
Show file tree
Hide file tree
Showing 11 changed files with 1,101 additions and 129 deletions.
19 changes: 15 additions & 4 deletions Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ PWD := $(shell pwd)
GOPATH := $(shell $(GO) env GOPATH)
PROTOC := $(shell which protoc)
PROTOC_VER := $(shell protoc --version)
INSTALL_PATH := $(PWD)/bin

all: check-protoc-version

Expand All @@ -24,20 +25,30 @@ else
@echo "using $(shell which protoc)"
endif

install-tool:
@mkdir -p $(INSTALL_PATH)
@$(INSTALL_PATH)/golangci-lint --version 2>&1 1>/dev/null || (echo "Installing golangci-lint into ./bin/" && curl -sSfL https://raw.githubusercontent.com/golangci/golangci-lint/master/install.sh | sh -s -- -b $(INSTALL_PATH) v1.46.2)
@$(INSTALL_PATH)/mockery --version 2>&1 1>/dev/null || (echo "Installing mockery v2.16.0 to ./bin/" && GOBIN=$(INSTALL_PATH)/ go install github.com/vektra/mockery/v2@v2.16.0)


check-protoc-version: check-protoc
@(env bash $(PWD)/scripts/check_protoc_version.sh)

generate-proto: check-protoc-version
@which protoc-gen-go 1>/dev/null || (echo "Installing protoc-gen-go" && go get github.com/golang/protobuf/protoc-gen-go@v1.5.2)
@(env bash $(PWD)/scripts/proto_gen_go.sh)

generate-mockery: install-tool
@echo "generating mockery milvus service server"
@$(INSTALL_PATH)/mockery --srcpkg=github.com/milvus-io/milvus-proto/go-api/milvuspb --name=MilvusServiceServer --output=mocks --outpkg=mocks --with-expecter

static-check:
@echo "Running $@ check:"
@golangci-lint cache clean
@golangci-lint run --timeout=30m --config ./.golangci.yml ./entity/...
@golangci-lint run --timeout=30m --config ./.golangci.yml ./client/...
@golangci-lint run --timeout=30m --config ./.golangci.yml ./internal/...
@golangci-lint run --timeout=30m --config ./.golangci.yml ./tests/...
@$(INSTALL_PATH)/golangci-lint run --timeout=30m --config ./.golangci.yml ./entity/...
@$(INSTALL_PATH)/golangci-lint run --timeout=30m --config ./.golangci.yml ./client/...
@$(INSTALL_PATH)/golangci-lint run --timeout=30m --config ./.golangci.yml ./internal/...
@$(INSTALL_PATH)/golangci-lint run --timeout=30m --config ./.golangci.yml ./test/...

test-go:
@echo "Running unit tests:"
Expand Down
13 changes: 13 additions & 0 deletions client/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -170,6 +170,19 @@ type Client interface {
// GetLoadState get the collection or partitions load state
GetLoadState(ctx context.Context, collectionName string, partitionNames []string) (entity.LoadState, error)

// ListResourceGroups returns list of resource group names in current Milvus instance.
ListResourceGroups(ctx context.Context) ([]string, error)
// CreateResourceGroup creates a resource group with provided name.
CreateResourceGroup(ctx context.Context, rgName string) error
// DescribeResourceGroup returns resource groups information.
DescribeResourceGroup(ctx context.Context, rgName string) (*entity.ResourceGroup, error)
// DropResourceGroup drops the resource group with provided name.
DropResourceGroup(ctx context.Context, rgName string) error
// TransferNode transfers querynodes between resource groups.
TransferNode(ctx context.Context, sourceRg, targetRg string, nodesNum int32) error
// TransferReplica transfer collection replicas between source,target resource group.
TransferReplica(ctx context.Context, sourceRg, targetRg string, collectionName string, replicaNum int64) error

// GetVersion get milvus version
GetVersion(ctx context.Context) (string, error)
}
Expand Down
142 changes: 142 additions & 0 deletions client/client_grpc_rg.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,142 @@
// Copyright (C) 2019-2021 Zilliz. All rights reserved.
//
// Licensed 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 client

import (
"context"

server "github.com/milvus-io/milvus-proto/go-api/milvuspb"
"github.com/milvus-io/milvus-sdk-go/v2/entity"
)

// ListResourceGroups returns list of resource group names in current Milvus instance.
func (c *GrpcClient) ListResourceGroups(ctx context.Context) ([]string, error) {
if c.Service == nil {
return nil, ErrClientNotReady
}

req := &server.ListResourceGroupsRequest{}

resp, err := c.Service.ListResourceGroups(ctx, req)
if err != nil {
return nil, err
}
if err = handleRespStatus(resp.GetStatus()); err != nil {
return nil, err
}

return resp.GetResourceGroups(), nil
}

// CreateResourceGroup creates a resource group with provided name.
func (c *GrpcClient) CreateResourceGroup(ctx context.Context, rgName string) error {
if c.Service == nil {
return ErrClientNotReady
}

req := &server.CreateResourceGroupRequest{
ResourceGroup: rgName,
}

resp, err := c.Service.CreateResourceGroup(ctx, req)
if err != nil {
return err
}
return handleRespStatus(resp)
}

// DescribeResourceGroup returns resource groups information.
func (c *GrpcClient) DescribeResourceGroup(ctx context.Context, rgName string) (*entity.ResourceGroup, error) {
if c.Service == nil {
return nil, ErrClientNotReady
}

req := &server.DescribeResourceGroupRequest{
ResourceGroup: rgName,
}

resp, err := c.Service.DescribeResourceGroup(ctx, req)
if err != nil {
return nil, err
}
if err = handleRespStatus(resp.GetStatus()); err != nil {
return nil, err
}

rg := resp.GetResourceGroup()
result := &entity.ResourceGroup{
Name: rg.GetName(),
Capacity: rg.GetCapacity(),
AvailableNodesNumber: rg.GetNumAvailableNode(),
LoadedReplica: rg.GetNumLoadedReplica(),
OutgoingNodeNum: rg.GetNumOutgoingNode(),
IncomingNodeNum: rg.GetNumIncomingNode(),
}

return result, nil
}

// DropResourceGroup drops the resource group with provided name.
func (c *GrpcClient) DropResourceGroup(ctx context.Context, rgName string) error {
if c.Service == nil {
return ErrClientNotReady
}

req := &server.DropResourceGroupRequest{
ResourceGroup: rgName,
}

resp, err := c.Service.DropResourceGroup(ctx, req)
if err != nil {
return err
}
return handleRespStatus(resp)
}

// TransferNode transfers querynodes between resource groups.
func (c *GrpcClient) TransferNode(ctx context.Context, sourceRg, targetRg string, nodesNum int32) error {
if c.Service == nil {
return ErrClientNotReady
}

req := &server.TransferNodeRequest{
SourceResourceGroup: sourceRg,
TargetResourceGroup: targetRg,
NumNode: nodesNum,
}

resp, err := c.Service.TransferNode(ctx, req)
if err != nil {
return err
}
return handleRespStatus(resp)
}

// TransferReplica transfer collection replicas between source,target resource group.
func (c *GrpcClient) TransferReplica(ctx context.Context, sourceRg, targetRg string, collectionName string, replicaNum int64) error {
if c.Service == nil {
return ErrClientNotReady
}

req := &server.TransferReplicaRequest{
SourceResourceGroup: sourceRg,
TargetResourceGroup: targetRg,
CollectionName: collectionName,
NumReplica: replicaNum,
}

resp, err := c.Service.TransferReplica(ctx, req)
if err != nil {
return err
}
return handleRespStatus(resp)
}
Loading

0 comments on commit 01b5dd8

Please sign in to comment.