From 87cacb427fa35e3ab541b22976436c5259bd9a17 Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Mon, 8 Oct 2018 16:09:41 -0700
Subject: [PATCH] parse devices

---
 api/resources.go                   | 11 +++++++-
 command/agent/job_endpoint.go      |  2 +-
 command/agent/job_endpoint_test.go |  4 +--
 jobspec/parse.go                   | 44 ++++++++++++++++++++++++++++--
 jobspec/parse_test.go              | 10 +++++++
 jobspec/test-fixtures/basic.hcl    |  6 ++++
 6 files changed, 71 insertions(+), 6 deletions(-)

diff --git a/api/resources.go b/api/resources.go
index a3c6da8c7..6cd4b483b 100644
--- a/api/resources.go
+++ b/api/resources.go
@@ -29,6 +29,9 @@ func (r *Resources) Canonicalize() {
 	for _, n := range r.Networks {
 		n.Canonicalize()
 	}
+	for _, d := range r.Devices {
+		d.Canonicalize()
+	}
 }
 
 // DefaultResources is a small resources object that contains the
@@ -117,5 +120,11 @@ type RequestedDevice struct {
 	Name string
 
 	// Count is the number of requested devices
-	Count uint64
+	Count *uint64
+}
+
+func (d *RequestedDevice) Canonicalize() {
+	if d.Count == nil {
+		d.Count = helper.Uint64ToPtr(1)
+	}
 }
diff --git a/command/agent/job_endpoint.go b/command/agent/job_endpoint.go
index ba763745a..0e4fcfe7f 100644
--- a/command/agent/job_endpoint.go
+++ b/command/agent/job_endpoint.go
@@ -934,7 +934,7 @@ func ApiResourcesToStructs(in *api.Resources) *structs.Resources {
 		for i, d := range in.Devices {
 			out.Devices[i] = &structs.RequestedDevice{
 				Name:  d.Name,
-				Count: d.Count,
+				Count: *d.Count,
 			}
 		}
 	}
diff --git a/command/agent/job_endpoint_test.go b/command/agent/job_endpoint_test.go
index 7119902ca..8654740e4 100644
--- a/command/agent/job_endpoint_test.go
+++ b/command/agent/job_endpoint_test.go
@@ -1419,11 +1419,11 @@ func TestJobs_ApiJobToStructsJob(t *testing.T) {
 							Devices: []*api.RequestedDevice{
 								{
 									Name:  "nvidia/gpu",
-									Count: 4,
+									Count: helper.Uint64ToPtr(4),
 								},
 								{
 									Name:  "gpu",
-									Count: 1,
+									Count: nil,
 								},
 							},
 						},
diff --git a/jobspec/parse.go b/jobspec/parse.go
index 0bb3949e0..9837cf0b3 100644
--- a/jobspec/parse.go
+++ b/jobspec/parse.go
@@ -755,7 +755,7 @@ func parseSpread(result *[]*api.Spread, list *ast.ObjectList) error {
 		// Parse spread target
 		if o := listVal.Filter("target"); len(o.Items) > 0 {
 			if err := parseSpreadTarget(&s.SpreadTarget, o); err != nil {
-				return multierror.Prefix(err, fmt.Sprintf("error parsing spread target"))
+				return multierror.Prefix(err, fmt.Sprintf("target ->"))
 			}
 		}
 
@@ -766,9 +766,11 @@ func parseSpread(result *[]*api.Spread, list *ast.ObjectList) error {
 }
 
 func parseSpreadTarget(result *[]*api.SpreadTarget, list *ast.ObjectList) error {
-
 	seen := make(map[string]struct{})
 	for _, item := range list.Items {
+		if len(item.Keys) != 1 {
+			return fmt.Errorf("missing spread target")
+		}
 		n := item.Keys[0].Token.Value().(string)
 
 		// Make sure we haven't already found this
@@ -1413,6 +1415,7 @@ func parseResources(result *api.Resources, list *ast.ObjectList) error {
 		"disk",
 		"memory",
 		"network",
+		"device",
 	}
 	if err := helper.CheckHCLKeys(listVal, valid); err != nil {
 		return multierror.Prefix(err, "resources ->")
@@ -1423,6 +1426,7 @@ func parseResources(result *api.Resources, list *ast.ObjectList) error {
 		return err
 	}
 	delete(m, "network")
+	delete(m, "device")
 
 	if err := mapstructure.WeakDecode(m, result); err != nil {
 		return err
@@ -1465,6 +1469,42 @@ func parseResources(result *api.Resources, list *ast.ObjectList) error {
 		result.Networks = []*api.NetworkResource{&r}
 	}
 
+	// Parse the device resources
+	if o := listVal.Filter("device"); len(o.Items) > 0 {
+		result.Devices = make([]*api.RequestedDevice, len(o.Items))
+		for idx, do := range o.Items {
+			if l := len(do.Keys); l == 0 {
+				return multierror.Prefix(fmt.Errorf("missing device name"), fmt.Sprintf("resources, device[%d]->", idx))
+			} else if l > 1 {
+				return multierror.Prefix(fmt.Errorf("only one name may be specified"), fmt.Sprintf("resources, device[%d]->", idx))
+			}
+			name := do.Keys[0].Token.Value().(string)
+
+			// Check for invalid keys
+			valid := []string{
+				"name",
+				"count",
+			}
+			if err := helper.CheckHCLKeys(do.Val, valid); err != nil {
+				return multierror.Prefix(err, fmt.Sprintf("resources, device[%d]->", idx))
+			}
+
+			// Set the name
+			var r api.RequestedDevice
+			r.Name = name
+
+			var m map[string]interface{}
+			if err := hcl.DecodeObject(&m, do.Val); err != nil {
+				return err
+			}
+			if err := mapstructure.WeakDecode(m, &r); err != nil {
+				return err
+			}
+
+			result.Devices[idx] = &r
+		}
+	}
+
 	return nil
 }
 
diff --git a/jobspec/parse_test.go b/jobspec/parse_test.go
index 394270bfe..721dcef1b 100644
--- a/jobspec/parse_test.go
+++ b/jobspec/parse_test.go
@@ -230,6 +230,16 @@ func TestParse(t *testing.T) {
 											DynamicPorts:  []api.Port{{Label: "http", Value: 0}, {Label: "https", Value: 0}, {Label: "admin", Value: 0}},
 										},
 									},
+									Devices: []*api.RequestedDevice{
+										{
+											Name:  "nvidia/gpu",
+											Count: helper.Uint64ToPtr(10),
+										},
+										{
+											Name:  "intel/gpu",
+											Count: nil,
+										},
+									},
 								},
 								KillTimeout:   helper.TimeToPtr(22 * time.Second),
 								ShutdownDelay: 11 * time.Second,
diff --git a/jobspec/test-fixtures/basic.hcl b/jobspec/test-fixtures/basic.hcl
index 3b6748cb8..c9ffd8e14 100644
--- a/jobspec/test-fixtures/basic.hcl
+++ b/jobspec/test-fixtures/basic.hcl
@@ -196,6 +196,12 @@ job "binstore-storagelocker" {
           port "admin" {
           }
         }
+
+        device "nvidia/gpu" {
+            count = 10
+        }
+        
+        device "intel/gpu" {}
       }
 
       kill_timeout = "22s"