feature: Add configurable labels for Argo Workflows

metaflow/metaflow_config.py

@@ -442,6 +442,7 @@
 KUBERNETES_CONDA_ARCH = from_conf("KUBERNETES_CONDA_ARCH")
 ARGO_WORKFLOWS_KUBERNETES_SECRETS = from_conf("ARGO_WORKFLOWS_KUBERNETES_SECRETS", "")
 ARGO_WORKFLOWS_ENV_VARS_TO_SKIP = from_conf("ARGO_WORKFLOWS_ENV_VARS_TO_SKIP", "")
+ARGO_WORKFLOWS_LABELS = from_conf("ARGO_WORKFLOWS_LABELS", "")
 
 KUBERNETES_JOBSET_GROUP = from_conf("KUBERNETES_JOBSET_GROUP", "jobset.x-k8s.io")
 KUBERNETES_JOBSET_VERSION = from_conf("KUBERNETES_JOBSET_VERSION", "v1alpha2")

metaflow/plugins/argo/argo_workflows.py

@@ -27,6 +27,7 @@
     ARGO_WORKFLOWS_CAPTURE_ERROR_SCRIPT,
     ARGO_WORKFLOWS_ENV_VARS_TO_SKIP,
     ARGO_WORKFLOWS_KUBERNETES_SECRETS,
+    ARGO_WORKFLOWS_LABELS,
     ARGO_WORKFLOWS_UI_URL,
     AWS_SECRETS_MANAGER_DEFAULT_REGION,
     AZURE_KEY_VAULT_PREFIX,
@@ -54,7 +55,11 @@
 from metaflow.metaflow_config_funcs import config_values
 from metaflow.mflog import BASH_SAVE_LOGS, bash_capture_logs, export_mflog_env_vars
 from metaflow.parameters import deploy_time_eval
-from metaflow.plugins.kubernetes.kube_utils import qos_requests_and_limits
+from metaflow.plugins.kubernetes.kube_utils import (
+    qos_requests_and_limits,
+    parse_kube_keyvalue_list,
+    validate_kube_labels,
+)
 
 from metaflow.plugins.kubernetes.kubernetes_jobsets import KubernetesArgoJobSet
 from metaflow.unbounded_foreach import UBF_CONTROL, UBF_TASK
