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

metaflow/plugins/kubernetes/kubernetes_job.py

@@ -1,17 +1,16 @@
-import copy
 import json
 import math
 import random
-import sys
 import time
 
 from metaflow.exception import MetaflowException
 from metaflow.metaflow_config import KUBERNETES_SECRETS
 from metaflow.tracing import inject_tracing_vars
-from metaflow.unbounded_foreach import UBF_CONTROL, UBF_TASK
 from metaflow.metaflow_config_funcs import init_config
 
 CLIENT_REFRESH_INTERVAL_SECONDS = 300
+
+from .kube_utils import qos_requests_and_limits
 from .kubernetes_jobsets import (
     KubernetesJobSet,
 )  # We need this import for Kubernetes Client.
@@ -75,8 +74,14 @@ class KubernetesJob(object):
             if self._kwargs["shared_memory"]
             else None
         )
+        qos_requests, qos_limits = qos_requests_and_limits(
+            self._kwargs["qos"],
+            self._kwargs["cpu"],
+            self._kwargs["memory"],
+            self._kwargs["disk"],
+        )
         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 KubernetesJobException(
                     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."
@@ -84,8 +89,8 @@ class KubernetesJob(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"
             ],
         }
 
@@ -176,20 +181,18 @@ class KubernetesJob(object):
                             image_pull_policy=self._kwargs["image_pull_policy"],
                             name=self._kwargs["step_name"].replace("_", "-"),
                             resources=client.V1ResourceRequirements(
-                                requests={
-                                    "cpu": str(self._kwargs["cpu"]),
-                                    "memory": "%sM" % str(self._kwargs["memory"]),
-                                    "ephemeral-storage": "%sM"
-                                    % str(self._kwargs["disk"]),
-                                },
+                                requests=qos_requests,
                                 limits={
-                                    "%s.com/gpu".lower()
-                                    % self._kwargs["gpu_vendor"]: str(
-                                        self._kwargs["gpu"]
-                                    )
-                                    for k in [0]
-                                    # Don't set GPU limits if gpu isn't specified.
-                                    if self._kwargs["gpu"] is not None
+                                    **qos_limits,
+                                    **{
+                                        "%s.com/gpu".lower()
+                                        % self._kwargs["gpu_vendor"]: str(
+                                            self._kwargs["gpu"]
+                                        )
+                                        for k in [0]
+                                        # Don't set GPU limits if gpu isn't specified.
+                                        if self._kwargs["gpu"] is not None
+                                    },
                                 },
                             ),
                             volume_mounts=(
@@ -445,7 +448,7 @@ class RunningJob(object):
         def best_effort_kill():
             try:
                 self.kill()
-            except Exception as ex:
+            except Exception:
                 pass
 
         atexit.register(best_effort_kill)

metaflow/plugins/kubernetes/kubernetes_jobsets.py

@@ -1,4 +1,3 @@
-import copy
 import json
 import math
 import random
@@ -7,7 +6,8 @@ from collections import namedtuple
 from metaflow.exception import MetaflowException
 from metaflow.metaflow_config import KUBERNETES_JOBSET_GROUP, KUBERNETES_JOBSET_VERSION
 from metaflow.tracing import inject_tracing_vars
-from metaflow.metaflow_config import KUBERNETES_SECRETS
+
+from .kube_utils import qos_requests_and_limits
 
 
 class KubernetesJobsetException(MetaflowException):
@@ -255,7 +255,7 @@ class RunningJobSet(object):
         def best_effort_kill():
             try:
                 self.kill()
-            except Exception as ex:
+            except Exception:
                 pass
 
         atexit.register(best_effort_kill)
@@ -340,7 +340,7 @@ class RunningJobSet(object):
                 stdout=True,
                 tty=False,
             )
-        except Exception as e:
+        except Exception:
             with client.ApiClient() as api_client:
                 # If we are unable to kill the control pod then
                 # Delete the jobset to kill the subsequent pods.
@@ -554,7 +554,12 @@ class JobSetSpec(object):
             if self._kwargs["shared_memory"]
             else None
         )
