test/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2024-12-02 20:09:57 +08:00
Code Issues Actions 12 Packages Projects Releases Wiki Activity
b31d1a1eb5
milvus/internal/datacoord/policy_test.go
XuanYang-cn 623939c9f5
enhance: Remove not in use policies (#29448) 
THe results don't meet our requirements, and the code hasn't been
maintained for a long time.

See also: #29447

Signed-off-by: yangxuan <xuan.yang@zilliz.com>
2023-12-28 10:38:46 +08:00

564 lines
13 KiB
Go
Raw Blame History

// 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 datacoord
import (
"testing"
"time"
"github.com/samber/lo"
"github.com/stretchr/testify/assert"
memkv "github.com/milvus-io/milvus/internal/kv/mem"
)
func TestBufferChannelAssignPolicy(t *testing.T) {
kv := memkv.NewMemoryKV()
channels := []RWChannel{getChannel("chan1", 1)}
store := &ChannelStore{
store: kv,
channelsInfo: map[int64]*NodeChannelInfo{
1: {1, []RWChannel{}},
bufferID: {bufferID, channels},
},
}
updates := BufferChannelAssignPolicy(store, 1).Collect()
assert.NotNil(t, updates)
assert.Equal(t, 2, len(updates))
assert.ElementsMatch(t,
NewChannelOpSet(
NewAddOp(1, channels...),
NewDeleteOp(bufferID, channels...),
).Collect(),
updates)
}
func getChannel(name string, collID int64) *channelMeta {
return &channelMeta{Name: name, CollectionID: collID}
}
func getChannels(ch2Coll map[string]int64) []RWChannel {
return lo.MapToSlice(ch2Coll, func(name string, coll int64) RWChannel {
return &channelMeta{Name: name, CollectionID: coll}
})
}
func TestAverageAssignPolicy(t *testing.T) {
type args struct {
store ROChannelStore
channels []RWChannel
}
tests := []struct {
name string
args args
want *ChannelOpSet
}{
{
"test assign empty cluster",
args{
&ChannelStore{
memkv.NewMemoryKV(),
map[int64]*NodeChannelInfo{},
},
[]RWChannel{getChannel("chan1", 1)},
},
NewChannelOpSet(NewAddOp(bufferID, getChannel("chan1", 1))),
},
{
"test watch same channel",
args{
&ChannelStore{
memkv.NewMemoryKV(),
map[int64]*NodeChannelInfo{
1: {1, []RWChannel{getChannel("chan1", 1)}},
},
},
[]RWChannel{getChannel("chan1", 1)},
},
NewChannelOpSet(),
},
{
"test normal assign",
args{
&ChannelStore{
memkv.NewMemoryKV(),
map[int64]*NodeChannelInfo{
1: {1, []RWChannel{getChannel("chan1", 1), getChannel("chan2", 1)}},
2: {2, []RWChannel{getChannel("chan3", 1)}},
},
},
[]RWChannel{getChannel("chan4", 1)},
},
NewChannelOpSet(NewAddOp(2, getChannel("chan4", 1))),
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := AverageAssignPolicy(tt.args.store, tt.args.channels)
assert.EqualValues(t, tt.want.Collect(), got.Collect())
})
}
}
func TestAvgAssignUnregisteredChannels(t *testing.T) {
type args struct {
store ROChannelStore
nodeID int64
}
tests := []struct {
name string
args args
want *ChannelOpSet
}{
{
"test deregister the last node",
args{
&ChannelStore{
memkv.NewMemoryKV(),
map[int64]*NodeChannelInfo{
1: {1, []RWChannel{getChannel("chan1", 1)}},
},
},
1,
},
NewChannelOpSet(
NewDeleteOp(1, getChannel("chan1", 1)),
NewAddOp(bufferID, getChannel("chan1", 1)),
),
},
{
"test rebalance channels after deregister",
args{
&ChannelStore{
memkv.NewMemoryKV(),
map[int64]*NodeChannelInfo{
1: {1, []RWChannel{getChannel("chan1", 1)}},
2: {2, []RWChannel{getChannel("chan2", 1)}},
3: {3, []RWChannel{}},
},
},
2,
},
NewChannelOpSet(
NewDeleteOp(2, getChannel("chan2", 1)),
NewAddOp(3, getChannel("chan2", 1)),
),
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := AvgAssignUnregisteredChannels(tt.args.store, tt.args.nodeID)
assert.EqualValues(t, tt.want.Collect(), got.Collect())
})
}
}
func TestBgCheckForChannelBalance(t *testing.T) {
type args struct {
channels []*NodeChannelInfo
timestamp time.Time
}
tests := []struct {
name string
args args
want []*NodeChannelInfo
wantErr error
}{
{
"test even distribution",
args{
[]*NodeChannelInfo{
{1, []RWChannel{getChannel("chan1", 1), getChannel("chan2", 1)}},
{2, []RWChannel{getChannel("chan1", 2), getChannel("chan2", 2)}},
{3, []RWChannel{getChannel("chan1", 3), getChannel("chan2", 3)}},
},
time.Now(),
},
// there should be no reallocate
[]*NodeChannelInfo{},
nil,
},
{
"test uneven with conservative effect",
args{
[]*NodeChannelInfo{
{1, []RWChannel{getChannel("chan1", 1), getChannel("chan2", 1)}},
{2, []RWChannel{}},
},
time.Now(),
},
// as we deem that the node having only one channel more than average as even, so there's no reallocation
// for this test case
[]*NodeChannelInfo{},
nil,
},
{
"test uneven with zero",
args{
[]*NodeChannelInfo{
{1, []RWChannel{
getChannel("chan1", 1),
getChannel("chan2", 1),
getChannel("chan3", 1),
}},
{2, []RWChannel{}},
},
time.Now(),
},
[]*NodeChannelInfo{{1, []RWChannel{getChannel("chan1", 1)}}},
nil,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
policy := BgBalanceCheck
got, err := policy(tt.args.channels, tt.args.timestamp)
assert.Equal(t, tt.wantErr, err)
assert.EqualValues(t, tt.want, got)
})
}
}
func TestAvgReassignPolicy(t *testing.T) {
type args struct {
store ROChannelStore
reassigns []*NodeChannelInfo
}
tests := []struct {
name string
args args
want *ChannelOpSet
}{
{
"test_only_one_node",
args{
&ChannelStore{
memkv.NewMemoryKV(),
map[int64]*NodeChannelInfo{
1: {1, []RWChannel{getChannel("chan1", 1)}},
},
},
[]*NodeChannelInfo{{1, []RWChannel{getChannel("chan1", 1)}}},
},
// as there's no available nodes except the input node, there's no reassign plan generated
NewChannelOpSet(),
},
{
"test_zero_avg",
args{
&ChannelStore{
memkv.NewMemoryKV(),
map[int64]*NodeChannelInfo{
1: {1, []RWChannel{getChannel("chan1", 1)}},
2: {2, []RWChannel{}},
3: {2, []RWChannel{}},
4: {2, []RWChannel{}},
},
},
[]*NodeChannelInfo{{1, []RWChannel{getChannel("chan1", 1)}}},
},
// as we use ceil to calculate the wanted average number, there should be one reassign
// though the average num less than 1
NewChannelOpSet(
NewDeleteOp(1, getChannel("chan1", 1)),
NewAddOp(2, getChannel("chan1", 1)),
),
},
{
"test_normal_reassigning_for_one_available_nodes",
args{
&ChannelStore{
memkv.NewMemoryKV(),
map[int64]*NodeChannelInfo{
1: {1, []RWChannel{getChannel("chan1", 1), getChannel("chan2", 1)}},
2: {2, []RWChannel{}},
},
},
[]*NodeChannelInfo{{1, []RWChannel{getChannel("chan1", 1), getChannel("chan2", 1)}}},
},
NewChannelOpSet(
NewDeleteOp(1, getChannel("chan1", 1), getChannel("chan2", 1)),
NewAddOp(2, getChannel("chan1", 1), getChannel("chan2", 1)),
),
},
{
"test_normal_reassigning_for_multiple_available_nodes",
args{
&ChannelStore{
memkv.NewMemoryKV(),
map[int64]*NodeChannelInfo{
1: {1, []RWChannel{
getChannel("chan1", 1),
getChannel("chan2", 1),
getChannel("chan3", 1),
getChannel("chan4", 1),
}},
2: {2, []RWChannel{}},
3: {3, []RWChannel{}},
4: {4, []RWChannel{}},
},
},
[]*NodeChannelInfo{{1, []RWChannel{
getChannel("chan1", 1),
getChannel("chan2", 1),
getChannel("chan3", 1),
}}},
},
NewChannelOpSet(
NewDeleteOp(1, []RWChannel{
getChannel("chan1", 1),
getChannel("chan2", 1),
getChannel("chan3", 1),
}...),
NewAddOp(2, getChannel("chan1", 1)),
NewAddOp(3, getChannel("chan2", 1)),
NewAddOp(4, getChannel("chan3", 1)),
),
},
{
"test_reassigning_for_extreme_case",
args{
&ChannelStore{
memkv.NewMemoryKV(),
map[int64]*NodeChannelInfo{
1: {1, []RWChannel{
getChannel("chan1", 1),
getChannel("chan2", 1),
getChannel("chan3", 1),
getChannel("chan4", 1),
getChannel("chan5", 1),
getChannel("chan6", 1),
getChannel("chan7", 1),
getChannel("chan8", 1),
getChannel("chan9", 1),
getChannel("chan10", 1),
getChannel("chan11", 1),
getChannel("chan12", 1),
}},
2: {2, []RWChannel{
getChannel("chan13", 1),
getChannel("chan14", 1),
}},
3: {3, []RWChannel{getChannel("chan15", 1)}},
4: {4, []RWChannel{}},
},
},
[]*NodeChannelInfo{{1, []RWChannel{
getChannel("chan1", 1),
getChannel("chan2", 1),
getChannel("chan3", 1),
getChannel("chan4", 1),
getChannel("chan5", 1),
getChannel("chan6", 1),
getChannel("chan7", 1),
getChannel("chan8", 1),
getChannel("chan9", 1),
getChannel("chan10", 1),
getChannel("chan11", 1),
getChannel("chan12", 1),
}}},
},
NewChannelOpSet(
NewDeleteOp(1, []RWChannel{
getChannel("chan1", 1),
getChannel("chan2", 1),
getChannel("chan3", 1),
getChannel("chan4", 1),
getChannel("chan5", 1),
getChannel("chan6", 1),
getChannel("chan7", 1),
getChannel("chan8", 1),
getChannel("chan9", 1),
getChannel("chan10", 1),
getChannel("chan11", 1),
getChannel("chan12", 1),
}...),
NewAddOp(4, []RWChannel{
getChannel("chan1", 1),
getChannel("chan2", 1),
getChannel("chan3", 1),
getChannel("chan4", 1),
getChannel("chan5", 1),
}...),
NewAddOp(3, []RWChannel{
getChannel("chan6", 1),
getChannel("chan7", 1),
getChannel("chan8", 1),
getChannel("chan9", 1),
}...),
NewAddOp(2, []RWChannel{
getChannel("chan10", 1),
getChannel("chan11", 1),
getChannel("chan12", 1),
}...),
),
},
}
for _, tt := range tests {
if tt.name == "test_reassigning_for_extreme_case" ||
tt.name == "test_normal_reassigning_for_multiple_available_nodes" {
continue
}
t.Run(tt.name, func(t *testing.T) {
got := AverageReassignPolicy(tt.args.store, tt.args.reassigns)
assert.ElementsMatch(t, tt.want.Collect(), got.Collect())
})
}
}
func TestAvgBalanceChannelPolicy(t *testing.T) {
type args struct {
store ROChannelStore
}
tests := []struct {
name string
args args
want *ChannelOpSet
}{
{
"test_only_one_node",
args{
&ChannelStore{
memkv.NewMemoryKV(),
map[int64]*NodeChannelInfo{
1: {
1, []RWChannel{
getChannel("chan1", 1),
getChannel("chan2", 1),
getChannel("chan3", 1),
getChannel("chan4", 1),
},
},
2: {2, []RWChannel{}},
},
},
},
NewChannelOpSet(NewAddOp(1, getChannel("chan1", 1))),
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := AvgBalanceChannelPolicy(tt.args.store, time.Now())
assert.EqualValues(t, tt.want.Collect(), got.Collect())
})
}
}
func TestAvgAssignRegisterPolicy(t *testing.T) {
type args struct {
store ROChannelStore
nodeID int64
}
tests := []struct {
name string
args args
bufferedUpdates *ChannelOpSet
balanceUpdates *ChannelOpSet
}{
{
"test empty",
args{
&ChannelStore{
memkv.NewMemoryKV(),
map[int64]*NodeChannelInfo{
1: {NodeID: 1, Channels: make([]RWChannel, 0)},
},
},
1,
},
NewChannelOpSet(),
NewChannelOpSet(),
},
{
"test with buffer channel",
args{
&ChannelStore{
memkv.NewMemoryKV(),
map[int64]*NodeChannelInfo{
bufferID: {bufferID, []RWChannel{getChannel("ch1", 1)}},
1: {NodeID: 1, Channels: []RWChannel{}},
},
},
1,
},
NewChannelOpSet(
NewDeleteOp(bufferID, getChannel("ch1", 1)),
NewAddOp(1, getChannel("ch1", 1)),
),
NewChannelOpSet(),
},
{
"test with avg assign",
args{
&ChannelStore{
memkv.NewMemoryKV(),
map[int64]*NodeChannelInfo{
1: {1, []RWChannel{getChannel("ch1", 1), getChannel("ch2", 1)}},
3: {3, []RWChannel{}},
},
},
3,
},
NewChannelOpSet(),
NewChannelOpSet(NewAddOp(1, getChannel("ch1", 1))),
},
{
"test with avg equals to zero",
args{
&ChannelStore{
memkv.NewMemoryKV(),
map[int64]*NodeChannelInfo{
1: {1, []RWChannel{getChannel("ch1", 1)}},
2: {2, []RWChannel{getChannel("ch3", 1)}},
3: {3, []RWChannel{}},
},
},
3,
},
NewChannelOpSet(),
NewChannelOpSet(),
},
{
"test node with empty channel",
args{
&ChannelStore{
memkv.NewMemoryKV(),
map[int64]*NodeChannelInfo{
1: {1, []RWChannel{getChannel("ch1", 1), getChannel("ch2", 1), getChannel("ch3", 1)}},
2: {2, []RWChannel{}},
3: {3, []RWChannel{}},
},
},
3,
},
NewChannelOpSet(),
NewChannelOpSet(NewAddOp(1, getChannel("ch1", 1))),
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
bufferedUpdates, balanceUpdates := AvgAssignRegisterPolicy(tt.args.store, tt.args.nodeID)
assert.EqualValues(t, tt.bufferedUpdates.Collect(), bufferedUpdates.Collect())
assert.EqualValues(t, tt.balanceUpdates.Collect(), balanceUpdates.Collect())
})
}
}
Reference in New Issue View Git Blame Copy Permalink