enhance: add configurable memory index load predict memory usage fact…

…or (#30563) pr: #30561 related pr: #30475 Signed-off-by: chyezh <[email protected]>
milvus-io · Feb 6, 2024 · be1bd96 · be1bd96
1 parent e1fb4f6
commit be1bd96
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 3 deletions.
diff --git a/internal/querynodev2/segments/segment_loader.go b/internal/querynodev2/segments/segment_loader.go
@@ -918,7 +918,7 @@ func GetIndexResourceUsage(indexInfo *querypb.FieldIndexInfo) (uint64, uint64, e
 		return uint64(neededMemSize), uint64(neededDiskSize), nil
 	}
 
-	factor := uint64(1)
+	factor := float64(1)
 
 	var isLoadWithDisk bool
 	GetDynamicPool().Submit(func() (any, error) {
@@ -930,10 +930,10 @@ func GetIndexResourceUsage(indexInfo *querypb.FieldIndexInfo) (uint64, uint64, e
 	}).Await()
 
 	if !isLoadWithDisk {
-		factor = 2
+		factor = paramtable.Get().QueryNodeCfg.MemoryIndexLoadPredictMemoryUsageFactor.GetAsFloat()
 	}
 
-	return uint64(indexInfo.IndexSize) * factor, 0, nil
+	return uint64(float64(indexInfo.IndexSize) * factor), 0, nil
 }
 
 // checkSegmentSize checks whether the memory & disk is sufficient to load the segments with given concurrency,

diff --git a/pkg/util/paramtable/component_param.go b/pkg/util/paramtable/component_param.go
@@ -1807,6 +1807,8 @@ type queryNodeConfig struct {
 	CGOPoolSizeRatio ParamItem `refreshable:"false"`
 
 	EnableWorkerSQCostMetrics ParamItem `refreshable:"true"`
+
+	MemoryIndexLoadPredictMemoryUsageFactor ParamItem `refreshable:"true"`
 }
 
 func (p *queryNodeConfig) init(base *BaseTable) {
@@ -2214,6 +2216,14 @@ Max read concurrency must greater than or equal to 1, and less than or equal to
 		Doc:          "whether use worker's cost to measure delegator's workload",
 	}
 	p.EnableWorkerSQCostMetrics.Init(base.mgr)
+
+	p.MemoryIndexLoadPredictMemoryUsageFactor = ParamItem{
+		Key:          "queryNode.memoryIndexLoadPredictMemoryUsageFactor",
+		Version:      "2.3.8",
+		DefaultValue: "2.5", // HNSW index needs more memory to load.
+		Doc:          "memory usage prediction factor for memory index loaded",
+	}
+	p.MemoryIndexLoadPredictMemoryUsageFactor.Init(base.mgr)
 }
 
 // /////////////////////////////////////////////////////////////////////////////

diff --git a/pkg/util/paramtable/component_param_test.go b/pkg/util/paramtable/component_param_test.go
@@ -360,6 +360,10 @@ func TestComponentParam(t *testing.T) {
 
 		assert.Equal(t, false, Params.EnableWorkerSQCostMetrics.GetAsBool())
 
+		assert.Equal(t, 2.5, Params.MemoryIndexLoadPredictMemoryUsageFactor.GetAsFloat())
+		params.Save("queryNode.memoryIndexLoadPredictMemoryUsageFactor", "2.0")
+		assert.Equal(t, 2.0, Params.MemoryIndexLoadPredictMemoryUsageFactor.GetAsFloat())
+
 		params.Save("querynode.gracefulStopTimeout", "100")
 		assert.Equal(t, 100*time.Second, Params.GracefulStopTimeout.GetAsDuration(time.Second))
 	})