update preempt and TDM to match VictimTasks function

Signed-off-by: Thor-wl <13164644535@163.com>
volcano-sh · Apr 16, 2022 · 56d6dcd · 56d6dcd
1 parent f464537
commit 56d6dcd
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 4 deletions.
diff --git a/pkg/scheduler/actions/preempt/preempt.go b/pkg/scheduler/actions/preempt/preempt.go
@@ -273,7 +273,12 @@ func preempt(
 
 func victimTasks(ssn *framework.Session) {
 	stmt := framework.NewStatement(ssn)
-	victimTasks := ssn.VictimTasks()
+	tasks := make([]*api.TaskInfo, 0)
+	victimTasksMap := ssn.VictimTasks(tasks)
+	victimTasks := make([]*api.TaskInfo, 0)
+	for task, _ := range victimTasksMap {
+		victimTasks = append(victimTasks, task)
+	}
 	for _, victim := range victimTasks {
 		if err := stmt.Evict(victim.Clone(), "evict"); err != nil {
 			klog.Errorf("Failed to evict Task <%s/%s>: %v",

diff --git a/pkg/scheduler/plugins/tdm/tdm.go b/pkg/scheduler/plugins/tdm/tdm.go
@@ -228,7 +228,7 @@ func (tp *tdmPlugin) OnSessionOpen(ssn *framework.Session) {
 		return victims, tutil.Permit
 	}
 
-	victimsFn := func() []*api.TaskInfo {
+	victimsFn := func([]*api.TaskInfo) []*api.TaskInfo {
 		if lastEvictAt.Add(tp.evictPeriod).After(time.Now()) {
 			klog.V(4).Infof("TDM next evict time at %v", lastEvictAt)
 			return nil
@@ -291,10 +291,12 @@ func (tp *tdmPlugin) OnSessionOpen(ssn *framework.Session) {
 		return len(jobInfo.TaskStatusIndex[api.Pending]) > 0
 	}
 
+	victimsFns := make([]api.VictimTasksFn, 0)
+	victimsFns = append(victimsFns, victimsFn)
 	ssn.AddPredicateFn(tp.Name(), predicateFn)
 	ssn.AddNodeOrderFn(tp.Name(), nodeOrderFn)
 	ssn.AddPreemptableFn(tp.Name(), preemptableFn)
-	ssn.AddVictimTasksFns(tp.Name(), victimsFn)
+	ssn.AddVictimTasksFns(tp.Name(), victimsFns)
 	ssn.AddJobOrderFn(tp.Name(), jobOrderFn)
 	ssn.AddJobPipelinedFn(tp.Name(), jobPipelinedFn)
 	ssn.AddJobStarvingFns(tp.Name(), jobStarvingFn)

diff --git a/pkg/scheduler/plugins/tdm/tdm_test.go b/pkg/scheduler/plugins/tdm/tdm_test.go
@@ -716,7 +716,8 @@ func Test_TDM_victimsFn(t *testing.T) {
 
 			d, _ := time.ParseDuration(test.args[evictPeriodLabel].(string))
 			time.Sleep(d)
-			res := ssn.VictimTasks()
+			tasks := make([]*api.TaskInfo, 0)
+			res := ssn.VictimTasks(tasks)
 			if len(res) != test.want {
 				t.Errorf("want %v, got %v", test.want, len(res))
 				return