enhance: update cagra index params in config and add params check (#29045)

issue:https://github.com/milvus-io/milvus/issues/29230
this pr do two things about cagra index:
 a.milvus yaml config support gpu memory settings

 b.add cagra-params check

Signed-off-by: cqy123456 <qianya.cheng@zilliz.com>
Co-authored-by: yusheng.ma <yusheng.ma@zilliz.com>

configs/milvus.yaml

@@ -724,3 +724,7 @@ trace:
 autoIndex:
   params:
     build: '{"M": 18,"efConstruction": 240,"index_type": "HNSW", "metric_type": "IP"}'
+
+gpu:
+  initMemSize: 0
+  maxMemSize: 0

internal/core/src/config/ConfigKnowhere.cpp

@@ -88,6 +88,16 @@ KnowhereInitSearchThreadPool(const uint32_t num_threads) {
     }
 }
 
+void
+KnowhereInitGPUMemoryPool(const uint32_t init_size, const uint32_t max_size) {
+    if (init_size == 0 && max_size == 0) {
+        knowhere::KnowhereConfig::SetRaftMemPool();
+        return;
+    }
+    knowhere::KnowhereConfig::SetRaftMemPool(size_t{init_size},
+                                             size_t{max_size});
+}
+
 int32_t
 GetMinimalIndexVersion() {
     return knowhere::Version::GetMinimalVersion().VersionNumber();

internal/core/src/config/ConfigKnowhere.h

@@ -37,4 +37,7 @@ GetMinimalIndexVersion();
 int32_t
 GetCurrentIndexVersion();
 
+void
+KnowhereInitGPUMemoryPool(const uint32_t init_size, const uint32_t max_size);
+
 }  // namespace milvus::config

internal/core/src/segcore/segcore_init_c.cpp

@@ -62,6 +62,12 @@ SegcoreSetKnowhereSearchThreadPoolNum(const uint32_t num_threads) {
     milvus::config::KnowhereInitSearchThreadPool(num_threads);
 }
 
