chore: change some files/ai suggestions

Author: Ethereal-O

vermeer/apps/master/bl/grpc_handlers.go

@@ -100,11 +100,15 @@ func (h *ServerHandler) SayHelloMaster(ctx context.Context, req *pb.HelloMasterR
 	}
 
 	_, err = workerMgr.AddWorker(reqWorker)
-	Scheduler.ChangeWorkerStatus(reqWorker.Name, schedules.WorkerOngoingStatusIdle)
 	if err != nil {
 		logrus.Errorf("failed to add a WorkerClient to the WorkerManager, error: %s", err)
 		return &pb.HelloMasterResp{}, err
 	}
+	_, err = Scheduler.ChangeWorkerStatus(reqWorker.Name, schedules.WorkerOngoingStatusIdle)
+	if err != nil {
+		logrus.Errorf("failed to change worker status to idle, error: %s", err)
+		return &pb.HelloMasterResp{}, err
+	}
 
 	logrus.Infof("worker say hello name: %s and set to workgroup: %s, client: %s", reqWorker.Name, reqWorker.Group, p.Addr.String())
 

vermeer/apps/master/bl/scheduler_bl.go

@@ -113,22 +113,24 @@ func (s *ScheduleBl) startTicker() {
 }
 
 // this make scheduler manager try to schedule next tasks
