ob-metaflow-extensions 1.1.84__tar.gz → 1.1.88__tar.gz

{ob-metaflow-extensions-1.1.84 → ob-metaflow-extensions-1.1.88}/PKG-INFO

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

{ob-metaflow-extensions-1.1.84 → ob-metaflow-extensions-1.1.88}/metaflow_extensions/outerbounds/config/__init__.py

@@ -14,6 +14,13 @@ DEFAULT_GCP_CLIENT_PROVIDER = "obp"
 FAST_BAKERY_URL = from_conf("FAST_BAKERY_URL", None)
 
 
+###
+# NVCF configuration
+###
+# Maximum number of consecutive heartbeats that can be missed.
+NVIDIA_HEARTBEAT_THRESHOLD = from_conf("NVIDIA_HEARTBEAT_THRESHOLD", "3")
+
+
 ###
 # Snowpark configuration
 ###

{ob-metaflow-extensions-1.1.84 → ob-metaflow-extensions-1.1.88}/metaflow_extensions/outerbounds/plugins/__init__.py

@@ -32,6 +32,126 @@ def hide_access_keys(*args, **kwds):
             os.environ["AWS_SESSION_TOKEN"] = AWS_SESSION_TOKEN
 
 
+# This is a special placeholder value that can be passed as role_arn to
+# get_boto3_session() which makes it use the CSPR role, if its set.
+USE_CSPR_ROLE_ARN_IF_SET = "__cspr__"
+
+
+def get_boto3_session(role_arn=None, session_vars=None):
+    import boto3
+    import botocore
+    from metaflow_extensions.outerbounds.plugins.auth_server import get_token
+
+    from hashlib import sha256
+    from metaflow.util import get_username
+
+    user = get_username()
+
+    token_info = get_token("/generate/aws")
+
+    # Write token to a file. The file name is derived from the user name
+    # so it works with multiple users on the same machine.
+    #
+    # We hash the user name so we don't have to deal with special characters
+    # in the file name and the file name is not exposed to the user
+    # anyways, so it doesn't matter that its a little ugly.
+    token_file = "/tmp/obp_token." + sha256(user.encode("utf-8")).hexdigest()[:16]
+
+    # Write to a temp file then rename to avoid a situation when someone
+    # tries to read the file after it was open for writing (and truncated)
+    # but before the token was written to it.
+    with tempfile.NamedTemporaryFile("w", delete=False) as f:
+        f.write(token_info["token"])
+        tmp_token_file = f.name
+    os.rename(tmp_token_file, token_file)
+
+    cspr_role = None
+    if token_info.get("cspr_role_arn"):
+        cspr_role = token_info["cspr_role_arn"]
+
+    if cspr_role:
+        # If CSPR role is set, we set it as the default role to assume
+        # for the AWS SDK. We do this by writing an AWS config file
+        # with two profiles. One to get credentials for the task role
+        # in exchange for the OIDC token, and second to assume the
+        # CSPR role using the task role credentials.
+        import configparser
+        from io import StringIO
+
+        aws_config = configparser.ConfigParser()
+
+        # Task role profile
+        aws_config["profile task"] = {
+            "role_arn": token_info["role_arn"],
+            "web_identity_token_file": token_file,
+        }
+
+        # CSPR role profile (default)
+        aws_config["profile cspr"] = {
+            "role_arn": cspr_role,
+            "source_profile": "task",
+        }
+
+        aws_config_string = StringIO()
+        aws_config.write(aws_config_string)
+        aws_config_file = (
+            "/tmp/aws_config." + sha256(user.encode("utf-8")).hexdigest()[:16]
+        )
+        with tempfile.NamedTemporaryFile(
+            "w", delete=False, dir=os.path.dirname(aws_config_file)
+        ) as f:
+            f.write(aws_config_string.getvalue())
+            tmp_aws_config_file = f.name
+        os.rename(tmp_aws_config_file, aws_config_file)
+        os.environ["AWS_CONFIG_FILE"] = aws_config_file
+        os.environ["AWS_PROFILE"] = "cspr"
+    else:
+        os.environ["AWS_WEB_IDENTITY_TOKEN_FILE"] = token_file
+        os.environ["AWS_ROLE_ARN"] = token_info["role_arn"]
+
+    # Enable regional STS endpoints. This is the new recommended way
+    # by AWS [1] and is the more performant way.
+    # [1] https://docs.aws.amazon.com/sdkref/latest/guide/feature-sts-regionalized-endpoints.html
+    os.environ["AWS_STS_REGIONAL_ENDPOINTS"] = "regional"
+    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)
+        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):
     name = "obp"
 
