enhance: Optimize the task scheduling logic in DataCoord

configs/milvus.yaml

@@ -633,6 +633,9 @@ dataCoord:
     clusteringCompactionUsage: 16 # slot usage of clustering compaction job.
     mixCompactionUsage: 8 # slot usage of mix compaction job.
     l0DeleteCompactionUsage: 8 # slot usage of l0 compaction job.
+    indexTaskSlotUsage: 16 # slot usage of index task
+    statsTaskSlotUsage: 4 # slot usage of stats task
+    analyzeTaskSlotUsage: 16 # slot usage of analyze task
   ip:  # TCP/IP address of dataCoord. If not specified, use the first unicastable address
   port: 13333 # TCP port of dataCoord
   grpc:

internal/datacoord/session/indexnode_manager.go

@@ -44,11 +44,18 @@
 	RemoveNode(nodeID typeutil.UniqueID)
 	StoppingNode(nodeID typeutil.UniqueID)
 	PickClient() (typeutil.UniqueID, types.IndexNodeClient)
+	QuerySlots() []WorkerSlots
 	ClientSupportDisk() bool
 	GetAllClients() map[typeutil.UniqueID]types.IndexNodeClient
 	GetClientByID(nodeID typeutil.UniqueID) (types.IndexNodeClient, bool)
 }
 
+type WorkerSlots struct {
+	NodeID         int64
+	TotalSlots     int64
+	AvailableSlots int64
+}
+
 // IndexNodeManager is used to manage the client of IndexNode.
 type IndexNodeManager struct {
 	nodeClients      map[typeutil.UniqueID]types.IndexNodeClient
@@ -115,6 +122,46 @@
 	return nil
 }
 
+func (nm *IndexNodeManager) QuerySlots() []WorkerSlots {
+	nm.lock.Lock()
+	defer nm.lock.Unlock()
+
+	ctx, cancel := context.WithTimeout(context.Background(), querySlotTimeout)
+	defer cancel()
+
+	nodeSlots := make([]WorkerSlots, 0)
+	mu := &sync.Mutex{}
+	wg := &sync.WaitGroup{}
+	for nodeID, client := range nm.nodeClients {
+		if _, ok := nm.stoppingNodes[nodeID]; !ok {
+			wg.Add(1)
+			go func(nodeID int64) {
+				defer wg.Done()
+				resp, err := client.GetJobStats(ctx, &workerpb.GetJobStatsRequest{})
+				if err != nil {
+					log.Warn("get IndexNode slots failed", zap.Int64("nodeID", nodeID), zap.Error(err))
+					return
+				}
+				if resp.GetStatus().GetErrorCode() != commonpb.ErrorCode_Success {
+					log.Warn("get IndexNode slots failed", zap.Int64("nodeID", nodeID),
+						zap.String("reason", resp.GetStatus().GetReason()))
+					return
+				}
+				mu.Lock()
+				defer mu.Unlock()
+				nodeSlots = append(nodeSlots, WorkerSlots{
+					NodeID:         nodeID,
+					TotalSlots:     resp.GetTotalSlots(),
+					AvailableSlots: resp.GetTaskSlots(),
+				})
+			}(nodeID)
+		}
+	}
+	wg.Wait()
+	log.Ctx(context.TODO()).Debug("query slot done", zap.Any("nodeSlots", nodeSlots))
+	return nodeSlots
+}
+
 func (nm *IndexNodeManager) PickClient() (typeutil.UniqueID, types.IndexNodeClient) {
 	nm.lock.Lock()
 	defer nm.lock.Unlock()

internal/datacoord/task_analyze.go

@@ -133,12 +133,12 @@
 	return nil
 }
 
