Add support for Kubernetes arbitrary node selector

iterative/kubernetes/provider.go

@@ -51,7 +51,6 @@ func ResourceMachineCreate(ctx context.Context, d *terraform_schema.ResourceData
 	}
 
 	// Define the accelerator settings (i.e. GPU type, model, ...)
-	jobNodeSelector := map[string]string{}
 	jobAccelerator := instanceType["accelerator"]["model"]
 	jobGPUType := instanceType["accelerator"]["type"]
 	jobGPUCount := instanceType["accelerator"]["count"]
@@ -63,6 +62,27 @@ func ResourceMachineCreate(ctx context.Context, d *terraform_schema.ResourceData
 	if diskAmount := d.Get("instance_hdd_size").(int); diskAmount > 0 {
 		jobLimits[kubernetes_core.ResourceName("ephemeral-storage")] = kubernetes_resource.MustParse(strconv.Itoa(diskAmount) + "G")
 	}
+	if jobGPUCount > "0" && jobGPUType != "" {
+		jobLimits[kubernetes_core.ResourceName(jobGPUType)] = kubernetes_resource.MustParse(jobGPUCount)
+	}
+
+	// Get the node selector defined by the user
+	kubernetesNodeSelector := d.Get("kubernetes_node_selector").(map[string]interface{})
+
+	// Set the default key value if none is defined
+	if len(kubernetesNodeSelector) == 0 {
+		kubernetesNodeSelector["accelerator"] = "infer"
+	}
+
+	// Define the node selector
+	jobNodeSelector := map[string]string{}
+	for selector, value := range kubernetesNodeSelector {
+		if value.(string) != "infer" {
+			jobNodeSelector[selector] = value.(string)
+		} else if jobGPUCount > "0" && jobAccelerator != "" {
+			jobNodeSelector[selector] = jobAccelerator
+		}
+	}
 
 	// Use the default CML Docker image unless specified otherwise.
 	jobImageName := jobName
@@ -77,14 +97,6 @@ func ResourceMachineCreate(ctx context.Context, d *terraform_schema.ResourceData
 		return err
 	}
 
-	// If the resource requires GPU provisioning, determine how many GPUs and the kind of GPU it needs.
-	if jobGPUCount > "0" {
-		jobLimits[kubernetes_core.ResourceName(jobGPUType)] = kubernetes_resource.MustParse(jobGPUCount)
-		if jobAccelerator != "" {
-			jobNodeSelector = map[string]string{"accelerator": jobAccelerator}
-		}
-	}
-
 	// Lookup service account if set.
 	svcAccount, svcTokenAutomount, err := getServiceAccount(ctx, conn, namespace, d.Get("instance_permission_set").(string))
 	if err != nil {

iterative/resource_machine.go

@@ -134,6 +134,14 @@ func machineSchema() *map[string]*schema.Schema {
 			Optional: true,
 			Default:  "",
 		},
+		"kubernetes_node_selector": &schema.Schema{
+			Type:     schema.TypeMap,
+			ForceNew: true,
+			Optional: true,
+			Elem: &schema.Schema{
+				Type: schema.TypeString,
+			},
+		},
 		"metadata": &schema.Schema{
 			Type:     schema.TypeMap,
 			ForceNew: true,

iterative/resource_runner.go

@@ -175,6 +175,14 @@ func resourceRunner() *schema.Resource {
 				Optional: true,
 				Default:  "",
 			},
+			"kubernetes_node_selector": &schema.Schema{
+				Type:     schema.TypeMap,
+				ForceNew: true,
+				Optional: true,
+				Elem: &schema.Schema{
+					Type: schema.TypeString,
+				},
+			},
 			"metadata": &schema.Schema{
 				Type:     schema.TypeMap,
 				ForceNew: true,