@@ -42,67 +162,13 @@ class ObpAuthProvider(object):
         if client_params is None:
             client_params = {}
 
-        import boto3
-        import botocore
         from botocore.exceptions import ClientError
-        from metaflow_extensions.outerbounds.plugins.auth_server import get_token
-
-        from hashlib import sha256
-        from metaflow.util import get_username
-
-        user = get_username()
-
-        token_info = get_token("/generate/aws")
 
-        # Write token to a file. The file name is derived from the user name
-        # so it works with multiple users on the same machine.
-        #
-        # We hash the user name so we don't have to deal with special characters
-        # in the file name and the file name is not exposed to the user
-        # anyways, so it doesn't matter that its a little ugly.
-        token_file = "/tmp/obp_token." + sha256(user.encode("utf-8")).hexdigest()[:16]
-
-        # Write to a temp file then rename to avoid a situation when someone
-        # tries to read the file after it was open for writing (and truncated)
-        # but before the token was written to it.
-        with tempfile.NamedTemporaryFile("w", delete=False) as f:
-            f.write(token_info["token"])
-            tmp_token_file = f.name
-        os.rename(tmp_token_file, token_file)
-
-        os.environ["AWS_WEB_IDENTITY_TOKEN_FILE"] = token_file
-        os.environ["AWS_ROLE_ARN"] = token_info["role_arn"]
-
-        # Enable regional STS endpoints. This is the new recommended way
-        # by AWS [1] and is the more performant way.
-        # [1] https://docs.aws.amazon.com/sdkref/latest/guide/feature-sts-regionalized-endpoints.html
-        os.environ["AWS_STS_REGIONAL_ENDPOINTS"] = "regional"
-        if token_info.get("region"):
-            os.environ["AWS_DEFAULT_REGION"] = token_info["region"]
-
-        with hide_access_keys():
-            if role_arn:
-                session = boto3.session.Session()
-                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
-                session = boto3.session.Session(botocore_session=botocore_session)
-                if with_error:
-                    return session.client(module, **client_params), ClientError
-                else:
-                    return session.client(module, **client_params)
-            if with_error:
-                return boto3.client(module, **client_params), ClientError
-            else:
-                return boto3.client(module, **client_params)
+        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")]

{ob-metaflow-extensions-1.1.84 → ob-metaflow-extensions-1.1.88}/metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py

@@ -1,30 +1,30 @@
 import hashlib
 import json
 import os
-
+import threading
+import time
+import uuid
 from concurrent.futures import ThreadPoolExecutor
 from typing import Dict
+
 from metaflow.exception import MetaflowException
-from metaflow.metaflow_config import (
-    FAST_BAKERY_URL,
-    get_pinned_conda_libs,
-)
+from metaflow.metaflow_config import FAST_BAKERY_URL, get_pinned_conda_libs
 from metaflow.metaflow_environment import MetaflowEnvironment
-from metaflow.plugins.pypi.conda_environment import CondaEnvironment
-from .fast_bakery import FastBakery, FastBakeryApiResponse, FastBakeryException
 from metaflow.plugins.aws.batch.batch_decorator import BatchDecorator
 from metaflow.plugins.kubernetes.kubernetes_decorator import KubernetesDecorator
 from metaflow.plugins.pypi.conda_decorator import CondaStepDecorator
