Add support for Kubernetes arbitrary node selector

bin/cml/runner/launch.js

@@ -135,6 +135,7 @@ const runCloud = async (opts) => {
       cloudStartupScript: startupScript,
       cloudAwsSecurityGroup: awsSecurityGroup,
       cloudAwsSubnet: awsSubnet,
+      cloudKubernetesNodeSelector: kubernetesNodeSelector,
       cloudImage: image,
       workdir
     } = opts;
@@ -172,6 +173,7 @@ const runCloud = async (opts) => {
       startupScript,
       awsSecurityGroup,
       awsSubnet,
+      kubernetesNodeSelector,
       image,
       dockerVolumes
     });
@@ -588,6 +590,20 @@ exports.options = kebabcaseKeys({
     description: 'Specifies the subnet to use within AWS',
     alias: 'cloud-aws-subnet-id'
   },
+  cloudKubernetesNodeSelector: {
+    type: 'string',
+    default: 'accelerator=inferred',
+    description:
+      'Specifies the node selector to use within Kubernetes. By default, the node selector is inferred from the GPU configuration',
+    coerce: (val) => {
+      const [key, value] = val.split(/=(.+)/);
+
+      return {
+        key,
+        value
+      };
+    }
+  },
   cloudImage: {
     type: 'string',
     description: 'Custom machine/container image',

src/terraform.js

@@ -115,7 +115,11 @@ const iterativeCmlRunnerTpl = (opts = {}) => {
           ...(opts.sshPrivate && { ssh_private: opts.sshPrivate }),
           ...(opts.startupScript && { startup_script: opts.startupScript }),
           ...(opts.token && { token: opts.token }),
-          ...(opts.type && { instance_type: opts.type })
+          ...(opts.type && { instance_type: opts.type }),
+          ...(opts.kubernetesNodeSelector && {
+            kubernetes_node_selector_key: opts.kubernetesNodeSelector.key,
+            kubernetes_node_selector_value: opts.kubernetesNodeSelector.value
+          })
         }
       }
     }