Add kubernetes labels

metaflow/metaflow_config.py

@@ -248,6 +248,8 @@
 KUBERNETES_NODE_SELECTOR = from_conf("KUBERNETES_NODE_SELECTOR", "")
 KUBERNETES_TOLERATIONS = from_conf("KUBERNETES_TOLERATIONS", "")
 KUBERNETES_SECRETS = from_conf("KUBERNETES_SECRETS", "")
+# Default labels for kubernetes pods
+KUBERNETES_LABELS = from_conf("KUBERNETES_LABELS", "")
 # Default GPU vendor to use by K8S jobs created by Metaflow (supports nvidia, amd)
 KUBERNETES_GPU_VENDOR = from_conf("KUBERNETES_GPU_VENDOR", "nvidia")
 # Default container image for K8S

metaflow/plugins/argo/argo_workflows.py

@@ -894,11 +894,14 @@ def _container_templates(self):
                     # the field 'task-id' in 'parameters'
                     # .annotation("metaflow/task_id", ...)
                     .annotation("metaflow/attempt", retry_count)
+                    # Set labels
+                    .labels(resources.get("labels"))
                 )
                 # Set emptyDir volume for state management
                 .empty_dir_volume("out")
                 # Set node selectors
                 .node_selectors(resources.get("node_selector"))
+                # Set tolerations
                 .tolerations(resources.get("tolerations"))
                 # Set container
                 .container(

metaflow/plugins/kubernetes/kubernetes.py

@@ -9,6 +9,8 @@
 from metaflow.metaflow_config import (
     SERVICE_HEADERS,
     SERVICE_INTERNAL_URL,
+    CARD_AZUREROOT,
+    CARD_GSROOT,
     CARD_S3ROOT,
     DATASTORE_SYSROOT_S3,
     DATATOOLS_S3ROOT,
@@ -28,8 +30,8 @@
     BASH_SAVE_LOGS,
     bash_capture_logs,
     export_mflog_env_vars,
-    tail_logs,
     get_log_tailer,
+    tail_logs,
 )
 
 from .kubernetes_client import KubernetesClient
@@ -151,6 +153,7 @@ def create_job(
         run_time_limit=None,
         env=None,
         tolerations=None,
+        labels=None,
     ):
 
         if env is None:
@@ -184,6 +187,7 @@ def create_job(
                 retries=0,
                 step_name=step_name,
                 tolerations=tolerations,
+                labels=labels,
             )
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -3,7 +3,7 @@
 import time
 import traceback
 
-from metaflow import util, JSONTypeClass
+from metaflow import JSONTypeClass, util
 from metaflow._vendor import click
 from metaflow.exception import METAFLOW_EXIT_DISALLOW_RETRY, CommandException
 from metaflow.metadata.util import sync_local_metadata_from_datastore
@@ -91,6 +91,12 @@ def kubernetes():
     type=JSONTypeClass(),
     multiple=False,
 )
+@click.option(
+    "--labels",
+    multiple=True,
+    default=None,
+    help="Labels for Kubernetes pod.",
+)
 @click.pass_context
 def step(
     ctx,
@@ -110,6 +116,7 @@ def step(
     gpu_vendor=None,
     run_time_limit=None,
     tolerations=None,
+    labels=None,
     **kwargs
 ):
     def echo(msg, stream="stderr", job_id=None):
@@ -175,7 +182,13 @@ def echo(msg, stream="stderr", job_id=None):
     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)
+    node_selector = KubernetesDecorator.parse_kube_keyvalue_list(node_selector)
+
+    # `labels` is a tuple of strings or a tuple with a single comma separated string
+    # convert it to a dict
+    labels = KubernetesDecorator.clean_kube_labels(
+        KubernetesDecorator.parse_kube_keyvalue_list(labels, False)
+    )
 
     def _sync_metadata():
         if ctx.obj.metadata.TYPE == "local":
