metaflow 2.12.9__py2.py3-none-any.whl → 2.12.10__py2.py3-none-any.whl

metaflow/plugins/argo/argo_workflows.py

@@ -117,6 +117,7 @@ class ArgoWorkflows(object):
         notify_on_success=False,
         notify_slack_webhook_url=None,
         notify_pager_duty_integration_key=None,
+        enable_heartbeat_daemon=True,
     ):
         # Some high-level notes -
         #
@@ -164,6 +165,7 @@ class ArgoWorkflows(object):
         self.notify_on_success = notify_on_success
         self.notify_slack_webhook_url = notify_slack_webhook_url
         self.notify_pager_duty_integration_key = notify_pager_duty_integration_key
+        self.enable_heartbeat_daemon = enable_heartbeat_daemon
 
         self.parameters = self._process_parameters()
         self.triggers, self.trigger_options = self._process_triggers()
@@ -853,6 +855,8 @@ class ArgoWorkflows(object):
                 .templates(self._container_templates())
                 # Exit hook template(s)
                 .templates(self._exit_hook_templates())
+                # Sidecar templates (Daemon Containers)
+                .templates(self._daemon_templates())
             )
         )
 
@@ -1265,7 +1269,13 @@ class ArgoWorkflows(object):
                     "Argo Workflows." % (node.type, node.name)
                 )
 
-        templates, _ = _visit(node=self.graph["start"])
+        # Generate daemon tasks
+        daemon_tasks = [
+            DAGTask("%s-task" % daemon_template.name).template(daemon_template.name)
+            for daemon_template in self._daemon_templates()
+        ]
+
+        templates, _ = _visit(node=self.graph["start"], dag_tasks=daemon_tasks)
         return templates
 
     # Visit every node and yield ContainerTemplates.
@@ -2015,6 +2025,7 @@ class ArgoWorkflows(object):
                         kubernetes_sdk.V1Container(
                             name=self._sanitize(node.name),
                             command=cmds,
+                            termination_message_policy="FallbackToLogsOnError",
                             ports=[kubernetes_sdk.V1ContainerPort(container_port=port)]
                             if port
                             else None,
@@ -2067,9 +2078,11 @@ class ArgoWorkflows(object):
                                 for k in list(
                                     []
                                     if not resources.get("secrets")
-                                    else [resources.get("secrets")]
-                                    if isinstance(resources.get("secrets"), str)
-                                    else resources.get("secrets")
+                                    else (
+                                        [resources.get("secrets")]
+                                        if isinstance(resources.get("secrets"), str)
+                                        else resources.get("secrets")
+                                    )
                                 )
                                 + KUBERNETES_SECRETS.split(",")
                                 + ARGO_WORKFLOWS_KUBERNETES_SECRETS.split(",")
@@ -2121,6 +2134,13 @@ class ArgoWorkflows(object):
                 )
             )
 
+    # Return daemon container templates for workflow execution notifications.
+    def _daemon_templates(self):
+        templates = []
+        if self.enable_heartbeat_daemon:
+            templates.append(self._heartbeat_daemon_template())
+        return templates
+
     # Return exit hook templates for workflow execution notifications.
     def _exit_hook_templates(self):
         templates = []
@@ -2327,6 +2347,117 @@ class ArgoWorkflows(object):
             Http("POST").url(self.notify_slack_webhook_url).body(json.dumps(payload))
         )
 
