ob-metaflow-extensions 1.1.89__py2.py3-none-any.whl → 1.1.91__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/__init__.py

@@ -116,40 +116,39 @@ def get_boto3_session(role_arn=None, session_vars=None):
     if token_info.get("region"):
         os.environ["AWS_DEFAULT_REGION"] = token_info["region"]
 
-    with hide_access_keys():
-        if cspr_role:
-            #  The generated AWS config will be used here since we set the
-            #  AWS_CONFIG_FILE environment variable above.
-            if role_arn == USE_CSPR_ROLE_ARN_IF_SET:
-                # Otherwise start from the default profile, assuming CSPR role
-                session = boto3.session.Session(profile_name="cspr")
-            else:
-                session = boto3.session.Session(profile_name="task")
-        else:
-            # Not using AWS config, just AWS_WEB_IDENTITY_TOKEN_FILE + AWS_ROLE_ARN
-            session = boto3.session.Session()
-
-        if role_arn and role_arn != USE_CSPR_ROLE_ARN_IF_SET:
-            # If the user provided a role_arn, we assume that role
-            # using the task role credentials. CSPR role is not used.
-            fetcher = botocore.credentials.AssumeRoleCredentialFetcher(
-                client_creator=session._session.create_client,
-                source_credentials=session._session.get_credentials(),
-                role_arn=role_arn,
-                extra_args={},
-            )
-            creds = botocore.credentials.DeferredRefreshableCredentials(
-                method="assume-role", refresh_using=fetcher.fetch_credentials
-            )
-            botocore_session = botocore.session.Session(session_vars=session_vars)
-            botocore_session._credentials = creds
-            return boto3.session.Session(botocore_session=botocore_session)
+    if cspr_role:
+        #  The generated AWS config will be used here since we set the
+        #  AWS_CONFIG_FILE environment variable above.
+        if role_arn == USE_CSPR_ROLE_ARN_IF_SET:
+            # Otherwise start from the default profile, assuming CSPR role
+            session = boto3.session.Session(profile_name="cspr")
         else:
-            # If the user didn't provide a role_arn, or if the role_arn
-            # is set to USE_CSPR_ROLE_ARN_IF_SET, we return the default
-            # session which would use the CSPR role if it is set on the
-            # server, and the task role otherwise.
-            return session
+            session = boto3.session.Session(profile_name="task")
+    else:
+        # Not using AWS config, just AWS_WEB_IDENTITY_TOKEN_FILE + AWS_ROLE_ARN
+        session = boto3.session.Session()
+
+    if role_arn and role_arn != USE_CSPR_ROLE_ARN_IF_SET:
+        # If the user provided a role_arn, we assume that role
+        # using the task role credentials. CSPR role is not used.
+        fetcher = botocore.credentials.AssumeRoleCredentialFetcher(
+            client_creator=session._session.create_client,
+            source_credentials=session._session.get_credentials(),
+            role_arn=role_arn,
+            extra_args={},
+        )
+        creds = botocore.credentials.DeferredRefreshableCredentials(
+            method="assume-role", refresh_using=fetcher.fetch_credentials
+        )
+        botocore_session = botocore.session.Session(session_vars=session_vars)
+        botocore_session._credentials = creds
+        return boto3.session.Session(botocore_session=botocore_session)
+    else:
+        # If the user didn't provide a role_arn, or if the role_arn
+        # is set to USE_CSPR_ROLE_ARN_IF_SET, we return the default
+        # session which would use the CSPR role if it is set on the
+        # server, and the task role otherwise.
+        return session
 
 
 class ObpAuthProvider(object):
@@ -164,11 +163,12 @@ class ObpAuthProvider(object):
 
         from botocore.exceptions import ClientError
 
-        session = get_boto3_session(role_arn, session_vars)
-        if with_error:
-            return session.client(module, **client_params), ClientError
-        else:
-            return session.client(module, **client_params)
+        with hide_access_keys():
+            session = get_boto3_session(role_arn, session_vars)
+            if with_error:
+                return session.client(module, **client_params), ClientError
+            else:
+                return session.client(module, **client_params)
 
 
 AWS_CLIENT_PROVIDERS_DESC = [("obp", ".ObpAuthProvider")]

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,
+    )

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()
 

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(

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.89.dist-info → ob_metaflow_extensions-1.1.91.dist-info}/METADATA

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

{ob_metaflow_extensions-1.1.89.dist-info → ob_metaflow_extensions-1.1.91.dist-info}/RECORD

@@ -1,7 +1,7 @@
 metaflow_extensions/outerbounds/__init__.py,sha256=TRGvIUMjkfneWtYUFSWoubu_Kf2ekAL4WLbV3IxOj9k,499
 metaflow_extensions/outerbounds/remote_config.py,sha256=Zpfpjgz68_ZgxlXezjzlsDLo4840rkWuZgwDB_5H57U,4059
 metaflow_extensions/outerbounds/config/__init__.py,sha256=JsQGRuGFz28fQWjUvxUgR8EKBLGRdLUIk_buPLJplJY,1225