+extern "C" void
+SegcoreSetKnowhereGpuMemoryPoolSize(const uint32_t init_size,
+                                    const uint32_t max_size) {
+    milvus::config::KnowhereInitGPUMemoryPool(init_size, max_size);
+}
+
 // return value must be freed by the caller
 extern "C" char*
 SegcoreSetSimdType(const char* value) {

internal/core/src/segcore/segcore_init_c.h

@@ -40,6 +40,10 @@ SegcoreSetKnowhereBuildThreadPoolNum(const uint32_t num_threads);
 void
 SegcoreSetKnowhereSearchThreadPoolNum(const uint32_t num_threads);
 
+void
+SegcoreSetKnowhereGpuMemoryPoolSize(const uint32_t init_size,
+                                    const uint32_t max_size);
+
 void
 SegcoreCloseGlog();
 

internal/core/unittest/test_init.cpp

@@ -37,3 +37,9 @@ TEST(Init, KnowhereThreadPoolInit) {
 #endif
     milvus::config::KnowhereInitSearchThreadPool(8);
 }
+
+TEST(Init, KnowhereGPUMemoryPoolInit) {
+#ifdef MILVUS_GPU_VERSION
+    ASSERT_NO_THROW(milvus::config::KnowhereInitGPUMemoryPool(0, 0));
+#endif
+}

internal/indexnode/indexnode.go

@@ -183,6 +183,9 @@ func (i *IndexNode) initSegcore() {
 
 	localDataRootPath := filepath.Join(Params.LocalStorageCfg.Path.GetValue(), typeutil.IndexNodeRole)
 	initcore.InitLocalChunkManager(localDataRootPath)
+	cGpuMemoryPoolInitSize := C.uint32_t(paramtable.Get().GpuConfig.InitSize.GetAsUint32())
+	cGpuMemoryPoolMaxSize := C.uint32_t(paramtable.Get().GpuConfig.MaxSize.GetAsUint32())
+	C.SegcoreSetKnowhereGpuMemoryPoolSize(cGpuMemoryPoolInitSize, cGpuMemoryPoolMaxSize)
 }
 
 func (i *IndexNode) CloseSegcore() {

internal/querynodev2/server.go

@@ -226,6 +226,10 @@ func (node *QueryNode) InitSegcore() error {
 	cExprBatchSize := C.int64_t(paramtable.Get().QueryNodeCfg.ExprEvalBatchSize.GetAsInt64())
 	C.InitDefaultExprEvalBatchSize(cExprBatchSize)
 
+	cGpuMemoryPoolInitSize := C.uint32_t(paramtable.Get().GpuConfig.InitSize.GetAsUint32())
+	cGpuMemoryPoolMaxSize := C.uint32_t(paramtable.Get().GpuConfig.MaxSize.GetAsUint32())
+	C.SegcoreSetKnowhereGpuMemoryPoolSize(cGpuMemoryPoolInitSize, cGpuMemoryPoolMaxSize)
+
 	localDataRootPath := filepath.Join(paramtable.Get().LocalStorageCfg.Path.GetValue(), typeutil.QueryNodeRole)
 	initcore.InitLocalChunkManager(localDataRootPath)
 

pkg/util/indexparamcheck/cagra_checker.go

@@ -0,0 +1,51 @@
+package indexparamcheck
+
+import (
+	"fmt"
+	"strconv"
+)
+
+// diskannChecker checks if an diskann index can be built.
+type cagraChecker struct {
+	floatVectorBaseChecker
+}
+
+func (c *cagraChecker) CheckTrain(params map[string]string) error {
+	err := c.baseChecker.CheckTrain(params)
+	if err != nil {
+		return err
+	}
+	interDegree := int(0)
+	graphDegree := int(0)
+	interDegreeStr, interDegreeExist := params[CagraInterDegree]
+	if interDegreeExist {
+		interDegree, err = strconv.Atoi(interDegreeStr)
+		if err != nil {
+			return fmt.Errorf("invalid cagra inter degree: %s", interDegreeStr)
+		}
+	}
+	graphDegreeStr, graphDegreeExist := params[CagraGraphDegree]
+	if graphDegreeExist {
+		graphDegree, err = strconv.Atoi(graphDegreeStr)
+		if err != nil {
+			return fmt.Errorf("invalid cagra graph degree: %s", graphDegreeStr)
+		}
+	}
+	if graphDegreeExist && interDegreeExist && interDegree < graphDegree {
+		return fmt.Errorf("Graph degree cannot be larger than intermediate graph degree")
+	}
+
+	if !CheckStrByValues(params, Metric, CagraMetrics) {
+		return fmt.Errorf("metric type not found or not supported, supported: %v", CagraMetrics)
+	}
+
+	return nil
+}
+
+func (c cagraChecker) StaticCheck(params map[string]string) error {
+	return c.staticCheck(params)
+}
+
+func newCagraChecker() IndexChecker {
+	return &cagraChecker{}
+}

pkg/util/indexparamcheck/cagra_checker_test.go

@@ -0,0 +1,90 @@
+package indexparamcheck
+
+import (
+	"strconv"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/milvus-io/milvus/pkg/util/metric"
+)
+
+func Test_cagraChecker_CheckTrain(t *testing.T) {
+	p1 := map[string]string{
+		DIM:    strconv.Itoa(128),
+		Metric: metric.L2,
+	}
+	p2 := map[string]string{
+		DIM:    strconv.Itoa(128),
+		Metric: metric.IP,
+	}
+	p3 := map[string]string{
+		DIM:              strconv.Itoa(128),
+		Metric:           metric.L2,
+		CagraInterDegree: strconv.Itoa(20),
+	}
+
+	p4 := map[string]string{
+		DIM:              strconv.Itoa(128),
+		Metric:           metric.L2,
+		CagraGraphDegree: strconv.Itoa(20),
+	}
+	p5 := map[string]string{
+		DIM:              strconv.Itoa(128),
+		Metric:           metric.L2,
+		CagraInterDegree: strconv.Itoa(60),
+		CagraGraphDegree: strconv.Itoa(20),
+	}
+	p6 := map[string]string{
+		DIM:              strconv.Itoa(128),
+		Metric:           metric.L2,
+		CagraInterDegree: strconv.Itoa(20),
+		CagraGraphDegree: strconv.Itoa(60),
+	}
+	p7 := map[string]string{
+		DIM:    strconv.Itoa(128),
+		Metric: metric.SUPERSTRUCTURE,
+	}
+	p8 := map[string]string{
+		DIM:              strconv.Itoa(128),
+		Metric:           metric.L2,
+		CagraInterDegree: "error",
+		CagraGraphDegree: strconv.Itoa(20),
+	}
+	p9 := map[string]string{
+		DIM:              strconv.Itoa(128),
+		Metric:           metric.L2,
+		CagraInterDegree: strconv.Itoa(20),
+		CagraGraphDegree: "error",
+	}
+	p10 := map[string]string{
+		DIM:    strconv.Itoa(0),
+		Metric: metric.L2,
+	}
+
+	cases := []struct {
+		params   map[string]string
+		errIsNil bool
+	}{
+		{p1, true},
+		{p2, false},
+		{p3, true},
+		{p4, true},
+		{p5, true},
+		{p6, false},
+		{p7, false},
+		{p8, false},
+		{p9, false},
+		{p10, false},
+	}
+
+	c := newCagraChecker()
+	for _, test := range cases {
+		err := c.CheckTrain(test.params)
+		if test.errIsNil {
+			assert.NoError(t, err)
+		} else {
+			assert.Error(t, err)
+		}
+	}
+}

pkg/util/indexparamcheck/conf_adapter_mgr.go

@@ -45,6 +45,7 @@ func (mgr *indexCheckerMgrImpl) GetChecker(indexType string) (IndexChecker, erro
 func (mgr *indexCheckerMgrImpl) registerIndexChecker() {
 	mgr.checkers[IndexRaftIvfFlat] = newIVFBaseChecker()
 	mgr.checkers[IndexRaftIvfPQ] = newRaftIVFPQChecker()
+	mgr.checkers[IndexRaftCagra] = newCagraChecker()
 	mgr.checkers[IndexFaissIDMap] = newFlatChecker()
 	mgr.checkers[IndexFaissIvfFlat] = newIVFBaseChecker()
 	mgr.checkers[IndexFaissIvfPQ] = newIVFPQChecker()

pkg/util/indexparamcheck/constraints.go

@@ -36,6 +36,10 @@ const (
 
 	EFConstruction = "efConstruction"
 	HNSWM          = "M"
+
+	// Cagra Train Param
+	CagraInterDegree = "intermediate_graph_degree"
+	CagraGraphDegree = "graph_degree"
 )
 
 // METRICS is a set of all metrics types supported for float vector.
@@ -46,6 +50,7 @@ var (
 	BinIDMapMetrics           = []string{metric.HAMMING, metric.JACCARD, metric.SUBSTRUCTURE, metric.SUPERSTRUCTURE} // const
 	BinIvfMetrics             = []string{metric.HAMMING, metric.JACCARD}                                             // const
 	HnswMetrics               = []string{metric.L2, metric.IP, metric.COSINE, metric.HAMMING, metric.JACCARD}        // const
+	CagraMetrics              = []string{metric.L2}                                                                  // const
 	supportDimPerSubQuantizer = []int{32, 28, 24, 20, 16, 12, 10, 8, 6, 4, 3, 2, 1}                                  // const
 	supportSubQuantizer       = []int{96, 64, 56, 48, 40, 32, 28, 24, 20, 16, 12, 8, 4, 3, 2, 1}                     // const
 )

pkg/util/indexparamcheck/index_type.go

@@ -18,6 +18,7 @@ type IndexType = string
 const (
 	IndexRaftIvfFlat     IndexType = "GPU_IVF_FLAT"
 	IndexRaftIvfPQ       IndexType = "GPU_IVF_PQ"
+	IndexRaftCagra       IndexType = "GPU_CAGRA"
 	IndexFaissIDMap      IndexType = "FLAT" // no index is built.
 	IndexFaissIvfFlat    IndexType = "IVF_FLAT"
 	IndexFaissIvfPQ      IndexType = "IVF_PQ"

pkg/util/paramtable/component_param.go

@@ -58,6 +58,7 @@ type ComponentParam struct {
 	CommonCfg       commonConfig
 	QuotaConfig     quotaConfig
 	AutoIndexConfig autoIndexConfig
+	GpuConfig       gpuConfig
 	TraceCfg        traceConfig
 
 	RootCoordCfg  rootCoordConfig
@@ -118,6 +119,7 @@ func (p *ComponentParam) init(bt *BaseTable) {
 	p.HTTPCfg.init(bt)
 	p.LogCfg.init(bt)
 	p.RoleCfg.init(bt)
+	p.GpuConfig.init(bt)
 
 	p.RootCoordGrpcServerCfg.Init("rootCoord", bt)
 	p.ProxyGrpcServerCfg.Init("proxy", bt)
@@ -672,6 +674,29 @@ like the old password verification when updating the credential`,
 	p.TraceLogMode.Init(base.mgr)
 }
 
+type gpuConfig struct {
+	InitSize ParamItem `refreshable:"false"`
+	MaxSize  ParamItem `refreshable:"false"`
+}
+
+func (t *gpuConfig) init(base *BaseTable) {
+	t.InitSize = ParamItem{
+		Key:     "gpu.initMemSize",
+		Version: "2.3.4",
+		Doc:     `Gpu Memory Pool init size`,
+		Export:  true,
+	}
+	t.InitSize.Init(base.mgr)
+
+	t.MaxSize = ParamItem{
+		Key:     "gpu.maxMemSize",
+		Version: "2.3.4",
+		Doc:     `Gpu Memory Pool Max size`,
+		Export:  true,
+	}
+	t.MaxSize.Init(base.mgr)
+}
+
 type traceConfig struct {
 	Exporter       ParamItem `refreshable:"false"`
 	SampleFraction ParamItem `refreshable:"false"`