ob-metaflow 2.12.30.2__py2.py3-none-any.whl → 2.13.6.1__py2.py3-none-any.whl

metaflow/plugins/kubernetes/kube_utils.py

@@ -1,7 +1,14 @@
-from metaflow.exception import CommandException
+import re
+from typing import Dict, List, Optional
+from metaflow.exception import CommandException, MetaflowException
 from metaflow.util import get_username, get_latest_run_id
 
 
+# avoid circular import by having the exception class contained here
+class KubernetesException(MetaflowException):
+    headline = "Kubernetes error"
+
+
 def parse_cli_options(flow_name, run_id, user, my_runs, echo):
     if user and my_runs:
         raise CommandException("--user and --my-runs are mutually exclusive.")
@@ -23,3 +30,79 @@ def parse_cli_options(flow_name, run_id, user, my_runs, echo):
             raise CommandException("A previous run id was not found. Specify --run-id.")
 
     return flow_name, run_id, user
+
+
+def qos_requests_and_limits(qos: str, cpu: int, memory: int, storage: int):
+    "return resource requests and limits for the kubernetes pod based on the given QoS Class"
+    # case insensitive matching for QoS class
+    qos = qos.lower()
+    # Determine the requests and limits to define chosen QoS class
+    qos_limits = {}
+    qos_requests = {}
+    if qos == "guaranteed":
+        # Guaranteed - has both cpu/memory limits. requests not required, as these will be inferred.
+        qos_limits = {
+            "cpu": str(cpu),
+            "memory": "%sM" % str(memory),
+            "ephemeral-storage": "%sM" % str(storage),
+        }
+        # NOTE: Even though Kubernetes will produce matching requests for the specified limits, this happens late in the lifecycle.
+        # We specify them explicitly here to make some K8S tooling happy, in case they rely on .resources.requests being present at time of submitting the job.
+        qos_requests = qos_limits
+    else:
+        # Burstable - not Guaranteed, and has a memory/cpu limit or request
+        qos_requests = {
+            "cpu": str(cpu),
+            "memory": "%sM" % str(memory),
+            "ephemeral-storage": "%sM" % str(storage),
+        }
+    # TODO: Add support for BestEffort once there is a use case for it.
+    # BestEffort - no limit or requests for cpu/memory
+    return qos_requests, qos_limits
+
+
+def validate_kube_labels(
+    labels: Optional[Dict[str, Optional[str]]],
+) -> bool:
+    """Validate label values.
+
+    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:
+
+    https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/#syntax-and-character-set
+    """
+
+    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):
+            raise KubernetesException(
+                'Invalid value: "%s"\n'
+                "A valid label must be an empty string or one that\n"
+                "  - Consist of alphanumeric, '-', '_' or '.' characters\n"
+                "  - Begins and ends with an alphanumeric character\n"
+                "  - Is at most 63 characters" % s
+            )
+        return True
+
+    return all([validate_label(v) for v in labels.values()]) if labels else True
+
+
+def parse_kube_keyvalue_list(items: List[str], requires_both: bool = True):
+    try:
+        ret = {}
+        for item_str in items:
+            item = item_str.split("=", 1)
+            if requires_both:
+                item[1]  # raise IndexError
+            if str(item[0]) in ret:
+                raise KubernetesException("Duplicate key found: %s" % str(item[0]))
+            ret[str(item[0])] = str(item[1]) if len(item) > 1 else None
+        return ret
+    except KubernetesException as e:
+        raise e
+    except (AttributeError, IndexError):
+        raise KubernetesException("Unable to parse kubernetes list: %s" % items)

metaflow/plugins/kubernetes/kubernetes.py

@@ -1,11 +1,8 @@
-import copy
 import json
 import math
 import os
-import re
 import shlex
 import time
-from typing import Dict, List, Optional
 from uuid import uuid4
 
 from metaflow import current, util
