From 670c565c00992976b70733cef60f5a64fbad1c34 Mon Sep 17 00:00:00 2001 From: nhatle Date: Thu, 28 Jul 2022 15:58:49 -0700 Subject: [PATCH] check role for dedicated constraints --- offer.go | 13 ++++++++++--- realis_e2e_test.go | 42 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 3 deletions(-) diff --git a/offer.go b/offer.go index e70088a..6f5346f 100644 --- a/offer.go +++ b/offer.go @@ -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 { diff --git a/realis_e2e_test.go b/realis_e2e_test.go index 6ff2e49..bf95f72 100644 --- a/realis_e2e_test.go +++ b/realis_e2e_test.go @@ -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)