Skip to content

Commit 7e1c3a6

Browse files
committed
feat(sdl): parse amd gpu
refs akash-network/support#142 Signed-off-by: Artur Troian <troian.ap@gmail.com>
1 parent f26d00c commit 7e1c3a6

File tree

3 files changed

+263
-120
lines changed

3 files changed

+263
-120
lines changed

sdl/gpu.go

Lines changed: 42 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package sdl
22

33
import (
4+
"errors"
45
"fmt"
56
"sort"
67

@@ -9,25 +10,27 @@ import (
910
types "github.com/akash-network/akash-api/go/node/types/v1beta3"
1011
)
1112

12-
type v2GPUNvidia struct {
13+
var (
14+
ErrResourceGPUEmptyVendors = errors.New("sdl: invalid GPU attributes. at least one vendor must be set")
15+
)
16+
17+
type v2GPU struct {
1318
Model string `yaml:"model"`
1419
RAM *memoryQuantity `yaml:"ram,omitempty"`
1520
}
1621

17-
func (sdl *v2GPUNvidia) String() string {
22+
func (sdl *v2GPU) String() string {
1823
key := sdl.Model
1924
if sdl.RAM != nil {
20-
key += "/" + sdl.RAM.StringWithSuffix("Gi")
25+
key += "/ram/" + sdl.RAM.StringWithSuffix("Gi")
2126
}
2227

2328
return key
2429
}
2530

26-
type v2GPUsNvidia []v2GPUNvidia
31+
type v2GPUs []v2GPU
2732

28-
type gpuVendor struct {
29-
Nvidia v2GPUsNvidia `yaml:"nvidia,omitempty"`
30-
}
33+
type gpuVendors map[string]v2GPUs
3134

3235
type v2GPUAttributes types.Attributes
3336

@@ -66,37 +69,54 @@ func (sdl *v2ResourceGPU) UnmarshalYAML(node *yaml.Node) error {
6669
func (sdl *v2GPUAttributes) UnmarshalYAML(node *yaml.Node) error {
6770
var res types.Attributes
6871

69-
var vendor *gpuVendor
72+
vendors := make(gpuVendors)
7073

7174
for i := 0; i < len(node.Content); i += 2 {
7275
switch node.Content[i].Value {
7376
case "vendor":
74-
if err := node.Content[i+1].Decode(&vendor); err != nil {
77+
if err := node.Content[i+1].Decode(&vendors); err != nil {
7578
return err
7679
}
7780
default:
7881
return fmt.Errorf("sdl: unsupported attribute (%s) for GPU resource", node.Content[i].Value)
7982
}
8083
}
8184

82-
if vendor == nil {
83-
return fmt.Errorf("sdl: invalid GPU attributes. at least one vendor must be set")
85+
if len(vendors) == 0 {
86+
return ErrResourceGPUEmptyVendors
8487
}
8588

86-
res = make(types.Attributes, 0, len(vendor.Nvidia))
89+
resPrealloc := 0
8790

88-
for _, model := range vendor.Nvidia {
89-
res = append(res, types.Attribute{
90-
Key: fmt.Sprintf("vendor/nvidia/model/%s", model.String()),
91-
Value: "true",
92-
})
91+
for _, models := range vendors {
92+
if len(models) == 0 {
93+
resPrealloc += 1
94+
} else {
95+
resPrealloc += len(models)
96+
}
9397
}
9498

95-
if len(res) == 0 {
96-
res = append(res, types.Attribute{
97-
Key: "vendor/nvidia/model/*",
98-
Value: "true",
99-
})
99+
for vendor, models := range vendors {
100+
switch vendor {
101+
case "nvidia":
102+
case "amd":
103+
default:
104+
return fmt.Errorf("sdl: unsupported GPU vendor (%s)", vendor)
105+
}
106+
107+
for _, model := range models {
108+
res = append(res, types.Attribute{
109+
Key: fmt.Sprintf("vendor/%s/model/%s", vendor, model.String()),
110+
Value: "true",
111+
})
112+
}
113+
114+
if len(models) == 0 {
115+
res = append(res, types.Attribute{
116+
Key: fmt.Sprintf("vendor/%s/model/*", vendor),
117+
Value: "true",
118+
})
119+
}
100120
}
101121

102122
sort.Sort(res)

0 commit comments

Comments
 (0)