+from metaflow.plugins.pypi.conda_environment import CondaEnvironment
 from metaflow.plugins.pypi.pypi_decorator import PyPIStepDecorator
 
+from .fast_bakery import FastBakery, FastBakeryApiResponse, FastBakeryException
+
 BAKERY_METAFILE = ".imagebakery-cache"
 
+import fcntl
 import json
 import os
-import fcntl
-from functools import wraps
 from concurrent.futures import ThreadPoolExecutor
-
+from functools import wraps
 
 # TODO - ensure that both @conda/@pypi are not assigned to the same step
 
@@ -36,6 +36,9 @@ def cache_request(cache_file):
             call_args = kwargs.copy()
             call_args.update(zip(func.__code__.co_varnames, args))
             call_args.pop("self", None)
+            call_args.pop("ref", None)
+            # invalidate cache when moving from one deployment to another
+            call_args.update({"fast_bakery_url": FAST_BAKERY_URL})
             cache_key = hashlib.md5(
                 json.dumps(call_args, sort_keys=True).encode("utf-8")
             ).hexdigest()
@@ -79,7 +82,7 @@ def cache_request(cache_file):
 
 
 class DockerEnvironmentException(MetaflowException):
-    headline = "Ran into an error while setting up the environment"
+    headline = "Ran into an error while baking image"
 
     def __init__(self, msg):
         super(DockerEnvironmentException, self).__init__(msg)
@@ -93,8 +96,8 @@ class DockerEnvironment(MetaflowEnvironment):
         self.skipped_steps = set()
         self.flow = flow
 
-        self.bakery = FastBakery(url=FAST_BAKERY_URL)
         self.results = {}
+        self.images_baked = 0
 
     def set_local_root(self, local_root):
         self.local_root = local_root
@@ -102,15 +105,31 @@ class DockerEnvironment(MetaflowEnvironment):
     def decospecs(self):
         return ("conda", "fast_bakery_internal") + super().decospecs()
 
-    def validate_environment(self, echo, datastore_type):
+    def validate_environment(self, logger, datastore_type):
         self.datastore_type = datastore_type
-        self.echo = echo
+        self.logger = logger
 
         # Avoiding circular imports.
         from metaflow.plugins import DATASTORES
 
         self.datastore = [d for d in DATASTORES if d.TYPE == self.datastore_type][0]
 
+        # Mixing @pypi/@conda in a single step is not supported yet
+        for step in self.flow:
+            if (
+                sum(
+                    1
+                    for deco in step.decorators
+                    if isinstance(deco, (PyPIStepDecorator, CondaStepDecorator))
+                )
+                > 1
+            ):
+                raise MetaflowException(
+                    "Mixing and matching PyPI packages and Conda packages within a\n"
+                    "step is not yet supported. Use one of @pypi or @conda only for the *%s* step."
+                    % step.name
+                )
+
     def init_environment(self, echo):
         self.skipped_steps = {
             step.name
@@ -125,14 +144,21 @@ class DockerEnvironment(MetaflowEnvironment):
             step for step in self.flow if step.name not in self.skipped_steps
         ]
         if steps_to_bake:
-            echo("Baking container image(s) ...")
-            self.results = self._bake(steps_to_bake, echo)
+            self.logger("🚀 Baking container image(s) ...")
+            start_time = time.time()
+            self.results = self._bake(steps_to_bake)
             for step in self.flow:
                 for d in step.decorators:
                     if isinstance(d, (BatchDecorator, KubernetesDecorator)):
                         d.attributes["image"] = self.results[step.name].container_image
                         d.attributes["executable"] = self.results[step.name].python_path