+    def _heartbeat_daemon_template(self):
+        # Use all the affordances available to _parameters task
+        executable = self.environment.executable("_parameters")
+        run_id = "argo-{{workflow.name}}"
+        entrypoint = [executable, "-m metaflow.plugins.argo.daemon"]
+        heartbeat_cmds = "{entrypoint} --flow_name {flow_name} --run_id {run_id} {tags} heartbeat".format(
+            entrypoint=" ".join(entrypoint),
+            flow_name=self.flow.name,
+            run_id=run_id,
+            tags=" ".join(["--tag %s" % t for t in self.tags]) if self.tags else "",
+        )
+
+        # TODO: we do not really need MFLOG logging for the daemon at the moment, but might be good for the future.
+        # Consider if we can do without this setup.
+        # Configure log capture.
+        mflog_expr = export_mflog_env_vars(
+            datastore_type=self.flow_datastore.TYPE,
+            stdout_path="$PWD/.logs/mflog_stdout",
+            stderr_path="$PWD/.logs/mflog_stderr",
+            flow_name=self.flow.name,
+            run_id=run_id,
+            step_name="_run_heartbeat_daemon",
+            task_id="1",
+            retry_count="0",
+        )
+        # TODO: Can the init be trimmed down?
+        # Can we do without get_package_commands fetching the whole code package?
+        init_cmds = " && ".join(
+            [
+                # For supporting sandboxes, ensure that a custom script is executed
+                # before anything else is executed. The script is passed in as an
+                # env var.
+                '${METAFLOW_INIT_SCRIPT:+eval \\"${METAFLOW_INIT_SCRIPT}\\"}',
+                "mkdir -p $PWD/.logs",
+                mflog_expr,
+            ]
+            + self.environment.get_package_commands(
+                self.code_package_url, self.flow_datastore.TYPE
+            )[:-1]
+            # Replace the line 'Task in starting'
+            # FIXME: this can be brittle.
+            + ["mflog 'Heartbeat daemon is starting.'"]
+        )
+
+        cmd_str = " && ".join([init_cmds, heartbeat_cmds])
+        cmds = shlex.split('bash -c "%s"' % cmd_str)
+
+        # TODO: Check that this is the minimal env.
+        # Env required for sending heartbeats to the metadata service, nothing extra.
+        env = {
+            # These values are needed by Metaflow to set it's internal
+            # state appropriately.
+            "METAFLOW_CODE_URL": self.code_package_url,
+            "METAFLOW_CODE_SHA": self.code_package_sha,
+            "METAFLOW_CODE_DS": self.flow_datastore.TYPE,
+            "METAFLOW_SERVICE_URL": SERVICE_INTERNAL_URL,
+            "METAFLOW_SERVICE_HEADERS": json.dumps(SERVICE_HEADERS),
+            "METAFLOW_USER": "argo-workflows",
+            "METAFLOW_DEFAULT_DATASTORE": self.flow_datastore.TYPE,
+            "METAFLOW_DEFAULT_METADATA": DEFAULT_METADATA,
+            "METAFLOW_OWNER": self.username,
+        }
+        # support Metaflow sandboxes
+        env["METAFLOW_INIT_SCRIPT"] = KUBERNETES_SANDBOX_INIT_SCRIPT
+
+        # cleanup env values
+        env = {
+            k: v
+            for k, v in env.items()
+            if v is not None
+            and k not in set(ARGO_WORKFLOWS_ENV_VARS_TO_SKIP.split(","))
+        }
+
+        # We want to grab the base image used by the start step, as this is known to be pullable from within the cluster,
+        # and it might contain the required libraries, allowing us to start up faster.
+        start_step = next(step for step in self.flow if step.name == "start")
+        resources = dict(
+            [deco for deco in start_step.decorators if deco.name == "kubernetes"][
+                0
+            ].attributes
+        )
+        from kubernetes import client as kubernetes_sdk
+
+        return DaemonTemplate("heartbeat-daemon").container(
+            to_camelcase(
+                kubernetes_sdk.V1Container(
+                    name="main",
+                    # TODO: Make the image configurable
+                    image=resources["image"],
+                    command=cmds,
+                    env=[
+                        kubernetes_sdk.V1EnvVar(name=k, value=str(v))
+                        for k, v in env.items()
+                    ],
+                    resources=kubernetes_sdk.V1ResourceRequirements(
+                        # NOTE: base resources for this are kept to a minimum to save on running costs.
+                        # This has an adverse effect on startup time for the daemon, which can be completely
+                        # alleviated by using a base image that has the required dependencies pre-installed
+                        requests={
+                            "cpu": "200m",
+                            "memory": "100Mi",
+                        },
+                        limits={
+                            "cpu": "200m",
+                            "memory": "100Mi",
+                        },
+                    ),
+                )
+            )
+        )
+
     def _compile_sensor(self):
         # This method compiles a Metaflow @trigger decorator into Argo Events Sensor.
         #
@@ -2899,6 +3030,25 @@ class Metadata(object):
         return json.dumps(self.to_json(), indent=4)
 
 
+class DaemonTemplate(object):
+    def __init__(self, name):
+        tree = lambda: defaultdict(tree)
+        self.name = name
+        self.payload = tree()
+        self.payload["daemon"] = True
+        self.payload["name"] = name
+
+    def container(self, container):
+        self.payload["container"] = container
+        return self
+
+    def to_json(self):
+        return self.payload
+
+    def __str__(self):
+        return json.dumps(self.payload, indent=4)
+
+
 class Template(object):
     # https://argoproj.github.io/argo-workflows/fields/#template
 

metaflow/plugins/argo/argo_workflows_cli.py

@@ -167,6 +167,12 @@ def argo_workflows(obj, name=None):
     default="",
     help="PagerDuty Events API V2 Integration key for workflow success/failure notifications.",
 )
