Add GPU support to compose service

Author: efekarakus

ecs-cli/modules/utils/compose/convert_task_definition_test.go

@@ -1807,6 +1807,41 @@ task_definition:
 	}
 }
 
+func TestConvertToTaskDefinitionWithECSParams_Gpu(t *testing.T) {
+	expectedGpuValue := "2"
+	content := `version: 1
+task_definition:
+  services:
+    web:
+      gpu: ` + expectedGpuValue
+	ecsParams, err := createTempECSParamsForTest(t, content)
+
+	containerConfig := &adapter.ContainerConfig{
+		Name:  "web",
+		Image: "wordpress",
+	}
+	containerConfigs := []adapter.ContainerConfig{*containerConfig}
+	taskDefinition, err := convertToTaskDefinitionForTest(t, containerConfigs, "", "", ecsParams, nil)
+
+	containerDefs := taskDefinition.ContainerDefinitions
+	web := findContainerByName("web", containerDefs)
+
+	resourceType := ecs.ResourceTypeGpu
+	expectedResourceRequirements := []*ecs.ResourceRequirement{
+		{
+			Type:  &resourceType,
+			Value: &expectedGpuValue,
+		},
+	}
+
+	if assert.NoError(t, err) {
+		assert.ElementsMatch(t,
+			expectedResourceRequirements,
+			web.ResourceRequirements,
+			"Expected resourceRequirements to match")
+	}
+}
+
 ///////////////////////
 // helper functions //
 //////////////////////
@@ -1842,3 +1877,23 @@ func findContainerByName(name string, containerDefs []*ecs.ContainerDefinition)
 	}
 	return nil
 }
+
+func createTempECSParamsForTest(t *testing.T, content string) (*ECSParams, error) {
+	b := []byte(content)
+
+	f, err := ioutil.TempFile("", "ecs-params")
+	assert.NoError(t, err, "Could not create ecs params tempfile")
+
+	defer os.Remove(f.Name())
+
+	_, err = f.Write(b)
+	assert.NoError(t, err, "Could not write data to ecs params tempfile")
+
+	err = f.Close()
+	assert.NoError(t, err, "Could not close tempfile")
+
+	ecsParamsFileName := f.Name()
+	ecsParams, err := ReadECSParams(ecsParamsFileName)
+	assert.NoError(t, err, "Could not read ECS Params file")
+	return ecsParams, err
+}

ecs-cli/modules/utils/compose/ecs_params_reader.go

@@ -67,6 +67,7 @@ type ContainerDef struct {
 	MemoryReservation libYaml.MemStringorInt `yaml:"mem_reservation"`
 	HealthCheck       *HealthCheck           `yaml:"healthcheck"`
 	Secrets           []Secret               `yaml:"secrets"`
+	GPU               string                 `yaml:"gpu"`
 }
 
 type DockerVolume struct {

ecs-cli/modules/utils/compose/reconcile_container_def.go

@@ -15,7 +15,6 @@ package utils
 
 import (
 	"errors"
-
 	"github.com/aws/amazon-ecs-cli/ecs-cli/modules/cli/compose/adapter"
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ecs"
@@ -96,6 +95,7 @@ func reconcileContainerDef(inputCfg *adapter.ContainerConfig, ecsConDef *Contain
 	memLimit := inputCfg.Memory
 	memRes := inputCfg.MemoryReservation
 	healthCheck := inputCfg.HealthCheck
+	var resourceRequirements []*ecs.ResourceRequirement
 
 	if ecsConDef != nil {
 		outputContDef.Essential = aws.Bool(ecsConDef.Essential)
@@ -126,6 +126,15 @@ func reconcileContainerDef(inputCfg *adapter.ContainerConfig, ecsConDef *Contain
 		if err != nil {
 			return nil, err
 		}
+
+		if ecsConDef.GPU != "" {
+			resourceType := ecs.ResourceTypeGpu
+			resourceRequirement := ecs.ResourceRequirement{
+				Type:  &resourceType,
+				Value: &ecsConDef.GPU,
+			}
+			resourceRequirements = append(resourceRequirements, &resourceRequirement)
+		}
 	}
 
 	// At least one memory value is required to register a task definition.
@@ -153,5 +162,9 @@ func reconcileContainerDef(inputCfg *adapter.ContainerConfig, ecsConDef *Contain
 		outputContDef.SetHealthCheck(healthCheck)
 	}
 
+	if len(resourceRequirements) > 0 {
+		outputContDef.SetResourceRequirements(resourceRequirements)
+	}
+
 	return outputContDef, nil
 }