-func (s *ScheduleBl) TryScheduleNextTasks() {
+func (s *ScheduleBl) TryScheduleNextTasks(noLock ...bool) {
 	defer func() {
 		if err := recover(); err != nil {
 			logrus.Errorln("TryScheduleNextTasks() has been recovered:", err)
 		}
 	}()
 
-	if err := s.tryScheduleInner(s.softSchedule); err != nil {
+	if err := s.tryScheduleInner(s.softSchedule, noLock...); err != nil {
 		logrus.Errorf("do scheduling error:%v", err)
 	}
 }
 
 // Main routine to schedule tasks
-func (s *ScheduleBl) tryScheduleInner(softSchedule bool) error {
+func (s *ScheduleBl) tryScheduleInner(softSchedule bool, noLock ...bool) error {
 	// Implement logic to get the next task in the queue for the given space
-	defer s.Unlock(s.Lock())
+	if !(len(noLock) > 0 && noLock[0]) {
+		defer s.Unlock(s.Lock())
+	}
 
 	// step 1: make sure all tasks have alloc to a worker group
 	// This is done by the TaskManager, which assigns a worker group to each task
@@ -264,6 +266,8 @@ func (s *ScheduleBl) BatchQueueTask(taskInfos []*structure.TaskInfo) ([]bool, []
 // ******** CloseCurrent ********
 
 func (s *ScheduleBl) CloseCurrent(taskId int32, removeWorkerName ...string) error {
+	defer s.Unlock(s.Lock())
+
 	// trace tasks need these workers, check if these tasks are available
 	s.taskManager.RemoveTask(taskId)
 	// release the worker group
@@ -282,24 +286,27 @@ func (s *ScheduleBl) CloseCurrent(taskId int32, removeWorkerName ...string) erro
 	}
 
 	logrus.Infof("invoke dispatch when task '%d' is closed", taskId)
-	s.TryScheduleNextTasks()
+	s.TryScheduleNextTasks(true)
 	return nil
 }
 
-func (s *ScheduleBl) ChangeWorkerStatus(workerName string, status schedules.WorkerOngoingStatus) {
+func (s *ScheduleBl) ChangeWorkerStatus(workerName string, status schedules.WorkerOngoingStatus) (bool, error) {
+	defer s.Unlock(s.Lock())
 	s.resourceManager.ChangeWorkerStatus(workerName, status)
 
 	logrus.Infof("worker '%s' status changed to '%s'", workerName, status)
 	// After changing the worker status, we may need to reschedule tasks
-	s.TryScheduleNextTasks()
+	s.TryScheduleNextTasks(true)
+
+	return true, nil
 }
 
 // ******** START TASK ********
 func (s *ScheduleBl) waitingStartedTask() {
 	for taskInfo := range s.startChan {
 		if taskInfo == nil {
 			logrus.Warnf("recieved a nil task from startChan")
-			return
+			continue
 		}
 
 		logrus.Infof("chan received task '%d' to start", taskInfo.ID)
@@ -388,6 +395,8 @@ func (s *ScheduleBl) CancelTask(taskInfo *structure.TaskInfo) error {
 		return errors.New("the argument `taskInfo` is nil")
 	}
 
+	defer s.Unlock(s.Lock())
+
 	isHeadTask := s.taskManager.IsTaskOngoing(taskInfo.ID)
 	task := s.taskManager.RemoveTask(taskInfo.ID)
 	s.cronManager.DeleteTask(taskInfo.ID)

vermeer/apps/master/bl/task_bl.go

@@ -71,7 +71,7 @@ func (tb *TaskBl) CreateTaskInfo(
 	// taskInfo.Exclusive = true
 	if params != nil {
 		if priority, ok := params["priority"]; ok {
-			if p, err := strconv.Atoi(priority); err == nil {
+			if p, err := strconv.ParseInt(priority, 10, 32); err == nil {
 				taskInfo.Priority = int32(p)
 			} else {
 				logrus.Warnf("priority convert to int32 error:%v", err)
@@ -80,7 +80,7 @@ func (tb *TaskBl) CreateTaskInfo(
 		if preorders, ok := params["preorders"]; ok {
 			preorderList := strings.Split(preorders, ",")
 			for _, preorder := range preorderList {
-				if pid, err := strconv.Atoi(preorder); err == nil {
+				if pid, err := strconv.ParseInt(preorder, 10, 32); err == nil {
 					taskInfo.Preorders = append(taskInfo.Preorders, int32(pid))
 				} else {
 					logrus.Warnf("preorder convert to int32 error:%v", err)

vermeer/apps/master/schedules/scheduler_algorithm_manager.go

@@ -1,6 +1,7 @@
 package schedules
 
 import (
+	"slices"
 	"sort"
 	"strconv"
 	"time"
@@ -23,13 +24,13 @@ type SchedulerAlgorithm interface {
 
 type SchedulerAlgorithmManager struct {
 	filteredSchedulerAlgorithms  map[string]SchedulerAlgorithm
-	schuduledSchedulerAlgorithms map[string]SchedulerAlgorithm
+	scheduledSchedulerAlgorithms map[string]SchedulerAlgorithm
 	dispatchPaused               bool
 }
 
 func (am *SchedulerAlgorithmManager) Init() {
 	am.filteredSchedulerAlgorithms = make(map[string]SchedulerAlgorithm)
-	am.schuduledSchedulerAlgorithms = make(map[string]SchedulerAlgorithm)
+	am.scheduledSchedulerAlgorithms = make(map[string]SchedulerAlgorithm)
 	am.dispatchPaused = false
 	// Register filter and schedule algorithms
 	am.RegisterFilterAlgorithm(&DependsSchedulerAlgorithm{})
@@ -38,33 +39,33 @@ func (am *SchedulerAlgorithmManager) Init() {
 	am.RegisterSchedulerAlgorithm(&PriorityElderSchedulerAlgorithm{})
 }
 
-func (am *SchedulerAlgorithmManager) RegisterSchedulerAlgorithm(SchedulerAlgorithm SchedulerAlgorithm) {
-	if SchedulerAlgorithm == nil {
+func (am *SchedulerAlgorithmManager) RegisterSchedulerAlgorithm(schedulerAlgorithm SchedulerAlgorithm) {
+	if schedulerAlgorithm == nil {
 		return
 	}
-	name := SchedulerAlgorithm.Name()
-	if _, exists := am.schuduledSchedulerAlgorithms[name]; exists {
+	name := schedulerAlgorithm.Name()
+	if _, exists := am.scheduledSchedulerAlgorithms[name]; exists {
 		return // SchedulerAlgorithm already registered
 	}
 
 	// only support one scheduling algorithm for now
-	if len(am.schuduledSchedulerAlgorithms) > 0 {
+	if len(am.scheduledSchedulerAlgorithms) > 0 {
 		return // Only one scheduling algorithm can be registered
 	}
-	SchedulerAlgorithm.Init()
-	am.schuduledSchedulerAlgorithms[name] = SchedulerAlgorithm
+	schedulerAlgorithm.Init()
+	am.scheduledSchedulerAlgorithms[name] = schedulerAlgorithm
 }
 
-func (am *SchedulerAlgorithmManager) RegisterFilterAlgorithm(FilterAlgorithm SchedulerAlgorithm) {
-	if FilterAlgorithm == nil {
+func (am *SchedulerAlgorithmManager) RegisterFilterAlgorithm(filterAlgorithm SchedulerAlgorithm) {
+	if filterAlgorithm == nil {
 		return
 	}
-	name := FilterAlgorithm.Name()
+	name := filterAlgorithm.Name()
 	if _, exists := am.filteredSchedulerAlgorithms[name]; exists {
 		return // SchedulerAlgorithm already registered
 	}
-	FilterAlgorithm.Init()
-	am.filteredSchedulerAlgorithms[name] = FilterAlgorithm
+	filterAlgorithm.Init()
+	am.filteredSchedulerAlgorithms[name] = filterAlgorithm
 }
 
 func (am *SchedulerAlgorithmManager) IsDispatchPaused() bool {
@@ -98,7 +99,7 @@ func (am *SchedulerAlgorithmManager) ScheduleNextTasks(allTasks []*structure.Tas
 
 	// only support one scheduling algorithm for now
 	// get first algorithm
-	for _, algorithm := range am.schuduledSchedulerAlgorithms {
+	for _, algorithm := range am.scheduledSchedulerAlgorithms {
 		tasks, err := algorithm.ScheduleNextTasks(filteredTasks, taskToWorkerGroupMap, idleWorkerGroups, concurrentWorkerGroups, softSchedule)
 		if err != nil {
 			return nil, err
@@ -268,7 +269,12 @@ func (p *PriorityElderSchedulerAlgorithm) CalculateTaskEmergency(task *structure
 	priorityCost := priorityParam * int64(task.Priority)
 	// step 3: resource cost
 	graph := structure.GraphManager.GetGraphByName(task.SpaceName, task.GraphName)
-	resourceCost := resourceParam / max(1, graph.VertexCount+graph.EdgeCount) // Avoid division by zero, ensure at least 1
+	resourceCost := int64(0)
+	if graph == nil {
+		resourceCost = resourceParam // if graph not found, use max resource cost
+	} else {
+		resourceCost = resourceParam / max(1, graph.VertexCount+graph.EdgeCount) // Avoid division by zero, ensure at least 1
+	}
 	// step 4: some random value
 	randomValue := int64(randomValueParam) // Placeholder for any random value logic
 	if printValue {
@@ -282,13 +288,19 @@ func (p *PriorityElderSchedulerAlgorithm) ScheduleNextTasks(allTasks []*structur
 		return nil, nil // No tasks to schedule
 	}
 
+	// calculate emergency value for each task
+	taskEmergencies := make(map[int32]int64)
+	for _, task := range allTasks {
+		taskEmergencies[task.ID] = p.CalculateTaskEmergency(task, taskToWorkerGroupMap, false)
+	}
+
 	// Sort tasks by priority (higher priority first)
 	sort.Slice(allTasks, func(i, j int) bool {
-		return p.CalculateTaskEmergency(allTasks[i], taskToWorkerGroupMap, false) > p.CalculateTaskEmergency(allTasks[j], taskToWorkerGroupMap, false)
+		return taskEmergencies[allTasks[i].ID] > taskEmergencies[allTasks[j].ID]
 	})
 
 	for _, task := range allTasks {
-		logrus.Debugf("Task %d: Emergency Value: %d", task.ID, p.CalculateTaskEmergency(task, taskToWorkerGroupMap, true))
+		logrus.Debugf("Task %d: Emergency Value: %d", task.ID, taskEmergencies[task.ID])
 	}
 
 	for _, task := range allTasks {
@@ -427,19 +439,17 @@ func (d *DependsSchedulerAlgorithm) ScheduleNextTasks(allTasks []*structure.Task
 		// Check if all dependencies are satisfied
 		allDepsSatisfied := true
 		for _, dep := range depends {
-			if depTask, exists := allTaskIDs[dep]; !exists || depTask.State != structure.TaskStateWaiting {
+			if depTask, exists := allTaskIDs[dep]; exists && depTask.State != structure.TaskStateComplete {
 				allDepsSatisfied = false
 				break
 			}
 		}
 		if allDepsSatisfied {
 			if group, exists := taskToWorkerGroupMap[task.ID]; exists && group != "" {
 				// only support idle worker groups for now
-				for _, idleGroup := range idleWorkerGroups {
-					if group == idleGroup {
-						logrus.Debugf("Task %d is assigned to worker group %s", task.ID, group)
-						return []*structure.TaskInfo{task}, nil // Return the first task that can be scheduled
-					}
+				if slices.Contains(idleWorkerGroups, group) {
+					logrus.Debugf("Task %d is assigned to worker group %s", task.ID, group)
+					return []*structure.TaskInfo{task}, nil // Return the first task that can be scheduled
 				}
 			}
 		}

vermeer/apps/master/schedules/scheduler_cron_manager.go

@@ -43,7 +43,6 @@ func (t *SchedulerCronManager) AddCronTask(taskInfo *structure.TaskInfo) error {
 	}
 
 	// add to cron tasks
-	t.cronTasks[taskInfo.ID] = append(t.cronTasks[taskInfo.ID], taskInfo)
 	cronJob := cron.New()
 	_, err := cronJob.AddFunc(taskInfo.CronExpr, func() {
 		if taskInfo == nil {
@@ -63,6 +62,7 @@ func (t *SchedulerCronManager) AddCronTask(taskInfo *structure.TaskInfo) error {
 		logrus.Errorf("Failed to add cron job for task %d: %v", taskInfo.ID, err)
 		return err
 	}
+	t.cronTasks[taskInfo.ID] = append(t.cronTasks[taskInfo.ID], taskInfo)
 	t.crons[taskInfo.ID] = append(t.crons[taskInfo.ID], cronJob)
 	cronJob.Start()
 	logrus.Infof("Added cron task for task ID %d with expression %s", taskInfo.ID, taskInfo.CronExpr)

vermeer/apps/master/schedules/scheduler_task_manager.go

@@ -8,6 +8,7 @@ import (
 )
 
 type SchedulerTaskManager struct {
+	structure.MutexLocker
 	// This struct is responsible for managing tasks in the scheduling system.
 	// A map from task ID to TaskInfo can be used to track tasks.
 	allTaskMap     map[int32]*structure.TaskInfo
@@ -32,6 +33,8 @@ func (t *SchedulerTaskManager) QueueTask(taskInfo *structure.TaskInfo) (bool, er
 		return false, errors.New("the property `SpaceName` of taskInfo is empty")
 	}
 
+	defer t.Unlock(t.Lock())
+
 	// Add the task to the task map
 	t.allTaskMap[taskInfo.ID] = taskInfo
 	t.allTaskQueue = append(t.allTaskQueue, taskInfo)
@@ -53,6 +56,13 @@ func (t *SchedulerTaskManager) RemoveTask(taskID int32) error {
 		return errors.New("task not found")
 	}
 	delete(t.allTaskMap, taskID)
+	// remove from queue
+	for i, task := range t.allTaskQueue {
+		if task.ID == taskID {
+			t.allTaskQueue = append(t.allTaskQueue[:i], t.allTaskQueue[i+1:]...)
+			break
+		}
+	}
 	delete(t.taskToworkerGroupMap, taskID)
 	return nil
 }

vermeer/config/master.ini

@@ -14,7 +14,7 @@
 ; limitations under the License.
 
 [default]
-log_level=debug
+log_level=info
 debug_mode=release
 http_peer=0.0.0.0:6688
 grpc_peer=0.0.0.0:6689

vermeer/test/functional/compute_base.go

@@ -140,7 +140,7 @@ func (ctb *ComputeTaskBase) SendComputeReqAsyncBatchPriority(params []map[string
 		//若成功启动Compute Task，开始轮询tasksGet，解析response，得到状态为完成时break。
 		var taskResp *client.TaskResponse
 		var err error
-		for i := 0; i < ctb.waitSecond; i++ {
+		for attempt := 0; attempt < ctb.waitSecond; attempt++ {
 			ctb.healthCheck.DoHealthCheck()
 			taskResp, err = ctb.masterHttp.GetTask(ctb.taskID)
 			require.NoError(ctb.t, err)

vermeer/test/functional/load_local.go

@@ -47,6 +47,8 @@ func (lt *LoadTaskLocal) TaskLoadBody() map[string]string {
 	}
 }
 
+// TaskLoadBodyWithNum creates load configuration with specified number of files.
+// If num <= 10, it will be automatically adjusted to 30 to ensure minimum test coverage.
 func (lt *LoadTaskLocal) TaskLoadBodyWithNum(num int) map[string]string {
 	vertexBackends := []string{"db", "mem"}
 

vermeer/vermeer_test.go

@@ -111,9 +111,9 @@ func testScheduler(t *testing.T) {
 
 func testAlgorithms(t *testing.T) {
 	// todo: 增加算法名称
-	// var computeTasks = []string{"pagerank", "lpa", "wcc", "degree_out", "degree_in", "degree_both", "triangle_count",
-	// 	"sssp", "closeness_centrality", "betweenness_centrality", "kcore", "jaccard", "ppr", "clustering_coefficient", "scc", "louvain"}
-	var computeTasks = []string{"pagerank"}
+	var computeTasks = []string{"pagerank", "lpa", "wcc", "degree_out", "degree_in", "degree_both", "triangle_count",
+		"sssp", "closeness_centrality", "betweenness_centrality", "kcore", "jaccard", "ppr", "clustering_coefficient", "scc", "louvain"}
+	// var computeTasks = []string{"pagerank"}
 
 	startTime := time.Now()
 	expectRes, err := functional.GetExpectRes(expectResPath)