ob-metaflow 2.11.13.1__py2.py3-none-any.whl → 2.19.7.1rc0__py2.py3-none-any.whl

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -7,26 +7,31 @@ import time
 from metaflow import current
 from metaflow.decorators import StepDecorator
 from metaflow.exception import MetaflowException
-from metaflow.metadata import MetaDatum
-from metaflow.metadata.util import sync_local_metadata_to_datastore
+from metaflow.metadata_provider import MetaDatum
+from metaflow.metadata_provider.util import sync_local_metadata_to_datastore
 from metaflow.metaflow_config import (
     DATASTORE_LOCAL_DIR,
+    FEAT_ALWAYS_UPLOAD_CODE_PACKAGE,
     KUBERNETES_CONTAINER_IMAGE,
     KUBERNETES_CONTAINER_REGISTRY,
+    KUBERNETES_CPU,
+    KUBERNETES_DISK,
     KUBERNETES_FETCH_EC2_METADATA,
-    KUBERNETES_IMAGE_PULL_POLICY,
     KUBERNETES_GPU_VENDOR,
+    KUBERNETES_IMAGE_PULL_POLICY,
+    KUBERNETES_IMAGE_PULL_SECRETS,
+    KUBERNETES_MEMORY,
+    KUBERNETES_LABELS,
+    KUBERNETES_ANNOTATIONS,
     KUBERNETES_NAMESPACE,
     KUBERNETES_NODE_SELECTOR,
     KUBERNETES_PERSISTENT_VOLUME_CLAIMS,
-    KUBERNETES_TOLERATIONS,
-    KUBERNETES_SERVICE_ACCOUNT,
     KUBERNETES_PORT,
+    KUBERNETES_SERVICE_ACCOUNT,
     KUBERNETES_SHARED_MEMORY,
-    KUBERNETES_PORT,
-    KUBERNETES_CPU,
-    KUBERNETES_MEMORY,
-    KUBERNETES_DISK,
+    KUBERNETES_TOLERATIONS,
+    KUBERNETES_QOS,
+    KUBERNETES_CONDA_ARCH,
 )
 from metaflow.plugins.resources_decorator import ResourcesDecorator
 from metaflow.plugins.timeout_decorator import get_run_time_limit_for_task
@@ -34,7 +39,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
 
@@ -44,6 +50,8 @@ except NameError:
     unicode = str
     basestring = str
 
+SUPPORTED_KUBERNETES_QOS_CLASSES = ["Guaranteed", "Burstable"]
+
 
 class KubernetesDecorator(StepDecorator):
     """
@@ -68,12 +76,21 @@ class KubernetesDecorator(StepDecorator):
         not, a default Docker image mapping to the current version of Python is used.
     image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
         If given, the imagePullPolicy to be applied to the Docker image of the step.
+    image_pull_secrets: List[str], default []
+        The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
+        Kubernetes image pull secrets to use when pulling container images
+        in Kubernetes.
     service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
         Kubernetes service account to use when launching pod in Kubernetes.
     secrets : List[str], optional, default None
         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.
+    node_selector: Union[Dict[str,str], str], optional, default None
+        Kubernetes node selector(s) to apply to the pod running the task.
+        Can be passed in as a comma separated string of values e.g.
+        'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
+        {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
     namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
         Kubernetes namespace to use when launching pod in Kubernetes.
     gpu : int, optional, default None
@@ -81,9 +98,13 @@ class KubernetesDecorator(StepDecorator):
         the scheduled node should not have GPUs.
     gpu_vendor : str, default KUBERNETES_GPU_VENDOR
         The vendor of the GPUs to be used for this step.
-    tolerations : List[str], default []
+    tolerations : List[Dict[str,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
@@ -101,6 +122,22 @@ class KubernetesDecorator(StepDecorator):
         Shared memory size (in MiB) required for this step
     port: int, optional
         Port number to specify in the Kubernetes job object
+    compute_pool : str, optional, default None
+        Compute pool to be used for for this step.
+        If not specified, any accessible compute pool within the perimeter is used.
+    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
+
+    security_context: Dict[str, Any], optional, default None
+        Container security context. Applies to the task container. Allows the following keys:
+        - privileged: bool, optional, default None
+        - allow_privilege_escalation: bool, optional, default None
+        - run_as_user: int, optional, default None
+        - run_as_group: int, optional, default None
+        - run_as_non_root: bool, optional, default None
     """
 
     name = "kubernetes"
