Skip to content

Commit

Permalink
Add GPU support to compose service
Browse files Browse the repository at this point in the history
  • Loading branch information
efekarakus committed Mar 18, 2019
1 parent 252549f commit d696be9
Show file tree
Hide file tree
Showing 3 changed files with 70 additions and 1 deletion.
55 changes: 55 additions & 0 deletions ecs-cli/modules/utils/compose/convert_task_definition_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 //
//////////////////////
Expand Down Expand Up @@ -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
}
1 change: 1 addition & 0 deletions ecs-cli/modules/utils/compose/ecs_params_reader.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
15 changes: 14 additions & 1 deletion ecs-cli/modules/utils/compose/reconcile_container_def.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -153,5 +162,9 @@ func reconcileContainerDef(inputCfg *adapter.ContainerConfig, ecsConDef *Contain
outputContDef.SetHealthCheck(healthCheck)
}

if len(resourceRequirements) > 0 {
outputContDef.SetResourceRequirements(resourceRequirements)
}

return outputContDef, nil
}

0 comments on commit d696be9

Please sign in to comment.