-func (at *analyzeTask) PreCheck(ctx context.Context, dependency *taskScheduler) bool {
+func (at *analyzeTask) PreCheck(ctx context.Context, dependency *taskScheduler) (bool, int64) {
 	t := dependency.meta.analyzeMeta.GetTask(at.GetTaskID())
 	if t == nil {
 		log.Ctx(ctx).Info("task is nil, delete it", zap.Int64("taskID", at.GetTaskID()))
 		at.SetState(indexpb.JobState_JobStateNone, "analyze task is nil")
-		return false
+		return false, 0
 	}
 
 	at.req = &workerpb.AnalyzeRequest{
@@ -170,7 +170,7 @@
 			log.Ctx(ctx).Warn("analyze stats task is processing, but segment is nil, delete the task",
 				zap.Int64("taskID", at.GetTaskID()), zap.Int64("segmentID", segID))
 			at.SetState(indexpb.JobState_JobStateFailed, fmt.Sprintf("segmentInfo with ID: %d is nil", segID))
-			return false
+			return false, 0
 		}
 
 		totalSegmentsRows += info.GetNumOfRows()
@@ -188,7 +188,7 @@
 		log.Ctx(ctx).Warn("analyze task get collection info failed", zap.Int64("collectionID",
 			segments[0].GetCollectionID()), zap.Error(err))
 		at.SetState(indexpb.JobState_JobStateInit, err.Error())
-		return false
+		return false, 0
 	}
 
 	schema := collInfo.Schema
@@ -203,7 +203,7 @@
 	dim, err := storage.GetDimFromParams(field.TypeParams)
 	if err != nil {
 		at.SetState(indexpb.JobState_JobStateInit, err.Error())
-		return false
+		return false, 0
 	}
 	at.req.Dim = int64(dim)
 
@@ -212,7 +212,7 @@
 	if numClusters < Params.DataCoordCfg.ClusteringCompactionMinCentroidsNum.GetAsInt64() {
 		log.Ctx(ctx).Info("data size is too small, skip analyze task", zap.Float64("raw data size", totalSegmentsRawDataSize), zap.Int64("num clusters", numClusters), zap.Int64("minimum num clusters required", Params.DataCoordCfg.ClusteringCompactionMinCentroidsNum.GetAsInt64()))
 		at.SetState(indexpb.JobState_JobStateFinished, "")
-		return false
+		return false, 0
 	}
 	if numClusters > Params.DataCoordCfg.ClusteringCompactionMaxCentroidsNum.GetAsInt64() {
 		numClusters = Params.DataCoordCfg.ClusteringCompactionMaxCentroidsNum.GetAsInt64()
@@ -223,8 +223,10 @@
 	at.req.MinClusterSizeRatio = Params.DataCoordCfg.ClusteringCompactionMinClusterSizeRatio.GetAsFloat()
 	at.req.MaxClusterSizeRatio = Params.DataCoordCfg.ClusteringCompactionMaxClusterSizeRatio.GetAsFloat()
 	at.req.MaxClusterSize = Params.DataCoordCfg.ClusteringCompactionMaxClusterSize.GetAsSize()
+	taskSlot := Params.DataCoordCfg.AnalyzeTaskSlotUsage.GetAsInt64()
+	at.req.TaskSlot = taskSlot
 
-	return true
+	return true, taskSlot
 }
 
 func (at *analyzeTask) AssignTask(ctx context.Context, client types.IndexNodeClient) bool {

internal/datacoord/task_index.go

@@ -132,19 +132,19 @@
 	return meta.indexMeta.BuildIndex(it.taskID)
 }
 
-func (it *indexBuildTask) PreCheck(ctx context.Context, dependency *taskScheduler) bool {
+func (it *indexBuildTask) PreCheck(ctx context.Context, dependency *taskScheduler) (bool, int64) {
 	segIndex, exist := dependency.meta.indexMeta.GetIndexJob(it.taskID)
 	if !exist || segIndex == nil {
 		log.Ctx(ctx).Info("index task has not exist in meta table, remove task", zap.Int64("taskID", it.taskID))
 		it.SetState(indexpb.JobState_JobStateNone, "index task has not exist in meta table")
-		return false
+		return false, 0
 	}
 
 	segment := dependency.meta.GetSegment(ctx, segIndex.SegmentID)
 	if !isSegmentHealthy(segment) || !dependency.meta.indexMeta.IsIndexExist(segIndex.CollectionID, segIndex.IndexID) {
 		log.Ctx(ctx).Info("task is no need to build index, remove it", zap.Int64("taskID", it.taskID))
 		it.SetState(indexpb.JobState_JobStateNone, "task is no need to build index")
-		return false
+		return false, 0
 	}
 	indexParams := dependency.meta.indexMeta.GetIndexParams(segIndex.CollectionID, segIndex.IndexID)
 	indexType := GetIndexType(indexParams)
@@ -154,7 +154,7 @@
 		it.SetStartTime(time.Now())
 		it.SetEndTime(time.Now())
 		it.SetState(indexpb.JobState_JobStateFinished, "fake finished index success")
-		return false
+		return false, 0
 	}
 
 	typeParams := dependency.meta.indexMeta.GetTypeParams(segIndex.CollectionID, segIndex.IndexID)
@@ -170,7 +170,7 @@
 		if ret != nil {
 			log.Ctx(ctx).Warn("failed to update index build params defined in yaml", zap.Int64("taskID", it.taskID), zap.Error(ret))
 			it.SetState(indexpb.JobState_JobStateInit, ret.Error())
-			return false
+			return false, 0
 		}
 	}
 
@@ -180,14 +180,14 @@
 		if err != nil {
 			log.Ctx(ctx).Warn("failed to append index build params", zap.Int64("taskID", it.taskID), zap.Error(err))
 			it.SetState(indexpb.JobState_JobStateInit, err.Error())
-			return false
+			return false, 0
 		}
 	}
 
 	collectionInfo, err := dependency.handler.GetCollection(ctx, segment.GetCollectionID())
 	if err != nil {
 		log.Ctx(ctx).Info("index builder get collection info failed", zap.Int64("collectionID", segment.GetCollectionID()), zap.Error(err))
-		return false
+		return false, 0
 	}
 
 	schema := collectionInfo.Schema
