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

metaflow/plugins/cards/card_modules/test_cards.py

@@ -38,7 +38,7 @@ class TestEditableCard(MetaflowCard):
 
     ALLOW_USER_COMPONENTS = True
 
-    def __init__(self, options={}, components=[], graph=None):
+    def __init__(self, components=[], **kwargs):
         self._components = components
 
     def render(self, task):
@@ -52,7 +52,7 @@ class TestEditableCard2(MetaflowCard):
 
     ALLOW_USER_COMPONENTS = True
 
-    def __init__(self, options={}, components=[], graph=None):
+    def __init__(self, components=[], **kwargs):
         self._components = components
 
     def render(self, task):
@@ -64,7 +64,7 @@ class TestNonEditableCard(MetaflowCard):
 
     seperator = "$&#!!@*"
 
-    def __init__(self, options={}, components=[], graph=None):
+    def __init__(self, components=[], **kwargs):
         self._components = components
 
     def render(self, task):
@@ -193,7 +193,7 @@ class TestRefreshComponentCard(MetaflowCard):
 
     type = "test_component_refresh_card"
 
-    def __init__(self, options={}, components=[], graph=None):
+    def __init__(self, components=[], **kwargs):
         self._components = components
 
     def render(self, task) -> str:

metaflow/plugins/cards/component_serializer.py

@@ -1,8 +1,8 @@
 from .card_modules import MetaflowCardComponent
+from .card_modules.card import create_component_id
 from .card_modules.basic import ErrorComponent, SectionComponent
 from .card_modules.components import (
     UserComponent,
-    create_component_id,
     StubComponent,
 )
 from .exception import ComponentOverwriteNotSupportedException

metaflow/plugins/events_decorator.py