@@ -35,7 +32,6 @@ from metaflow.metaflow_config import (
     DEFAULT_SECRETS_BACKEND_TYPE,
     GCP_SECRET_MANAGER_PREFIX,
     KUBERNETES_FETCH_EC2_METADATA,
-    KUBERNETES_LABELS,
     KUBERNETES_SANDBOX_INIT_SCRIPT,
     OTEL_ENDPOINT,
     S3_ENDPOINT_URL,
@@ -193,9 +189,11 @@ class Kubernetes(object):
         persistent_volume_claims=None,
         tolerations=None,
         labels=None,
+        annotations=None,
         shared_memory=None,
         port=None,
         num_parallel=None,
+        qos=None,
     ):
         name = "js-%s" % str(uuid4())[:6]
         jobset = (
@@ -228,6 +226,7 @@ class Kubernetes(object):
                 shared_memory=shared_memory,
                 port=port,
                 num_parallel=num_parallel,
+                qos=qos,
             )
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)
@@ -302,17 +301,13 @@ class Kubernetes(object):
             # see get_datastore_root_from_config in datastore/local.py).
         )
 
-        _labels = self._get_labels(labels)
-        for k, v in _labels.items():
-            jobset.label(k, v)
-
         for k in list(
             [] if not secrets else [secrets] if isinstance(secrets, str) else secrets
         ) + KUBERNETES_SECRETS.split(","):
             jobset.secret(k)
 
         initial_configs = init_config()
