ob-metaflow-extensions 1.1.90__tar.gz → 1.1.91__tar.gz

{ob-metaflow-extensions-1.1.90 → ob-metaflow-extensions-1.1.91}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.1.90
+Version: 1.1.91
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial

{ob-metaflow-extensions-1.1.90 → ob-metaflow-extensions-1.1.91}/metaflow_extensions/outerbounds/plugins/nvcf/heartbeat_store.py

@@ -8,26 +8,27 @@ from datetime import datetime, timezone
 from metaflow.exception import MetaflowException
 
 
-def kill_process_and_descendants(pid, termination_timeout=5):
-    try:
-        subprocess.check_call(["pkill", "-TERM", "-P", str(pid)])
-        subprocess.check_call(["kill", "-TERM", str(pid)])
-    except subprocess.CalledProcessError:
-        pass
+def kill_process_and_descendants(pid, termination_timeout=1, iterations=20, delay=0.5):
+    for i in range(iterations):
+        try:
+            subprocess.check_call(["pkill", "-TERM", "-P", str(pid)])
+            subprocess.check_call(["kill", "-TERM", str(pid)])
+
+            time.sleep(termination_timeout)
 
-    time.sleep(termination_timeout)
+            subprocess.check_call(["pkill", "-KILL", "-P", str(pid)])
+            subprocess.check_call(["kill", "-KILL", str(pid)])
+        except subprocess.CalledProcessError:
+            pass
 
-    try:
-        subprocess.check_call(["pkill", "-KILL", "-P", str(pid)])
-        subprocess.check_call(["kill", "-KILL", str(pid)])
-    except subprocess.CalledProcessError:
-        pass
+        # Don't delay after the last iteration
+        if i < iterations - 1:
+            time.sleep(delay)
 
 
 class HeartbeatStore(object):
     def __init__(
         self,
-        heartbeat_prefix,
         main_pid=None,
         storage_backend=None,
         emit_frequency=30,
@@ -35,7 +36,6 @@ class HeartbeatStore(object):
         monitor_frequency=15,
         max_missed_heartbeats=3,
     ) -> None:
-        self.heartbeat_prefix = heartbeat_prefix
         self.main_pid = main_pid
         self.storage_backend = storage_backend
         self.emit_frequency = emit_frequency
@@ -44,8 +44,8 @@ class HeartbeatStore(object):
         self.max_missed_heartbeats = max_missed_heartbeats
         self.missed_heartbeats = 0
 
-    def emit_heartbeat(self, folder_name=None):
-        heartbeat_key = f"{self.heartbeat_prefix}/heartbeat"
+    def emit_heartbeat(self, heartbeat_prefix: str, folder_name=None):
+        heartbeat_key = f"{heartbeat_prefix}/heartbeat"
         if folder_name:
             heartbeat_key = f"{folder_name}/{heartbeat_key}"
 
@@ -63,8 +63,8 @@ class HeartbeatStore(object):
 
             time.sleep(self.emit_frequency)
 
-    def emit_tombstone(self, folder_name=None):
-        tombstone_key = f"{self.heartbeat_prefix}/tombstone"
+    def emit_tombstone(self, tombstone_prefix: str, folder_name=None):
+        tombstone_key = f"{tombstone_prefix}/tombstone"
         if folder_name:
             tombstone_key = f"{folder_name}/{tombstone_key}"
 
@@ -113,12 +113,12 @@ class HeartbeatStore(object):
             return False
         return True
 
-    def monitor(self, folder_name=None):
-        heartbeat_key = f"{self.heartbeat_prefix}/heartbeat"
+    def monitor(self, heartbeat_prefix: str, tombstone_prefix: str, folder_name=None):
+        heartbeat_key = f"{heartbeat_prefix}/heartbeat"
         if folder_name:
             heartbeat_key = f"{folder_name}/{heartbeat_key}"
 
-        tombstone_key = f"{self.heartbeat_prefix}/tombstone"
+        tombstone_key = f"{tombstone_prefix}/tombstone"
         if folder_name:
             tombstone_key = f"{folder_name}/{tombstone_key}"
 
@@ -162,12 +162,17 @@ if __name__ == "__main__":
     storage = datastores[0](datastore_sysroot)
 
     heartbeat_prefix = f"{os.getenv('MF_PATHSPEC')}/{os.getenv('MF_ATTEMPT')}"
+    flow_name, run_id, _, _ = os.getenv("MF_PATHSPEC").split("/")
+    tombstone_prefix = f"{flow_name}/{run_id}"
 
     store = HeartbeatStore(
-        heartbeat_prefix=heartbeat_prefix,
         main_pid=int(main_pid),
         storage_backend=storage,
         max_missed_heartbeats=int(NVIDIA_HEARTBEAT_THRESHOLD),
     )
 
-    store.monitor(folder_name=folder_name)
+    store.monitor(
+        heartbeat_prefix=heartbeat_prefix,
+        tombstone_prefix=tombstone_prefix,
+        folder_name=folder_name,
+    )

{ob-metaflow-extensions-1.1.90 → ob-metaflow-extensions-1.1.91}/metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py