@@ -188,7 +193,10 @@ def __init__(
         self.triggers, self.trigger_options = self._process_triggers()
         self._schedule, self._timezone = self._get_schedule()
 
+        # _workflow_labels (unlike _base_labels) includes user-supplied ARGO_WORKFLOWS_LABELS
+        # and must stay scoped to the WorkflowTemplate/Workflow level, not per-task resources.
         self._base_labels = self._base_kubernetes_labels()
+        self._workflow_labels = self._base_argo_labels()
         self._base_annotations = self._base_kubernetes_annotations()
         self._workflow_template = self._compile_workflow_template()
         self._sensor = self._compile_sensor()
@@ -403,12 +411,35 @@ def trigger(cls, name, parameters=None):
 
     def _base_kubernetes_labels(self):
         """
-        Get shared Kubernetes labels for Argo resources.
+        Get shared Kubernetes labels for all resources.
         """
-        # TODO: Add configuration through an environment variable or Metaflow config in the future if required.
-        labels = {"app.kubernetes.io/part-of": "metaflow"}
+        return {"app.kubernetes.io/part-of": "metaflow"}
+
+    def _custom_argo_labels(self):
+        """
+        Parse and validate custom labels from the METAFLOW_ARGO_WORKFLOWS_LABELS
+        env var. Format: comma-separated key=value pairs (e.g., "team=ml,env=prod").
+
+        Returns an empty dict if the env var is unset.
+        """
+        if not ARGO_WORKFLOWS_LABELS:
+            return {}
 
-        return labels
+        env_labels = parse_kube_keyvalue_list(
+            ARGO_WORKFLOWS_LABELS.split(","), requires_both=True
+        )
+        validate_kube_labels(env_labels, validate_keys=True)
+        return env_labels
+
+    def _base_argo_labels(self):
+        """
+        Get Kubernetes labels for WorkflowTemplate/Workflow-level Argo resources.
+
+        Merges custom labels from METAFLOW_ARGO_WORKFLOWS_LABELS with base
+        Kubernetes labels, with base (internal) labels taking precedence so
+        that they cannot be overridden by user-supplied custom labels.
+        """
+        return {**self._custom_argo_labels(), **self._base_kubernetes_labels()}
 
     def _base_kubernetes_annotations(self):
         """
@@ -903,7 +934,7 @@ def _compile_workflow_template(self):
                 .namespace(KUBERNETES_NAMESPACE)
                 .annotations(annotations)
                 .annotations(self._base_annotations)
-                .labels(self._base_labels)
+                .labels(self._workflow_labels)
                 .label("app.kubernetes.io/name", "metaflow-flow")
                 .annotations(dag_annotation)
             )
@@ -935,7 +966,7 @@ def _compile_workflow_template(self):
                 # Set workflow metadata
                 .workflow_metadata(
                     Metadata()
-                    .labels(self._base_labels)
+                    .labels(self._workflow_labels)
                     .label("app.kubernetes.io/name", "metaflow-run")
                     .annotations(
                         {
@@ -982,6 +1013,7 @@ def _compile_workflow_template(self):
                     )
                 )
                 # Set common pod metadata.
+                # internal labels only
                 .pod_metadata(
                     Metadata()
                     .labels(self._base_labels)
@@ -2718,6 +2750,7 @@ def _container_templates(self):
                     "metaflow/argo-workflows-name": "{{workflow.name}}",
                     "workflows.argoproj.io/workflow": "{{workflow.name}}",
                 }
+                # internal labels only
                 jobset.labels(
                     {
                         **resources["labels"],
@@ -3917,6 +3950,7 @@ def _compile_sensor(self):
             Sensor()
             .metadata(
                 # Sensor metadata.
+                # internal labels only
                 ObjectMeta()
                 .name(ArgoWorkflows._sensor_name(self.name))
                 .namespace(ARGO_EVENTS_SENSOR_NAMESPACE)

metaflow/plugins/kubernetes/kube_utils.py

@@ -63,22 +63,27 @@ def qos_requests_and_limits(qos: str, cpu: int, memory: int, storage: int):
 
 def validate_kube_labels(
     labels: Optional[Dict[str, Optional[str]]],
+    validate_keys: bool = False,
 ) -> bool:
-    """Validate label values.
+    """Validate label values, and optionally keys.
 
-    This validates the kubernetes label values.  It does not validate the keys.
-    Ideally, keys should be static and also the validation rules for keys are
-    more complex than those for values.  For full validation rules, see:
+    This validates the kubernetes label values.  By default, it does not
+    validate the keys, since keys have historically been static/internal and
+    the validation rules for keys are more complex than those for values. Set
+    validate_keys=True to also validate label keys, e.g. when keys are
+    user-supplied. For full validation rules, see:
 
     https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/#syntax-and-character-set
     """
 
+    # shared with the "name" segment of a label key
+    segment_regex = r"[A-Za-z0-9]([-A-Za-z0-9_.]{0,61}[A-Za-z0-9])?"
+
     def validate_label(s: Optional[str]):
-        regex_match = r"^(([A-Za-z0-9][-A-Za-z0-9_.]{0,61})?[A-Za-z0-9])?$"
         if not s:
             # allow empty label
             return True
-        if not re.search(regex_match, s):
+        if not re.search(r"^(%s)?$" % segment_regex, s):
             raise KubernetesException(
                 'Invalid value: "%s"\n'
                 "A valid label must be an empty string or one that\n"
@@ -88,7 +93,33 @@ def validate_label(s: Optional[str]):
             )
         return True
 
-    return all([validate_label(v) for v in labels.values()]) if labels else True
+    def validate_label_key(key: str):
+        prefix, _, name = key.rpartition("/")
+        if prefix:
+            prefix_regex = r"^[A-Za-z0-9]([-A-Za-z0-9.]{0,251}[A-Za-z0-9])?$"
+            if not re.search(prefix_regex, prefix):
+                raise KubernetesException(
+                    'Invalid key: "%s"\n'
+                    "The prefix of a label key, if present, must be a DNS\n"
+                    "subdomain: a series of DNS labels separated by '.',\n"
+                    "not longer than 253 characters in total" % key
+                )
+        if not re.search(r"^%s$" % segment_regex, name):
+            raise KubernetesException(
+                'Invalid key: "%s"\n'
+                "The name segment of a label key must be non-empty and\n"
+                "  - Consist of alphanumeric, '-', '_' or '.' characters\n"
+                "  - Begin and end with an alphanumeric character\n"
+                "  - Be at most 63 characters" % key
+            )
+        return True
+
+    if not labels:
+        return True
+    if validate_keys:
+        for key in labels:
+            validate_label_key(key)
+    return all([validate_label(v) for v in labels.values()])
 
 
 def parse_kube_keyvalue_list(items: List[str], requires_both: bool = True):

test/unit/test_argo_workflows_labels.py

@@ -0,0 +1,64 @@
+import pytest
+
+from metaflow.plugins.argo.argo_workflows import ArgoWorkflows
+from metaflow.plugins.kubernetes.kube_utils import KubernetesException
+
+
+@pytest.fixture
+def argo_workflows():
+    return ArgoWorkflows.__new__(ArgoWorkflows)
+
+
+@pytest.mark.parametrize(
+    ("configured_labels", "expected"),
+    [
+        (
+            "",
+            {"app.kubernetes.io/part-of": "metaflow"},
+        ),
+        (
+            "team=ml,env=prod",
+            {
+                "app.kubernetes.io/part-of": "metaflow",
+                "team": "ml",
+                "env": "prod",
+            },
+        ),
+        (
+            "app.kubernetes.io/part-of=custom,team=ml",
+            {
+                "app.kubernetes.io/part-of": "metaflow",
+                "team": "ml",
+            },
+        ),
+    ],
+    ids=["default", "custom-labels", "protected-label"],
+)
+def test_base_argo_labels(mocker, argo_workflows, configured_labels, expected):
+    mocker.patch(
+        "metaflow.plugins.argo.argo_workflows.ARGO_WORKFLOWS_LABELS",
+        configured_labels,
+    )
+
+    assert argo_workflows._base_argo_labels() == expected
+
+
+@pytest.mark.parametrize(
+    "configured_labels",
+    [
+        "missing-value",
+        "team=value with spaces",
+        "team=%s" % ("a" * 64),
+    ],
+    ids=["missing-equals", "invalid-value", "value-too-long"],
+)
+def test_base_argo_labels_rejects_invalid_configuration(
+    mocker, argo_workflows, configured_labels
+):
+    mocker.patch(
+        "metaflow.plugins.argo.argo_workflows.ARGO_WORKFLOWS_LABELS",
+        configured_labels,
+    )
+
+    with pytest.raises(KubernetesException):
+        argo_workflows._base_argo_labels()

test/ux/core/test_argo_compilation.py

@@ -64,6 +64,55 @@ def test_argo_only_json_exposes_workflow_template(
     assert workflow_template["spec"]["templates"]
 
 
+def test_configured_labels_are_emitted_on_argo_workflows(
+    exec_mode, decospecs, compute_env, tag, scheduler_config
+):
+    if exec_mode != "deployer":
+        pytest.skip("Argo compilation tests require deployer mode")
+    if scheduler_config.scheduler_type != "argo-workflows":
+        pytest.skip("Argo compilation tests require the argo-workflows scheduler")
+
+    from metaflow import Deployer
+
+    from .test_utils import _resolve_flow_path, prepare_runner_deployer_args
+
+    env = dict(compute_env)
+    env["METAFLOW_ARGO_WORKFLOWS_LABELS"] = (
+        "team=ml-platform,environment=test,app.kubernetes.io/name=custom"
+    )
+
+    deployed_flow = (
+        Deployer(
+            flow_file=_resolve_flow_path("basic/helloworld.py"),
+            show_output=False,
+            **prepare_runner_deployer_args({"decospecs": decospecs, "env": env}),
+        )
+        .argo_workflows()
+        .create(
+            only_json=True,
+            tags=tag + ["test_configured_argo_labels"],
+            **(scheduler_config.deploy_args or {}),
+        )
+    )
+
+    workflow_template = deployed_flow.workflow_template
+    template_labels = workflow_template["metadata"]["labels"]
+    workflow_labels = workflow_template["spec"]["workflowMetadata"]["labels"]
+
+    assert template_labels["team"] == "ml-platform"
+    assert template_labels["environment"] == "test"
+    assert template_labels["app.kubernetes.io/name"] == "metaflow-flow"
+
+    assert workflow_labels["team"] == "ml-platform"
+    assert workflow_labels["environment"] == "test"
+    assert workflow_labels["app.kubernetes.io/name"] == "metaflow-run"
+
+    # Custom Argo labels are intentionally workflow-level only.
+    pod_labels = workflow_template["spec"]["podMetadata"]["labels"]
+    assert "team" not in pod_labels
+    assert "environment" not in pod_labels
+
+
 def test_foreach_split_switch_join_task_names_are_deduplicated(
     exec_mode, decospecs, tag, scheduler_config
 ):