-        for entry in ["OBP_PERIMETER", "OBP_INTEGRATIONS_SECRETS_METADATA_URL"]:
+        for entry in ["OBP_PERIMETER", "OBP_INTEGRATIONS_URL"]:
             if entry not in initial_configs:
                 raise KubernetesException(
                     f"{entry} was not found in metaflow config. Please make sure to run `outerbounds configure <...>` command which can be found on the Ourebounds UI or reach out to your Outerbounds support team."
@@ -320,8 +315,8 @@ class Kubernetes(object):
 
         additional_obp_configs = {
             "OBP_PERIMETER": initial_configs["OBP_PERIMETER"],
-            "OBP_INTEGRATIONS_SECRETS_METADATA_URL": initial_configs[
-                "OBP_INTEGRATIONS_SECRETS_METADATA_URL"
+            "OBP_INTEGRATIONS_URL": initial_configs[
+                "OBP_INTEGRATIONS_URL"
             ],
         }
         for k, v in additional_obp_configs.items():
@@ -393,13 +388,16 @@ class Kubernetes(object):
         for name, value in env.items():
             jobset.environment_variable(name, value)
 
-        annotations = {
+        system_annotations = {
             "metaflow/user": user,
             "metaflow/flow_name": flow_name,
             "metaflow/control-task-id": task_id,
+            "metaflow/run_id": run_id,
+            "metaflow/step_name": step_name,
+            "metaflow/attempt": attempt,
         }
         if current.get("project_name"):
-            annotations.update(
+            system_annotations.update(
                 {
                     "metaflow/project_name": current.project_name,
                     "metaflow/branch_name": current.branch_name,
@@ -407,15 +405,15 @@ class Kubernetes(object):
                 }
             )
 
-        for name, value in annotations.items():
-            jobset.annotation(name, value)
+        system_labels = {
+            "app.kubernetes.io/name": "metaflow-task",
+            "app.kubernetes.io/part-of": "metaflow",
+        }
+
+        jobset.labels({**({} if not labels else labels), **system_labels})
 
-        (
-            jobset.annotation("metaflow/run_id", run_id)
-            .annotation("metaflow/step_name", step_name)
-            .annotation("metaflow/attempt", attempt)
-            .label("app.kubernetes.io/name", "metaflow-task")
-            .label("app.kubernetes.io/part-of", "metaflow")
+        jobset.annotations(
+            {**({} if not annotations else annotations), **system_annotations}
         )
         # We need this task-id set so that all the nodes are aware of the control
         # task's task-id. These "MF_" variables populate the `current.parallel` namedtuple
@@ -504,6 +502,8 @@ class Kubernetes(object):
         shared_memory=None,
         port=None,
         name_pattern=None,
+        qos=None,
+        annotations=None,
     ):
         if env is None:
             env = {}
@@ -536,7 +536,8 @@ class Kubernetes(object):
                 retries=0,
                 step_name=step_name,
                 tolerations=tolerations,
-                labels=self._get_labels(labels),
+                labels=labels,
+                annotations=annotations,
                 use_tmpfs=use_tmpfs,
                 tmpfs_tempdir=tmpfs_tempdir,
                 tmpfs_size=tmpfs_size,
@@ -544,6 +545,7 @@ class Kubernetes(object):
                 persistent_volume_claims=persistent_volume_claims,
                 shared_memory=shared_memory,
                 port=port,
+                qos=qos,
             )
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)
@@ -654,13 +656,25 @@ class Kubernetes(object):
 
         for name, value in env.items():
             job.environment_variable(name, value)
+        # Add job specific labels
+        system_labels = {
+            "app.kubernetes.io/name": "metaflow-task",
+            "app.kubernetes.io/part-of": "metaflow",
+        }
+        for name, value in system_labels.items():
+            job.label(name, value)
 
-        annotations = {
-            "metaflow/user": user,
+        # Add job specific annotations not set in the decorator.
+        system_annotations = {
             "metaflow/flow_name": flow_name,
+            "metaflow/run_id": run_id,
+            "metaflow/step_name": step_name,
+            "metaflow/task_id": task_id,
+            "metaflow/attempt": attempt,
+            "metaflow/user": user,
         }
         if current.get("project_name"):
-            annotations.update(
+            system_annotations.update(
                 {
                     "metaflow/project_name": current.project_name,
                     "metaflow/branch_name": current.branch_name,
@@ -668,18 +682,9 @@ class Kubernetes(object):
                 }
             )
 
-        for name, value in annotations.items():
+        for name, value in system_annotations.items():
             job.annotation(name, value)
 
-        (
-            job.annotation("metaflow/run_id", run_id)
-            .annotation("metaflow/step_name", step_name)
-            .annotation("metaflow/task_id", task_id)
-            .annotation("metaflow/attempt", attempt)
-            .label("app.kubernetes.io/name", "metaflow-task")
-            .label("app.kubernetes.io/part-of", "metaflow")
-        )
-
         return job
 
     def create_k8sjob(self, job):
@@ -787,60 +792,3 @@ class Kubernetes(object):
             "stderr",
             job_id=self._job.id,
         )
-
-    @staticmethod
-    def _get_labels(extra_labels=None):
-        if extra_labels is None:
-            extra_labels = {}
-        env_labels = KUBERNETES_LABELS.split(",") if KUBERNETES_LABELS else []
-        env_labels = parse_kube_keyvalue_list(env_labels, False)
-        labels = {**env_labels, **extra_labels}
-        validate_kube_labels(labels)
-        return labels
-
-
-def validate_kube_labels(
-    labels: Optional[Dict[str, Optional[str]]],
-) -> bool:
-    """Validate label values.
-
-    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:
-
-    https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/#syntax-and-character-set
-    """
-
-    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):
-            raise KubernetesException(
-                'Invalid value: "%s"\n'
-                "A valid label must be an empty string or one that\n"
-                "  - Consist of alphanumeric, '-', '_' or '.' characters\n"
-                "  - Begins and ends with an alphanumeric character\n"
-                "  - Is at most 63 characters" % s
-            )
-        return True
-
-    return all([validate_label(v) for v in labels.values()]) if labels else True
-
-
-def parse_kube_keyvalue_list(items: List[str], requires_both: bool = True):
-    try:
-        ret = {}
-        for item_str in items:
-            item = item_str.split("=", 1)
-            if requires_both:
-                item[1]  # raise IndexError
-            if str(item[0]) in ret:
-                raise KubernetesException("Duplicate key found: %s" % str(item[0]))
-            ret[str(item[0])] = str(item[1]) if len(item) > 1 else None
-        return ret
-    except KubernetesException as e:
-        raise e
-    except (AttributeError, IndexError):
-        raise KubernetesException("Unable to parse kubernetes list: %s" % items)

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -3,14 +3,17 @@ import sys
 import time
 import traceback
 
-from metaflow.plugins.kubernetes.kube_utils import parse_cli_options
+from metaflow.plugins.kubernetes.kube_utils import (
+    parse_cli_options,
+    parse_kube_keyvalue_list,
+)
 from metaflow.plugins.kubernetes.kubernetes_client import KubernetesClient
 import metaflow.tracing as tracing
 from metaflow import JSONTypeClass, util
 from metaflow._vendor import click
 from metaflow.exception import METAFLOW_EXIT_DISALLOW_RETRY, MetaflowException
 from metaflow.metadata_provider.util import sync_local_metadata_from_datastore
-from metaflow.metaflow_config import DATASTORE_LOCAL_DIR, KUBERNETES_LABELS
+from metaflow.metaflow_config import DATASTORE_LOCAL_DIR
 from metaflow.mflog import TASK_LOG_SOURCE
 from metaflow.unbounded_foreach import UBF_CONTROL, UBF_TASK
 
@@ -18,9 +21,7 @@ from .kubernetes import (
     Kubernetes,
     KubernetesException,
     KubernetesKilledException,
-    parse_kube_keyvalue_list,
 )
-from .kubernetes_decorator import KubernetesDecorator
 
 
 @click.group()
@@ -33,12 +34,12 @@ def kubernetes():
     pass
 
 
-@tracing.cli_entrypoint("kubernetes/step")
 @kubernetes.command(
     help="Execute a single task on Kubernetes. This command calls the top-level step "
     "command inside a Kubernetes pod with the given options. Typically you do not call "
     "this command directly; it is used internally by Metaflow."
 )
+@tracing.cli_entrypoint("kubernetes/step")
 @click.argument("step-name")
 @click.argument("code-package-sha")
 @click.argument("code-package-url")
@@ -126,6 +127,24 @@ def kubernetes():
     type=int,
     help="Number of parallel nodes to run as a multi-node job.",
 )
+@click.option(
+    "--qos",
+    default=None,
+    type=str,
+    help="Quality of Service class for the Kubernetes pod",
+)
+@click.option(
+    "--labels",
+    default=None,
+    type=JSONTypeClass(),
+    multiple=False,
+)
+@click.option(
+    "--annotations",
+    default=None,
+    type=JSONTypeClass(),
+    multiple=False,
+)
 @click.pass_context
 def step(
     ctx,
@@ -154,6 +173,9 @@ def step(
     shared_memory=None,
     port=None,
     num_parallel=None,
+    qos=None,
+    labels=None,
+    annotations=None,
     **kwargs
 ):
     def echo(msg, stream="stderr", job_id=None, **kwargs):
@@ -168,7 +190,7 @@ def step(
     executable = ctx.obj.environment.executable(step_name, executable)
 
     # Set environment
-    env = {}
+    env = {"METAFLOW_FLOW_FILENAME": os.path.basename(sys.argv[0])}
     env_deco = [deco for deco in node.decorators if deco.name == "environment"]
     if env_deco:
         env = env_deco[0].attributes["vars"]
@@ -294,8 +316,11 @@ def step(
                 shared_memory=shared_memory,
                 port=port,
                 num_parallel=num_parallel,
+                qos=qos,
+                labels=labels,
+                annotations=annotations,
             )
-    except Exception as e:
+    except Exception:
         traceback.print_exc(chain=False)
         _sync_metadata()
         sys.exit(METAFLOW_EXIT_DISALLOW_RETRY)

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -19,6 +19,8 @@ from metaflow.metaflow_config import (
     KUBERNETES_GPU_VENDOR,
     KUBERNETES_IMAGE_PULL_POLICY,
     KUBERNETES_MEMORY,
+    KUBERNETES_LABELS,
+    KUBERNETES_ANNOTATIONS,
     KUBERNETES_NAMESPACE,
     KUBERNETES_NODE_SELECTOR,
     KUBERNETES_PERSISTENT_VOLUME_CLAIMS,
@@ -26,6 +28,7 @@ from metaflow.metaflow_config import (
     KUBERNETES_SERVICE_ACCOUNT,
     KUBERNETES_SHARED_MEMORY,
     KUBERNETES_TOLERATIONS,
+    KUBERNETES_QOS,
 )
 from metaflow.plugins.resources_decorator import ResourcesDecorator
 from metaflow.plugins.timeout_decorator import get_run_time_limit_for_task
@@ -33,7 +36,8 @@ from metaflow.sidecar import Sidecar
 from metaflow.unbounded_foreach import UBF_CONTROL
 
 from ..aws.aws_utils import get_docker_registry, get_ec2_instance_metadata
-from .kubernetes import KubernetesException, parse_kube_keyvalue_list
+from .kubernetes import KubernetesException
+from .kube_utils import validate_kube_labels, parse_kube_keyvalue_list
 
 from metaflow.metaflow_config import MAX_MEMORY_PER_TASK, MAX_CPU_PER_TASK
 
@@ -43,6 +47,8 @@ except NameError:
     unicode = str
     basestring = str
 
+SUPPORTED_KUBERNETES_QOS_CLASSES = ["Guaranteed", "Burstable"]
+
 
 class KubernetesDecorator(StepDecorator):
     """
@@ -88,6 +94,10 @@ class KubernetesDecorator(StepDecorator):
     tolerations : List[str], default []
         The default is extracted from 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.
+    annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
+        Kubernetes annotations to use when launching pod in Kubernetes.
     use_tmpfs : bool, default False
         This enables an explicit tmpfs mount for this step.
     tmpfs_tempdir : bool, default True
@@ -111,6 +121,8 @@ class KubernetesDecorator(StepDecorator):
     hostname_resolution_timeout: int, default 10 * 60
         Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
         Only applicable when @parallel is used.
+    qos: str, default: Burstable
+        Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
     """
 
     name = "kubernetes"
@@ -128,6 +140,8 @@ class KubernetesDecorator(StepDecorator):
         "gpu_vendor": None,
         "tolerations": None,  # e.g., [{"key": "arch", "operator": "Equal", "value": "amd"},
         #                              {"key": "foo", "operator": "Equal", "value": "bar"}]
+        "labels": None,  # e.g. {"test-label": "value", "another-label":"value2"}
+        "annotations": None,  # e.g. {"note": "value", "another-note": "value2"}
         "use_tmpfs": None,
         "tmpfs_tempdir": True,
         "tmpfs_size": None,
@@ -138,6 +152,7 @@ class KubernetesDecorator(StepDecorator):
         "compute_pool": None,
         "executable": None,
         "hostname_resolution_timeout": 10 * 60,
+        "qos": KUBERNETES_QOS,
     }
     package_url = None
     package_sha = None
@@ -147,8 +162,8 @@ class KubernetesDecorator(StepDecorator):
     supports_conda_environment = True
     target_platform = "linux-64"
 
-    def __init__(self, attributes=None, statically_defined=False):
-        super(KubernetesDecorator, self).__init__(attributes, statically_defined)
+    def init(self):
+        super(KubernetesDecorator, self).init()
 
         if not self.attributes["namespace"]:
             self.attributes["namespace"] = KUBERNETES_NAMESPACE
@@ -213,6 +228,36 @@ class KubernetesDecorator(StepDecorator):
             self.attributes["memory"] = KUBERNETES_MEMORY
         if self.attributes["disk"] == self.defaults["disk"] and KUBERNETES_DISK:
             self.attributes["disk"] = KUBERNETES_DISK
+        # Label source precedence (decreasing):
+        # - System labels (set outside of decorator)
+        # - Decorator labels: @kubernetes(labels={})
+        # - Environment variable labels: METAFLOW_KUBERNETES_LABELS=
+        deco_labels = {}
+        if self.attributes["labels"] is not None:
+            deco_labels = self.attributes["labels"]
+
+        env_labels = {}
+        if KUBERNETES_LABELS:
+            env_labels = parse_kube_keyvalue_list(KUBERNETES_LABELS.split(","), False)
+
+        self.attributes["labels"] = {**env_labels, **deco_labels}
+
+        # Annotations
+        # annotation precedence (decreasing):
+        # - System annotations (set outside of decorator)
+        # - Decorator annotations: @kubernetes(annotations={})
+        # - Environment annotations: METAFLOW_KUBERNETES_ANNOTATIONS=
+        deco_annotations = {}
+        if self.attributes["annotations"] is not None:
+            deco_annotations = self.attributes["annotations"]
+
+        env_annotations = {}
+        if KUBERNETES_ANNOTATIONS:
+            env_annotations = parse_kube_keyvalue_list(
+                KUBERNETES_ANNOTATIONS.split(","), False
+            )
+
+        self.attributes["annotations"] = {**env_annotations, **deco_annotations}
 
         # If no docker image is explicitly specified, impute a default image.
         if not self.attributes["image"]:
@@ -261,6 +306,17 @@ class KubernetesDecorator(StepDecorator):
         self.step = step
         self.flow_datastore = flow_datastore
 
+        if (
+            self.attributes["qos"] is not None
+            # case insensitive matching.
+            and self.attributes["qos"].lower()
+            not in [c.lower() for c in SUPPORTED_KUBERNETES_QOS_CLASSES]
+        ):
+            raise MetaflowException(
+                "*%s* is not a valid Kubernetes QoS class. Choose one of the following: %s"
+                % (self.attributes["qos"], ", ".join(SUPPORTED_KUBERNETES_QOS_CLASSES))
+            )
+
         if any([deco.name == "batch" for deco in decos]):
             raise MetaflowException(
                 "Step *{step}* is marked for execution both on AWS Batch and "
@@ -369,6 +425,9 @@ class KubernetesDecorator(StepDecorator):
                     )
                 )
 
+        validate_kube_labels(self.attributes["labels"])
+        # TODO: add validation to annotations as well?
+
     def package_init(self, flow, step_name, environment):
         try:
             # Kubernetes is a soft dependency.
@@ -424,7 +483,12 @@ class KubernetesDecorator(StepDecorator):
                         "=".join([key, str(val)]) if val else key
                         for key, val in v.items()
                     ]
-                elif k in ["tolerations", "persistent_volume_claims"]:
+                elif k in [
+                    "tolerations",
+                    "persistent_volume_claims",
+                    "labels",
+                    "annotations",
+                ]:
                     cli_args.command_options[k] = json.dumps(v)
                 else:
                     cli_args.command_options[k] = v
@@ -498,6 +562,13 @@ class KubernetesDecorator(StepDecorator):
             self._save_logs_sidecar = Sidecar("save_logs_periodically")
             self._save_logs_sidecar.start()
 
+            # Start spot termination monitor sidecar.
+            current._update_env(
+                {"spot_termination_notice": "/tmp/spot_termination_notice"}
+            )
+            self._spot_monitor_sidecar = Sidecar("spot_termination_monitor")
+            self._spot_monitor_sidecar.start()
+
         num_parallel = None
         if hasattr(flow, "_parallel_ubf_iter"):
             num_parallel = flow._parallel_ubf_iter.num_parallel
@@ -556,6 +627,7 @@ class KubernetesDecorator(StepDecorator):
 
         try:
             self._save_logs_sidecar.terminate()
+            self._spot_monitor_sidecar.terminate()
         except:
             # Best effort kill
             pass