-            echo("Container image(s) baked!")
+            if self.images_baked > 0:
+                bake_time = time.time() - start_time
+                self.logger(
+                    f"🎉 All container image(s) baked in {bake_time:.2f} seconds!"
+                )
+            else:
+                self.logger("🎉 All container image(s) baked!")
 
         if self.skipped_steps:
             self.delegate = CondaEnvironment(self.flow)
@@ -140,29 +166,54 @@ class DockerEnvironment(MetaflowEnvironment):
             self.delegate.validate_environment(echo, self.datastore_type)
             self.delegate.init_environment(echo, self.skipped_steps)
 
-    def _bake(self, steps, echo) -> Dict[str, FastBakeryApiResponse]:
+    def _bake(self, steps) -> Dict[str, FastBakeryApiResponse]:
         metafile_path = get_fastbakery_metafile_path(self.local_root, self.flow.name)
+        logger_lock = threading.Lock()
 
         @cache_request(metafile_path)
         def _cached_bake(
-            python=None, pypi_packages=None, conda_packages=None, base_image=None
+            ref=None,
+            python=None,
+            pypi_packages=None,
+            conda_packages=None,
+            base_image=None,
         ):
-            self.bakery._reset_payload()
-            self.bakery.python_version(python)
-            self.bakery.pypi_packages(pypi_packages)
-            self.bakery.conda_packages(conda_packages)
-            self.bakery.base_image(base_image)
-            # self.bakery.ignore_cache()
+            bakery = FastBakery(url=FAST_BAKERY_URL)
+            bakery._reset_payload()
+            bakery.python_version(python)
+            bakery.pypi_packages(pypi_packages)
+            bakery.conda_packages(conda_packages)
+            bakery.base_image(base_image)
+            # bakery.ignore_cache()
+
+            with logger_lock:
+                self.logger(f"🍳 Baking [{ref}] ...")
+                self.logger(f"     🐍 Python: {python}")
+
+                if pypi_packages:
+                    self.logger(f"     📦 PyPI packages:")
+                    for package, version in pypi_packages.items():
+                        self.logger(f"        🔧 {package}: {version}")
+
+                if conda_packages:
+                    self.logger(f"     📦 Conda packages:")
+                    for package, version in conda_packages.items():
+                        self.logger(f"        🔧 {package}: {version}")
+
+                self.logger(f"     🏗️  Base image: {base_image}")
+
+            start_time = time.time()
             try:
-                res = self.bakery.bake()
-                if res.baking_stats:
-                    echo(
-                        "baked image in: %s milliseconds"
-                        % res.baking_stats.solver_stats.duration_ms
-                    )
+                res = bakery.bake()
+                # TODO: Get actual bake time from bakery
+                bake_time = time.time() - start_time
+
+                with logger_lock:
+                    self.logger(f"🏁 Baked [{ref}] in {bake_time:.2f} seconds!")
+                self.images_baked += 1
                 return res
             except FastBakeryException as ex:
-                raise DockerEnvironmentException(str(ex))
+                raise DockerEnvironmentException(f"Bake [{ref}] failed: {str(ex)}")
 
         def prepare_step(step):
             base_image = next(
@@ -216,10 +267,15 @@ class DockerEnvironment(MetaflowEnvironment):
             }
 
         with ThreadPoolExecutor() as executor:
-            return {
-                step.name: _cached_bake(**args)
-                for step, args in zip(steps, executor.map(prepare_step, steps))
-            }
+            prepared_args = list(executor.map(prepare_step, steps))
+            for i, args in enumerate(prepared_args, 1):
+                args["ref"] = f"#{i:02d}"
+            futures = [executor.submit(_cached_bake, **args) for args in prepared_args]
+            results = {}
+            for step, future in zip(steps, futures):
+                results[step.name] = future.result()
+
+            return results
 
     def executable(self, step_name, default=None):
         if step_name in self.skipped_steps:

{ob-metaflow-extensions-1.1.84 → ob-metaflow-extensions-1.1.88}/metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery.py

@@ -1,5 +1,6 @@
 from typing import Dict, Optional
 import requests