+@click.option(
+    "--enable-heartbeat-daemon/--no-enable-heartbeat-daemon",
+    default=False,
+    show_default=True,
+    help="Use a daemon container to broadcast heartbeats.",
+)
 @click.option(
     "--deployer-attribute-file",
     default=None,
@@ -192,6 +198,7 @@ def create(
     notify_on_success=False,
     notify_slack_webhook_url=None,
     notify_pager_duty_integration_key=None,
+    enable_heartbeat_daemon=True,
     deployer_attribute_file=None,
 ):
     validate_tags(tags)
@@ -240,6 +247,7 @@ def create(
         notify_on_success,
         notify_slack_webhook_url,
         notify_pager_duty_integration_key,
+        enable_heartbeat_daemon,
     )
 
     if only_json:
@@ -412,6 +420,7 @@ def make_flow(
     notify_on_success,
     notify_slack_webhook_url,
     notify_pager_duty_integration_key,
+    enable_heartbeat_daemon,
 ):
     # TODO: Make this check less specific to Amazon S3 as we introduce
     #       support for more cloud object stores.
@@ -474,6 +483,7 @@ def make_flow(
         notify_on_success=notify_on_success,
         notify_slack_webhook_url=notify_slack_webhook_url,
         notify_pager_duty_integration_key=notify_pager_duty_integration_key,
+        enable_heartbeat_daemon=enable_heartbeat_daemon,
     )
 
 

metaflow/plugins/argo/daemon.py

@@ -0,0 +1,59 @@
+from collections import namedtuple
+from time import sleep
+from metaflow.metaflow_config import DEFAULT_METADATA
+from metaflow.metaflow_environment import MetaflowEnvironment
+from metaflow.plugins import METADATA_PROVIDERS
+from metaflow._vendor import click
+
+
+class CliState:
+    pass
+
+
+@click.group()
+@click.option("--flow_name", required=True)
+@click.option("--run_id", required=True)
+@click.option(
+    "--tag",
+    "tags",
+    multiple=True,
+    default=None,
+    help="Annotate all objects produced by Argo Workflows runs "
+    "with the given tag. You can specify this option multiple "
+    "times to attach multiple tags.",
+)
+@click.pass_context
+def cli(ctx, flow_name, run_id, tags=None):
+    ctx.obj = CliState()
+    ctx.obj.flow_name = flow_name
+    ctx.obj.run_id = run_id
+    ctx.obj.tags = tags
+    # Use a dummy flow to initialize the environment and metadata service,
+    # as we only need a name for the flow object.
+    flow = namedtuple("DummyFlow", "name")
+    dummyflow = flow(flow_name)
+
+    # Initialize a proper metadata service instance
+    environment = MetaflowEnvironment(dummyflow)
+
+    ctx.obj.metadata = [m for m in METADATA_PROVIDERS if m.TYPE == DEFAULT_METADATA][0](
+        environment, dummyflow, None, None
+    )
+
+
+@cli.command(help="start heartbeat process for a run")
+@click.pass_obj
+def heartbeat(obj):
+    # Try to register a run in case the start task has not taken care of it yet.
+    obj.metadata.register_run_id(obj.run_id, obj.tags)
+    # Start run heartbeat
+    obj.metadata.start_run_heartbeat(obj.flow_name, obj.run_id)
+    # Keepalive loop
+    while True:
+        # Do not pollute daemon logs with anything unnecessary,
+        # as they might be extremely long running.
+        sleep(10)
+
+
+if __name__ == "__main__":
+    cli()

metaflow/plugins/kubernetes/kubernetes_job.py

@@ -98,6 +98,7 @@ class KubernetesJob(object):
                     containers=[
                         client.V1Container(
                             command=self._kwargs["command"],
+                            termination_message_policy="FallbackToLogsOnError",
                             ports=[]
                             if self._kwargs["port"] is None
                             else [

metaflow/plugins/kubernetes/kubernetes_jobsets.py

@@ -586,6 +586,7 @@ class JobSetSpec(object):
                                 containers=[
                                     client.V1Container(
                                         command=self._kwargs["command"],
+                                        termination_message_policy="FallbackToLogsOnError",
                                         ports=[]
                                         if self._kwargs["port"] is None
                                         else [

metaflow/version.py

@@ -1 +1 @@
-metaflow_version = "2.12.9"
+metaflow_version = "2.12.10"

{metaflow-2.12.9.dist-info → metaflow-2.12.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: metaflow
-Version: 2.12.9
+Version: 2.12.10
 Summary: Metaflow: More Data Science, Less Engineering
 Author: Metaflow Developers
 Author-email: help@metaflow.org
@@ -26,7 +26,7 @@ License-File: LICENSE
 Requires-Dist: requests
 Requires-Dist: boto3
 Provides-Extra: stubs
-Requires-Dist: metaflow-stubs ==2.12.9 ; extra == 'stubs'
+Requires-Dist: metaflow-stubs ==2.12.10 ; extra == 'stubs'
 
 ![Metaflow_Logo_Horizontal_FullColor_Ribbon_Dark_RGB](https://user-images.githubusercontent.com/763451/89453116-96a57e00-d713-11ea-9fa6-82b29d4d6eff.png)
 

{metaflow-2.12.9.dist-info → metaflow-2.12.10.dist-info}/RECORD

@@ -35,7 +35,7 @@ metaflow/tuple_util.py,sha256=_G5YIEhuugwJ_f6rrZoelMFak3DqAR2tt_5CapS1XTY,830
 metaflow/unbounded_foreach.py,sha256=p184WMbrMJ3xKYHwewj27ZhRUsSj_kw1jlye5gA9xJk,387
 metaflow/util.py,sha256=m5womQ7y-jXehuMyHPfByDbZ4HwTJxzs869cPOlMR8s,13057
 metaflow/vendor.py,sha256=FchtA9tH22JM-eEtJ2c9FpUdMn8sSb1VHuQS56EcdZk,5139
-metaflow/version.py,sha256=TkQOTBK44Rf1jdRra0slObbAYISL1VgE5Ar3nFR7PC0,28
+metaflow/version.py,sha256=t4RI_U0iAZBohPGuAobz1uO1iHJdG0OWNq9m857GkRw,29
 metaflow/_vendor/__init__.py,sha256=y_CiwUD3l4eAKvTVDZeqgVujMy31cAM1qjAB-HfI-9s,353
 metaflow/_vendor/typing_extensions.py,sha256=0nUs5p1A_UrZigrAVBoOEM6TxU37zzPDUtiij1ZwpNc,110417
 metaflow/_vendor/zipp.py,sha256=ajztOH-9I7KA_4wqDYygtHa6xUBVZgFpmZ8FE74HHHI,8425
@@ -174,10 +174,11 @@ metaflow/plugins/airflow/sensors/s3_sensor.py,sha256=iDReG-7FKnumrtQg-HY6cCUAAqN
 metaflow/plugins/argo/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow/plugins/argo/argo_client.py,sha256=MKKhMCbWOPzf6z5zQQiyDRHHkAXcO7ipboDZDqAAvOk,15849
 metaflow/plugins/argo/argo_events.py,sha256=_C1KWztVqgi3zuH57pInaE9OzABc2NnncC-zdwOMZ-w,5909
-metaflow/plugins/argo/argo_workflows.py,sha256=sWS-uG1csYkUGEv1LMex8tWhoF0GzIoQdXRdusbIJLw,154446
-metaflow/plugins/argo/argo_workflows_cli.py,sha256=FnjnXPXas-ANfZ5m9s4YCHjejrkF7SbOS87J7z6339c,35036
+metaflow/plugins/argo/argo_workflows.py,sha256=a5el_pkK6B-XymAbFyxBdPjuy8OvtijHuSvWYLqe8iA,160728
+metaflow/plugins/argo/argo_workflows_cli.py,sha256=N_QMeFypWLFBtBzQ68WadCzOqdaA1FQQ4HHw5lizImE,35370
 metaflow/plugins/argo/argo_workflows_decorator.py,sha256=A23cRd8IBDUFJRmtQ0U8kik4gk0lg5s8UKrSO7obSEM,8475
 metaflow/plugins/argo/argo_workflows_deployer.py,sha256=yMIXAVoAuBLHCqQyFriV_Wc_Lp5D041Ay83R5pYNoXE,8066
+metaflow/plugins/argo/daemon.py,sha256=dJOS_UUISXBYffi3oGVKPwq4Pa4P_nGBGL15piPaPto,1776
 metaflow/plugins/argo/generate_input_paths.py,sha256=loYsI6RFX9LlFsHb7Fe-mzlTTtRdySoOu7sYDy-uXK0,881
 metaflow/plugins/argo/jobset_input_paths.py,sha256=_JhZWngA6p9Q_O2fx3pdzKI0WE-HPRHz_zFvY2pHPTQ,525
 metaflow/plugins/aws/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -281,8 +282,8 @@ metaflow/plugins/kubernetes/kubernetes.py,sha256=0UjKZy5_5KOGYC-k0DZqpv-ca_tZBkw
 metaflow/plugins/kubernetes/kubernetes_cli.py,sha256=qBDdr1Lvtt-RO9pB-9_HTOPdzAmDvvJ0aiQ1OoCcrMU,10892
 metaflow/plugins/kubernetes/kubernetes_client.py,sha256=GKg-gT3qhXMRQV-sG1YyoOf3Z32NXr_wwEN2ytMVSEg,2471
 metaflow/plugins/kubernetes/kubernetes_decorator.py,sha256=DDc6N1Q0Cmcl44U-uZVBVr0tS4yL9N4GcAnf4eeX3Bk,24623
-metaflow/plugins/kubernetes/kubernetes_job.py,sha256=1uQmEHsLyJ5jewmIA5smrFtSO8MymOqPDb-YIS65u50,31614
-metaflow/plugins/kubernetes/kubernetes_jobsets.py,sha256=Cr5M7m0lKShv6o8xjRJap8hSRgf51fz4RR0THlCSM08,40780
+metaflow/plugins/kubernetes/kubernetes_job.py,sha256=QYmb2VIEniJa0xZlRBofyoxpuhQjbpDbSjBdB0X6N70,31694
+metaflow/plugins/kubernetes/kubernetes_jobsets.py,sha256=60331lf8UE95zJrkM0jtTV9vxJvre2SU63j-_i2e-8I,40872
 metaflow/plugins/metadata/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 metaflow/plugins/metadata/local.py,sha256=YhLJC5zjVJrvQFIyQ92ZBByiUmhCC762RUX7ITX12O8,22428
 metaflow/plugins/metadata/service.py,sha256=ihq5F7KQZlxvYwzH_-jyP2aWN_I96i2vp92j_d697s8,20204
@@ -342,9 +343,9 @@ metaflow/tutorials/07-worldview/README.md,sha256=5vQTrFqulJ7rWN6r20dhot9lI2sVj9W
 metaflow/tutorials/07-worldview/worldview.ipynb,sha256=ztPZPI9BXxvW1QdS2Tfe7LBuVzvFvv0AToDnsDJhLdE,2237
 metaflow/tutorials/08-autopilot/README.md,sha256=GnePFp_q76jPs991lMUqfIIh5zSorIeWznyiUxzeUVE,1039
 metaflow/tutorials/08-autopilot/autopilot.ipynb,sha256=DQoJlILV7Mq9vfPBGW-QV_kNhWPjS5n6SJLqePjFYLY,3191
-metaflow-2.12.9.dist-info/LICENSE,sha256=nl_Lt5v9VvJ-5lWJDT4ddKAG-VZ-2IaLmbzpgYDz2hU,11343
-metaflow-2.12.9.dist-info/METADATA,sha256=VVvmqX24fHz8PbJ483_p9Tq7i42DuZkfQFQ1T96alYo,5906
-metaflow-2.12.9.dist-info/WHEEL,sha256=ANi2y8tYx-p70pj7MSdqQMJNFJEUgAOyfPpHzqz0w84,109
-metaflow-2.12.9.dist-info/entry_points.txt,sha256=IKwTN1T3I5eJL3uo_vnkyxVffcgnRdFbKwlghZfn27k,57
-metaflow-2.12.9.dist-info/top_level.txt,sha256=v1pDHoWaSaKeuc5fKTRSfsXCKSdW1zvNVmvA-i0if3o,9
-metaflow-2.12.9.dist-info/RECORD,,
+metaflow-2.12.10.dist-info/LICENSE,sha256=nl_Lt5v9VvJ-5lWJDT4ddKAG-VZ-2IaLmbzpgYDz2hU,11343
+metaflow-2.12.10.dist-info/METADATA,sha256=KgfU7CsRqUnD_STSlhUclySsag0ieBOT8qH_vbH5ISE,5908
+metaflow-2.12.10.dist-info/WHEEL,sha256=XRxW4r1PNiVhMpP4bT9oWtu3HyndxpJ84SkubFgzp_Y,109
+metaflow-2.12.10.dist-info/entry_points.txt,sha256=IKwTN1T3I5eJL3uo_vnkyxVffcgnRdFbKwlghZfn27k,57
+metaflow-2.12.10.dist-info/top_level.txt,sha256=v1pDHoWaSaKeuc5fKTRSfsXCKSdW1zvNVmvA-i0if3o,9
+metaflow-2.12.10.dist-info/RECORD,,

{metaflow-2.12.9.dist-info → metaflow-2.12.10.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (71.1.0)
+Generator: setuptools (72.1.0)
 Root-Is-Purelib: true
 Tag: py2-none-any
 Tag: py3-none-any