@@ -215,7 +215,7 @@
 		if collectionInfo == nil {
 			log.Ctx(ctx).Warn("get collection failed", zap.Int64("collID", segIndex.CollectionID), zap.Error(err))
 			it.SetState(indexpb.JobState_JobStateInit, err.Error())
-			return true
+			return true, 0
 		}
 		partitionKeyField, _ := typeutil.GetPartitionKeyFieldSchema(schema)
 		if partitionKeyField != nil && typeutil.IsFieldDataTypeSupportMaterializedView(partitionKeyField) {
@@ -235,6 +235,8 @@
 		}
 	}
 
+	taskSlot := calculateIndexTaskSlot(segment.getSegmentSize())
+
 	it.req = &workerpb.CreateJobRequest{
 		ClusterID:                 Params.CommonCfg.ClusterPrefix.GetValue(),
 		IndexFilePrefix:           path.Join(dependency.chunkManager.RootPath(), common.SegmentIndexPath),
@@ -257,11 +259,23 @@
 		OptionalScalarFields:      optionalFields,
 		Field:                     field,
 		PartitionKeyIsolation:     partitionKeyIsolation,
+		TaskSlot:                  taskSlot,
 	}
 
 	log.Ctx(ctx).Info("index task pre check successfully", zap.Int64("taskID", it.GetTaskID()),
 		zap.Int64("segID", segment.GetID()))
-	return true
+	return true, taskSlot
+}
+
+func calculateIndexTaskSlot(segmentSize int64) int64 {
+	if segmentSize > 500*1024*1024 {
+		return max(Params.DataCoordCfg.IndexTaskSlotUsage.GetAsInt64(), 1)
+	} else if segmentSize > 100*1024*1024 {
+		return max(Params.DataCoordCfg.IndexTaskSlotUsage.GetAsInt64()/2, 1)
+	} else if segmentSize > 10*1024*1024 {
+		return max(Params.DataCoordCfg.IndexTaskSlotUsage.GetAsInt64()/4, 1)
+	}
+	return max(Params.DataCoordCfg.IndexTaskSlotUsage.GetAsInt64()/8, 1)
 }
 
 func (it *indexBuildTask) AssignTask(ctx context.Context, client types.IndexNodeClient) bool {