@@ -218,6 +231,7 @@ def _sync_metadata():
                 run_time_limit=run_time_limit,
                 env=env,
                 tolerations=tolerations,
+                labels=labels,
             )
     except Exception as e:
         traceback.print_exc(chain=False)

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -1,8 +1,11 @@
+import hashlib
 import json
 import os
 import platform
+import re
 import requests
 import sys
+from typing import Dict, List, Optional, Union
 
 from metaflow.decorators import StepDecorator
 from metaflow.exception import MetaflowException
@@ -13,18 +16,18 @@
     KUBERNETES_CONTAINER_IMAGE,
     KUBERNETES_CONTAINER_REGISTRY,
     KUBERNETES_GPU_VENDOR,
+    KUBERNETES_LABELS,
     KUBERNETES_NAMESPACE,
     KUBERNETES_NODE_SELECTOR,
-    KUBERNETES_TOLERATIONS,
-    KUBERNETES_SERVICE_ACCOUNT,
     KUBERNETES_SECRETS,
+    KUBERNETES_SERVICE_ACCOUNT,
+    KUBERNETES_TOLERATIONS,
 )
 from metaflow.plugins.resources_decorator import ResourcesDecorator
 from metaflow.plugins.timeout_decorator import get_run_time_limit_for_task
 from metaflow.sidecar import Sidecar
 
 from ..aws.aws_utils import get_docker_registry
-
 from .kubernetes import KubernetesException
 
 try:
@@ -65,6 +68,8 @@ class KubernetesDecorator(StepDecorator):
         in Metaflow configuration.
     tolerations : List[str], default: METAFLOW_KUBERNETES_TOLERATIONS
         Kubernetes tolerations to use when launching pod in Kubernetes.