@@ -362,11 +362,7 @@ class TriggerOnFinishDecorator(FlowDecorator):
     """
 
     name = "trigger_on_finish"
-    defaults = {
-        "flow": None,  # flow_name or project_flow_name
-        "flows": [],  # flow_names or project_flow_names
-        "options": {},
-    }
+
     options = {
         "trigger": dict(
             multiple=True,
@@ -374,6 +370,14 @@ class TriggerOnFinishDecorator(FlowDecorator):
             help="Specify run pathspec for testing @trigger_on_finish locally.",
         ),
     }
+    defaults = {
+        "flow": None,  # flow_name or project_flow_name
+        "flows": [],  # flow_names or project_flow_names
+        "options": {},
+        # Re-enable if you want to support TL options directly in the decorator like
+        # for @project decorator
+        #    **{k: v["default"] for k, v in options.items()},
+    }
 
     def flow_init(
         self,
@@ -394,111 +398,23 @@ class TriggerOnFinishDecorator(FlowDecorator):
             )
         elif self.attributes["flow"]:
             # flow supports the format @trigger_on_finish(flow='FooFlow')
-            if is_stringish(self.attributes["flow"]):
-                self.triggers.append(
-                    {
-                        "fq_name": self.attributes["flow"],
-                    }
-                )
-            elif isinstance(self.attributes["flow"], dict):
-                if "name" not in self.attributes["flow"]:
-                    raise MetaflowException(
-                        "The *flow* attribute for *@trigger_on_finish* is missing the "
-                        "*name* key."
-                    )
-                flow_name = self.attributes["flow"]["name"]
-
-                if not is_stringish(flow_name) or "." in flow_name:
-                    raise MetaflowException(
-                        "The *name* attribute of the *flow* is not a valid string"
-                    )
-                result = {"fq_name": flow_name}
-                if "project" in self.attributes["flow"]:
-                    if is_stringish(self.attributes["flow"]["project"]):
-                        result["project"] = self.attributes["flow"]["project"]
-                    else:
-                        raise MetaflowException(
-                            "The *project* attribute of the *flow* is not a string"
-                        )
-                if "project_branch" in self.attributes["flow"]:
-                    if is_stringish(self.attributes["flow"]["project_branch"]):
-                        result["branch"] = self.attributes["flow"]["project_branch"]
-                    else:
-                        raise MetaflowException(
-                            "The *project_branch* attribute of the *flow* is not a string"
-                        )
-                self.triggers.append(result)
-            elif callable(self.attributes["flow"]) and not isinstance(
+            flow = self.attributes["flow"]
+            if callable(flow) and not isinstance(
                 self.attributes["flow"], DeployTimeField
             ):
-                trig = DeployTimeField(
-                    "fq_name", [str, dict], None, self.attributes["flow"], False
-                )
+                trig = DeployTimeField("fq_name", [str, dict], None, flow, False)
                 self.triggers.append(trig)
             else:
-                raise MetaflowException(
-                    "Incorrect type for *flow* attribute in *@trigger_on_finish* "
-                    " decorator. Supported type is string or Dict[str, str] - \n"
-                    "@trigger_on_finish(flow='FooFlow') or "
-                    "@trigger_on_finish(flow={'name':'FooFlow', 'project_branch': 'branch'})"
-                )
+                self.triggers.extend(self._parse_static_triggers([flow]))
         elif self.attributes["flows"]:
             # flows attribute supports the following formats -
             #     1. flows=['FooFlow', 'BarFlow']
-            if isinstance(self.attributes["flows"], list):
-                for flow in self.attributes["flows"]:
-                    if is_stringish(flow):
-                        self.triggers.append(
-                            {
-                                "fq_name": flow,
-                            }
-                        )
-                    elif isinstance(flow, dict):
-                        if "name" not in flow:
-                            raise MetaflowException(
-                                "One or more flows in the *flows* attribute for "
-                                "*@trigger_on_finish* is missing the "
-                                "*name* key."
-                            )
-                        flow_name = flow["name"]
-
-                        if not is_stringish(flow_name) or "." in flow_name:
-                            raise MetaflowException(
-                                "The *name* attribute '%s' is not a valid string"
-                                % str(flow_name)
-                            )
-                        result = {"fq_name": flow_name}
-                        if "project" in flow:
-                            if is_stringish(flow["project"]):
-                                result["project"] = flow["project"]
-                            else:
-                                raise MetaflowException(
-                                    "The *project* attribute of the *flow* '%s' is not "
-                                    "a string" % flow_name
-                                )
-                        if "project_branch" in flow:
-                            if is_stringish(flow["project_branch"]):
-                                result["branch"] = flow["project_branch"]
-                            else:
-                                raise MetaflowException(
-                                    "The *project_branch* attribute of the *flow* %s "
-                                    "is not a string" % flow_name
-                                )
-                        self.triggers.append(result)
-                    else:
-                        raise MetaflowException(
-                            "One or more flows in *flows* attribute in "
-                            "*@trigger_on_finish* decorator have an incorrect type. "
-                            "Supported type is string or Dict[str, str]- \n"
-                            "@trigger_on_finish(flows=['FooFlow', 'BarFlow']"
-                        )
-            elif callable(self.attributes["flows"]) and not isinstance(
-                self.attributes["flows"], DeployTimeField
-            ):
-                trig = DeployTimeField(
-                    "flows", list, None, self.attributes["flows"], False
-                )
+            flows = self.attributes["flows"]
+            if callable(flows) and not isinstance(flows, DeployTimeField):
+                trig = DeployTimeField("flows", list, None, flows, False)
                 self.triggers.append(trig)
+            elif isinstance(flows, list):
+                self.triggers.extend(self._parse_static_triggers(flows))
             else:
                 raise MetaflowException(
                     "Incorrect type for *flows* attribute in *@trigger_on_finish* "
@@ -515,37 +431,48 @@ class TriggerOnFinishDecorator(FlowDecorator):
         for trigger in self.triggers:
             if isinstance(trigger, DeployTimeField):
                 continue
-            if trigger["fq_name"].count(".") == 0:
-                # fully qualified name is just the flow name
-                trigger["flow"] = trigger["fq_name"]
-            elif trigger["fq_name"].count(".") >= 2:
-                # fully qualified name is of the format - project.branch.flow_name
-                trigger["project"], tail = trigger["fq_name"].split(".", maxsplit=1)
-                trigger["branch"], trigger["flow"] = tail.rsplit(".", maxsplit=1)
-            else:
-                raise MetaflowException(
-                    "Incorrect format for *flow* in *@trigger_on_finish* "
-                    "decorator. Specify either just the *flow_name* or a fully "
-                    "qualified name like *project_name.branch_name.flow_name*."
-                )
-            # TODO: Also sanity check project and branch names
-            if not re.match(r"^[A-Za-z0-9_]+$", trigger["flow"]):
-                raise MetaflowException(
-                    "Invalid flow name *%s* in *@trigger_on_finish* "
-                    "decorator. Only alphanumeric characters and "
-                    "underscores(_) are allowed." % trigger["flow"]
-                )
+            self._parse_fq_name(trigger)
 
         self.options = self.attributes["options"]
 
         # Handle scenario for local testing using --trigger.
+
+        # Re-enable this code if you want to support passing trigger directly in the
+        # decorator in a way similar to how production and branch are passed in the
+        # project decorator.
+
+        # # This is overkill since default is None for all options but adding this code
+        # # to make it safe if other non None-default options are added in the future.
+        # 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 @trigger_on_finish decorator." % op
+        #         )
+
+        # if "trigger" in self._user_defined_attributes:
+        #    trigger_option = self.attributes["trigger"]
+        # else:
+        trigger_option = options["trigger"]
+
         self._option_values = options
-        if options["trigger"]:
+        if trigger_option:
             from metaflow import Run
             from metaflow.events import Trigger
 
             run_objs = []
-            for run_pathspec in options["trigger"]:
+            for run_pathspec in trigger_option:
                 if len(run_pathspec.split("/")) != 2:
                     raise MetaflowException(
                         "Incorrect format for run pathspec for *--trigger*. "
@@ -559,9 +486,67 @@ class TriggerOnFinishDecorator(FlowDecorator):
                 run_objs.append(run_obj)
             current._update_env({"trigger": Trigger.from_runs(run_objs)})
 
+    @staticmethod
+    def _parse_static_triggers(flows):
+        results = []
+        for flow in flows:
+            if is_stringish(flow):
+                results.append(
+                    {
+                        "fq_name": flow,
+                    }
+                )
+            elif isinstance(flow, dict):
+                if "name" not in flow:
+                    if len(flows) > 1:
+                        raise MetaflowException(
+                            "One or more flows in the *flows* attribute for "
+                            "*@trigger_on_finish* is missing the "
+                            "*name* key."
+                        )
+                    raise MetaflowException(
+                        "The *flow* attribute for *@trigger_on_finish* is missing the "
+                        "*name* key."
+                    )
+                flow_name = flow["name"]
+
+                if not is_stringish(flow_name) or "." in flow_name:
+                    raise MetaflowException(
+                        f"The *name* attribute of the *flow* {flow_name} is not a valid string"
+                    )
+                result = {"fq_name": flow_name}
+                if "project" in flow:
+                    if is_stringish(flow["project"]):
+                        result["project"] = flow["project"]
+                    else:
+                        raise MetaflowException(
+                            f"The *project* attribute of the *flow* {flow_name} is not a string"
+                        )
+                if "project_branch" in flow:
+                    if is_stringish(flow["project_branch"]):
+                        result["branch"] = flow["project_branch"]
+                    else:
+                        raise MetaflowException(
+                            f"The *project_branch* attribute of the *flow* {flow_name} is not a string"
+                        )
+                results.append(result)
+            else:
+                if len(flows) > 1:
+                    raise MetaflowException(
+                        "One or more flows in the *flows* attribute for "
+                        "*@trigger_on_finish* decorator have an incorrect type. "
+                        "Supported type is string or Dict[str, str]- \n"
+                        "@trigger_on_finish(flows=['FooFlow', 'BarFlow']"
+                    )
+                raise MetaflowException(
+                    "Incorrect type for *flow* attribute in *@trigger_on_finish* "
+                    " decorator. Supported type is string or Dict[str, str] - \n"
+                    "@trigger_on_finish(flow='FooFlow') or "
+                    "@trigger_on_finish(flow={'name':'FooFlow', 'project_branch': 'branch'})"
+                )
+        return results
+
     def _parse_fq_name(self, trigger):
-        if isinstance(trigger, DeployTimeField):
-            trigger["fq_name"] = deploy_time_eval(trigger["fq_name"])
         if trigger["fq_name"].count(".") == 0:
             # fully qualified name is just the flow name
             trigger["flow"] = trigger["fq_name"]
@@ -581,32 +566,18 @@ class TriggerOnFinishDecorator(FlowDecorator):
                 "decorator. Only alphanumeric characters and "
                 "underscores(_) are allowed." % trigger["flow"]
             )
-        return trigger
 
     def format_deploytime_value(self):
-        for trigger in self.triggers:
-            # Case were trigger is a function that returns a list
-            # Need to do this bc we need to iterate over list and process
-            if isinstance(trigger, DeployTimeField):
-                deploy_value = deploy_time_eval(trigger)
-                if isinstance(deploy_value, list):
-                    self.triggers = deploy_value
+        if len(self.triggers) == 1 and isinstance(self.triggers[0], DeployTimeField):
+            deploy_value = deploy_time_eval(self.triggers[0])
+            if isinstance(deploy_value, list):
+                self.triggers = deploy_value
             else:
-                break
-        for trigger in self.triggers:
-            # Entire trigger is a function (returns either string or dict)
-            old_trig = trigger
-            if isinstance(trigger, DeployTimeField):
-                trigger = deploy_time_eval(trigger)
-            if isinstance(trigger, dict):
-                trigger["fq_name"] = trigger.get("name")
-                trigger["project"] = trigger.get("project")
-                trigger["branch"] = trigger.get("project_branch")
-            # We also added this bc it won't be formatted yet
-            if isinstance(trigger, str):
-                trigger = {"fq_name": trigger}
-                trigger = self._parse_fq_name(trigger)
-            self.triggers[self.triggers.index(old_trig)] = trigger
+                self.triggers = [deploy_value]
+            triggers = self._parse_static_triggers(self.triggers)
+            for trigger in triggers:
+                self._parse_fq_name(trigger)
+            self.triggers = triggers
 
     def get_top_level_options(self):
         return list(self._option_values.items())

metaflow/plugins/kubernetes/kubernetes.py

@@ -685,15 +685,6 @@ class Kubernetes(object):
         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):

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -190,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"]

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -562,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
@@ -620,6 +627,7 @@ class KubernetesDecorator(StepDecorator):
 
         try:
             self._save_logs_sidecar.terminate()
+            self._spot_monitor_sidecar.terminate()
         except:
             # Best effort kill
             pass

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/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,
             }
         )