Add support for Kubernetes arbitrary node selector

[ludelafo]

Hello,

We have an on-premises Kubernetes cluster with multiple nodes. Some of these nodes have dedicated GPUs.

We would like to use CML to train our models on those specific nodes. Kubernetes supports many ways to specify on which node a pod should be run:

• Assigning Pods to Nodes
• Assigning Pods to Nodes #nodeSelector
• Assigning Pods to Nodes #Affinity and anti-affinity

Examples

nodeSelector

Assigning a pod to a node with nodeSelector is as simple as labeling a node with a certain tag and add a nodeSelector to the pod configuration. Example inspired by the Assign Pods to Nodes [with nodeSelector] documentation:

# Label a node with `disktype=ssd`
kubectl label nodes <your-node-name> disktype=ssd

# pod-nginx.yaml (note the `nodeSelector`)
apiVersion: v1
kind: Pod
metadata:
  name: nginx
spec:
  containers:
    - name: nginx
      image: nginx
      imagePullPolicy: IfNotPresent
  nodeSelector:
    disktype: ssd

# Apply the Pod configuration
kubectl apply -f pod-nginx.yaml

The pod will be created on the node with the label disktype=ssd.

Bonus: Assign a pod to a node by its name

You can also assign a pod to a node using the node's name. Example inspired by the Assign Pods to Nodes [with nodeSelector] #Create a pod that gets scheduled to specific node documentation:

# pod-nginx.yaml (note the `nodeName`)
apiVersion: v1
kind: Pod
metadata:
  name: nginx
spec:
  nodeName: the-node-name
  containers:
    - name: nginx
      image: nginx
      imagePullPolicy: IfNotPresent

Affinity and anti-affinity

Another way to assign a pod to a node is using the affinity and anti-affinity feature. Example inspired by the Assigning Pods to Nodes #Node affinity documentation:

# pod-nginx.yaml (note the `affinity`)
apiVersion: v1
kind: Pod
metadata:
  name: nginx
spec:
  affinity:
    nodeAffinity:
      requiredDuringSchedulingIgnoredDuringExecution:
        nodeSelectorTerms:
          - matchExpressions:
              - key: disktype
                operator: In
                values:
                  - ssd
  containers:
    - name: nginx
      image: nginx
      imagePullPolicy: IfNotPresent

The requiredDuringSchedulingIgnoredDuringExecution means that the pod will be scheduled on a node that has the label disktype=ssd. If no node has the label, the pod will not be scheduled.

The preferredDuringSchedulingIgnoredDuringExecution allows to specify a preference for a node that has the label disktype=ssd. If no node has the label, the pod will be scheduled on any node.

The affinity and anti-affinity feature is more powerful than the nodeSelector feature. It allows to specify more complex rules. For example, you can specify that a pod should be scheduled on a node that has the label disktype=ssd and the label gputype=k80.

CML and Kubernetes abritary node selector

As discussed with @0x2b3bfa0 on Discord, the only node selector supported by CML at the moment is the accelerator node selector (https://github.com/iterative/terraform-provider-iterative/blob/ce4f3bec2300b3a15d615f3456575794f829f72b/iterative/kubernetes/provider.go#L84).

Of my understanding, the accelerator node selector is set by Cloud providers when you use a GPU instance. We could set those accelerator labels on our nodes so that CML could then use when using the --cloud-gpu flag. It could work but it seems a bit hacky. Allowing a --cloud-k8s-node-selector flag would be more flexible as any labels could be set and used.

It seems to me that the implementation of the nodeSelector feature would be rather simple (I have not thought about the implementation for the affinity feature yet as it seems more complex):

• Add the necessary code to the Iterative Terraform Provider codebase (https://github.com/iterative/terraform-provider-iterative/blob/ce4f3bec2300b3a15d615f3456575794f829f72b/iterative/kubernetes/provider.go#L31) to set the nodeSelector field for the Kubernetes pod.
• Add the necessary code to CML to support the --cloud-k8s-node-selector flag (

  cml/src/terraform.js

  Line 84 in 1be24ed

  const iterativeCmlRunnerTpl = (opts = {}) => {

  ,

  cml/bin/cml/runner/launch.js

  Line 113 in 1be24ed

  const runTerraform = async (opts) => {

  ,

  cml/bin/cml/runner/launch.js

  Line 457 in 1be24ed

  exports.options = kebabcaseKeys({

  ) to set up the Iterative Terraform Provider configuration.

I would like to know if you are interested in this and if you have any feedback on this proposal.

I would be willing to add support for this feature in CML. Any resources to start implementing it (your contribution guide, your usual workflows, etc.) would be helpful to get started. I'm available on Discord if you want to discuss this further with the same username.

[0x2b3bfa0]

Hopefully closed with the upcoming CML release, v0.19.0.

Thank you very much, @ludelafo! And... cough it would be awesome if you could test it on your side and tell us houw it goes. 😅

[ludelafo]

Great news! Thanks!

I'll test the new release in the following days and report back to you.

[ludelafo]

Seems to be working well from my testing, thank you again!

[dacbd]

Awesome, thanks for your patience.