Skip to content

Commit

Permalink
fix enable node device score plugin
Browse files Browse the repository at this point in the history
  • Loading branch information
coldzerofear committed Jul 10, 2024
1 parent 94944b5 commit f2cf893
Show file tree
Hide file tree
Showing 3 changed files with 81 additions and 7 deletions.
19 changes: 19 additions & 0 deletions pkg/scheduler/framework/arguments.go
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,25 @@ func (a Arguments) GetBool(ptr *bool, key string) {
*ptr = value
}

func (a Arguments) GetString(ptr *string, key string) {
if ptr == nil {
return
}

argv, ok := a[key]
if !ok {
return
}

value, ok := argv.(string)
if !ok {
klog.Warningf("Could not parse argument: %v for key %s to string", argv, key)
return
}

*ptr = value
}

// GetArgOfActionFromConf return argument of action reading from configuration of schedule
func GetArgOfActionFromConf(configurations []conf.Configuration, actionName string) Arguments {
for _, c := range configurations {
Expand Down
57 changes: 57 additions & 0 deletions pkg/scheduler/framework/arguments_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -187,3 +187,60 @@ func TestGetArgOfActionFromConf(t *testing.T) {
}
}
}

func TestArgumentsGetString(t *testing.T) {
key1 := "stringkey"

cases := []struct {
name string
arg Arguments
key string
baseValue string
expectValue string
}{
{
name: "key not exist",
arg: Arguments{
"anotherKey": "test",
},
key: key1,
baseValue: "test1",
expectValue: "test1",
},
{
name: "key exist",
arg: Arguments{
key1: "test1",
},
key: key1,
baseValue: "test",
expectValue: "test1",
},
{
name: "value of key invalid",
arg: Arguments{
key1: 0,
},
key: key1,
baseValue: "test",
expectValue: "test",
},
{
name: "value of key null",
arg: Arguments{
key1: nil,
},
key: key1,
baseValue: "test",
expectValue: "test",
},
}

for index, c := range cases {
baseValue := c.baseValue
c.arg.GetString(&baseValue, c.key)
if baseValue != c.expectValue {
t.Errorf("index %d, case %s, value should be %v, but not %v", index, c.name, c.expectValue, baseValue)
}
}
}
12 changes: 5 additions & 7 deletions pkg/scheduler/plugins/deviceshare/deviceshare.go
Original file line number Diff line number Diff line change
Expand Up @@ -73,10 +73,7 @@ func enablePredicate(dsp *deviceSharePlugin) {
args.GetBool(&gpushare.NodeLockEnable, NodeLockEnable)
args.GetBool(&vgpu.VGPUEnable, VGPUEnable)

_, ok := args[SchedulePolicyArgument]
if ok {
dsp.schedulePolicy = args[SchedulePolicyArgument].(string)
}
args.GetString(&dsp.schedulePolicy, SchedulePolicyArgument)
args.GetInt(&dsp.scheduleWeight, ScheduleWeight)

if gpushare.GpuSharingEnable && gpushare.GpuNumberEnable {
Expand Down Expand Up @@ -149,18 +146,19 @@ func (dp *deviceSharePlugin) OnSessionOpen(ssn *framework.Session) {

ssn.AddNodeOrderFn(dp.Name(), func(task *api.TaskInfo, node *api.NodeInfo) (float64, error) {
// DeviceScore
if len(dp.schedulePolicy) > 0 {
nodeScore := float64(0)
if dp.scheduleWeight > 0 {
score, status := getDeviceScore(context.TODO(), task.Pod, node, dp.schedulePolicy)
if !status.IsSuccess() {
klog.Warningf("Node: %s, Calculate Device Score Failed because of Error: %v", node.Name, status.AsError())
return 0, status.AsError()
}

// TODO: we should use a seperate plugin for devices, and seperate them from predicates and nodeOrder plugin.
nodeScore := float64(score) * float64(dp.scheduleWeight)
nodeScore = float64(score) * float64(dp.scheduleWeight)
klog.V(5).Infof("Node: %s, task<%s/%s> Device Score weight %d, score: %f", node.Name, task.Namespace, task.Name, dp.scheduleWeight, nodeScore)
}
return 0, nil
return nodeScore, nil
})
}

Expand Down

0 comments on commit f2cf893

Please sign in to comment.