@@ -110,6 +147,7 @@ class KubernetesDecorator(StepDecorator):
         "disk": "10240",
         "image": None,
         "image_pull_policy": None,
+        "image_pull_secrets": None,  # e.g., ["regcred"]
         "service_account": None,
         "secrets": None,  # e.g., mysecret
         "node_selector": None,  # e.g., kubernetes.io/os=linux
@@ -118,6 +156,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,
@@ -125,14 +165,22 @@ class KubernetesDecorator(StepDecorator):
         "persistent_volume_claims": None,  # e.g., {"pvc-name": "/mnt/vol", "another-pvc": "/mnt/vol2"}
         "shared_memory": None,
         "port": None,
+        "compute_pool": None,
+        "executable": None,
+        "hostname_resolution_timeout": 10 * 60,
+        "qos": KUBERNETES_QOS,
+        "security_context": None,
     }
+    package_metadata = None
     package_url = None
     package_sha = None
     run_time_limit = None
 
-    def __init__(self, attributes=None, statically_defined=False):
-        super(KubernetesDecorator, self).__init__(attributes, statically_defined)
+    # Conda environment support
+    supports_conda_environment = True
+    target_platform = KUBERNETES_CONDA_ARCH or "linux-64"
 
+    def init(self):
         if not self.attributes["namespace"]:
             self.attributes["namespace"] = KUBERNETES_NAMESPACE
         if not self.attributes["service_account"]:
@@ -152,11 +200,21 @@ class KubernetesDecorator(StepDecorator):
             )
         if not self.attributes["image_pull_policy"] and KUBERNETES_IMAGE_PULL_POLICY:
             self.attributes["image_pull_policy"] = KUBERNETES_IMAGE_PULL_POLICY
+        if not self.attributes["image_pull_secrets"] and KUBERNETES_IMAGE_PULL_SECRETS:
+            self.attributes["image_pull_secrets"] = json.loads(
+                KUBERNETES_IMAGE_PULL_SECRETS
+            )
 
         if isinstance(self.attributes["node_selector"], str):
             self.attributes["node_selector"] = parse_kube_keyvalue_list(
                 self.attributes["node_selector"].split(",")
             )
+        if self.attributes["compute_pool"]:
+            if self.attributes["node_selector"] is None:
+                self.attributes["node_selector"] = {}
+            self.attributes["node_selector"].update(
+                {"outerbounds.co/compute-pool": self.attributes["compute_pool"]}
+            )
 
         if self.attributes["tolerations"]:
             try:
@@ -190,6 +248,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"]:
@@ -238,12 +326,33 @@ 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 "
                 "Kubernetes. Please use one or the other.".format(step=step)
             )
 
+        if any([deco.name == "parallel" for deco in decos]) and any(
+            [deco.name == "catch" for deco in decos]
+        ):
+            raise MetaflowException(
+                "Step *{step}* contains a @parallel decorator "
+                "with the @catch decorator. @catch is not supported with @parallel on Kubernetes.".format(
+                    step=step
+                )
+            )
+
         # Set run time limit for the Kubernetes job.
         self.run_time_limit = get_run_time_limit_for_task(decos)
         if self.run_time_limit < 60:
