Port task instance types to runner

iterative/kubernetes/provider.go

@@ -9,6 +9,7 @@ import (
 	"log"
 	"net/http"
 	"os"
+	"regexp"
 	"strconv"
 	"time"
 
@@ -391,9 +392,13 @@ func getInstanceType(instanceType string, instanceGPU string) (map[string]map[st
 		"memory": {"amount": "512Gi"},
 	}
 
-	if val, ok := instanceTypes[instanceType+"+"+instanceGPU]; ok {
-		return val, nil
-	} else if val, ok := instanceTypes[instanceType]; ok && instanceGPU == "" {
+	size := instanceType
+
+	if instanceGPU != "" {
+		size += "+"+instanceGPU
+	}
+
+	if val, ok := instanceTypes[size]; ok {
 		return val, nil
 	} else if val, ok := instanceTypes[instanceType]; ok {
 		// Allow users to specify custom accelerator selectors.
@@ -408,6 +413,23 @@ func getInstanceType(instanceType string, instanceGPU string) (map[string]map[st
 		}, nil
 	}
 
+
+	if match := regexp.MustCompile(`^(\d+)-(\d+)(?:\+([^*]+)\*([1-9]\d*))?$`).FindStringSubmatch(size); match != nil {
+		return map[string]map[string]string{
+			"accelerator": {
+				"count": match[4],
+				"model": match[3],
+				"type":  "nvidia.com/gpu",
+			},
+			"cores":  {
+				"count": match[1],
+			},
+			"memory": {
+				"amount": match[2] + "M",
+			},
+		}, nil
+	}
+
 	return nil, fmt.Errorf("invalid instance type")
 }