admin/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2026-02-02 01:06:41 +08:00
Code Issues Packages Projects Releases Wiki Activity

Seperate reading build and load params for diskann (#23722)

Browse Source 
Signed-off-by: chasingegg <chao.gao@zilliz.com>
This commit is contained in:
Gao 2023-04-27 10:54:36 +08:00 committed by GitHub
parent 1a773b8db1
commit 21762d9c23
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 96 additions and 52 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
internal/querynode/load_index_info.go
View File

@ -73,7 +73,7 @@ func (li *LoadIndexInfo) appendLoadIndexInfo(bytesIndex [][]byte, indexInfo *que
// some build params also exist in indexParams, which are useless during loading process
indexParams := funcutil.KeyValuePair2Map(indexInfo.IndexParams)
if indexParams["index_type"] == indexparamcheck.IndexDISKANN {
err = indexparams.SetDiskIndexLoadParams(indexParams, indexInfo.GetNumRows())
err = indexparams.SetDiskIndexLoadParams(&Params, indexParams, indexInfo.GetNumRows())
if err != nil {
return err
}

47
internal/util/indexparams/disk_index_params.go
View File

@ -59,10 +59,6 @@ func FillDiskIndexParams(params *paramtable.ComponentParam, indexParams map[stri
pqCodeBudgetGBRatio := params.CommonCfg.PQCodeBudgetGBRatio
buildNumThreadsRatio := params.CommonCfg.BuildNumThreadsRatio
searchCacheBudgetGBRatio := params.CommonCfg.SearchCacheBudgetGBRatio
loadNumThreadRatio := params.CommonCfg.LoadNumThreadRatio
beamWidthRatio := params.CommonCfg.BeamWidthRatio
if params.AutoIndexConfig.Enable {
var ok bool
maxDegree, ok = params.AutoIndexConfig.IndexParams[MaxDegreeKey]
@ -75,18 +71,12 @@ func FillDiskIndexParams(params *paramtable.ComponentParam, indexParams map[stri
}
pqCodeBudgetGBRatio = params.AutoIndexConfig.BigDataExtraParams.PQCodeBudgetGBRatio
buildNumThreadsRatio = params.AutoIndexConfig.BigDataExtraParams.BuildNumThreadsRatio
searchCacheBudgetGBRatio = params.AutoIndexConfig.BigDataExtraParams.SearchCacheBudgetGBRatio
loadNumThreadRatio = params.AutoIndexConfig.BigDataExtraParams.LoadNumThreadRatio
beamWidthRatio = params.AutoIndexConfig.BigDataExtraParams.BeamWidthRatio
}
indexParams[MaxDegreeKey] = maxDegree
indexParams[SearchListSizeKey] = searchListSize
indexParams[PQCodeBudgetRatioKey] = fmt.Sprintf("%f", pqCodeBudgetGBRatio)
indexParams[NumBuildThreadRatioKey] = fmt.Sprintf("%f", buildNumThreadsRatio)
indexParams[SearchCacheBudgetRatioKey] = fmt.Sprintf("%f", searchCacheBudgetGBRatio)
indexParams[NumLoadThreadRatioKey] = fmt.Sprintf("%f", loadNumThreadRatio)
indexParams[BeamWidthRatioKey] = fmt.Sprintf("%f", beamWidthRatio)
return nil
}
@ -120,7 +110,7 @@ func SetDiskIndexBuildParams(indexParams map[string]string, fieldDataSize int64)
// SetDiskIndexLoadParams set disk index load params with ratio params on queryNode
// QueryNode cal load params with ratio params ans cpu count...
func SetDiskIndexLoadParams(indexParams map[string]string, numRows int64) error {
func SetDiskIndexLoadParams(params *paramtable.ComponentParam, indexParams map[string]string, numRows int64) error {
dimStr, ok := indexParams["dim"]
if !ok {
// type param dim has been put into index params before build index
@ -131,29 +121,18 @@ func SetDiskIndexLoadParams(indexParams map[string]string, numRows int64) error
return err
}
searchCacheBudgetGBRatioStr, ok := indexParams[SearchCacheBudgetRatioKey]
if !ok {
return fmt.Errorf("index param searchCacheBudgetGBRatio not exist")
}
searchCacheBudgetGBRatio, err := strconv.ParseFloat(searchCacheBudgetGBRatioStr, 64)
if err != nil {
return err
}
loadNumThreadRatioStr, ok := indexParams[NumLoadThreadRatioKey]
if !ok {
return fmt.Errorf("index param loadNumThreadRatio not exist")
}
loadNumThreadRatio, err := strconv.ParseFloat(loadNumThreadRatioStr, 64)
if err != nil {
return err
}
beamWidthRatioStr, ok := indexParams[BeamWidthRatioKey]
if !ok {
return fmt.Errorf("index param beamWidthRatio not exist")
}
beamWidthRatio, err := strconv.ParseFloat(beamWidthRatioStr, 64)
if err != nil {
return err
var searchCacheBudgetGBRatio float64
var loadNumThreadRatio float64
var beamWidthRatio float64
if params.AutoIndexConfig.Enable {
searchCacheBudgetGBRatio = params.AutoIndexConfig.BigDataExtraParams.SearchCacheBudgetGBRatio
loadNumThreadRatio = params.AutoIndexConfig.BigDataExtraParams.LoadNumThreadRatio
beamWidthRatio = params.AutoIndexConfig.BigDataExtraParams.BeamWidthRatio
} else {
searchCacheBudgetGBRatio = params.CommonCfg.SearchCacheBudgetGBRatio
loadNumThreadRatio = params.CommonCfg.LoadNumThreadRatio
beamWidthRatio = params.CommonCfg.BeamWidthRatio
}
indexParams[SearchCacheBudgetKey] = fmt.Sprintf("%f",

99
internal/util/indexparams/disk_index_params_test.go
View File

@ -17,22 +17,25 @@
package indexparams
import (
"fmt"
"strconv"
"testing"
"github.com/stretchr/testify/assert"
"github.com/milvus-io/milvus/internal/util/autoindex"
"github.com/milvus-io/milvus/internal/util/hardware"
"github.com/milvus-io/milvus/internal/util/paramtable"
)
func TestDiskIndexParams(t *testing.T) {
t.Run("fill index params with auto index", func(t *testing.T) {
var params paramtable.ComponentParam
params.Init()
params.AutoIndexConfig.Enable = true
mapString := make(map[string]string)
mapString[autoindex.BuildRatioKey] = "{\"pg_code_budget_gb\": 0.125, \"num_threads\": 1}"
mapString[autoindex.BuildRatioKey] = "{\"pq_code_budget_gb\": 0.325, \"num_threads\": 2}"
mapString[autoindex.PrepareRatioKey] = "{\"search_cache_budget_gb\": 0.225, \"num_threads\": 4}"
extraParams, err := autoindex.NewBigDataExtraParamsFromMap(mapString)
assert.NoError(t, err)
@ -48,23 +51,29 @@ func TestDiskIndexParams(t *testing.T) {
pqCodeBudgetGBRatio, err := strconv.ParseFloat(indexParams[PQCodeBudgetRatioKey], 64)
assert.NoError(t, err)
assert.Equal(t, 0.125, pqCodeBudgetGBRatio)
assert.Equal(t, 0.325, pqCodeBudgetGBRatio)
assert.NotEqual(t, params.CommonCfg.PQCodeBudgetGBRatio, pqCodeBudgetGBRatio)
buildNumThreadsRatio, err := strconv.ParseFloat(indexParams[NumBuildThreadRatioKey], 64)
assert.NoError(t, err)
assert.Equal(t, 1.0, buildNumThreadsRatio)
assert.Equal(t, 2.0, buildNumThreadsRatio)
assert.NotEqual(t, params.CommonCfg.BuildNumThreadsRatio, buildNumThreadsRatio)
searchCacheBudgetGBRatio, err := strconv.ParseFloat(indexParams[SearchCacheBudgetRatioKey], 64)
// disable autoindex, use default config
params.AutoIndexConfig.Enable = false
indexParams = make(map[string]string)
err = FillDiskIndexParams(&params, indexParams)
assert.NoError(t, err)
assert.Equal(t, 0.225, searchCacheBudgetGBRatio)
loadNumThreadRatio, err := strconv.ParseFloat(indexParams[NumLoadThreadRatioKey], 64)
pqCodeBudgetGBRatio, err = strconv.ParseFloat(indexParams[PQCodeBudgetRatioKey], 64)
assert.NoError(t, err)
assert.Equal(t, 4.0, loadNumThreadRatio)
assert.NotEqual(t, 0.325, pqCodeBudgetGBRatio)
assert.Equal(t, params.CommonCfg.PQCodeBudgetGBRatio, pqCodeBudgetGBRatio)
beamWidthRatio, err := strconv.ParseFloat(indexParams[BeamWidthRatioKey], 64)
buildNumThreadsRatio, err = strconv.ParseFloat(indexParams[NumBuildThreadRatioKey], 64)
assert.NoError(t, err)
assert.Equal(t, 4.0, beamWidthRatio)
assert.NotEqual(t, 2.0, buildNumThreadsRatio)
assert.Equal(t, params.CommonCfg.BuildNumThreadsRatio, buildNumThreadsRatio)
})
t.Run("set disk index build params", func(t *testing.T) {
@ -84,23 +93,79 @@ func TestDiskIndexParams(t *testing.T) {
})
t.Run("set disk index load params", func(t *testing.T) {
indexParams := make(map[string]string)
indexParams[SearchCacheBudgetRatioKey] = "0.125"
indexParams[NumLoadThreadRatioKey] = "8.0"
indexParams[BeamWidthRatioKey] = "4.0"
var params paramtable.ComponentParam
params.Init()
params.AutoIndexConfig.Enable = true
err := SetDiskIndexLoadParams(indexParams, 100)
mapString := make(map[string]string)
mapString[autoindex.BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}"
mapString[autoindex.PrepareRatioKey] = "{\"search_cache_budget_gb\": 0.15, \"num_threads\": 4}"
extraParams, err := autoindex.NewBigDataExtraParamsFromMap(mapString)
assert.NoError(t, err)
params.AutoIndexConfig.BigDataExtraParams = extraParams
params.AutoIndexConfig.IndexParams = make(map[string]string)
params.AutoIndexConfig.IndexParams["max_degree"] = "56"
params.AutoIndexConfig.IndexParams["search_list_size"] = "100"
params.AutoIndexConfig.IndexParams["index_type"] = "DISKANN"
indexParams := make(map[string]string)
err = SetDiskIndexLoadParams(&params, indexParams, 100)
assert.Error(t, err)
indexParams["dim"] = "128"
err = SetDiskIndexLoadParams(indexParams, 100)
err = SetDiskIndexLoadParams(&params, indexParams, 100)
assert.NoError(t, err)
_, ok := indexParams[SearchCacheBudgetKey]
assert.True(t, ok)
_, ok = indexParams[NumLoadThreadKey]
searchCacheBudget, ok := indexParams[SearchCacheBudgetKey]
assert.True(t, ok)
_, ok = indexParams[BeamWidthKey]
assert.Equal(t, fmt.Sprintf("%f", float32(getRowDataSizeOfFloatVector(100, 128))*float32(extraParams.SearchCacheBudgetGBRatio)/(1<<30)), searchCacheBudget)
numLoadThread, ok := indexParams[NumLoadThreadKey]
assert.True(t, ok)
expectedNumLoadThread := int(float32(hardware.GetCPUNum()) * float32(extraParams.LoadNumThreadRatio))
if expectedNumLoadThread > MaxLoadThread {
expectedNumLoadThread = MaxLoadThread
}
indexParams[NumLoadThreadKey] = strconv.Itoa(expectedNumLoadThread)
assert.Equal(t, strconv.Itoa(expectedNumLoadThread), numLoadThread)
beamWidth, ok := indexParams[BeamWidthKey]
assert.True(t, ok)
expectedBeamWidth := int(float32(hardware.GetCPUNum()) * float32(extraParams.BeamWidthRatio))
if expectedBeamWidth > MaxBeamWidth {
expectedBeamWidth = MaxBeamWidth
}
indexParams[BeamWidthKey] = strconv.Itoa(expectedBeamWidth)
assert.Equal(t, strconv.Itoa(expectedBeamWidth), beamWidth)
// disable autoindex, use default config
params.AutoIndexConfig.Enable = false
err = SetDiskIndexLoadParams(&params, indexParams, 100)
assert.NoError(t, err)
searchCacheBudget, ok = indexParams[SearchCacheBudgetKey]
assert.True(t, ok)
assert.Equal(t, fmt.Sprintf("%f", float32(getRowDataSizeOfFloatVector(100, 128))*float32(params.CommonCfg.SearchCacheBudgetGBRatio)/(1<<30)), searchCacheBudget)
numLoadThread, ok = indexParams[NumLoadThreadKey]
assert.True(t, ok)
expectedNumLoadThread = int(float32(hardware.GetCPUNum()) * float32(params.CommonCfg.LoadNumThreadRatio))
if expectedNumLoadThread > MaxLoadThread {
expectedNumLoadThread = MaxLoadThread
}
indexParams[NumLoadThreadKey] = strconv.Itoa(expectedNumLoadThread)
assert.Equal(t, strconv.Itoa(expectedNumLoadThread), numLoadThread)
beamWidth, ok = indexParams[BeamWidthKey]
assert.True(t, ok)
expectedBeamWidth = int(float32(hardware.GetCPUNum()) * float32(params.CommonCfg.BeamWidthRatio))
if expectedBeamWidth > MaxBeamWidth {
expectedBeamWidth = MaxBeamWidth
}
indexParams[BeamWidthKey] = strconv.Itoa(expectedBeamWidth)
assert.Equal(t, strconv.Itoa(expectedBeamWidth), beamWidth)
})
}