diff --git a/internal/querynode/load_index_info.go b/internal/querynode/load_index_info.go index c3a21ed495..1edf01067e 100644 --- a/internal/querynode/load_index_info.go +++ b/internal/querynode/load_index_info.go @@ -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 } diff --git a/internal/util/indexparams/disk_index_params.go b/internal/util/indexparams/disk_index_params.go index cd339048bb..47ca78bd67 100644 --- a/internal/util/indexparams/disk_index_params.go +++ b/internal/util/indexparams/disk_index_params.go @@ -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", diff --git a/internal/util/indexparams/disk_index_params_test.go b/internal/util/indexparams/disk_index_params_test.go index b55deaf821..6ab6d2671c 100644 --- a/internal/util/indexparams/disk_index_params_test.go +++ b/internal/util/indexparams/disk_index_params_test.go @@ -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(¶ms, 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(¶ms, indexParams, 100) assert.Error(t, err) indexParams["dim"] = "128" - err = SetDiskIndexLoadParams(indexParams, 100) + err = SetDiskIndexLoadParams(¶ms, 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(¶ms, 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) }) }