From fef73b8af380b3d1bd07ac36f92d82399c810f1a Mon Sep 17 00:00:00 2001
From: aiordache <anca.iordache@docker.com>
Date: Tue, 5 Jan 2021 16:40:52 +0100
Subject: [PATCH] Parse reservation.devices for GPU requests

Signed-off-by: aiordache <anca.iordache@docker.com>
---
 docs/ecs-compose-examples.md |  6 +++---
 ecs/compatibility.go         |  8 ++++++++
 ecs/convert.go               | 19 +++++++++++++++++++
 ecs/gpu.go                   | 19 ++++++++++++++++++-
 4 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/docs/ecs-compose-examples.md b/docs/ecs-compose-examples.md
index ce5d32a95..f0e4e7f47 100644
--- a/docs/ecs-compose-examples.md
+++ b/docs/ecs-compose-examples.md
@@ -108,9 +108,9 @@ services:
           memory: 32Gb
           cpus: "32"
           generic_resources:
-          - discrete_resource_spec:
-            kind: gpus
-            value: 2
+            - discrete_resource_spec:
+                kind: gpus
+                value: 2
 ```
 
 
diff --git a/ecs/compatibility.go b/ecs/compatibility.go
index 8cc9b3503..814c239bd 100644
--- a/ecs/compatibility.go
+++ b/ecs/compatibility.go
@@ -66,7 +66,9 @@ var compatibleComposeAttributes = []string{
 	"services.deploy.resources.reservations.cpus",
 	"services.deploy.resources.reservations.memory",
 	"services.deploy.resources.reservations.devices",
+	"services.deploy.resources.reservations.devices.capabilities",
 	"services.deploy.resources.reservations.devices.count",
+	"services.deploy.resources.reservations.devices.driver",
 	"services.deploy.resources.reservations.generic_resources",
 	"services.deploy.resources.reservations.generic_resources.discrete_resource_spec",
 	"services.deploy.update_config",
@@ -167,3 +169,9 @@ func (c *fargateCompatibilityChecker) CheckDeployResourcesDevicesCapabilities(s
 		}
 	}
 }
+
+func (c *fargateCompatibilityChecker) CheckDeployResourcesDevicesDriver(s string, r types.DeviceRequest) {
+	if r.Driver != "" && r.Driver != "nvidia" {
+		c.Unsupported("services.deploy.resources.%s.devices.driver = %s", s, r.Driver)
+	}
+}
diff --git a/ecs/convert.go b/ecs/convert.go
index 220f9670f..fc23a9618 100644
--- a/ecs/convert.go
+++ b/ecs/convert.go
@@ -181,6 +181,25 @@ func toTaskResourceRequirements(reservations *types.Resource) []ecs.TaskDefiniti
 			})
 		}
 	}
+	for _, r := range reservations.Devices {
+		hasGpuCap := false
+		for _, c := range r.Capabilities {
+			if c == "gpu" {
+				hasGpuCap = true
+				break
+			}
+		}
+		if hasGpuCap {
+			count := r.Count
+			if count <= 0 {
+				count = 1
+			}
+			requirements = append(requirements, ecs.TaskDefinition_ResourceRequirement{
+				Type:  ecsapi.ResourceTypeGpu,
+				Value: fmt.Sprint(count),
+			})
+		}
+	}
 	return requirements
 }
 
diff --git a/ecs/gpu.go b/ecs/gpu.go
index e38667634..a5f0386f0 100644
--- a/ecs/gpu.go
+++ b/ecs/gpu.go
@@ -132,7 +132,7 @@ type resourceRequirements struct {
 func getResourceRequirements(project *types.Project) (*resourceRequirements, error) {
 	return toResourceRequirementsSlice(project).
 		filter(func(requirements *resourceRequirements) bool {
-			return requirements.gpus != 0
+			return requirements != nil && requirements.gpus != 0
 		}).
 		max()
 }
@@ -184,6 +184,23 @@ func toResourceRequirements(service types.ServiceConfig) (*resourceRequirements,
 		}
 	}
 
+	for _, r := range reservations.Devices {
+		requiresGpus := false
+		for _, c := range r.Capabilities {
+			if c == "gpu" {
+				requiresGpus = true
+				break
+			}
+		}
+		if requiresGpus {
+			requiredGPUs = r.Count
+			if requiredGPUs <= 0 {
+				requiredGPUs = 1
+			}
+			break
+		}
+	}
+
 	var nanocpu float64
 	if reservations.NanoCPUs != "" {
 		v, err := strconv.ParseFloat(reservations.NanoCPUs, 64)