metaflow 2.13.4__py2.py3-none-any.whl → 2.13.5__py2.py3-none-any.whl

metaflow/plugins/__init__.py

@@ -16,6 +16,7 @@ CLIS_DESC = [
     ("argo-workflows", ".argo.argo_workflows_cli.cli"),
     ("card", ".cards.card_cli.cli"),
     ("tag", ".tag_cli.cli"),
+    ("spot-metadata", ".kubernetes.spot_metadata_cli.cli"),
     ("logs", ".logs_cli.cli"),
 ]
 
@@ -104,6 +105,10 @@ SIDECARS_DESC = [
         "save_logs_periodically",
         "..mflog.save_logs_periodically.SaveLogsPeriodicallySidecar",
     ),
+    (
+        "spot_termination_monitor",
+        ".kubernetes.spot_monitor_sidecar.SpotTerminationMonitorSidecar",
+    ),
     ("heartbeat", "metaflow.metadata_provider.heartbeat.MetadataHeartBeat"),
 ]
 

metaflow/plugins/argo/argo_workflows.py

@@ -1705,6 +1705,7 @@ class ArgoWorkflows(object):
                     },
                     **{
                         # Some optional values for bookkeeping
+                        "METAFLOW_FLOW_FILENAME": os.path.basename(sys.argv[0]),
                         "METAFLOW_FLOW_NAME": self.flow.name,
                         "METAFLOW_STEP_NAME": node.name,
                         "METAFLOW_RUN_ID": run_id,

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

@@ -547,6 +547,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
@@ -605,6 +612,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/version.py

@@ -1 +1 @@
-metaflow_version = "2.13.4"
+metaflow_version = "2.13.5"

{metaflow-2.13.4.dist-info → metaflow-2.13.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: metaflow
-Version: 2.13.4
+Version: 2.13.5
 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.13.4; extra == "stubs"
+Requires-Dist: metaflow-stubs==2.13.5; extra == "stubs"
 Dynamic: author
 Dynamic: author-email
 Dynamic: classifier

{metaflow-2.13.4.dist-info → metaflow-2.13.5.dist-info}/RECORD

@@ -36,7 +36,7 @@ metaflow/tuple_util.py,sha256=_G5YIEhuugwJ_f6rrZoelMFak3DqAR2tt_5CapS1XTY,830
 metaflow/unbounded_foreach.py,sha256=p184WMbrMJ3xKYHwewj27ZhRUsSj_kw1jlye5gA9xJk,387
 metaflow/util.py,sha256=hKjHl6NYJkKBSU2tzdVbddfOX1zWK73T4GCO42A0XB4,14666
 metaflow/vendor.py,sha256=FchtA9tH22JM-eEtJ2c9FpUdMn8sSb1VHuQS56EcdZk,5139
-metaflow/version.py,sha256=drr-g04woqZf71pNYGX7pz8CSPausyn9M8cgyFIvKlE,28
+metaflow/version.py,sha256=3DskisNYGb1X_79clFayeGjmYoYe6Nydsy4YBI-w0Ew,28
 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
@@ -149,7 +149,7 @@ metaflow/mflog/mflog.py,sha256=VebXxqitOtNAs7VJixnNfziO_i_urG7bsJ5JiB5IXgY,4370
 metaflow/mflog/save_logs.py,sha256=ZBAF4BMukw4FMAC7odpr9OI2BC_2petPtDX0ca6srC4,2352
 metaflow/mflog/save_logs_periodically.py,sha256=2Uvk9hi-zlCqXxOQoXmmjH1SCugfw6eG6w70WgfI-ho,1256
 metaflow/mflog/tee.py,sha256=wTER15qeHuiRpCkOqo-bd-r3Gj-EVlf3IvWRCA4beW4,887
-metaflow/plugins/__init__.py,sha256=NXlwhFvhLYhAVhjCyRJZMIpTBBBJlzFupM7MgDKNYv0,7872
+metaflow/plugins/__init__.py,sha256=Lr7i7ssJI_-czorJYjMFcRhGspqArobNoXUl9T1p3MY,8055
 metaflow/plugins/catch_decorator.py,sha256=UOM2taN_OL2RPpuJhwEOA9ZALm0-hHD0XS2Hn2GUev0,4061
 metaflow/plugins/debug_logger.py,sha256=mcF5HYzJ0NQmqCMjyVUk3iAP-heroHRIiVWQC6Ha2-I,879
 metaflow/plugins/debug_monitor.py,sha256=Md5X_sDOSssN9pt2D8YcaIjTK5JaQD55UAYTcF6xYF0,1099
@@ -181,7 +181,7 @@ 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=PS_cYGnPw9h4X7TP_plObDH3clMw4reOsBLkkGPTd0Y,16282
 metaflow/plugins/argo/argo_events.py,sha256=_C1KWztVqgi3zuH57pInaE9OzABc2NnncC-zdwOMZ-w,5909
-metaflow/plugins/argo/argo_workflows.py,sha256=CdZoBZU8aSkza1wRw60VogJv9QdF4CLbfD4o8XPAf5o,175510
+metaflow/plugins/argo/argo_workflows.py,sha256=HgreJyYibFiWScq9mvd0p0bM8NJPX49n0gdnhUVGoHI,175591
 metaflow/plugins/argo/argo_workflows_cli.py,sha256=11_8l4IrtkwviKsijInTZPt7YK5TZzClREnw_Cf4D5o,36706
 metaflow/plugins/argo/argo_workflows_decorator.py,sha256=ogCSBmwsC2C3eusydrgjuAJd4qK18f1sI4jJwA4Fd-o,7800
 metaflow/plugins/argo/argo_workflows_deployer.py,sha256=6kHxEnYXJwzNCM9swI8-0AckxtPWqwhZLerYkX8fxUM,4444
@@ -289,11 +289,13 @@ metaflow/plugins/gcp/includefile_support.py,sha256=OQO0IVWv4ObboL0VqEZwcDOyj9ORL
 metaflow/plugins/kubernetes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow/plugins/kubernetes/kube_utils.py,sha256=jdFMGbEmIow-oli26v31W9CmbZXigx06b3D_xIobpk0,4140
 metaflow/plugins/kubernetes/kubernetes.py,sha256=7yaa1TL3TcC-Js6_kAi0HGFLbXesMw3WiKWPlN9yIxo,30028
-metaflow/plugins/kubernetes/kubernetes_cli.py,sha256=A6hI6KZ6sadPAOAyGhjwITMfnabr6voBXLRlDDxylcg,13874
+metaflow/plugins/kubernetes/kubernetes_cli.py,sha256=o_o0BDEJFpTuga7txRmkvZH8OIuTb5kI4UaG6xbzf84,13929
 metaflow/plugins/kubernetes/kubernetes_client.py,sha256=tuvXP-QKpdeSmzVolB2R_TaacOr5DIb0j642eKcjsiM,6491
-metaflow/plugins/kubernetes/kubernetes_decorator.py,sha256=5NfrCZaGf2a2oQK4CeJExcizbojynCnEXzXqSN5Hoz0,30500
+metaflow/plugins/kubernetes/kubernetes_decorator.py,sha256=OwIuB9MAeO_fmTv3_IurDnbL_szXH7et3TwEva4PCfc,30853
 metaflow/plugins/kubernetes/kubernetes_job.py,sha256=pO9ExyAVCDoAoWFn9oFcos2aa0MQk4_D61O-T4E10E8,31826
 metaflow/plugins/kubernetes/kubernetes_jobsets.py,sha256=9kU43eE5IvIa7y-POzBdxnJOazWsedKhwQ51Tu1HN_A,42471
+metaflow/plugins/kubernetes/spot_metadata_cli.py,sha256=an0nWCxgflmqIPBCBrlb4m3DereDFFJBLt-KKhqcHc8,1670
+metaflow/plugins/kubernetes/spot_monitor_sidecar.py,sha256=zrWU-smQwPnL6MBHmzTxWyEA00R6iKKQbhhy50xFwQ8,3832
 metaflow/plugins/metadata_providers/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 metaflow/plugins/metadata_providers/local.py,sha256=9UAxe9caN6kU1lkSlIoJbRGgTqsMa62cBTnyMwhqiaA,22446
 metaflow/plugins/metadata_providers/service.py,sha256=NKZfFMamx6upP6aFRJfXlfYIhySgFNzz6kbp1yPD7LA,20222
@@ -358,9 +360,9 @@ metaflow/user_configs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3h
 metaflow/user_configs/config_decorators.py,sha256=Tj0H88UT8Q6pylXxHXgiA6cqnNlw4d3mR7M8J9g3ZUg,20139
 metaflow/user_configs/config_options.py,sha256=Knpiax_YGmYAdR3zKmaepN8puW1MyL9g6-eMGAkcylo,20942
 metaflow/user_configs/config_parameters.py,sha256=T0Zz18o9zKEV7mMcKotFWvXixhJpotLRBVrKx6ENErQ,15416
-metaflow-2.13.4.dist-info/LICENSE,sha256=nl_Lt5v9VvJ-5lWJDT4ddKAG-VZ-2IaLmbzpgYDz2hU,11343
-metaflow-2.13.4.dist-info/METADATA,sha256=PBmkd2ZuaamBlNTJCyBWZWWVO8xs8J1anqV4YV_P4kQ,6121
-metaflow-2.13.4.dist-info/WHEEL,sha256=9Hm2OB-j1QcCUq9Jguht7ayGIIZBRTdOXD1qg9cCgPM,109
-metaflow-2.13.4.dist-info/entry_points.txt,sha256=IKwTN1T3I5eJL3uo_vnkyxVffcgnRdFbKwlghZfn27k,57
-metaflow-2.13.4.dist-info/top_level.txt,sha256=v1pDHoWaSaKeuc5fKTRSfsXCKSdW1zvNVmvA-i0if3o,9
-metaflow-2.13.4.dist-info/RECORD,,
+metaflow-2.13.5.dist-info/LICENSE,sha256=nl_Lt5v9VvJ-5lWJDT4ddKAG-VZ-2IaLmbzpgYDz2hU,11343
+metaflow-2.13.5.dist-info/METADATA,sha256=8_zAjFu6yf75apsgNeoWmncpzpEhuji4C-JYe3w8HBg,6121
+metaflow-2.13.5.dist-info/WHEEL,sha256=9Hm2OB-j1QcCUq9Jguht7ayGIIZBRTdOXD1qg9cCgPM,109
+metaflow-2.13.5.dist-info/entry_points.txt,sha256=IKwTN1T3I5eJL3uo_vnkyxVffcgnRdFbKwlghZfn27k,57
+metaflow-2.13.5.dist-info/top_level.txt,sha256=v1pDHoWaSaKeuc5fKTRSfsXCKSdW1zvNVmvA-i0if3o,9
+metaflow-2.13.5.dist-info/RECORD,,