Allow custom machine+gpu*count on runner cloud=gcp (#380)

iterative · Feb 2, 2022 · 610c449 · 610c449
1 parent 952f1a9
commit 610c449
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/iterative/gcp/provider.go b/iterative/gcp/provider.go
@@ -441,7 +441,18 @@ func getInstanceType(instanceType string, instanceGPU string) (map[string]map[st
 		},
 	}
 
-	if val, ok := instanceTypes[instanceType+"+"+instanceGPU]; ok {
+	match := regexp.MustCompile(`^([^+]+)\+([^*]+)\*([1-9]\d*)?$`).FindStringSubmatch(instanceType)
+	if match != nil {
+		return map[string]map[string]string{
+			"accelerator": {
+				"count": match[3],
+				"type":  match[2],
+			},
+			"machine": {
+				"type": match[1],
+			},
+		}, nil
+	} else if val, ok := instanceTypes[instanceType+"+"+instanceGPU]; ok {
 		return val, nil
 	} else if val, ok := instanceTypes[instanceType]; ok && instanceGPU == "" {
 		return val, nil