check role for dedicated constraints

This commit is contained in:
nhatle 2022-07-28 15:58:49 -07:00
parent 5605030fd9
commit 670c565c00
2 changed files with 52 additions and 3 deletions

View file

@ -20,6 +20,7 @@ import (
"encoding/json"
"fmt"
"net/http"
"strings"
"github.com/aurora-scheduler/gorealis/v2/gen-go/apache/aurora"
)
@ -176,6 +177,7 @@ func (c *Client) AvailOfferReport() (OfferReport, error) {
for _, a := range o.Attributes {
if a.Name == "dedicated" {
group = a.Text.Value
break
}
}
@ -229,6 +231,7 @@ func (c *Client) FitTasks(taskConfig *aurora.TaskConfig, offers []Offer) (int64,
for _, resVal := range request {
if resVal > 0 {
isValid = true
break
}
}
@ -300,7 +303,7 @@ func fitConstraints(taskConfig *aurora.TaskConfig,
numTasksPerOffer int64) int64 {
// check dedicated attributes vs. constraints
if !isDedicated(offer, taskConfig.Constraints) {
if !isDedicated(offer, taskConfig.Job.Role, taskConfig.Constraints) {
return 0
}
@ -345,7 +348,7 @@ func fitConstraints(taskConfig *aurora.TaskConfig,
return numTasksPerOffer
}
func isDedicated(offer *Offer, constraints []*aurora.Constraint) bool {
func isDedicated(offer *Offer, role string, constraints []*aurora.Constraint) bool {
// get all dedicated attributes of an offer
dedicatedAtts := map[string]bool{}
for _, a := range offer.Attributes {
@ -365,7 +368,7 @@ func isDedicated(offer *Offer, constraints []*aurora.Constraint) bool {
found := false
for _, v := range c.Constraint.Value.Values {
if dedicatedAtts[v] {
if dedicatedAtts[v] && strings.HasPrefix(v, fmt.Sprintf("%s/", role)) {
found = true
break
}
@ -382,6 +385,8 @@ func isDedicated(offer *Offer, constraints []*aurora.Constraint) bool {
return matched
}
// valueConstraint checks Value Contraints of task if the are matched by the offer.
// more details can be found here https://aurora.apache.org/documentation/latest/features/constraints/
func valueConstraint(offer *Offer, constraint *aurora.Constraint) bool {
matched := false
@ -405,6 +410,8 @@ func valueConstraint(offer *Offer, constraint *aurora.Constraint) bool {
return matched
}
// limitConstraint limits the number of pods on a group which has the same attribute.
// more details can be found here https://aurora.apache.org/documentation/latest/features/constraints/
func limitConstraint(offer *Offer, constraint *aurora.Constraint, limitCounts map[string]map[string]int64) int64 {
limit := int64(-1)
for _, a := range offer.Attributes {

View file

@ -1054,6 +1054,7 @@ func TestRealisClient_FitTasks(t *testing.T) {
tests := []struct {
message string
role string
request aurora.Resource
constraints []*aurora.Constraint
expected int64
@ -1061,6 +1062,7 @@ func TestRealisClient_FitTasks(t *testing.T) {
}{
{
message: "task with gpu request",
role: "vagrant",
request: aurora.Resource{
NumGpus: &gpu,
},
@ -1069,12 +1071,14 @@ func TestRealisClient_FitTasks(t *testing.T) {
},
{
message: "empty resource request",
role: "vagrant",
request: aurora.Resource{},
expected: -1,
isError: true,
},
{
message: "valid resource request",
role: "vagrant",
request: aurora.Resource{
NumCpus: &validCpu,
},
@ -1083,6 +1087,7 @@ func TestRealisClient_FitTasks(t *testing.T) {
},
{
message: "invalid cpu request",
role: "vagrant",
request: aurora.Resource{
NumCpus: &inValidCpu,
},
@ -1091,9 +1096,11 @@ func TestRealisClient_FitTasks(t *testing.T) {
},
{
message: "dedicated constraint",
role: "vagrant",
request: aurora.Resource{
NumCpus: &validCpu,
},
constraints: []*aurora.Constraint{
{
Name: "dedicated",
@ -1108,8 +1115,29 @@ func TestRealisClient_FitTasks(t *testing.T) {
expected: 2,
isError: false,
},
{
message: "dedicated constraint with unauthorized role",
role: "unauthorized",
request: aurora.Resource{
NumCpus: &validCpu,
},
constraints: []*aurora.Constraint{
{
Name: "dedicated",
Constraint: &aurora.TaskConstraint{
Value: &aurora.ValueConstraint{
Negated: false,
Values: []string{"vagrant/bar"},
},
},
},
},
expected: 0,
isError: false,
},
{
message: "value constraint on zone",
role: "vagrant",
request: aurora.Resource{
NumCpus: &validCpu,
},
@ -1129,6 +1157,7 @@ func TestRealisClient_FitTasks(t *testing.T) {
},
{
message: "negative value constraint on zone",
role: "vagrant",
request: aurora.Resource{
NumCpus: &validCpu,
},
@ -1148,6 +1177,7 @@ func TestRealisClient_FitTasks(t *testing.T) {
},
{
message: "negative value constraint on host",
role: "vagrant",
request: aurora.Resource{
NumCpus: &validCpu,
},
@ -1167,6 +1197,7 @@ func TestRealisClient_FitTasks(t *testing.T) {
},
{
message: "value constraint on unavailable zone",
role: "vagrant",
request: aurora.Resource{
NumCpus: &validCpu,
},
@ -1186,6 +1217,7 @@ func TestRealisClient_FitTasks(t *testing.T) {
},
{
message: "value constraint on unavailable attribute",
role: "vagrant",
request: aurora.Resource{
NumCpus: &validCpu,
},
@ -1205,6 +1237,7 @@ func TestRealisClient_FitTasks(t *testing.T) {
},
{
message: "1 value constraint with 2 values",
role: "vagrant",
request: aurora.Resource{
NumCpus: &validCpu,
},
@ -1224,6 +1257,7 @@ func TestRealisClient_FitTasks(t *testing.T) {
},
{
message: "2 value constraints",
role: "vagrant",
request: aurora.Resource{
NumCpus: &validCpu,
},
@ -1252,6 +1286,7 @@ func TestRealisClient_FitTasks(t *testing.T) {
},
{
message: "limit constraint on host",
role: "vagrant",
request: aurora.Resource{
NumCpus: &validCpu,
},
@ -1270,6 +1305,7 @@ func TestRealisClient_FitTasks(t *testing.T) {
},
{
message: "limit constraint on zone",
role: "vagrant",
request: aurora.Resource{
NumCpus: &validCpu,
},
@ -1288,6 +1324,7 @@ func TestRealisClient_FitTasks(t *testing.T) {
},
{
message: "limit constraint on zone & host",
role: "vagrant",
request: aurora.Resource{
NumCpus: &validCpu,
},
@ -1314,6 +1351,7 @@ func TestRealisClient_FitTasks(t *testing.T) {
},
{
message: "limit constraint on unavailable zone",
role: "vagrant",
request: aurora.Resource{
NumCpus: &validCpu,
},
@ -1332,6 +1370,7 @@ func TestRealisClient_FitTasks(t *testing.T) {
},
{
message: "limit & dedicated constraint",
role: "vagrant",
request: aurora.Resource{
NumCpus: &validCpu,
},
@ -1363,6 +1402,9 @@ func TestRealisClient_FitTasks(t *testing.T) {
task := aurora.NewTaskConfig()
task.Resources = []*aurora.Resource{&tc.request}
task.Constraints = tc.constraints
task.Job = &aurora.JobKey{
Role: tc.role,
}
numTasks, err := r.FitTasks(task, offers)