+import time
 
 
 class FastBakeryException(Exception):
@@ -140,6 +141,19 @@ class FastBakery:
             headers = {**self.headers, **(SERVICE_HEADERS or {})}
         except ImportError:
             headers = self.headers
+
+        retryable_status_codes = [409]
+
+        for attempt in range(2):  # 0 = initial attempt, 1-2 = retries
+            response = requests.post(self.url, json=payload, headers=headers)
+
+            if response.status_code not in retryable_status_codes:
+                break
+
+            if attempt < 2:  # Don't sleep after the last attempt
+                sleep_time = 0.5 * (attempt + 1)
+                time.sleep(sleep_time)
+
         response = requests.post(self.url, json=payload, headers=headers)
         self._handle_error_response(response)
         return FastBakeryApiResponse(response.json())

ob-metaflow-extensions-1.1.88/metaflow_extensions/outerbounds/plugins/nvcf/heartbeat_store.py

@@ -0,0 +1,157 @@
+import os
+import sys
+import time
+import signal
+from io import BytesIO
+from datetime import datetime, timezone
+
+from metaflow.exception import MetaflowException
+
+
+class HeartbeatStore(object):
+    def __init__(
+        self,
+        heartbeat_prefix,
+        main_pid=None,
+        storage_backend=None,
+        emit_frequency=30,
+        missed_heartbeat_timeout=60,
+        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
+        self.monitor_frequency = monitor_frequency
+        self.missed_heartbeat_timeout = missed_heartbeat_timeout
+        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"
+        if folder_name:
+            heartbeat_key = f"{folder_name}/{heartbeat_key}"
+
+        while True:
+            try:
+                epoch_string = str(datetime.now(timezone.utc).timestamp()).encode(
+                    "utf-8"
+                )
+                self.storage_backend.save_bytes(
+                    [(heartbeat_key, BytesIO(bytes(epoch_string)))], overwrite=True
+                )
+            except Exception as e:
+                print(f"Error writing heartbeat: {e}")
+                sys.exit(1)
+
+            time.sleep(self.emit_frequency)
+
+    def emit_tombstone(self, folder_name=None):
+        tombstone_key = f"{self.heartbeat_prefix}/tombstone"
+        if folder_name:
+            tombstone_key = f"{folder_name}/{tombstone_key}"
+
+        tombstone_string = "tombstone".encode("utf-8")
+        try:
+            self.storage_backend.save_bytes(
+                [(tombstone_key, BytesIO(bytes(tombstone_string)))], overwrite=True
+            )
+        except Exception as e:
+            print(f"Error writing tombstone: {e}")
+            sys.exit(1)
+
+    def __handle_tombstone(self, path):
+        if path is not None:
+            with open(path) as f:
+                contents = f.read()
+                if "tombstone" in contents:
+                    print("[Outerbounds] Tombstone detected. Terminating the task..")
+                    os.kill(self.main_pid, signal.SIGTERM)
+                    sys.exit(1)
+
+    def __handle_heartbeat(self, path):
+        if path is not None:
+            with open(path) as f:
+                contents = f.read()
+                current_timestamp = datetime.now(timezone.utc).timestamp()
+                if current_timestamp - float(contents) > self.missed_heartbeat_timeout:
+                    self.missed_heartbeats += 1
+                else:
+                    self.missed_heartbeats = 0
+        else:
+            self.missed_heartbeats += 1
+
+        if self.missed_heartbeats > self.max_missed_heartbeats:
+            print(
+                f"[Outerbounds] Missed {self.max_missed_heartbeats} consecutive heartbeats. Terminating the task.."
+            )
+            os.kill(self.main_pid, signal.SIGTERM)
+            sys.exit(1)
+
+    def is_main_process_running(self):
+        try:
+            # Check if the process is running
+            os.kill(self.main_pid, 0)
+        except ProcessLookupError:
+            return False
+        return True
+
+    def monitor(self, folder_name=None):
+        heartbeat_key = f"{self.heartbeat_prefix}/heartbeat"
+        if folder_name:
+            heartbeat_key = f"{folder_name}/{heartbeat_key}"
+
+        tombstone_key = f"{self.heartbeat_prefix}/tombstone"
+        if folder_name:
+            tombstone_key = f"{folder_name}/{tombstone_key}"
+
+        while self.is_main_process_running():
+            with self.storage_backend.load_bytes(
+                [heartbeat_key, tombstone_key]
+            ) as results:
+                for key, path, _ in results:
+                    if key == tombstone_key:
+                        self.__handle_tombstone(path)
+                    elif key == heartbeat_key:
+                        self.__handle_heartbeat(path)
+
+            time.sleep(self.monitor_frequency)
+
+
+if __name__ == "__main__":
+    from metaflow.plugins import DATASTORES
+    from metaflow.metaflow_config import NVIDIA_HEARTBEAT_THRESHOLD
+
+    if len(sys.argv) != 4:
+        print("Usage: heartbeat_store.py <main_pid> <datastore_type> <folder_name>")
+        sys.exit(1)
+    _, main_pid, datastore_type, folder_name = sys.argv
+
+    if datastore_type not in ("azure", "gs", "s3"):
+        print(f"Datastore unsupported for type: {datastore_type}")
+        sys.exit(1)
+
+    datastores = [d for d in DATASTORES if d.TYPE == datastore_type]
+    datastore_sysroot = datastores[0].get_datastore_root_from_config(
+        lambda *args, **kwargs: None
+    )
+    if datastore_sysroot is None:
+        raise MetaflowException(
+            msg="METAFLOW_DATASTORE_SYSROOT_{datastore_type} must be set!".format(
+                datastore_type=datastore_type.upper()
+            )
+        )
+
+    storage = datastores[0](datastore_sysroot)
+
+    heartbeat_prefix = f"{os.getenv('MF_PATHSPEC')}/{os.getenv('MF_ATTEMPT')}"
+
+    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)

