Use region field as comma-delimited node selector for k8s (#657)

* Use `region` field as comma-delimited node selector for k8s

* UPdate docs for k8s region/nodeSelectoR

* Apply suggestions from code review

Co-authored-by: Helio Machado <0x2b3bfa0+git@googlemail.com>

Author: sjawhar

docs/resources/task.md

@@ -54,7 +54,7 @@ resource "iterative_task" "example" {
 
 ### Optional
 
-- `region` - (Optional) [Cloud region/zone](#cloud-region) to run the task on.
+- `region` - (Optional) [Cloud region/zone](#cloud-region) to run the task on, or node selector labels for Kubernetes.
 - `machine` - (Optional) See [Machine Types](#machine-type) below.
 - `disk_size` - (Optional) Size of the ephemeral machine storage in GB. `-1`: automatic based on `image`.
 - `spot` - (Optional) Spot instance price. `-1`: disabled, `0`: automatic price, any other positive number: maximum bidding price in USD per hour (above which the instance is terminated until the price drops).
@@ -257,7 +257,25 @@ In addition to generic regions, it's possible to specify any cloud region suppor
 
 ### Kubernetes
 
-The `region` attribute is ignored.
+For Kubernetes, the `region` attribute can be used to set the [`nodeSelector`](https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/#nodeselector) field of the job specification. The value of `region` should be a comma-delimited list of `key=value` nodel label pairs.
+
+For example:
+```tf
+resource "iterative_task" "example" {
+  cloud  = "k8s"
+  region = "foo=bar,goo=baz"
+}
+```
+
+This will create a pod with a `nodeSelector` value like the following:
+```yaml
+apiVersion: V1
+kind: Pod
+spec:
+  nodeSelector:
+    foo: bar
+    goo: baz
+```
 
 ## Permission Set
 
@@ -285,10 +303,6 @@ A comma-separated list of [user-assigned identity](https://docs.microsoft.com/en
 
 ### Kubernetes
 
-#### Region attribute
-
-Setting the `region` attribute results in undefined behaviour.
-
 #### Directory storage
 
 Unlike public cloud providers, Kubernetes does not offer any portable way of persisting and sharing storage between pods. When specified, the `storage.workdir` attribute will create a `PersistentVolumeClaim` of the default `StorageClass`, with the same lifecycle as the task and the specified `disk_size`.

task/k8s/resources/resource_job.go

@@ -36,18 +36,26 @@ func NewJob(client *client.Client, identifier common.Identifier, persistentVolum
 	j.Dependencies.PermissionSet = permissionSet
 	j.Attributes.Task = task
 	j.Attributes.Parallelism = task.Parallelism
+	j.Attributes.NodeSelector = map[string]string{}
+	for _, selector := range strings.Split(string(client.Cloud.Region), ",") {
+		key, value, is_found := strings.Cut(selector, "=")
+		if is_found && len(value) > 0 {
+			j.Attributes.NodeSelector[key] = value
+		}
+	}
 	return j
 }
 
 type Job struct {
 	Client     *client.Client
 	Identifier string
 	Attributes struct {
-		Task        common.Task
-		Parallelism uint16
-		Addresses   []net.IP
-		Status      common.Status
-		Events      []common.Event
+		Task         common.Task
+		Parallelism  uint16
+		NodeSelector map[string]string
+		Addresses    []net.IP
+		Status       common.Status
+		Events       []common.Event
 	}
 	Dependencies struct {
 		*PersistentVolumeClaim
@@ -90,8 +98,12 @@ func (j *Job) Create(ctx context.Context) error {
 		return common.NotFoundError
 	}
 
-	// Define the accelerator settings (i.e. GPU type, model, ...)
 	jobNodeSelector := map[string]string{}
+	for selector, value := range j.Attributes.NodeSelector {
+		jobNodeSelector[selector] = value
+	}
+
+	// Define the accelerator settings (i.e. GPU type, model, ...)
 	jobAccelerator := match[3]
 	jobGPUType := "nvidia.com/gpu"
 	jobGPUCount := match[4]
@@ -108,7 +120,7 @@ func (j *Job) Create(ctx context.Context) error {
 	if jobGPUCount > "0" {
 		jobLimits[kubernetes_core.ResourceName(jobGPUType)] = kubernetes_resource.MustParse(jobGPUCount)
 		if jobAccelerator != "" {
-			jobNodeSelector = map[string]string{"accelerator": jobAccelerator}
+			jobNodeSelector["accelerator"] = jobAccelerator
 		}
 	}