-
+        qos_requests, qos_limits = qos_requests_and_limits(
+            self._kwargs["qos"],
+            self._kwargs["cpu"],
+            self._kwargs["memory"],
+            self._kwargs["disk"],
+        )
         return dict(
             name=self.name,
             template=client.api_client.ApiClient().sanitize_for_serialization(
@@ -653,21 +658,18 @@ class JobSetSpec(object):
                                             "_", "-"
                                         ),
                                         resources=client.V1ResourceRequirements(
-                                            requests={
-                                                "cpu": str(self._kwargs["cpu"]),
-                                                "memory": "%sM"
-                                                % str(self._kwargs["memory"]),
-                                                "ephemeral-storage": "%sM"
-                                                % str(self._kwargs["disk"]),
-                                            },
+                                            requests=qos_requests,
                                             limits={
-                                                "%s.com/gpu".lower()
-                                                % self._kwargs["gpu_vendor"]: str(
-                                                    self._kwargs["gpu"]
-                                                )
-                                                for k in [0]
-                                                # Don't set GPU limits if gpu isn't specified.
-                                                if self._kwargs["gpu"] is not None
+                                                **qos_limits,
+                                                **{
+                                                    "%s.com/gpu".lower()
+                                                    % self._kwargs["gpu_vendor"]: str(
+                                                        self._kwargs["gpu"]
+                                                    )
+                                                    for k in [0]
+                                                    # Don't set GPU limits if gpu isn't specified.
+                                                    if self._kwargs["gpu"] is not None
+                                                },
                                             },
                                         ),
                                         volume_mounts=(
@@ -858,6 +860,16 @@ class KubernetesJobSet(object):
         self._annotations = dict(self._annotations, **{name: value})
         return self
 
+    def labels(self, labels):
+        for k, v in labels.items():
+            self.label(k, v)
+        return self
+
+    def annotations(self, annotations):
+        for k, v in annotations.items():
+            self.annotation(k, v)
+        return self
+
     def secret(self, name):
         self.worker.secret(name)
         self.control.secret(name)
@@ -983,15 +995,24 @@ class KubernetesArgoJobSet(object):
         self._labels = dict(self._labels, **{name: value})
         return self
 
+    def labels(self, labels):
+        for k, v in labels.items():
+            self.label(k, v)
+        return self
+
     def annotation(self, name, value):
         self.worker.annotation(name, value)
         self.control.annotation(name, value)
         self._annotations = dict(self._annotations, **{name: value})
         return self
 
+    def annotations(self, annotations):
+        for k, v in annotations.items():
+            self.annotation(k, v)
+        return self
+
     def dump(self):
         client = self._kubernetes_sdk
-        import json
 
         data = json.dumps(
             client.ApiClient().sanitize_for_serialization(

metaflow/plugins/kubernetes/spot_metadata_cli.py

@@ -0,0 +1,69 @@
+from metaflow._vendor import click
+from datetime import datetime, timezone
+from metaflow.tagging_util import validate_tags
+from metaflow.metadata_provider import MetaDatum
+
+
+@click.group()
+def cli():
+    pass
+
+
+@cli.group(help="Commands related to spot metadata.")
+def spot_metadata():
+    pass
+
+
+@spot_metadata.command(help="Record spot termination metadata for a task.")
+@click.option(
+    "--run-id",
+    required=True,
+    help="Run ID for which metadata is to be recorded.",
+)
+@click.option(
+    "--step-name",
+    required=True,
+    help="Step Name for which metadata is to be recorded.",
+)
+@click.option(
+    "--task-id",
+    required=True,
+    help="Task ID for which metadata is to be recorded.",
+)
+@click.option(
+    "--termination-notice-time",
+    required=True,
+    help="Spot termination notice time.",
+)
+@click.option(
+    "--tag",
+    "tags",
+    multiple=True,
+    required=False,
+    default=None,
+    help="List of tags.",
+)
+@click.pass_obj
+def record(obj, run_id, step_name, task_id, termination_notice_time, tags=None):
+    validate_tags(tags)
+
+    tag_list = list(tags) if tags else []
+
+    entries = [
+        MetaDatum(
+            field="spot-termination-received-at",
+            value=datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
+            type="spot-termination-received-at",
+            tags=tag_list,
+        ),
+        MetaDatum(
+            field="spot-termination-time",
+            value=termination_notice_time,
+            type="spot-termination-time",
+            tags=tag_list,
+        ),
+    ]
+
+    obj.metadata.register_metadata(
+        run_id=run_id, step_name=step_name, task_id=task_id, metadata=entries
+    )

metaflow/plugins/kubernetes/spot_monitor_sidecar.py

@@ -0,0 +1,109 @@
+import os
+import sys
+import time
+import signal
+import requests
+import subprocess
+from multiprocessing import Process
+from datetime import datetime, timezone
+from metaflow.sidecar import MessageTypes
+
+
+class SpotTerminationMonitorSidecar(object):
+    EC2_TYPE_URL = "http://169.254.169.254/latest/meta-data/instance-life-cycle"
+    METADATA_URL = "http://169.254.169.254/latest/meta-data/spot/termination-time"
+    TOKEN_URL = "http://169.254.169.254/latest/api/token"
+    POLL_INTERVAL = 5  # seconds
+
+    def __init__(self):
+        self.is_alive = True
+        self._process = None
+        self._token = None
+        self._token_expiry = 0
+
+        if self._is_aws_spot_instance():
+            self._process = Process(target=self._monitor_loop)
+            self._process.start()
+
+    def process_message(self, msg):
+        if msg.msg_type == MessageTypes.SHUTDOWN:
+            self.is_alive = False
+            if self._process:
+                self._process.terminate()
+
+    @classmethod
+    def get_worker(cls):
+        return cls
+
+    def _get_imds_token(self):
+        current_time = time.time()
+        if current_time >= self._token_expiry - 60:  # Refresh 60s before expiry
+            try:
+                response = requests.put(
+                    url=self.TOKEN_URL,
+                    headers={"X-aws-ec2-metadata-token-ttl-seconds": "300"},
+                    timeout=1,
+                )
+                if response.status_code == 200:
+                    self._token = response.text
+                    self._token_expiry = current_time + 240  # Slightly less than TTL
+            except requests.exceptions.RequestException:
+                pass
+        return self._token
+
+    def _make_ec2_request(self, url, timeout):
+        token = self._get_imds_token()
+        headers = {"X-aws-ec2-metadata-token": token} if token else {}
+        response = requests.get(url=url, headers=headers, timeout=timeout)
+        return response
+
+    def _is_aws_spot_instance(self):
+        try:
+            response = self._make_ec2_request(url=self.EC2_TYPE_URL, timeout=1)
+            return response.status_code == 200 and response.text == "spot"
+        except (requests.exceptions.RequestException, requests.exceptions.Timeout):
+            return False
+
+    def _monitor_loop(self):
+        while self.is_alive:
+            try:
+                response = self._make_ec2_request(url=self.METADATA_URL, timeout=1)
+                if response.status_code == 200:
+                    termination_time = response.text
+                    self._emit_termination_metadata(termination_time)
+                    os.kill(os.getppid(), signal.SIGTERM)
+                    break
+            except (requests.exceptions.RequestException, requests.exceptions.Timeout):
+                pass
+            time.sleep(self.POLL_INTERVAL)
+
+    def _emit_termination_metadata(self, termination_time):
+        flow_filename = os.getenv("METAFLOW_FLOW_FILENAME")
+        pathspec = os.getenv("MF_PATHSPEC")
+        _, run_id, step_name, task_id = pathspec.split("/")
+        retry_count = os.getenv("MF_ATTEMPT")
+
+        with open("/tmp/spot_termination_notice", "w") as fp:
+            fp.write(termination_time)
+
+        command = [
+            sys.executable,
+            f"/metaflow/{flow_filename}",
+            "spot-metadata",
+            "record",
+            "--run-id",
+            run_id,
+            "--step-name",
+            step_name,
+            "--task-id",
+            task_id,
+            "--termination-notice-time",
+            termination_time,
+            "--tag",
+            "attempt_id:{}".format(retry_count),
+        ]
+
+        result = subprocess.run(command, capture_output=True, text=True)
+
+        if result.returncode != 0:
+            print(f"Failed to record spot termination metadata: {result.stderr}")

metaflow/plugins/parallel_decorator.py

@@ -45,6 +45,8 @@ class ParallelDecorator(StepDecorator):
         if ubf_context == UBF_CONTROL:
             num_parallel = cli_args.task.ubf_iter.num_parallel
             cli_args.command_options["num-parallel"] = str(num_parallel)
+            if os.environ.get("METAFLOW_RUNTIME_ENVIRONMENT", "local") == "local":
+                cli_args.command_options["split_index"] = "0"
 
     def step_init(
         self, flow, graph, step_name, decorators, environment, flow_datastore, logger
@@ -126,6 +128,8 @@ class ParallelDecorator(StepDecorator):
                     tags=["attempt_id:{0}".format(0)],
                 )
             ]
+            flow._control_task_is_mapper_zero = True
+
         metadata.register_metadata(run_id, step_name, task_id, task_metadata_list)
 
     def task_decorate(
@@ -221,7 +225,6 @@ def _local_multinode_control_task_step_func(
         "%s/%s/%s" % (run_id, step_name, mapper_task_id)
         for mapper_task_id in mapper_task_ids
     ]
-    flow._control_task_is_mapper_zero = True
 
     # run the step function ourselves
     os.environ["MF_PARALLEL_NODE_INDEX"] = "0"

metaflow/plugins/project_decorator.py

@@ -72,7 +72,6 @@ class ProjectDecorator(FlowDecorator):
     """
 
     name = "project"
-    defaults = {"name": None}
 
     options = {
         "production": dict(
@@ -91,19 +90,48 @@ class ProjectDecorator(FlowDecorator):
         ),
     }
 
+    defaults = {"name": None, **{k: v["default"] for k, v in options.items()}}
+
     def flow_init(
         self, flow, graph, environment, flow_datastore, metadata, logger, echo, options
     ):
         self._option_values = options
         project_name = self.attributes.get("name")
+        for op in options:
+            if (
+                op in self._user_defined_attributes
+                and options[op] != self.defaults[op]
+                and self.attributes[op] != options[op]
+            ):
+                # Exception if:
+                #  - the user provides a value in the attributes field
+                #  - AND the user provided a value in the command line (non default)
+                #  - AND the values are different
+                # Note that this won't raise an error if the user provided the default
+                # value in the command line and provided one in attribute but although
+                # slightly inconsistent, it is not incorrect.
+                raise MetaflowException(
+                    "You cannot pass %s as both a command-line argument and an attribute "
+                    "of the @project decorator." % op
+                )
+        if "branch" in self._user_defined_attributes:
+            project_branch = self.attributes["branch"]
+        else:
+            project_branch = options["branch"]
+
+        if "production" in self._user_defined_attributes:
+            project_production = self.attributes["production"]
+        else:
+            project_production = options["production"]
+
         project_flow_name, branch_name = format_name(
             flow.name,
             project_name,
-            options["production"],
-            options["branch"],
+            project_production,
+            project_branch,
             get_username(),
         )
-        is_user_branch = options["branch"] is None and not options["production"]
+        is_user_branch = project_branch is None and not project_production
         echo(
             "Project: *%s*, Branch: *%s*" % (project_name, branch_name),
             fg="magenta",
@@ -114,7 +142,7 @@ class ProjectDecorator(FlowDecorator):
                 "project_name": project_name,
                 "branch_name": branch_name,
                 "is_user_branch": is_user_branch,
-                "is_production": options["production"],
+                "is_production": project_production,
                 "project_flow_name": project_flow_name,
             }
         )