Add support for Kubernetes tolerations

metaflow/metaflow_config.py

@@ -242,6 +242,7 @@
 KUBERNETES_SERVICE_ACCOUNT = from_conf("KUBERNETES_SERVICE_ACCOUNT")
 # Default node selectors to use by K8S jobs created by Metaflow - foo=bar,baz=bab
 KUBERNETES_NODE_SELECTOR = from_conf("KUBERNETES_NODE_SELECTOR", "")
+KUBERNETES_TOLERATIONS = from_conf("KUBERNETES_TOLERATIONS", "")
 KUBERNETES_SECRETS = from_conf("KUBERNETES_SECRETS", "")
 # Default GPU vendor to use by K8S jobs created by Metaflow (supports nvidia, amd)
 KUBERNETES_GPU_VENDOR = from_conf("KUBERNETES_GPU_VENDOR", "nvidia")

metaflow/plugins/argo/argo_workflows.py

@@ -140,7 +140,9 @@ def _sanitize(name):
         return name.replace("_", "-")
 
     @classmethod
-    def trigger(cls, name, parameters={}):
+    def trigger(cls, name, parameters=None):
+        if parameters is None:
+            parameters = {}
         try:
             workflow_template = ArgoClient(
                 namespace=KUBERNETES_NAMESPACE
@@ -378,14 +380,18 @@ def _compile(self):
 
     # Visit every node and yield the uber DAGTemplate(s).
     def _dag_templates(self):
-        def _visit(node, exit_node=None, templates=[], dag_tasks=[]):
+        def _visit(node, exit_node=None, templates=None, dag_tasks=None):
             # Every for-each node results in a separate subDAG and an equivalent
             # DAGTemplate rooted at the child of the for-each node. Each DAGTemplate
             # has a unique name - the top-level DAGTemplate is named as the name of
             # the flow and the subDAG DAGTemplates are named after the (only) descendant
             # of the for-each node.
 
             # Emit if we have reached the end of the sub workflow
+            if dag_tasks is None:
+                dag_tasks = []
+            if templates is None:
+                templates = []
             if exit_node is not None and exit_node is node.name:
                 return templates, dag_tasks
 
@@ -859,7 +865,8 @@ def _container_templates(self):
                 .retry_strategy(
                     times=total_retries,
                     minutes_between_retries=minutes_between_retries,
-                ).metadata(
+                )
+                .metadata(
                     ObjectMeta().annotation("metaflow/step_name", node.name)
                     # Unfortunately, we can't set the task_id since it is generated
                     # inside the pod. However, it can be inferred from the annotation
@@ -871,19 +878,8 @@ def _container_templates(self):
                 # Set emptyDir volume for state management
                 .empty_dir_volume("out")
                 # Set node selectors
-                .node_selectors(
-                    {
-                        str(k.split("=", 1)[0]): str(k.split("=", 1)[1])
-                        for k in (
-                            resources.get("node_selector")
-                            or (
-                                KUBERNETES_NODE_SELECTOR.split(",")
-                                if KUBERNETES_NODE_SELECTOR
-                                else []
-                            )
-                        )
-                    }
-                )
+                .node_selectors(resources.get("node_selector"))
+                .tolerations(resources.get("tolerations"))
                 # Set container
                 .container(
                     # TODO: Unify the logic with kubernetes.py
@@ -1234,8 +1230,13 @@ def empty_dir_volume(self, name):
 
     def node_selectors(self, node_selectors):
         if "nodeSelector" not in self.payload:
-            self.payload["labels"] = {}
-        self.payload["nodeSelector"].update(node_selectors)
+            self.payload["nodeSelector"] = {}
+        if node_selectors:
+            self.payload["nodeSelector"].update(node_selectors)
+        return self
+
+    def tolerations(self, tolerations):
+        self.payload["tolerations"] = tolerations
         return self
 
     def to_json(self):

metaflow/plugins/kubernetes/kubernetes.py

@@ -147,9 +147,13 @@ def create_job(
         disk=None,
         memory=None,
         run_time_limit=None,
-        env={},
+        env=None,
+        tolerations=None,
     ):
 
+        if env is None:
+            env = {}
+
         job = (
             KubernetesClient()
             .job(
@@ -177,6 +181,7 @@ def create_job(
                 # Retries are handled by Metaflow runtime
                 retries=0,
                 step_name=step_name,
+                tolerations=tolerations,
             )
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -3,14 +3,15 @@
 import time
 import traceback
 
-from metaflow import util
+from metaflow import util, JSONTypeClass
 from metaflow._vendor import click
 from metaflow.exception import METAFLOW_EXIT_DISALLOW_RETRY, CommandException
 from metaflow.metadata.util import sync_local_metadata_from_datastore
 from metaflow.metaflow_config import DATASTORE_LOCAL_DIR
 from metaflow.mflog import TASK_LOG_SOURCE
 
 from .kubernetes import Kubernetes, KubernetesKilledException
+from .kubernetes_decorator import KubernetesDecorator
 
 
 @click.group()
@@ -84,6 +85,12 @@ def kubernetes():
     default=5 * 24 * 60 * 60,  # Default is set to 5 days
     help="Run time limit in seconds for Kubernetes pod.",
 )
+@click.option(
+    "--tolerations",
+    default=None,
+    type=JSONTypeClass(),
+    multiple=False,
+)
 @click.pass_context
 def step(
     ctx,
@@ -102,6 +109,7 @@ def step(
     gpu=None,
     gpu_vendor=None,
     run_time_limit=None,
+    tolerations=None,
     **kwargs
 ):
     def echo(msg, stream="stderr", job_id=None):
@@ -166,6 +174,9 @@ def echo(msg, stream="stderr", job_id=None):
     stdout_location = ds.get_log_location(TASK_LOG_SOURCE, "stdout")
     stderr_location = ds.get_log_location(TASK_LOG_SOURCE, "stderr")
 
+    # `node_selector` is a tuple of strings, convert it to a dictionary
+    node_selector = KubernetesDecorator.parse_node_selector(node_selector)
+
     def _sync_metadata():
         if ctx.obj.metadata.TYPE == "local":
             sync_local_metadata_from_datastore(
@@ -206,6 +217,7 @@ def _sync_metadata():
                 gpu_vendor=gpu_vendor,
                 run_time_limit=run_time_limit,
                 env=env,
+                tolerations=tolerations,
             )
     except Exception as e:
         traceback.print_exc(chain=False)

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -1,10 +1,8 @@
+import json
 import os
 import platform
 import sys
 
-import requests
-
-from metaflow import util
 from metaflow.decorators import StepDecorator
 from metaflow.exception import MetaflowException
 from metaflow.metadata import MetaDatum
@@ -16,6 +14,7 @@
     KUBERNETES_GPU_VENDOR,
     KUBERNETES_NAMESPACE,
     KUBERNETES_NODE_SELECTOR,
+    KUBERNETES_TOLERATIONS,
     KUBERNETES_SERVICE_ACCOUNT,
     KUBERNETES_SECRETS,
 )
@@ -24,6 +23,7 @@
 from metaflow.sidecar import Sidecar
 
 from ..aws.aws_utils import get_docker_registry
+
 from .kubernetes import KubernetesException
 
 try:
@@ -65,6 +65,10 @@ class KubernetesDecorator(StepDecorator):
         Kubernetes secrets to use when launching pod in Kubernetes. These
         secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
         in Metaflow configuration.
+    tolerations : List[str]
+        Kubernetes tolerations to use when launching pod in Kubernetes. If
+        not specified, the value of `METAFLOW_KUBERNETES_TOLERATIONS` is used
+        from Metaflow configuration.
     """
 
     name = "kubernetes"
@@ -79,6 +83,8 @@ class KubernetesDecorator(StepDecorator):
         "namespace": None,
         "gpu": None,  # value of 0 implies that the scheduled node should not have GPUs
         "gpu_vendor": None,
+        "tolerations": None,  # e.g., [{"key": "arch", "operator": "Equal", "value": "amd"},
+        #                              {"key": "foo", "operator": "Equal", "value": "bar"}]
     }
     package_url = None
     package_sha = None
@@ -93,9 +99,39 @@ def __init__(self, attributes=None, statically_defined=False):
             self.attributes["service_account"] = KUBERNETES_SERVICE_ACCOUNT
         if not self.attributes["gpu_vendor"]:
             self.attributes["gpu_vendor"] = KUBERNETES_GPU_VENDOR
+        if not self.attributes["node_selector"] and KUBERNETES_NODE_SELECTOR:
+            self.attributes["node_selector"] = KUBERNETES_NODE_SELECTOR
+        if not self.attributes["tolerations"] and KUBERNETES_TOLERATIONS:
+            self.attributes["tolerations"] = json.loads(KUBERNETES_TOLERATIONS)
+
+        if isinstance(self.attributes["node_selector"], str):
+            self.attributes["node_selector"] = self.parse_node_selector(
+                self.attributes["node_selector"].split(",")
+            )
 
-        # TODO: Handle node_selector in a better manner. Currently it is special
-        #       cased in kubernetes_client.py
+        if self.attributes["tolerations"]:
+            try:
+                from kubernetes.client import V1Toleration
+
+                for toleration in self.attributes["tolerations"]:
+                    try:
+                        invalid_keys = [
+                            k
+                            for k in toleration.keys()
+                            if k not in V1Toleration.attribute_map.keys()
+                        ]
+                        if len(invalid_keys) > 0:
+                            raise KubernetesException(
+                                "Tolerations parameter contains invalid keys: %s"
+                                % invalid_keys
+                            )
+                    except AttributeError:
+                        raise KubernetesException(
+                            "Unable to parse tolerations: %s"
+                            % self.attributes["tolerations"]
+                        )
+            except (NameError, ImportError):
+                pass
 
         # If no docker image is explicitly specified, impute a default image.
         if not self.attributes["image"]:
@@ -248,6 +284,12 @@ def runtime_step_cli(
             for k, v in self.attributes.items():
                 if k == "namespace":
                     cli_args.command_options["k8s_namespace"] = v
+                elif k == "node_selector" and v:
+                    cli_args.command_options[k] = ",".join(
+                        ["=".join([key, str(val)]) for key, val in v.items()]
+                    )
+                elif k == "tolerations":
+                    cli_args.command_options[k] = json.dumps(v)
                 else:
                     cli_args.command_options[k] = v
             cli_args.command_options["run-time-limit"] = self.run_time_limit
@@ -340,3 +382,15 @@ def _save_package_once(cls, flow_datastore, package):
             cls.package_url, cls.package_sha = flow_datastore.save_data(
                 [package.blob], len_hint=1
             )[0]
+
+    @staticmethod
+    def parse_node_selector(node_selector: list):
+        try:
+            return {
+                str(k.split("=", 1)[0]): str(k.split("=", 1)[1])
+                for k in node_selector or []
+            }
+        except (AttributeError, IndexError):
+            raise KubernetesException(
+                "Unable to parse node_selector: %s" % node_selector
+            )

metaflow/plugins/kubernetes/kubernetes_job.py

@@ -1,12 +1,10 @@
 import json
 import math
-import os
 import random
-import sys
 import time
 
 from metaflow.exception import MetaflowException
-from metaflow.metaflow_config import KUBERNETES_NODE_SELECTOR, KUBERNETES_SECRETS
+from metaflow.metaflow_config import KUBERNETES_SECRETS
 
 CLIENT_REFRESH_INTERVAL_SECONDS = 300
 
@@ -71,6 +69,7 @@ def create(self):
         # Note: This implementation ensures that there is only one unique Pod
         # (unique UID) per Metaflow task attempt.
         client = self._client.get()
+
         self._job = client.V1Job(
             api_version="batch/v1",
             kind="Job",
@@ -163,17 +162,7 @@ def create(self):
                                 ),
                             )
                         ],
-                        node_selector={
-                            str(k.split("=", 1)[0]): str(k.split("=", 1)[1])
-                            for k in (
-                                self._kwargs.get("node_selector")
-                                or (
-                                    KUBERNETES_NODE_SELECTOR.split(",")
-                                    if KUBERNETES_NODE_SELECTOR
-                                    else []
-                                )
-                            )
-                        },
+                        node_selector=self._kwargs.get("node_selector"),
                         # TODO (savin): Support image_pull_secrets
                         # image_pull_secrets=?,
                         # TODO (savin): Support preemption policies
@@ -188,8 +177,10 @@ def create(self):
                         service_account_name=self._kwargs["service_account"],
                         # Terminate the container immediately on SIGTERM
                         termination_grace_period_seconds=0,
-                        # TODO (savin): Enable tolerations for GPU scheduling.
-                        # tolerations=?,
+                        tolerations=[
+                            client.V1Toleration(**toleration)
+                            for toleration in self._kwargs.get("tolerations") or []
+                        ],
                         # volumes=?,
                         # TODO (savin): Set termination_message_policy
                     ),