chore: change getting all task to complete/ refresh group assign

Author: Ethereal-O

vermeer/apps/master/bl/compute_task.go

@@ -142,6 +142,8 @@ func (ctb *ComputeTaskBl) ComputeTaskStatus(
 				}
 			}
 			taskMgr.ForceState(computeTask.Task, structure.TaskStateComplete)
+			// for scheduler, mark task complete
+			Scheduler.taskManager.MarkTaskComplete(taskId)
 			graph.SubUsingNum()
 			computeTask.FreeMemory()
 			needQuery := options.GetInt(computeTask.Task.Params, "output.need_query") == 1

vermeer/apps/master/bl/load_task.go

@@ -204,6 +204,9 @@ func (lb *LoadTaskBl) LoadTaskStatus(taskId int32, state string, workerName stri
 			loadTask.Task.SetState(structure.TaskStateLoaded)
 			//TaskMgr.ForceState(loadTask.Task, structure.TaskStateLoaded)
 
+			// for scheduler, mark task complete
+			Scheduler.taskManager.MarkTaskComplete(taskId)
+
 			logrus.Infof("graph: %s, vertex: %d, edge: %d", graph.Name, graph.VertexCount, graph.EdgeCount)
 			for _, w := range graph.Workers {
 				logrus.Infof(

vermeer/apps/master/bl/scheduler_bl.go

@@ -134,19 +134,23 @@ func (s *ScheduleBl) tryScheduleInner(softSchedule bool, noLock ...bool) error {
 
 	// 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
+	s.taskManager.RefreshTaskToWorkerGroupMap()
 
 	// step 2: get available resources and tasks
 	logrus.Debugf("scheduling next tasks, softSchedule: %v", softSchedule)
 	idleWorkerGroups := s.resourceManager.GetIdleWorkerGroups()
 	concurrentWorkerGroups := s.resourceManager.GetConcurrentWorkerGroups()
-	allTasks := s.taskManager.GetAllTasks()
+	allTasks := s.taskManager.GetAllTasksNotComplete()
 	if len(allTasks) == 0 || (len(idleWorkerGroups) == 0 && len(concurrentWorkerGroups) == 0) {
 		logrus.Debugf("no available tasks or workerGroups, allTasks: %d, workerGroups: %d/%d",
 			len(allTasks), len(idleWorkerGroups), len(concurrentWorkerGroups))
 		return nil
 	}
 	logrus.Debugf("all tasks: %d, workerGroups: %d/%d", len(allTasks), len(idleWorkerGroups), len(concurrentWorkerGroups))
 
+	// TODO: NEED TO JUDGE IF THE TASK CAN CONCURRENTLY RUNNING
+	// NOT only by user setting, but also by scheduler setting
+
 	// step 3: return the task with the highest priority or small tasks which can be executed immediately
 	taskToWorkerGroupMap := s.taskManager.GetTaskToWorkerGroupMap()
 	nextTasks, err := s.algorithmManager.ScheduleNextTasks(allTasks, taskToWorkerGroupMap, idleWorkerGroups, concurrentWorkerGroups, softSchedule)
@@ -195,6 +199,19 @@ func (s *ScheduleBl) QueueTask(taskInfo *structure.TaskInfo) (bool, error) {
 
 	logrus.Debugf("queuing task %d with parameters: %+v", taskInfo.ID, taskInfo)
 
+	// check dependency if exists
+	if len(taskInfo.Preorders) > 0 {
+		for _, depTaskID := range taskInfo.Preorders {
+			depTask := taskMgr.GetTaskByID(depTaskID)
+			if depTask == nil {
+				err := errors.New("the dependency task with ID " + strconv.Itoa(int(depTaskID)) + " does not exist")
+				logrus.Error(err)
+				taskMgr.SetError(taskInfo, err.Error())
+				return false, err
+			}
+		}
+	}
+
 	// Notice: Ensure successful invocation.
 	// make sure all tasks have alloc to a worker group
 	ok, err := s.taskManager.QueueTask(taskInfo)

vermeer/apps/master/schedules/scheduler_task_manager.go

@@ -14,12 +14,15 @@ type SchedulerTaskManager struct {
 	allTaskMap     map[int32]*structure.TaskInfo
 	allTaskQueue   []*structure.TaskInfo
 	startTaskQueue []*structure.TaskInfo
+	// onGoingTasks
+	notCompleteTasks map[int32]*structure.TaskInfo
 	// A map from task ID to worker group can be used to track which worker group is handling which task.
 	taskToworkerGroupMap map[int32]string
 }
 
 func (t *SchedulerTaskManager) Init() *SchedulerTaskManager {
 	t.allTaskMap = make(map[int32]*structure.TaskInfo)
+	t.notCompleteTasks = make(map[int32]*structure.TaskInfo)
 	t.taskToworkerGroupMap = make(map[int32]string)
 	return t
 }
@@ -38,10 +41,23 @@ func (t *SchedulerTaskManager) QueueTask(taskInfo *structure.TaskInfo) (bool, er
 	// Add the task to the task map
 	t.allTaskMap[taskInfo.ID] = taskInfo
 	t.allTaskQueue = append(t.allTaskQueue, taskInfo)
+	t.notCompleteTasks[taskInfo.ID] = taskInfo
 	t.AssignGroup(taskInfo)
 	return true, nil
 }
 
+func (t *SchedulerTaskManager) RefreshTaskToWorkerGroupMap() {
+	defer t.Unlock(t.Lock())
+
+	for _, taskInfo := range t.GetAllTasksNotComplete() {
+		if taskInfo == nil {
+			continue
+		}
+		t.AssignGroup(taskInfo)
+		t.taskToworkerGroupMap[taskInfo.ID] = workerMgr.ApplyGroup(taskInfo.SpaceName, taskInfo.GraphName)
+	}
+}
+
 // Only for debug or test, get task start sequence
 func (t *SchedulerTaskManager) AddTaskStartSequence(taskID int32) error {
 	if _, exists := t.allTaskMap[taskID]; !exists {
@@ -64,6 +80,15 @@ func (t *SchedulerTaskManager) RemoveTask(taskID int32) error {
 		}
 	}
 	delete(t.taskToworkerGroupMap, taskID)
+	delete(t.notCompleteTasks, taskID)
+	return nil
+}
+
+func (t *SchedulerTaskManager) MarkTaskComplete(taskID int32) error {
+	if _, exists := t.allTaskMap[taskID]; !exists {
+		return errors.New("task not found")
+	}
+	delete(t.notCompleteTasks, taskID)
 	return nil
 }
 
@@ -106,9 +131,17 @@ func (t *SchedulerTaskManager) GetAllTasks() []*structure.TaskInfo {
 	return tasks
 }
 
+func (t *SchedulerTaskManager) GetAllTasksNotComplete() []*structure.TaskInfo {
+	tasks := make([]*structure.TaskInfo, 0, len(t.allTaskMap))
+	for _, task := range t.notCompleteTasks {
+		tasks = append(tasks, task)
+	}
+	return tasks
+}
+
 func (t *SchedulerTaskManager) GetAllTasksWaitng() []*structure.TaskInfo {
 	tasks := make([]*structure.TaskInfo, 0, len(t.allTaskMap))
-	for _, task := range t.allTaskMap {
+	for _, task := range t.GetAllTasksNotComplete() {
 		if task.State == structure.TaskStateWaiting {
 			tasks = append(tasks, task)
 		}
@@ -118,7 +151,7 @@ func (t *SchedulerTaskManager) GetAllTasksWaitng() []*structure.TaskInfo {
 
 func (t *SchedulerTaskManager) GetTasksInQueue(space string) []*structure.TaskInfo {
 	tasks := make([]*structure.TaskInfo, 0)
-	for _, task := range t.allTaskQueue {
+	for _, task := range t.GetAllTasksNotComplete() {
 		if task.SpaceName == space {
 			tasks = append(tasks, task)
 		}
@@ -153,9 +186,12 @@ func (t *SchedulerTaskManager) GetTaskStartSequence(queryTasks []int32) []*struc
 
 func (t *SchedulerTaskManager) GetTaskToWorkerGroupMap() map[int32]string {
 	// Return a copy of the worker group map to avoid external modifications
-	groupMap := make(map[int32]string, len(t.taskToworkerGroupMap))
-	for k, v := range t.taskToworkerGroupMap {
-		groupMap[k] = v
+	taskNotComplete := t.GetAllTasksNotComplete()
+	groupMap := make(map[int32]string, len(taskNotComplete))
+	for _, task := range taskNotComplete {
+		if group, exists := t.taskToworkerGroupMap[task.ID]; exists {
+			groupMap[task.ID] = group
+		}
 	}
 	return groupMap
 }