Skip to content
This repository was archived by the owner on Nov 27, 2023. It is now read-only.

Commit dbe87e2

Browse files
ndeloofndeloof
committed
Guess AWS machine type based on service resources reservations
Signed-off-by: Nicolas De Loof <[email protected]>
1 parent 387370d commit dbe87e2

File tree

3 files changed

+259
-0
lines changed

3 files changed

+259
-0
lines changed

‎ecs/compatibility.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,8 @@ var compatibleComposeAttributes = []string{
3838
"services.deploy.resources.reservations",
3939
"services.deploy.resources.reservations.cpus",
4040
"services.deploy.resources.reservations.memory",
41+
"services.deploy.resources.reservations.generic_resources",
42+
"services.deploy.resources.reservations.generic_resources.discrete_resource_spec",
4143
"services.deploy.update_config",
4244
"services.deploy.update_config.parallelism",
4345
"services.entrypoint",

‎ecs/gpu.go‎

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
/*
2+
Copyright 2020 Docker, Inc.
3+
4+
Licensed under the Apache License, Version 2.0 (the "License");
5+
you may not use this file except in compliance with the License.
6+
You may obtain a copy of the License at
7+
8+
http://www.apache.org/licenses/LICENSE-2.0
9+
10+
Unless required by applicable law or agreed to in writing, software
11+
distributed under the License is distributed on an "AS IS" BASIS,
12+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
See the License for the specific language governing permissions and
14+
limitations under the License.
15+
*/
16+
17+
package ecs
18+
19+
import (
20+
"fmt"
21+
"strconv"
22+
23+
"github.com/compose-spec/compose-go/types"
24+
"github.com/docker/go-units"
25+
)
26+
27+
type machine struct {
28+
id string
29+
cpus float64
30+
memory types.UnitBytes
31+
gpus int64
32+
}
33+
34+
type family []machine
35+
36+
var p3family = family{
37+
{
38+
id: "p3.2xlarge",
39+
cpus: 8,
40+
memory: 64 * units.GiB,
41+
gpus: 2,
42+
},
43+
{
44+
id: "p3.8xlarge",
45+
cpus: 32,
46+
memory: 244 * units.GiB,
47+
gpus: 4,
48+
},
49+
{
50+
id: "p3.16xlarge",
51+
cpus: 64,
52+
memory: 488 * units.GiB,
53+
gpus: 8,
54+
},
55+
}
56+
57+
type filterFn func(machine) bool
58+
59+
func (f family) filter(fn filterFn) family {
60+
var filtered family
61+
for _, machine := range f {
62+
if fn(machine) {
63+
filtered = append(filtered, machine)
64+
}
65+
}
66+
return filtered
67+
}
68+
69+
func (f family) firstOrError(msg string, args ...interface{}) (machine, error) {
70+
if len(f) == 0 {
71+
return machine{}, fmt.Errorf(msg, args...)
72+
}
73+
return f[0], nil
74+
}
75+
76+
func guessMachineType(project *types.Project) (string, error) {
77+
// we select a machine type to match all gpu-bound services requirements
78+
// once https://github.com/aws/containers-roadmap/issues/631 is implemented we can define dedicated CapacityProviders per service.
79+
minMemory, minCPU, minGPU, err := getResourceRequirements(project)
80+
if err != nil {
81+
return "", err
82+
}
83+
84+
instanceType, err := p3family.
85+
filter(func(m machine) bool {
86+
return m.memory >= minMemory
87+
}).
88+
filter(func(m machine) bool {
89+
return m.cpus >= minCPU
90+
}).
91+
filter(func(m machine) bool {
92+
return m.gpus >= minGPU
93+
}).
94+
firstOrError("none of the AWS p3 machines match requirement for memory:%d cpu:%f gpu:%d", minMemory, minCPU, minGPU)
95+
if err != nil {
96+
return "", err
97+
}
98+
return instanceType.id, nil
99+
}
100+
101+
func getResourceRequirements(project *types.Project) (types.UnitBytes, float64, int64, error) {
102+
var minMemory types.UnitBytes
103+
var minCPU float64
104+
var minGPU int64
105+
for _, service := range project.Services {
106+
if service.Deploy == nil {
107+
continue
108+
}
109+
reservations := service.Deploy.Resources.Reservations
110+
if reservations == nil {
111+
continue
112+
}
113+
114+
var requiredGPUs int64
115+
for _, r := range reservations.GenericResources {
116+
if r.DiscreteResourceSpec.Kind == "gpu" {
117+
requiredGPUs = r.DiscreteResourceSpec.Value
118+
break
119+
}
120+
}
121+
if requiredGPUs == 0 {
122+
continue
123+
}
124+
if requiredGPUs > minGPU {
125+
minGPU = requiredGPUs
126+
}
127+
128+
if reservations.MemoryBytes > minMemory {
129+
minMemory = reservations.MemoryBytes
130+
}
131+
if reservations.NanoCPUs != "" {
132+
nanocpu, err := strconv.ParseFloat(reservations.NanoCPUs, 64)
133+
if err != nil {
134+
return 0, 0, 0, err
135+
}
136+
if nanocpu > minCPU {
137+
minCPU = nanocpu
138+
}
139+
}
140+
}
141+
return minMemory, minCPU, minGPU, nil
142+
}

‎ecs/gpu_test.go‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
/*
2+
Copyright 2020 Docker, Inc.
3+
4+
Licensed under the Apache License, Version 2.0 (the "License");
5+
you may not use this file except in compliance with the License.
6+
You may obtain a copy of the License at
7+
8+
http://www.apache.org/licenses/LICENSE-2.0
9+
10+
Unless required by applicable law or agreed to in writing, software
11+
distributed under the License is distributed on an "AS IS" BASIS,
12+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
See the License for the specific language governing permissions and
14+
limitations under the License.
15+
*/
16+
17+
package ecs
18+
19+
import (
20+
"testing"
21+
)
22+
23+
func TestGuessMachineType(t *testing.T) {
24+
tests := []struct {
25+
name string
26+
yaml string
27+
want string
28+
wantErr bool
29+
}{
30+
{
31+
name: "1-gpu",
32+
yaml: `
33+
services:
34+
learning:
35+
image: tensorflow/tensorflow:latest-gpu
36+
deploy:
37+
resources:
38+
reservations:
39+
generic_resources:
40+
- discrete_resource_spec:
41+
kind: gpu
42+
value: 1
43+
`,
44+
want: "p3.2xlarge",
45+
wantErr: false,
46+
},
47+
{
48+
name: "4-gpu",
49+
yaml: `
50+
services:
51+
learning:
52+
image: tensorflow/tensorflow:latest-gpu
53+
deploy:
54+
resources:
55+
reservations:
56+
generic_resources:
57+
- discrete_resource_spec:
58+
kind: gpu
59+
value: 4
60+
`,
61+
want: "p3.8xlarge",
62+
wantErr: false,
63+
},
64+
{
65+
name: "1-gpu, high-memory",
66+
yaml: `
67+
services:
68+
learning:
69+
image: tensorflow/tensorflow:latest-gpu
70+
deploy:
71+
resources:
72+
reservations:
73+
memory: 300Gb
74+
generic_resources:
75+
- discrete_resource_spec:
76+
kind: gpu
77+
value: 2
78+
`,
79+
want: "p3.16xlarge",
80+
wantErr: false,
81+
},
82+
{
83+
name: "1-gpu, high-cpu",
84+
yaml: `
85+
services:
86+
learning:
87+
image: tensorflow/tensorflow:latest-gpu
88+
deploy:
89+
resources:
90+
reservations:
91+
memory: 32Gb
92+
cpus: "32"
93+
generic_resources:
94+
- discrete_resource_spec:
95+
kind: gpu
96+
value: 2
97+
`,
98+
want: "p3.8xlarge",
99+
wantErr: false,
100+
},
101+
}
102+
for _, tt := range tests {
103+
t.Run(tt.name, func(t *testing.T) {
104+
project := loadConfig(t, tt.yaml)
105+
got, err := guessMachineType(project)
106+
if (err != nil) != tt.wantErr {
107+
t.Errorf("guessMachineType() error = %v, wantErr %v", err, tt.wantErr)
108+
return
109+
}
110+
if got != tt.want {
111+
t.Errorf("guessMachineType() got = %v, want %v", got, tt.want)
112+
}
113+
})
114+
}
115+
}

0 commit comments

Comments
 (0)