Skip to content

Commit ce4f3be

Browse files
authored
Port task instance types to runner (#742)
* Port `task` instance types to `runner` * Remove stray parsing * Import regexo * Fix structure
1 parent c95cd28 commit ce4f3be

File tree

1 file changed

+25
-3
lines changed

1 file changed

+25
-3
lines changed

iterative/kubernetes/provider.go

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"log"
1010
"net/http"
1111
"os"
12+
"regexp"
1213
"strconv"
1314
"time"
1415

@@ -391,9 +392,13 @@ func getInstanceType(instanceType string, instanceGPU string) (map[string]map[st
391392
"memory": {"amount": "512Gi"},
392393
}
393394

394-
if val, ok := instanceTypes[instanceType+"+"+instanceGPU]; ok {
395-
return val, nil
396-
} else if val, ok := instanceTypes[instanceType]; ok && instanceGPU == "" {
395+
size := instanceType
396+
397+
if instanceGPU != "" {
398+
size += "+"+instanceGPU
399+
}
400+
401+
if val, ok := instanceTypes[size]; ok {
397402
return val, nil
398403
} else if val, ok := instanceTypes[instanceType]; ok {
399404
// Allow users to specify custom accelerator selectors.
@@ -408,6 +413,23 @@ func getInstanceType(instanceType string, instanceGPU string) (map[string]map[st
408413
}, nil
409414
}
410415

416+
417+
if match := regexp.MustCompile(`^(\d+)-(\d+)(?:\+([^*]+)\*([1-9]\d*))?$`).FindStringSubmatch(size); match != nil {
418+
return map[string]map[string]string{
419+
"accelerator": {
420+
"count": match[4],
421+
"model": match[3],
422+
"type": "nvidia.com/gpu",
423+
},
424+
"cores": {
425+
"count": match[1],
426+
},
427+
"memory": {
428+
"amount": match[2] + "M",
429+
},
430+
}, nil
431+
}
432+
411433
return nil, fmt.Errorf("invalid instance type")
412434
}
413435

0 commit comments

Comments
 (0)