{ob-metaflow-extensions-1.1.84 → ob-metaflow-extensions-1.1.88}/metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py

@@ -1,6 +1,7 @@
 import json
 import os
-import time
+import threading
+from urllib.parse import urlparse
 from urllib.request import HTTPError, Request, URLError, urlopen
 
 from metaflow import util
@@ -14,6 +15,7 @@ from metaflow.mflog import (
 )
 import requests
 from metaflow.metaflow_config_funcs import init_config
+from metaflow_extensions.outerbounds.plugins.nvcf.heartbeat_store import HeartbeatStore
 
 
 class NvcfException(MetaflowException):
@@ -58,10 +60,11 @@ class Nvcf(object):
             code_package_url, code_package_ds
         )
         init_expr = " && ".join(init_cmds)
+        heartbeat_expr = f'python -m metaflow_extensions.outerbounds.plugins.nvcf.heartbeat_store "$MAIN_PID" {code_package_ds} nvcf_heartbeats 1>> $MFLOG_STDOUT 2>> $MFLOG_STDERR'
         step_expr = bash_capture_logs(
             " && ".join(
                 self.environment.bootstrap_commands(step_name, code_package_ds)
-                + [step_cli]
+                + [step_cli + " & MAIN_PID=$!; " + heartbeat_expr]
             )
         )
 
@@ -87,7 +90,9 @@ class Nvcf(object):
             '${METAFLOW_INIT_SCRIPT:+eval \\"${METAFLOW_INIT_SCRIPT}\\"} && %s'
             % cmd_str
         )
-        self.job = Job('bash -c "%s"' % cmd_str, env)
+        self.job = Job(
+            'bash -c "%s"' % cmd_str, env, task_spec, self.datastore._storage_impl
+        )
         self.job.submit()
 
     def wait(self, stdout_location, stderr_location, echo=None):
