parse affinities and constraints on devices

api/resources.go

@@ -121,6 +121,14 @@ type RequestedDevice struct {
 
 	// Count is the number of requested devices
 	Count *uint64
+
+	// Constraints are a set of constraints to apply when selecting the device
+	// to use.
+	Constraints []*Constraint
+
+	// Affinities are a set of affinites to apply when selecting the device
+	// to use.
+	Affinities []*Affinity
 }
 
 func (d *RequestedDevice) Canonicalize() {

command/agent/job_endpoint.go

@@ -605,22 +605,8 @@ func ApiJobToStructJob(job *api.Job) *structs.Job {
 		Payload:     job.Payload,
 		Meta:        job.Meta,
 		VaultToken:  *job.VaultToken,
-	}
-
-	if l := len(job.Constraints); l != 0 {
-		j.Constraints = make([]*structs.Constraint, l)
-		for i, c := range job.Constraints {
-			con := &structs.Constraint{}
-			ApiConstraintToStructs(c, con)
-			j.Constraints[i] = con
-		}
-	}
-
-	if l := len(job.Affinities); l != 0 {
-		j.Affinities = make([]*structs.Affinity, l)
-		for i, a := range job.Affinities {
-			j.Affinities[i] = ApiAffinityToStructs(a)
-		}
+		Constraints: ApiConstraintsToStructs(job.Constraints),
+		Affinities:  ApiAffinitiesToStructs(job.Affinities),
 	}
 
 	// COMPAT: Remove in 0.7.0. Update has been pushed into the task groups
@@ -679,22 +665,8 @@ func ApiTgToStructsTG(taskGroup *api.TaskGroup, tg *structs.TaskGroup) {
 	tg.Name = *taskGroup.Name
 	tg.Count = *taskGroup.Count
 	tg.Meta = taskGroup.Meta
-
-	if l := len(taskGroup.Constraints); l != 0 {
-		tg.Constraints = make([]*structs.Constraint, l)
-		for k, constraint := range taskGroup.Constraints {
-			c := &structs.Constraint{}
-			ApiConstraintToStructs(constraint, c)
-			tg.Constraints[k] = c
-		}
-	}
-
-	if l := len(taskGroup.Affinities); l != 0 {
-		tg.Affinities = make([]*structs.Affinity, l)
-		for k, affinity := range taskGroup.Affinities {
-			tg.Affinities[k] = ApiAffinityToStructs(affinity)
-		}
-	}
+	tg.Constraints = ApiConstraintsToStructs(taskGroup.Constraints)
+	tg.Affinities = ApiAffinitiesToStructs(taskGroup.Affinities)
 
 	tg.RestartPolicy = &structs.RestartPolicy{
 		Attempts: *taskGroup.RestartPolicy.Attempts,
@@ -772,22 +744,8 @@ func ApiTaskToStructsTask(apiTask *api.Task, structsTask *structs.Task) {
 	structsTask.KillTimeout = *apiTask.KillTimeout
 	structsTask.ShutdownDelay = apiTask.ShutdownDelay
 	structsTask.KillSignal = apiTask.KillSignal
-
-	if l := len(apiTask.Constraints); l != 0 {
-		structsTask.Constraints = make([]*structs.Constraint, l)
-		for i, constraint := range apiTask.Constraints {
-			c := &structs.Constraint{}
-			ApiConstraintToStructs(constraint, c)
-			structsTask.Constraints[i] = c
-		}
-	}
-
-	if l := len(apiTask.Affinities); l != 0 {
-		structsTask.Affinities = make([]*structs.Affinity, l)
-		for i, a := range apiTask.Affinities {
-			structsTask.Affinities[i] = ApiAffinityToStructs(a)
-		}
-	}
+	structsTask.Constraints = ApiConstraintsToStructs(apiTask.Constraints)
+	structsTask.Affinities = ApiAffinitiesToStructs(apiTask.Affinities)
 
 	if l := len(apiTask.Services); l != 0 {
 		structsTask.Services = make([]*structs.Service, l)
@@ -935,6 +893,8 @@ func ApiResourcesToStructs(in *api.Resources) *structs.Resources {
 			out.Devices[i] = &structs.RequestedDevice{
 				Name:        d.Name,
 				Count:       *d.Count,
+				Constraints: ApiConstraintsToStructs(d.Constraints),
+				Affinities:  ApiAffinitiesToStructs(d.Affinities),
 			}
 		}
 	}
@@ -942,10 +902,42 @@ func ApiResourcesToStructs(in *api.Resources) *structs.Resources {
 	return out
 }
 
-func ApiConstraintToStructs(c1 *api.Constraint, c2 *structs.Constraint) {
-	c2.LTarget = c1.LTarget
-	c2.RTarget = c1.RTarget
-	c2.Operand = c1.Operand
+func ApiConstraintsToStructs(in []*api.Constraint) []*structs.Constraint {
+	if in == nil {
+		return nil
+	}
+
+	out := make([]*structs.Constraint, len(in))
+	for i, ac := range in {
+		out[i] = ApiConstraintToStructs(ac)
+	}
+
+	return out
+}
+
+func ApiConstraintToStructs(in *api.Constraint) *structs.Constraint {
+	if in == nil {
+		return nil
+	}
+
+	return &structs.Constraint{
+		LTarget: in.LTarget,
+		RTarget: in.RTarget,
+		Operand: in.Operand,
+	}
+}
+
+func ApiAffinitiesToStructs(in []*api.Affinity) []*structs.Affinity {
+	if in == nil {
+		return nil
+	}
+
+	out := make([]*structs.Affinity, len(in))
+	for i, ac := range in {
+		out[i] = ApiAffinityToStructs(ac)
+	}
+
+	return out
 }
 
 func ApiAffinityToStructs(a1 *api.Affinity) *structs.Affinity {

command/agent/job_endpoint_test.go

@@ -1420,6 +1420,21 @@ func TestJobs_ApiJobToStructsJob(t *testing.T) {
 								{
 									Name:  "nvidia/gpu",
 									Count: helper.Uint64ToPtr(4),
+									Constraints: []*api.Constraint{
+										{
+											LTarget: "x",
+											RTarget: "y",
+											Operand: "z",
+										},
+									},
+									Affinities: []*api.Affinity{
+										{
+											LTarget: "a",
+											RTarget: "b",
+											Operand: "c",
+											Weight:  50,
+										},
+									},
 								},
 								{
 									Name:  "gpu",
@@ -1704,6 +1719,21 @@ func TestJobs_ApiJobToStructsJob(t *testing.T) {
 								{
 									Name:  "nvidia/gpu",
 									Count: 4,
+									Constraints: []*structs.Constraint{
+										{
+											LTarget: "x",
+											RTarget: "y",
+											Operand: "z",
+										},
+									},
+									Affinities: []*structs.Affinity{
+										{
+											LTarget: "a",
+											RTarget: "b",
+											Operand: "c",
+											Weight:  50,
+										},
+									},
 								},
 								{
 									Name:  "gpu",

jobspec/parse.go

@@ -1480,10 +1480,20 @@ func parseResources(result *api.Resources, list *ast.ObjectList) error {
 			}
 			name := do.Keys[0].Token.Value().(string)
 
+			// Value should be an object
+			var listVal *ast.ObjectList
+			if ot, ok := do.Val.(*ast.ObjectType); ok {
+				listVal = ot.List
+			} else {
+				return fmt.Errorf("device should be an object")
+			}
+
 			// Check for invalid keys
 			valid := []string{
 				"name",
 				"count",
+				"affinity",
+				"constraint",
 			}
 			if err := helper.CheckHCLKeys(do.Val, valid); err != nil {
 				return multierror.Prefix(err, fmt.Sprintf("resources, device[%d]->", idx))
@@ -1497,10 +1507,28 @@ func parseResources(result *api.Resources, list *ast.ObjectList) error {
 			if err := hcl.DecodeObject(&m, do.Val); err != nil {
 				return err
 			}
+
+			delete(m, "constraint")
+			delete(m, "affinity")
+
 			if err := mapstructure.WeakDecode(m, &r); err != nil {
 				return err
 			}
 
+			// Parse constraints
+			if o := listVal.Filter("constraint"); len(o.Items) > 0 {
+				if err := parseConstraints(&r.Constraints, o); err != nil {
+					return multierror.Prefix(err, "constraint ->")
+				}
+			}
+
+			// Parse affinities
+			if o := listVal.Filter("affinity"); len(o.Items) > 0 {
+				if err := parseAffinities(&r.Affinities, o); err != nil {
+					return multierror.Prefix(err, "affinity ->")
+				}
+			}
+
 			result.Devices[idx] = &r
 		}
 	}

jobspec/parse_test.go

@@ -234,6 +234,21 @@ func TestParse(t *testing.T) {
 										{
 											Name:  "nvidia/gpu",
 											Count: helper.Uint64ToPtr(10),
+											Constraints: []*api.Constraint{
+												{
+													LTarget: "${driver.attr.memory}",
+													RTarget: "2GB",
+													Operand: ">",
+												},
+											},
+											Affinities: []*api.Affinity{
+												{
+													LTarget: "${driver.model}",
+													RTarget: "1080ti",
+													Operand: "=",
+													Weight:  50,
+												},
+											},
 										},
 										{
 											Name:  "intel/gpu",

jobspec/test-fixtures/basic.hcl

@@ -199,6 +199,17 @@ job "binstore-storagelocker" {
 
         device "nvidia/gpu" {
             count = 10
+            constraint {
+              attribute = "${driver.attr.memory}"
+              value = "2GB"
+              operator = ">"
+            }
+
+            affinity {
+              attribute = "${driver.model}"
+              value     = "1080ti"
+              weight = 50
+            }
         }
         
         device "intel/gpu" {}

nomad/structs/structs.go

@@ -2032,6 +2032,15 @@ type RequestedDevice struct {
 
 	// Count is the number of requested devices
 	Count uint64
+
+	// TODO validate
+	// Constraints are a set of constraints to apply when selecting the device
+	// to use.
+	Constraints []*Constraint
+
+	// Affinities are a set of affinites to apply when selecting the device
+	// to use.
+	Affinities []*Affinity
 }
 
 func (r *RequestedDevice) Copy() *RequestedDevice {
@@ -2040,6 +2049,9 @@ func (r *RequestedDevice) Copy() *RequestedDevice {
 	}
 
 	nr := *r
+	nr.Constraints = CopySliceConstraints(nr.Constraints)
+	nr.Affinities = CopySliceAffinities(nr.Affinities)
+
 	return &nr
 }