@@ -327,7 +436,7 @@ class KubernetesDecorator(StepDecorator):
 
         if self.attributes["shared_memory"]:
             if not (
-                isinstance(self.attributes["shared_memory"], (int, unicode, basestring))
+                isinstance(self.attributes["shared_memory"], int)
                 and int(self.attributes["shared_memory"]) > 0
             ):
                 raise KubernetesException(
@@ -336,6 +445,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.
@@ -374,12 +486,17 @@ class KubernetesDecorator(StepDecorator):
             # to execute on Kubernetes anymore. We can execute possible fallback
             # code locally.
             cli_args.commands = ["kubernetes", "step"]
+            cli_args.command_args.append(self.package_metadata)
             cli_args.command_args.append(self.package_sha)
             cli_args.command_args.append(self.package_url)
 
+            # skip certain keys as CLI arguments
+            _skip_keys = ["compute_pool", "hostname_resolution_timeout"]
             # --namespace is used to specify Metaflow namespace (a different
             # concept from k8s namespace).
             for k, v in self.attributes.items():
+                if k in _skip_keys:
+                    continue
                 if k == "namespace":
                     cli_args.command_options["k8s_namespace"] = v
                 elif k in {"node_selector"} and v:
@@ -387,7 +504,14 @@ 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 [
+                    "image_pull_secrets",
+                    "tolerations",
+                    "persistent_volume_claims",
+                    "labels",
+                    "annotations",
+                    "security_context",
+                ]:
                     cli_args.command_options[k] = json.dumps(v)
                 else:
                     cli_args.command_options[k] = v
@@ -422,8 +546,8 @@ class KubernetesDecorator(StepDecorator):
         # check for the existence of METAFLOW_KUBERNETES_WORKLOAD environment
         # variable.
 
+        meta = {}
         if "METAFLOW_KUBERNETES_WORKLOAD" in os.environ:
-            meta = {}
             meta["kubernetes-pod-name"] = os.environ["METAFLOW_KUBERNETES_POD_NAME"]
             meta["kubernetes-pod-namespace"] = os.environ[
                 "METAFLOW_KUBERNETES_POD_NAMESPACE"
@@ -434,10 +558,14 @@ class KubernetesDecorator(StepDecorator):
             ]
             meta["kubernetes-node-ip"] = os.environ["METAFLOW_KUBERNETES_NODE_IP"]
 
+            meta["kubernetes-jobset-name"] = os.environ.get(
+                "METAFLOW_KUBERNETES_JOBSET_NAME"
+            )
+
             # TODO (savin): Introduce equivalent support for Microsoft Azure and
             #               Google Cloud Platform
-            # TODO: Introduce a way to detect Cloud Provider, so unnecessary requests (and delays)
-            # can be avoided by not having to try out all providers.
+            # TODO: Introduce a way to detect Cloud Provider, so unnecessary requests
+            # (and delays) can be avoided by not having to try out all providers.
             if KUBERNETES_FETCH_EC2_METADATA:
                 instance_meta = get_ec2_instance_metadata()
                 meta.update(instance_meta)
@@ -453,38 +581,51 @@ class KubernetesDecorator(StepDecorator):
             #         "METAFLOW_KUBERNETES_POD_NAME"
             #     ].rpartition("-")[0]
 
-            entries = [
-                MetaDatum(field=k, value=v, type=k, tags=[])
-                for k, v in meta.items()
-                if v is not None
-            ]
-            # Register book-keeping metadata for debugging.
-            metadata.register_metadata(run_id, step_name, task_id, entries)
-
             # Start MFLog sidecar to collect task logs.
             self._save_logs_sidecar = Sidecar("save_logs_periodically")
             self._save_logs_sidecar.start()
 
-        num_parallel = int(os.environ.get("WORLD_SIZE", 0))
-        if num_parallel >= 1:
+            # 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
+
+        if num_parallel and num_parallel > 1:
+            _setup_multinode_environment(
+                ubf_context, self.attributes["hostname_resolution_timeout"]
+            )
+            # current.parallel.node_index will be correctly available over here.
+            meta.update({"parallel-node-index": current.parallel.node_index})
             if ubf_context == UBF_CONTROL:
-                control_task_id = current.task_id
-                top_task_id = control_task_id.replace("control-", "")
-                mapper_task_ids = [control_task_id] + [
-                    "%s-node-%d" % (top_task_id, node_idx)
-                    for node_idx in range(1, num_parallel)
-                ]
                 flow._control_mapper_tasks = [
-                    "%s/%s/%s" % (run_id, step_name, mapper_task_id)
-                    for mapper_task_id in mapper_task_ids
+                    "{}/{}/{}".format(run_id, step_name, task_id)
+                    for task_id in [task_id]
+                    + [
+                        "%s-worker-%d" % (task_id, idx)
+                        for idx in range(num_parallel - 1)
+                    ]
                 ]
                 flow._control_task_is_mapper_zero = True
-            else:
-                worker_job_rank = int(os.environ["RANK"])
-                os.environ["RANK"] = str(worker_job_rank + 1)
 
-        if num_parallel >= 1:
-            _setup_multinode_environment()
+        if len(meta) > 0:
+            entries = [
+                MetaDatum(
+                    field=k,
+                    value=v,
+                    type=k,
+                    tags=["attempt_id:{0}".format(retry_count)],
+                )
+                for k, v in meta.items()
+                if v is not None
+            ]
+            # Register book-keeping metadata for debugging.
+            metadata.register_metadata(run_id, step_name, task_id, entries)
 
     def task_finished(
         self, step_name, flow, graph, is_task_ok, retry_count, max_retries
@@ -497,10 +638,10 @@ class KubernetesDecorator(StepDecorator):
             # local file system after the user code has finished execution.
             # This happens via datastore as a communication bridge.
 
-            # TODO:  There is no guarantee that task_prestep executes before
-            #        task_finished is invoked. That will result in AttributeError:
-            #        'KubernetesDecorator' object has no attribute 'metadata' error.
-            if self.metadata.TYPE == "local":
+            # TODO:  There is no guarantee that task_pre_step executes before
+            #        task_finished is invoked.
+            # For now we guard against the missing metadata object in this case.
+            if hasattr(self, "metadata") and self.metadata.TYPE == "local":
                 # Note that the datastore is *always* Amazon S3 (see
                 # runtime_task_created function).
                 sync_local_metadata_to_datastore(
@@ -509,57 +650,74 @@ class KubernetesDecorator(StepDecorator):
 
         try:
             self._save_logs_sidecar.terminate()
+            self._spot_monitor_sidecar.terminate()
         except:
             # Best effort kill
             pass
 
-        if is_task_ok and len(getattr(flow, "_control_mapper_tasks", [])) > 1:
-            self._wait_for_mapper_tasks(flow, step_name)
+    @classmethod
+    def _save_package_once(cls, flow_datastore, package):
+        if cls.package_url is None:
+            if not FEAT_ALWAYS_UPLOAD_CODE_PACKAGE:
+                cls.package_url, cls.package_sha = flow_datastore.save_data(
+                    [package.blob], len_hint=1
+                )[0]
+                cls.package_metadata = package.package_metadata
+            else:
+                # Blocks until the package is uploaded
+                cls.package_url = package.package_url()
+                cls.package_sha = package.package_sha()
+                cls.package_metadata = package.package_metadata
+
+
+# TODO: Unify this method with the multi-node setup in @batch
+def _setup_multinode_environment(ubf_context, hostname_resolution_timeout):
+    import socket
 
-    def _wait_for_mapper_tasks(self, flow, step_name):
+    def _wait_for_hostname_resolution(max_wait_timeout=10 * 60):
         """
-        When launching multinode task with UBF, need to wait for the secondary
-        tasks to finish cleanly and produce their output before exiting the
-        main task. Otherwise, the main task finishing will cause secondary nodes
-        to terminate immediately, and possibly prematurely.
+        keep trying to resolve the hostname of the control task until the hostname is resolved
+        or the max_wait_timeout is reached. This is a workaround for the issue where the control
+        task is not scheduled before the worker task and the worker task fails because it cannot
+        resolve the hostname of the control task.
         """
-        from metaflow import Step  # avoid circular dependency
-
-        TIMEOUT = 600
-        last_completion_timeout = time.time() + TIMEOUT
-        print("Waiting for batch secondary tasks to finish")
-        while last_completion_timeout > time.time():
-            time.sleep(2)
+        start_time = time.time()
+        while True:
             try:
-                step_path = "%s/%s/%s" % (flow.name, current.run_id, step_name)
-                tasks = [task for task in Step(step_path)]
-                if len(tasks) == len(flow._control_mapper_tasks):
-                    if all(
-                        task.finished_at is not None for task in tasks
-                    ):  # for some reason task.finished fails
-                        return True
-                else:
-                    print(
-                        "Waiting for all parallel tasks to finish. Finished: {}/{}".format(
-                            len(tasks),
-                            len(flow._control_mapper_tasks),
+                return socket.gethostbyname(os.environ["MF_MASTER_ADDR"])
+            except socket.gaierror:
+                if time.time() - start_time > max_wait_timeout:
+                    raise MetaflowException(
+                        "Failed to get host by name for MF_MASTER_ADDR after waiting for {} seconds.".format(
+                            max_wait_timeout
                         )
                     )
-            except Exception as e:
-                pass
-        raise Exception(
-            "Batch secondary workers did not finish in %s seconds" % TIMEOUT
-        )
-
-    @classmethod
-    def _save_package_once(cls, flow_datastore, package):
-        if cls.package_url is None:
-            cls.package_url, cls.package_sha = flow_datastore.save_data(
-                [package.blob], len_hint=1
-            )[0]
+                time.sleep(1)
+
+    try:
+        # Even if Kubernetes may deploy control pods before worker pods, there is always a
+        # possibility that the worker pods may start before the control. In the case that this happens,
+        # the worker pods will not be able to resolve the control pod's IP address and this will cause
+        # the worker pods to fail. So if the worker pods are requesting a hostname resolution, we will
+        # make it wait for the name to be resolved within a reasonable timeout period.
+        if ubf_context != UBF_CONTROL:
+            os.environ["MF_PARALLEL_MAIN_IP"] = _wait_for_hostname_resolution(
+                hostname_resolution_timeout
+            )
+        else:
+            os.environ["MF_PARALLEL_MAIN_IP"] = socket.gethostbyname(
+                os.environ["MF_MASTER_ADDR"]
+            )
 
-def _setup_multinode_environment():
-    import socket
-    os.environ["MF_PARALLEL_MAIN_IP"] = socket.gethostbyname(os.environ["MASTER_ADDR"])
-    os.environ["MF_PARALLEL_NUM_NODES"] = os.environ["WORLD_SIZE"]
-    os.environ["MF_PARALLEL_NODE_INDEX"] = os.environ["RANK"]
+        os.environ["MF_PARALLEL_NUM_NODES"] = os.environ["MF_WORLD_SIZE"]
+        os.environ["MF_PARALLEL_NODE_INDEX"] = (
+            str(0)
+            if "MF_CONTROL_INDEX" in os.environ
+            else str(int(os.environ["MF_WORKER_REPLICA_INDEX"]) + 1)
+        )
+    except KeyError as e:
+        raise MetaflowException("Environment variable {} is missing.".format(e))
+    except socket.gaierror:
+        raise MetaflowException("Failed to get host by name for MF_MASTER_ADDR.")
+    except ValueError:
+        raise MetaflowException("Invalid value for MF_WORKER_REPLICA_INDEX.")