@@ -185,13 +185,17 @@ class Job(object):
         )
 
         store = HeartbeatStore(
-            heartbeat_prefix=heartbeat_prefix,
             main_pid=None,
             storage_backend=backend,
         )
 
         self.heartbeat_thread = threading.Thread(
-            target=store.emit_heartbeat, args=("nvcf_heartbeats",), daemon=True
+            target=store.emit_heartbeat,
+            args=(
+                heartbeat_prefix,
+                "nvcf_heartbeats",
+            ),
+            daemon=True,
         )
         self.heartbeat_thread.start()
 

{ob-metaflow-extensions-1.1.90 → ob-metaflow-extensions-1.1.91}/metaflow_extensions/outerbounds/plugins/nvcf/nvcf_cli.py

@@ -56,23 +56,21 @@ def list(ctx, run_id):
     flow_name = ctx.obj.flow.name
     run_obj = Run(pathspec=f"{flow_name}/{run_id}", _namespace_check=False)
     running_invocations = []
+
     for each_step in run_obj:
-        if (
-            not each_step.task.finished
-            and "nvcf-function-id" in each_step.task.metadata_dict
-        ):
-            task_pathspec = each_step.task.pathspec
-            attempt = each_step.task.metadata_dict.get("attempt")
-            flow_name, run_id, step_name, task_id = task_pathspec.split("/")
-            running_invocations.append(
-                f"Flow Name: {flow_name}, Run ID: {run_id}, Step Name: {step_name}, Task ID: {task_id}, Retry Count: {attempt}"
-            )
+        for each_task in each_step:
+            if not each_task.finished and "nvcf-function-id" in each_task.metadata_dict:
+
+                task_pathspec = each_task.pathspec
+                attempt = each_task.metadata_dict.get("attempt")
+                flow_name, run_id, step_name, task_id = task_pathspec.split("/")
+                running_invocations.append(
+                    f"Flow Name: {flow_name}, Run ID: {run_id}, Step Name: {step_name}, Task ID: {task_id}, Retry Count: {attempt}"
+                )
 
     if running_invocations:
         for each_invocation in running_invocations:
             ctx.obj.echo(each_invocation)
-    else:
-        ctx.obj.echo("No running @nvidia invocations for Run ID: %s" % run_id)
 
 
 @nvidia.command(help="Kill steps / tasks running as an nvidia job.")
@@ -88,29 +86,17 @@ def kill(ctx, run_id):
         HeartbeatStore,
     )
 
-    flow_name = ctx.obj.flow.name
-    run_obj = Run(pathspec=f"{flow_name}/{run_id}", _namespace_check=False)
-
-    for each_step in run_obj:
-        if (
-            not each_step.task.finished
-            and "nvcf-function-id" in each_step.task.metadata_dict
-        ):
-            task_pathspec = each_step.task.pathspec
-            attempt = each_step.task.metadata_dict.get("attempt")
-            heartbeat_prefix = "{task_pathspec}/{attempt}".format(
-                task_pathspec=task_pathspec, attempt=attempt
-            )
+    datastore_root = ctx.obj.datastore_impl.datastore_root
+    store = HeartbeatStore(
+        main_pid=None,
+        storage_backend=ctx.obj.datastore_impl(datastore_root),
+    )
 
-            datastore_root = ctx.obj.datastore_impl.datastore_root
-            store = HeartbeatStore(
-                heartbeat_prefix=heartbeat_prefix,
-                main_pid=None,
-                storage_backend=ctx.obj.datastore_impl(datastore_root),
-            )
-            store.emit_tombstone(folder_name="nvcf_heartbeats")
-        else:
-            ctx.obj.echo("No running @nvidia invocations for Run ID: %s" % run_id)
+    flow_name = ctx.obj.flow.name
+    tombstone_prefix = f"{flow_name}/{run_id}"
+    store.emit_tombstone(
+        tombstone_prefix=tombstone_prefix, folder_name="nvcf_heartbeats"
+    )
 
 
 @nvidia.command(

{ob-metaflow-extensions-1.1.90 → ob-metaflow-extensions-1.1.91}/metaflow_extensions/outerbounds/plugins/snowpark/snowpark.py

@@ -31,7 +31,7 @@ from metaflow.mflog import (
 )
 
 from .snowpark_client import SnowparkClient
-from .snowpark_exceptions import SnowparkException
+from .snowpark_exceptions import SnowparkException, SnowparkKilledException
 from .snowpark_job import SnowparkJob
 
 # Redirect structured logs to $PWD/.logs/
@@ -291,7 +291,7 @@ class Snowpark(object):
         else:
             if self.job.is_running:
                 # Kill the job if it is still running by throwing an exception.
-                raise SnowparkException("Task failed!")
+                raise SnowparkKilledException("Task failed!")
             echo(
                 "Task finished with message '%s'." % self.job.message,
                 "stderr",

{ob-metaflow-extensions-1.1.90 → ob-metaflow-extensions-1.1.91}/ob_metaflow_extensions.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.1.90
+Version: 1.1.91
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial

{ob-metaflow-extensions-1.1.90 → ob-metaflow-extensions-1.1.91}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_namespace_packages
 from pathlib import Path
 
 
-version = "1.1.90"
+version = "1.1.91"
 this_directory = Path(__file__).parent
 long_description = (this_directory / "README.md").read_text()