+    labels : Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
+        Kubernetes labels to use when launching pod in Kubernetes.
     """
 
     name = "kubernetes"
@@ -76,6 +81,7 @@ class KubernetesDecorator(StepDecorator):
         "service_account": None,
         "secrets": None,  # e.g., mysecret
         "node_selector": None,  # e.g., kubernetes.io/os=linux
+        "labels": None,  # e.g., my_label=my_value
         "namespace": None,
         "gpu": None,  # value of 0 implies that the scheduled node should not have GPUs
         "gpu_vendor": None,
@@ -99,9 +105,17 @@ def __init__(self, attributes=None, statically_defined=False):
             self.attributes["node_selector"] = KUBERNETES_NODE_SELECTOR
         if not self.attributes["tolerations"] and KUBERNETES_TOLERATIONS:
             self.attributes["tolerations"] = json.loads(KUBERNETES_TOLERATIONS)
+        if not self.attributes["labels"] and KUBERNETES_LABELS:
+            self.attributes["labels"] = KUBERNETES_LABELS
+
+        if isinstance(self.attributes["labels"], str):
+            self.attributes["labels"] = self.parse_kube_keyvalue_list(
+                self.attributes["labels"].split(","), False
+            )
+        self.attributes["labels"] = self.clean_kube_labels(self.attributes["labels"])
 
         if isinstance(self.attributes["node_selector"], str):
-            self.attributes["node_selector"] = self.parse_node_selector(
+            self.attributes["node_selector"] = self.parse_kube_keyvalue_list(
                 self.attributes["node_selector"].split(",")
             )
 
@@ -280,10 +294,11 @@ 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 in {"node_selector", "labels"} and v:
+                    cli_args.command_options[k] = [
+                        "=".join([key, str(val)]) if val else key
+                        for key, val in v.items()
+                    ]
                 elif k == "tolerations":
                     cli_args.command_options[k] = json.dumps(v)
                 else:
@@ -399,14 +414,61 @@ def _save_package_once(cls, flow_datastore, package):
                 [package.blob], len_hint=1
             )[0]
 
+    @classmethod
+    def _parse_decorator_spec(cls, deco_spec: str):
+        if not deco_spec:
+            return cls()
+
+        valid_options = "|".join(cls.defaults.keys())
+        deco_spec_parts = []
+        for part in re.split(f""",(?=[\s\w]+[{valid_options}]=)""", deco_spec):
+            name, val = part.split("=", 1)
+            if name in {"labels", "node_selector"}:
+                try:
+                    json.loads(val.strip().replace('\\"', '"'))
+                except json.JSONDecodeError:
+                    both = name == "node_selector"
+                    val = json.dumps(
+                        cls.parse_kube_keyvalue_list(val.split(","), both),
+                        separators=(",", ":"),
+                    )
+            deco_spec_parts.append("=".join([name, val]))
+        deco_spec_parsed = ",".join(deco_spec_parts)
+        return super()._parse_decorator_spec(deco_spec_parsed)
+
     @staticmethod
-    def parse_node_selector(node_selector: list):
+    def parse_kube_keyvalue_list(items: List[str], requires_both: bool = True):
         try:
-            return {
-                str(k.split("=", 1)[0]): str(k.split("=", 1)[1])
-                for k in node_selector or []
-            }
+            ret = {}
+            for item_str in items:
+                item = item_str.split("=", 1)
+                if requires_both:
+                    item[1]  # raise IndexError
+                ret[str(item[0])] = str(item[1]) if len(item) > 1 else None
+            return ret
         except (AttributeError, IndexError):
-            raise KubernetesException(
-                "Unable to parse node_selector: %s" % node_selector
-            )
+            raise KubernetesException("Unable to parse kubernetes list: %s" % items)
+
+    @staticmethod
+    def clean_kube_labels(
+        labels: Optional[Dict[str, Optional[str]]],
+        max_len: int = 63,
+        regex_sub: str = r"^[^a-z0-9A-Z]*|[^a-zA-Z0-9_\-\.]|[^a-z0-9A-Z]*$",
+    ):
+        """Inspired by apache airflow label cleaner."""
+
+        def clean_label(s: Optional[str]):
+            if not s:
+                # allow empty label
+                return s
+            s_clean = re.sub(regex_sub, "", s)
+            if len(s_clean) > max_len or s != s_clean:
+                s_hash = hashlib.blake2b(s.encode(), digest_size=5).hexdigest()
+                s_clean = f"{s_clean[: max_len - len(s_hash) - 1]}-{s_hash}"
+            # Could recursively call this function but a random hash label is most likely a mistake
+            assert not s_clean.startswith(
+                "-"
+            ), f"Label, {s}, contains no valid characters"
+            return s_clean
+
+        return {k: clean_label(v) for k, v in labels.items()} if labels else labels

test/unit/test_kubernetes_decorator.py

@@ -0,0 +1,33 @@
+import pytest
+
+from metaflow.plugins.kubernetes.kubernetes_decorator import KubernetesDecorator
+
+
+@pytest.mark.parametrize(
+    "labels,expected",
+    [
+        (None, None),
+        ({"label": "value"}, {"label": "value"}),
+        ({"label1": "val1", "label2": "val2"}, {"label1": "val1", "label2": "val2"}),
+        ({"label1": "val1", "label2": None}, {"label1": "val1", "label2": None}),
+        ({"label": "test()"}, {"label": "test-1ba883287c"}),
+        (
+            {
+                "label": "reallyreallyreallyreallyreallyreallyreallyreallyreallyreallyreallylong"
+            },
+            {
+                "label": "reallyreallyreallyreallyreallyreallyreallyreallyreal-768ebdbe62"
+            },
+        ),
+    ],
+)
+def test_kubernetes_decorator_clean_kube_labels(labels, expected):
+    cleaned_labels = KubernetesDecorator.clean_kube_labels(labels)
+    assert cleaned_labels == expected
+
+
+def test_kubernetes_decorator_clean_kube_labels_fail():
+    """Fail if label contains no valid characters"""
+    labels = {"labels": "(){}??"}
+    with pytest.raises(AssertionError):
+        KubernetesDecorator.clean_kube_labels(labels)