mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-02 11:59:00 +08:00
27cc9f2630
issue: #30633 Signed-off-by: Cai Zhang <cai.zhang@zilliz.com> Co-authored-by: chasingegg <chao.gao@zilliz.com>
320 lines
11 KiB
Go
320 lines
11 KiB
Go
// Licensed to the LF AI & Data foundation 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 indexnode
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
|
|
"github.com/cockroachdb/errors"
|
|
clientv3 "go.etcd.io/etcd/client/v3"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
|
|
"github.com/milvus-io/milvus/internal/proto/indexpb"
|
|
"github.com/milvus-io/milvus/internal/proto/internalpb"
|
|
"github.com/milvus-io/milvus/internal/types"
|
|
"github.com/milvus-io/milvus/pkg/util/hardware"
|
|
"github.com/milvus-io/milvus/pkg/util/merr"
|
|
"github.com/milvus-io/milvus/pkg/util/metricsinfo"
|
|
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
|
"github.com/milvus-io/milvus/pkg/util/typeutil"
|
|
)
|
|
|
|
// Mock is an alternative to IndexNode, it will return specific results based on specific parameters.
|
|
type Mock struct {
|
|
types.IndexNode
|
|
|
|
CallInit func() error
|
|
CallStart func() error
|
|
CallStop func() error
|
|
CallGetComponentStates func(ctx context.Context) (*milvuspb.ComponentStates, error)
|
|
CallGetStatisticsChannel func(ctx context.Context) (*milvuspb.StringResponse, error)
|
|
CallRegister func() error
|
|
|
|
CallSetAddress func(address string)
|
|
CallSetEtcdClient func(etcdClient *clientv3.Client)
|
|
CallUpdateStateCode func(stateCode commonpb.StateCode)
|
|
|
|
CallCreateJob func(ctx context.Context, req *indexpb.CreateJobRequest) (*commonpb.Status, error)
|
|
CallQueryJobs func(ctx context.Context, in *indexpb.QueryJobsRequest) (*indexpb.QueryJobsResponse, error)
|
|
CallDropJobs func(ctx context.Context, in *indexpb.DropJobsRequest) (*commonpb.Status, error)
|
|
CallGetJobStats func(ctx context.Context, in *indexpb.GetJobStatsRequest) (*indexpb.GetJobStatsResponse, error)
|
|
CallCreateJobV2 func(ctx context.Context, req *indexpb.CreateJobV2Request) (*commonpb.Status, error)
|
|
CallQueryJobV2 func(ctx context.Context, req *indexpb.QueryJobsV2Request) (*indexpb.QueryJobsV2Response, error)
|
|
CallDropJobV2 func(ctx context.Context, req *indexpb.DropJobsV2Request) (*commonpb.Status, error)
|
|
|
|
CallGetMetrics func(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error)
|
|
CallShowConfigurations func(ctx context.Context, req *internalpb.ShowConfigurationsRequest) (*internalpb.ShowConfigurationsResponse, error)
|
|
}
|
|
|
|
func NewIndexNodeMock() *Mock {
|
|
return &Mock{
|
|
CallInit: func() error {
|
|
return nil
|
|
},
|
|
CallStart: func() error {
|
|
return nil
|
|
},
|
|
CallRegister: func() error {
|
|
return nil
|
|
},
|
|
CallStop: func() error {
|
|
return nil
|
|
},
|
|
CallSetAddress: func(address string) {
|
|
},
|
|
CallSetEtcdClient: func(etcdClient *clientv3.Client) {
|
|
},
|
|
CallUpdateStateCode: func(stateCode commonpb.StateCode) {
|
|
},
|
|
CallGetComponentStates: func(ctx context.Context) (*milvuspb.ComponentStates, error) {
|
|
return &milvuspb.ComponentStates{
|
|
State: &milvuspb.ComponentInfo{
|
|
NodeID: 1,
|
|
Role: typeutil.IndexNodeRole,
|
|
StateCode: commonpb.StateCode_Healthy,
|
|
},
|
|
SubcomponentStates: nil,
|
|
Status: merr.Success(),
|
|
}, nil
|
|
},
|
|
CallGetStatisticsChannel: func(ctx context.Context) (*milvuspb.StringResponse, error) {
|
|
return &milvuspb.StringResponse{
|
|
Status: merr.Success(),
|
|
}, nil
|
|
},
|
|
CallCreateJob: func(ctx context.Context, req *indexpb.CreateJobRequest) (*commonpb.Status, error) {
|
|
return merr.Success(), nil
|
|
},
|
|
CallQueryJobs: func(ctx context.Context, in *indexpb.QueryJobsRequest) (*indexpb.QueryJobsResponse, error) {
|
|
indexInfos := make([]*indexpb.IndexTaskInfo, 0)
|
|
for _, buildID := range in.BuildIDs {
|
|
indexInfos = append(indexInfos, &indexpb.IndexTaskInfo{
|
|
BuildID: buildID,
|
|
State: commonpb.IndexState_Finished,
|
|
IndexFileKeys: []string{"file1", "file2"},
|
|
})
|
|
}
|
|
return &indexpb.QueryJobsResponse{
|
|
Status: merr.Success(),
|
|
ClusterID: in.ClusterID,
|
|
IndexInfos: indexInfos,
|
|
}, nil
|
|
},
|
|
CallDropJobs: func(ctx context.Context, in *indexpb.DropJobsRequest) (*commonpb.Status, error) {
|
|
return merr.Success(), nil
|
|
},
|
|
CallCreateJobV2: func(ctx context.Context, req *indexpb.CreateJobV2Request) (*commonpb.Status, error) {
|
|
return merr.Success(), nil
|
|
},
|
|
CallQueryJobV2: func(ctx context.Context, req *indexpb.QueryJobsV2Request) (*indexpb.QueryJobsV2Response, error) {
|
|
switch req.GetJobType() {
|
|
case indexpb.JobType_JobTypeIndexJob:
|
|
results := make([]*indexpb.IndexTaskInfo, 0)
|
|
for _, buildID := range req.GetTaskIDs() {
|
|
results = append(results, &indexpb.IndexTaskInfo{
|
|
BuildID: buildID,
|
|
State: commonpb.IndexState_Finished,
|
|
IndexFileKeys: []string{},
|
|
SerializedSize: 1024,
|
|
FailReason: "",
|
|
CurrentIndexVersion: 1,
|
|
IndexStoreVersion: 1,
|
|
})
|
|
}
|
|
return &indexpb.QueryJobsV2Response{
|
|
Status: merr.Success(),
|
|
ClusterID: req.GetClusterID(),
|
|
Result: &indexpb.QueryJobsV2Response_IndexJobResults{
|
|
IndexJobResults: &indexpb.IndexJobResults{
|
|
Results: results,
|
|
},
|
|
},
|
|
}, nil
|
|
case indexpb.JobType_JobTypeAnalyzeJob:
|
|
results := make([]*indexpb.AnalyzeResult, 0)
|
|
for _, taskID := range req.GetTaskIDs() {
|
|
results = append(results, &indexpb.AnalyzeResult{
|
|
TaskID: taskID,
|
|
State: indexpb.JobState_JobStateFinished,
|
|
CentroidsFile: fmt.Sprintf("%d/stats_file", taskID),
|
|
FailReason: "",
|
|
})
|
|
}
|
|
return &indexpb.QueryJobsV2Response{
|
|
Status: merr.Success(),
|
|
ClusterID: req.GetClusterID(),
|
|
Result: &indexpb.QueryJobsV2Response_AnalyzeJobResults{
|
|
AnalyzeJobResults: &indexpb.AnalyzeResults{
|
|
Results: results,
|
|
},
|
|
},
|
|
}, nil
|
|
default:
|
|
return &indexpb.QueryJobsV2Response{
|
|
Status: merr.Status(errors.New("unknown job type")),
|
|
ClusterID: req.GetClusterID(),
|
|
}, nil
|
|
}
|
|
},
|
|
CallDropJobV2: func(ctx context.Context, req *indexpb.DropJobsV2Request) (*commonpb.Status, error) {
|
|
return merr.Success(), nil
|
|
},
|
|
CallGetJobStats: func(ctx context.Context, in *indexpb.GetJobStatsRequest) (*indexpb.GetJobStatsResponse, error) {
|
|
return &indexpb.GetJobStatsResponse{
|
|
Status: merr.Success(),
|
|
TotalJobNum: 1,
|
|
EnqueueJobNum: 0,
|
|
InProgressJobNum: 1,
|
|
TaskSlots: 1,
|
|
JobInfos: []*indexpb.JobInfo{
|
|
{
|
|
NumRows: 1024,
|
|
Dim: 128,
|
|
StartTime: 1,
|
|
EndTime: 10,
|
|
PodID: 1,
|
|
},
|
|
},
|
|
}, nil
|
|
},
|
|
CallGetMetrics: func(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) {
|
|
return getMockSystemInfoMetrics(ctx, req, nil)
|
|
},
|
|
CallShowConfigurations: func(ctx context.Context, req *internalpb.ShowConfigurationsRequest) (*internalpb.ShowConfigurationsResponse, error) {
|
|
return &internalpb.ShowConfigurationsResponse{
|
|
Status: merr.Success(),
|
|
}, nil
|
|
},
|
|
}
|
|
}
|
|
|
|
func (m *Mock) Init() error {
|
|
return m.CallInit()
|
|
}
|
|
|
|
func (m *Mock) Start() error {
|
|
return m.CallStart()
|
|
}
|
|
|
|
func (m *Mock) Stop() error {
|
|
return m.CallStop()
|
|
}
|
|
|
|
func (m *Mock) GetComponentStates(ctx context.Context, req *milvuspb.GetComponentStatesRequest) (*milvuspb.ComponentStates, error) {
|
|
return m.CallGetComponentStates(ctx)
|
|
}
|
|
|
|
func (m *Mock) GetStatisticsChannel(ctx context.Context, req *internalpb.GetStatisticsChannelRequest) (*milvuspb.StringResponse, error) {
|
|
return m.CallGetStatisticsChannel(ctx)
|
|
}
|
|
|
|
func (m *Mock) Register() error {
|
|
return m.CallRegister()
|
|
}
|
|
|
|
func (m *Mock) SetAddress(address string) {
|
|
m.CallSetAddress(address)
|
|
}
|
|
|
|
func (m *Mock) GetAddress() string {
|
|
return ""
|
|
}
|
|
|
|
func (m *Mock) SetEtcdClient(etcdClient *clientv3.Client) {
|
|
}
|
|
|
|
func (m *Mock) UpdateStateCode(stateCode commonpb.StateCode) {
|
|
}
|
|
|
|
func (m *Mock) CreateJob(ctx context.Context, req *indexpb.CreateJobRequest) (*commonpb.Status, error) {
|
|
return m.CallCreateJob(ctx, req)
|
|
}
|
|
|
|
func (m *Mock) QueryJobs(ctx context.Context, req *indexpb.QueryJobsRequest) (*indexpb.QueryJobsResponse, error) {
|
|
return m.CallQueryJobs(ctx, req)
|
|
}
|
|
|
|
func (m *Mock) DropJobs(ctx context.Context, req *indexpb.DropJobsRequest) (*commonpb.Status, error) {
|
|
return m.CallDropJobs(ctx, req)
|
|
}
|
|
|
|
func (m *Mock) GetJobStats(ctx context.Context, req *indexpb.GetJobStatsRequest) (*indexpb.GetJobStatsResponse, error) {
|
|
return m.CallGetJobStats(ctx, req)
|
|
}
|
|
|
|
func (m *Mock) GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) {
|
|
return m.CallGetMetrics(ctx, req)
|
|
}
|
|
|
|
func (m *Mock) CreateJobV2(ctx context.Context, req *indexpb.CreateJobV2Request) (*commonpb.Status, error) {
|
|
return m.CallCreateJobV2(ctx, req)
|
|
}
|
|
|
|
func (m *Mock) QueryJobsV2(ctx context.Context, req *indexpb.QueryJobsV2Request) (*indexpb.QueryJobsV2Response, error) {
|
|
return m.CallQueryJobV2(ctx, req)
|
|
}
|
|
|
|
func (m *Mock) DropJobsV2(ctx context.Context, req *indexpb.DropJobsV2Request) (*commonpb.Status, error) {
|
|
return m.CallDropJobV2(ctx, req)
|
|
}
|
|
|
|
// ShowConfigurations returns the configurations of Mock indexNode matching req.Pattern
|
|
func (m *Mock) ShowConfigurations(ctx context.Context, req *internalpb.ShowConfigurationsRequest) (*internalpb.ShowConfigurationsResponse, error) {
|
|
return m.CallShowConfigurations(ctx, req)
|
|
}
|
|
|
|
func getMockSystemInfoMetrics(
|
|
ctx context.Context,
|
|
req *milvuspb.GetMetricsRequest,
|
|
node *Mock,
|
|
) (*milvuspb.GetMetricsResponse, error) {
|
|
// TODO(dragondriver): add more metrics
|
|
nodeInfos := metricsinfo.IndexNodeInfos{
|
|
BaseComponentInfos: metricsinfo.BaseComponentInfos{
|
|
Name: metricsinfo.ConstructComponentName(typeutil.IndexNodeRole, paramtable.GetNodeID()),
|
|
HardwareInfos: metricsinfo.HardwareMetrics{
|
|
CPUCoreCount: hardware.GetCPUNum(),
|
|
CPUCoreUsage: hardware.GetCPUUsage(),
|
|
Memory: 1000,
|
|
MemoryUsage: hardware.GetUsedMemoryCount(),
|
|
Disk: hardware.GetDiskCount(),
|
|
DiskUsage: hardware.GetDiskUsage(),
|
|
},
|
|
SystemInfo: metricsinfo.DeployMetrics{},
|
|
CreatedTime: paramtable.GetCreateTime().String(),
|
|
UpdatedTime: paramtable.GetUpdateTime().String(),
|
|
Type: typeutil.IndexNodeRole,
|
|
},
|
|
SystemConfigurations: metricsinfo.IndexNodeConfiguration{
|
|
MinioBucketName: Params.MinioCfg.BucketName.GetValue(),
|
|
SimdType: Params.CommonCfg.SimdType.GetValue(),
|
|
},
|
|
}
|
|
|
|
metricsinfo.FillDeployMetricsWithEnv(&nodeInfos.SystemInfo)
|
|
|
|
resp, _ := metricsinfo.MarshalComponentInfos(nodeInfos)
|
|
|
|
return &milvuspb.GetMetricsResponse{
|
|
Status: merr.Success(),
|
|
Response: resp,
|
|
ComponentName: metricsinfo.ConstructComponentName(typeutil.IndexNodeRole, paramtable.GetNodeID()),
|
|
}, nil
|
|
}
|