@@ -137,8 +142,7 @@ result_endpoint = f"{nvcf_url}/v2/nvcf/pexec/status"
 
 
 class Job(object):
-    def __init__(self, command, env):
-
+    def __init__(self, command, env, task_spec, backend):
         self._payload = {
             "command": command,
             "env": {k: v for k, v in env.items() if v is not None},
@@ -170,6 +174,27 @@ class Job(object):
             if f["model_key"] == "metaflow_task_executor":
                 self._function_id = f["id"]
 
+        flow_name = task_spec.get("flow_name")
+        run_id = task_spec.get("run_id")
+        step_name = task_spec.get("step_name")
+        task_id = task_spec.get("task_id")
+        retry_count = task_spec.get("retry_count")
+
+        heartbeat_prefix = "/".join(
+            (flow_name, str(run_id), step_name, str(task_id), str(retry_count))
+        )
+
+        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
+        )
+        self.heartbeat_thread.start()
+
     def submit(self):
         try:
             headers = {
@@ -224,7 +249,6 @@ class Job(object):
 
     def _poll(self):
         try:
-            invocation_id = self._invocation_id
             headers = {
                 "Authorization": f"Bearer {self._ngc_api_key}",
                 "Content-Type": "application/json",

{ob-metaflow-extensions-1.1.84 → ob-metaflow-extensions-1.1.88}/metaflow_extensions/outerbounds/plugins/nvcf/nvcf_cli.py

@@ -4,7 +4,7 @@ import sys
 import time
 import traceback
 
-from metaflow import util
+from metaflow import util, Run
 from metaflow._vendor import click
 from metaflow.exception import METAFLOW_EXIT_DISALLOW_RETRY
 from metaflow.metadata.util import sync_local_metadata_from_datastore
@@ -27,6 +27,7 @@ from metaflow.metaflow_config import (
     CARD_GSROOT,
     KUBERNETES_SANDBOX_INIT_SCRIPT,
     OTEL_ENDPOINT,
+    NVIDIA_HEARTBEAT_THRESHOLD,
 )
 from metaflow.mflog import TASK_LOG_SOURCE
 
@@ -43,6 +44,65 @@ def nvcf():
     pass
 
 
+@nvcf.command(help="List steps / tasks running as an NVCF job.")
+@click.argument("run-id")
+@click.pass_context
+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}"
+            )
+
+    if running_invocations:
+        for each_invocation in running_invocations:
+            print(each_invocation)
+    else:
+        print("No running NVCF invocations for Run ID: %s" % run_id)
+
+
+@nvcf.command(help="Kill steps / tasks running as an NVCF job.")
+@click.argument("run-id")
+@click.pass_context
+def kill(ctx, run_id):
+    from metaflow_extensions.outerbounds.plugins.nvcf.heartbeat_store import (
+        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(
+                heartbeat_prefix=heartbeat_prefix,
+                main_pid=None,
+                storage_backend=ctx.obj.datastore_impl(datastore_root),
+            )
+            store.emit_tombstone(folder_name="nvcf_heartbeats")
+        else:
+            print("No running NVCF invocations for Run ID: %s" % run_id)
+
+
 @nvcf.command(
     help="Execute a single task using NVCF. This command calls the "
     "top-level step command inside a NVCF job with the given options. "
@@ -139,6 +199,7 @@ def step(ctx, step_name, code_package_sha, code_package_url, **kwargs):
         "METAFLOW_CARD_GSROOT": CARD_GSROOT,
         "METAFLOW_INIT_SCRIPT": KUBERNETES_SANDBOX_INIT_SCRIPT,
         "METAFLOW_OTEL_ENDPOINT": OTEL_ENDPOINT,
+        "METAFLOW_NVIDIA_HEARTBEAT_THRESHOLD": str(NVIDIA_HEARTBEAT_THRESHOLD),
     }
 
     env_deco = [deco for deco in node.decorators if deco.name == "environment"]

ob-metaflow-extensions-1.1.88/metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py

@@ -0,0 +1,53 @@
+# These two fields will show up within `metaflow_version` task metadata.
+# Setting to major version of ob-metaflow-extensions only, so we don't keep trying
+# (and failing) to keep this in sync with setup.py
+# E.g. "2.7.22.1+ob(v1)"
+__version__ = "v1"
+__mf_extensions__ = "ob"
+
+
+# Must match the signature of metaflow.plugins.aws.aws_client.get_aws_client
+# This function is called by the "userland" code inside tasks. Metaflow internals
+# will call the function in metaflow.plugins.aws.aws_client.get_aws_client directly.
+#
+# Unlike the original function, this wrapper will use the CSPR role if both of the following
+# conditions are met:
+#
+#  1. CSPR is set
+#  2. user didn't provide a role to assume explicitly.
+#
+def get_aws_client(
+    module, with_error=False, role_arn=None, session_vars=None, client_params=None
+):
+    import metaflow.plugins.aws.aws_client
+
+    from metaflow_extensions.outerbounds.plugins import USE_CSPR_ROLE_ARN_IF_SET
+
+    return metaflow.plugins.aws.aws_client.get_aws_client(
+        module,
+        with_error=with_error,
+        role_arn=role_arn or USE_CSPR_ROLE_ARN_IF_SET,
+        session_vars=session_vars,
+        client_params=client_params,
+    )
+
+
+# This should match the signature of metaflow.plugins.datatools.s3.S3.
+#
+# This assumes that "userland" code inside tasks will call this, while Metaflow
+# internals will call metaflow.plugins.datatools.s3.S3 directly.
+#
+# This wrapper will make S3() use the CSPR role if its set, and user didn't provide
+# a role to assume explicitly.
+def S3(*args, **kwargs):
+    import sys
+    import metaflow.plugins.datatools.s3
+    from metaflow_extensions.outerbounds.plugins import USE_CSPR_ROLE_ARN_IF_SET
+
+    if "role" not in kwargs or kwargs["role"] is None:
+        kwargs["role"] = USE_CSPR_ROLE_ARN_IF_SET
+
+    return metaflow.plugins.datatools.s3.S3(*args, **kwargs)
+
+
+from .. import profilers

{ob-metaflow-extensions-1.1.84 → ob-metaflow-extensions-1.1.88}/ob_metaflow_extensions.egg-info/PKG-INFO

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

{ob-metaflow-extensions-1.1.84 → ob-metaflow-extensions-1.1.88}/ob_metaflow_extensions.egg-info/SOURCES.txt

@@ -16,6 +16,7 @@ metaflow_extensions/outerbounds/plugins/kubernetes/kubernetes_client.py
 metaflow_extensions/outerbounds/plugins/nim/__init__.py
 metaflow_extensions/outerbounds/plugins/nim/nim_manager.py
 metaflow_extensions/outerbounds/plugins/nvcf/__init__.py
+metaflow_extensions/outerbounds/plugins/nvcf/heartbeat_store.py
 metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py
 metaflow_extensions/outerbounds/plugins/nvcf/nvcf_cli.py
 metaflow_extensions/outerbounds/plugins/nvcf/nvcf_decorator.py

{ob-metaflow-extensions-1.1.84 → ob-metaflow-extensions-1.1.88}/setup.py

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

ob-metaflow-extensions-1.1.84/metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py

@@ -1,10 +0,0 @@
-# These two fields will show up within `metaflow_version` task metadata.
-# Setting to major version of ob-metaflow-extensions only, so we don't keep trying
-# (and failing) to keep this in sync with setup.py
-# E.g. "2.7.22.1+ob(v1)"
-__version__ = "v1"
-__mf_extensions__ = "ob"
-
-# To support "from metaflow import get_aws_client"
-from metaflow.plugins.aws.aws_client import get_aws_client
-from .. import profilers