-metaflow_extensions/outerbounds/plugins/__init__.py,sha256=9DCqKsb2bPfw1f7x-3EuB2Mqc9uKlLtGEG6yLJI8Xx0,12510
+metaflow_extensions/outerbounds/plugins/__init__.py,sha256=WwvFcN5kserbPwhrE4hXprnXzJxonEPT0Mlik2kmGMA,12406
 metaflow_extensions/outerbounds/plugins/auth_server.py,sha256=1v2GBqoMBxp5E7Lejz139w-jxJtPnLDvvHXP0HhEIHI,2361
 metaflow_extensions/outerbounds/plugins/perimeters.py,sha256=QXh3SFP7GQbS-RAIxUOPbhPzQ7KDFVxZkTdKqFKgXjI,2697
 metaflow_extensions/outerbounds/plugins/fast_bakery/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -14,12 +14,12 @@ metaflow_extensions/outerbounds/plugins/kubernetes/kubernetes_client.py,sha256=g
 metaflow_extensions/outerbounds/plugins/nim/__init__.py,sha256=GVnvSTjqYVj5oG2yh8KJFt7iZ33cEadDD5HbdmC9hJ0,1457
 metaflow_extensions/outerbounds/plugins/nim/nim_manager.py,sha256=SWieODDxtIaeZwdMYtObDi57Kjyfw2DUuE6pJtU750w,9206
 metaflow_extensions/outerbounds/plugins/nvcf/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-metaflow_extensions/outerbounds/plugins/nvcf/heartbeat_store.py,sha256=4aQZ0kpW2LlJbHx6AG4A9eaFLH9rWC_ENWnnfYNq1qk,5910
-metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py,sha256=NIt1kJHuYpnCF7n73A90ZITWsk5QWtsbiHfzvdVjgqk,8997
-metaflow_extensions/outerbounds/plugins/nvcf/nvcf_cli.py,sha256=9nQBwm6AYtaKIAxdb937MOnsut3INEXN3v5eSnXy4cg,9811
+metaflow_extensions/outerbounds/plugins/nvcf/heartbeat_store.py,sha256=pOWwm8LFQBbtku0zNBBwCyXxLK8U-hhC4naQcmU69nE,6217
+metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py,sha256=s9S_Ntm8Y23fgzQtN_hLMrQaSfOvKqNBUyL1K8Xo9vU,9039
+metaflow_extensions/outerbounds/plugins/nvcf/nvcf_cli.py,sha256=YZb5AvbVgUwUJVxRxQ4JqqP8e1RMJr6dZ9U4KkHE-M8,9134
 metaflow_extensions/outerbounds/plugins/nvcf/nvcf_decorator.py,sha256=E7h94ni8yW9BQkKSBUptPdGAaVmXpR9FlXkPWpLyPd0,6054
 metaflow_extensions/outerbounds/plugins/snowpark/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-metaflow_extensions/outerbounds/plugins/snowpark/snowpark.py,sha256=vzgpVLCKvHjzHNfJvmH0jcxefYNsVggw_vof_y_U_a8,10643
+metaflow_extensions/outerbounds/plugins/snowpark/snowpark.py,sha256=0R8aFN9MpgWraqiaI6ZF82YpLdFJ1f-3z_-BPRpZfxM,10674
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_cli.py,sha256=ErsVoCQLa33byiykOQzDEeEkRKk0mgffZme43f3jxn4,8747
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_client.py,sha256=JEW0EUxj_mNZXo9OFkJFmWfg-P7_CEgvNbgsMTCBTAE,4273
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_decorator.py,sha256=3A9LKg7EarWM8WQ0PTGLUetjxzemQeUiJivvv_4uzr0,9886
@@ -33,7 +33,7 @@ metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py,
 metaflow_extensions/outerbounds/toplevel/plugins/azure/__init__.py,sha256=WUuhz2YQfI4fz7nIcipwwWq781eaoHEk7n4GAn1npDg,63
 metaflow_extensions/outerbounds/toplevel/plugins/gcp/__init__.py,sha256=BbZiaH3uILlEZ6ntBLKeNyqn3If8nIXZFq_Apd7Dhco,70
 metaflow_extensions/outerbounds/toplevel/plugins/kubernetes/__init__.py,sha256=5zG8gShSj8m7rgF4xgWBZFuY3GDP5n1T0ktjRpGJLHA,69
-ob_metaflow_extensions-1.1.89.dist-info/METADATA,sha256=48w-ipRZMJOxQ7yZx8uqyDmiBQIirZ3MWYtygEXscNI,520
-ob_metaflow_extensions-1.1.89.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.1.89.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.1.89.dist-info/RECORD,,
+ob_metaflow_extensions-1.1.91.dist-info/METADATA,sha256=rTZDEBQQ5bWEZnqLoralSTjginYiJ-Q81f054qg3vLs,520
+ob_metaflow_extensions-1.1.91.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.1.91.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.1.91.dist-info/RECORD,,