ob-metaflow-extensions 1.1.102__py2.py3-none-any.whl → 1.1.103__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/__init__.py

@@ -318,6 +318,7 @@ STEP_DECORATORS_DESC = [
     ),
     ("snowpark", ".snowpark.snowpark_decorator.SnowparkDecorator"),
     ("tensorboard", ".tensorboard.TensorboardDecorator"),
+    ("gpu_profile", ".profilers.gpu_profile_decorator.GPUProfileDecorator"),
 ]
 FLOW_DECORATORS_DESC = [("nim", ".nim.NimDecorator")]
 TOGGLE_STEP_DECORATOR = [
@@ -331,3 +332,7 @@ TOGGLE_CLI = ["-batch", "-step-functions", "-airflow"]
 ENVIRONMENTS_DESC = [
     ("fast-bakery", ".fast_bakery.docker_environment.DockerEnvironment")
 ]
+
+SECRETS_PROVIDERS_DESC = [
+    ("outerbounds", ".secrets.secrets.OuterboundsSecretsProvider"),
+]

metaflow_extensions/outerbounds/plugins/profilers/deco_injector.py

@@ -0,0 +1,70 @@
+from metaflow.exception import MetaflowException
+from collections import defaultdict
+
+
+class CardDecoratorInjector:
+    """
+    Mixin Useful for injecting @card decorators from other first class Metaflow decorators.
+    """
+
+    _first_time_init = defaultdict(dict)
+
+    @classmethod
+    def _get_first_time_init_cached_value(cls, step_name, card_id):
+        return cls._first_time_init.get(step_name, {}).get(card_id, None)
+
+    @classmethod
+    def _set_first_time_init_cached_value(cls, step_name, card_id, value):
+        cls._first_time_init[step_name][card_id] = value
+
+    def _card_deco_already_attached(self, step, card_id):
+        for decorator in step.decorators:
+            if decorator.name == "card":
+                if decorator.attributes["id"] and card_id in decorator.attributes["id"]:
+                    return True
+        return False
+
+    def _get_step(self, flow, step_name):
+        for step in flow:
+            if step.name == step_name:
+                return step
+        return None
+
+    def _first_time_init_check(self, step_dag_node, card_id):
+        """ """
+        return not self._card_deco_already_attached(step_dag_node, card_id)
+
+    def attach_card_decorator(
+        self,
+        flow,
+        step_name,
+        card_id,
+        card_type,
+        refresh_interval=5,
+    ):
+        """
+        This method is called `step_init` in your StepDecorator code since
+        this class is used as a Mixin
+        """
+        from metaflow import decorators as _decorators
+
+        if not all([card_id, card_type]):
+            raise MetaflowException(
+                "`card_id` and `card_type` must be set in the `CardDecoratorInjector` Mixin"
+            )
+
+        step_dag_node = self._get_step(flow, step_name)
+        if (
+            self._get_first_time_init_cached_value(step_name, card_id) is None
+        ):  # First check class level setting.
+            if self._first_time_init_check(step_dag_node, card_id):
+                self._set_first_time_init_cached_value(step_name, card_id, True)
+                _decorators._attach_decorators_to_step(
+                    step_dag_node,
+                    [
+                        "card:type=%s,id=%s,refresh_interval=%s"
+                        % (card_type, card_id, str(refresh_interval))
+                    ],
+                )
+            else:
+                self._set_first_time_init_cached_value(step_name, card_id, False)

metaflow_extensions/outerbounds/plugins/profilers/gpu_profile_decorator.py

@@ -0,0 +1,88 @@
+from datetime import datetime
+from metaflow.decorators import StepDecorator
+from ...profilers.gpu import GPUProfiler  # Fix import
+from .deco_injector import CardDecoratorInjector
+import threading
+
+
+class GPUProfileDecorator(StepDecorator):
+    name = "gpu_profile"
+
+    defaults = {
+        "include_artifacts": True,
+        "artifact_prefix": "gpu_profile_",
+        "interval": 1,
+    }
+
+    def step_init(
+        self, flow, graph, step_name, decorators, environment, flow_datastore, logger
+    ):
+        self.deco_injector = CardDecoratorInjector()
+        self.deco_injector.attach_card_decorator(
+            flow,
+            step_name,
+            "gpu_profile",
+            "blank",
+            refresh_interval=self.attributes["interval"],
+        )
+
+    def task_pre_step(
+        self,
+        step_name,
+        task_datastore,
+        metadata,
+        run_id,
+        task_id,
+        flow,
+        graph,
+        retry_count,
+        max_user_code_retries,
+        ubf_context,
+        inputs,
+    ):
+        self._profiler = GPUProfiler(
+            interval=self.attributes["interval"],
+            artifact_name=self.attributes["artifact_prefix"] + "data",
+        )
+
+    def task_decorate(
+        self, step_func, flow, graph, retry_count, max_user_code_retries, ubf_context
+    ):
+        from metaflow import current
+        from metaflow.cards import Markdown
+
+        if self.attributes["include_artifacts"]:
+            setattr(
+                flow,
+                self.attributes["artifact_prefix"] + "num_gpus",
+                len(self._profiler.devices),
+            )
+
+        current.card["gpu_profile"].append(
+            Markdown("# GPU profile for `%s`" % current.pathspec)
+        )
+        current.card["gpu_profile"].append(
+            Markdown(
+                "_Started at: %s_"
+                % datetime.now().astimezone().strftime("%Y-%m-%dT%H:%M:%S %z")
+            )
+        )
+        self._profiler._setup_card()
+        current.card["gpu_profile"].refresh()
+        self._update_thread = threading.Thread(
+            target=self._profiler._update_card, daemon=True
+        )
+        self._update_thread.start()
+
+        def wrapped_step_func():
+            try:
+                step_func()
+            finally:
+                try:
+                    results = self._profiler.finish()
+                except:
+                    results = {"error": "couldn't read profiler results"}
+                if self.attributes["include_artifacts"]:
+                    setattr(flow, self.attributes["artifact_prefix"] + "data", results)
+
+        return wrapped_step_func

metaflow_extensions/outerbounds/plugins/secrets/secrets.py

@@ -0,0 +1,173 @@
+from metaflow.plugins.secrets import SecretsProvider
+from typing import Dict
+
+import base64
+import json
+import requests
+import time
+
+
+class OuterboundsSecretsException(Exception):
+    pass
+
+
+class OuterboundsSecretsApiResponse:
+    def __init__(self, response):
+        self.response = response
+
+    @property
+    def secret_resource_id(self):
+        return self.response["secret_resource_id"]
+
+    @property
+    def secret_backend_type(self):
+        return self.response["secret_backend_type"]
+
+
+class OuterboundsSecretsProvider(SecretsProvider):
+    TYPE = "outerbounds"
+
+    def get_secret_as_dict(self, secret_id, options={}, role=None):
+        """
+        Supports a special way of specifying secrets sources in outerbounds using the format:
+            @secrets(sources=["outerbounds.<integrations_name>"])
+
+        When invoked it makes a requests to the integrations secrets metadata endpoint on the
+        keywest server to get the cloud resource id for a secret. It then uses that to invoke
+        secrets manager on the core oss and returns the secrets.
+        """
+        headers = {"Content-Type": "application/json", "Connection": "keep-alive"}
+        perimeter, integrations_secrets_metadata_url = self._get_secret_configs()
+        integration_name = secret_id
+        request_payload = {
+            "perimeter_name": perimeter,
+            "integration_name": integration_name,
+        }
+        response = self._make_request(
+            integrations_secrets_metadata_url, headers, request_payload
+        )
+        secret_resource_id = response.secret_resource_id
+        secret_backend_type = response.secret_backend_type
+
+        from metaflow.plugins.secrets.secrets_decorator import (
+            get_secrets_backend_provider,
+        )
+
+        secrets_provider = get_secrets_backend_provider(secret_backend_type)
+        secret_dict = secrets_provider.get_secret_as_dict(
+            secret_resource_id, options={}, role=role
+        )
+
+        # Outerbounds stores secrets as binaries. Hence we expect the returned secret to be
+        # {<cloud-secret-name>: <base64 encoded full secret>}. We decode the secret here like:
+        # 1. decode the base64 encoded full secret
+        # 2. load the decoded secret as a json
+        # 3. decode the base64 encoded values in the dict
+        # 4. return the decoded dict
+        binary_secret = next(iter(secret_dict.values()))
+        return self._decode_secret(binary_secret)
+
+    def _is_base64_encoded(self, data):
+        try:
+            if isinstance(data, str):
+                # Check if the string can be base64 decoded
+                base64.b64decode(data).decode("utf-8")
+                return True
+            return False
+        except Exception:
+            return False
+
+    def _decode_secret(self, secret):
+        try:
+            result = {}
+            secret_str = secret
+            if self._is_base64_encoded(secret):
+                # we check if the secret string is base64 encoded because the returned secret from
+                # AWS secret manager is base64 encoded while the secret from GCP is not
+                secret_str = base64.b64decode(secret).decode("utf-8")
+
+            secret_dict = json.loads(secret_str)
+            for key, value in secret_dict.items():
+                result[key] = base64.b64decode(value).decode("utf-8")
+
+            return result
+        except Exception as e:
+            raise OuterboundsSecretsException(f"Error decoding secret: {e}")
+
+    def _get_secret_configs(self):
+        from metaflow_extensions.outerbounds.remote_config import init_config
+        from os import environ
+
+        conf = init_config()
+        if "OBP_PERIMETER" in conf:
+            perimeter = conf["OBP_PERIMETER"]
+        else:
+            # if the perimeter is not in metaflow config, try to get it from the environment
+            perimeter = environ.get("OBP_PERIMETER", "")
+
+        if "OBP_INTEGRATIONS_SECRETS_METADATA_URL" in conf:
+            integrations_secrets_metadata_url = conf[
+                "OBP_INTEGRATIONS_SECRETS_METADATA_URL"
+            ]
+        else:
+            # if the integrations secrets metadata url is not in metaflow config, try to get it from the environment
+            integrations_secrets_metadata_url = environ.get(
+                "OBP_INTEGRATIONS_SECRETS_METADATA_URL", ""
+            )
+
+        if not perimeter:
+            raise OuterboundsSecretsException(
+                "No perimeter set. Please make sure to run `outerbounds configure <...>` command which can be found on the Ourebounds UI or reach out to your Outerbounds support team."
+            )
+
+        if not integrations_secrets_metadata_url:
+            raise OuterboundsSecretsException(
+                "No integrations secrets metadata url set. Please notify your Outerbounds support team about this issue."
+            )
+
+        return perimeter, integrations_secrets_metadata_url
+
+    def _make_request(self, url, headers: Dict, payload: Dict):
+        try:
+            from metaflow.metaflow_config import SERVICE_HEADERS
+
+            request_headers = {**headers, **(SERVICE_HEADERS or {})}
+        except ImportError:
+            headers = self.headers
+
+        retryable_status_codes = [409]
+        json_payload = json.dumps(payload)
+        for attempt in range(2):  # 0 = initial attempt, 1-2 = retries
+            response = requests.get(url, data=json_payload, headers=request_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.get(url, data=json_payload, headers=request_headers)
+        self._handle_error_response(response)
+        return OuterboundsSecretsApiResponse(response.json())
+
+    @staticmethod
+    def _handle_error_response(response: requests.Response):
+        if response.status_code >= 500:
+            raise OuterboundsSecretsException(
+                f"Server error: {response.text}. Please reach out to your Outerbounds support team."
+            )
+
+        body = response.json()
+        status_code = body.get("error", {}).get("statusCode", response.status_code)
+        if status_code == 404:
+            raise OuterboundsSecretsException(f"Secret not found: {body}")
+
+        if status_code >= 400:
+            try:
+                raise OuterboundsSecretsException(
+                    f"status_code={status_code}\t*{body['error']['details']['kind']}*\n{body['error']['details']['message']}"
+                )
+            except KeyError:
+                raise OuterboundsSecretsException(
+                    f"status_code={status_code} Unexpected error: {body}"
+                )

metaflow_extensions/outerbounds/profilers/gpu.py

@@ -187,6 +187,9 @@ class GPUMonitor:
         all_readings = []
         if self._current_file is None:
             return None
+
+        if not os.path.exists(self._current_file):
+            return None
         # Extract everything from the CVS File and store it in a list of dictionaries
         all_fields = ["gpu_id"] + MONITOR_FIELDS
         with open(self._current_file, "r") as _monitor_out:
@@ -317,11 +320,17 @@ def _update_charts(results, md_dict):
 # This code is adapted from: https://github.com/outerbounds/monitorbench
 class GPUProfiler:
     def __init__(
-        self, interval=1, monitor_batch_duration=200, artifact_name="gpu_profile_data"
+        self,
+        interval=1,
+        monitor_batch_duration=200,
+        artifact_name="gpu_profile_data",
+        max_check_timeout=60,
     ):
         self._interval = interval
+        self.max_check_timeout = max_check_timeout
         self._monitor_batch_duration = monitor_batch_duration
         self.artifact_name = artifact_name
+        self._started_at = datetime.now().astimezone().strftime("%Y-%m-%dT%H:%M:%S %z")
         self._card_setup_finished = False
         self._card_comps = {
             "max_utilization": {},
@@ -334,14 +343,14 @@ class GPUProfiler:
 
     def _start_monitor(self):
         self.driver_ver, self.cuda_ver, self.error = self._read_versions()
-        (
-            self.interconnect_data,
-            self.interconnect_legend,
-        ) = self._read_multi_gpu_interconnect()
         if self.error:
             self.devices = []
             return
         else:
+            (
+                self.interconnect_data,
+                self.interconnect_legend,
+            ) = self._read_multi_gpu_interconnect()
             self.devices = self._read_devices()
             self._monitor = GPUMonitor(
                 interval=self._interval, duration=self._monitor_batch_duration
@@ -391,6 +400,12 @@ class GPUProfiler:
     def _update_card(self):
         if len(self.devices) == 0:
             current.card["gpu_profile"].clear()
+            current.card["gpu_profile"].append(
+                Markdown("# GPU profile for `%s`" % current.pathspec)
+            )
+            current.card["gpu_profile"].append(
+                Markdown("_Started at: %s_" % self._started_at)
+            )
             current.card["gpu_profile"].append(
                 Markdown("## GPU profile failed: %s" % self.error)
             )
@@ -398,14 +413,36 @@ class GPUProfiler:
 
             return
 
+        _check_time = 0
+        stop_checking = False
+        # Before writing anything to the card, we need to make sure that:
+        # 1. GPU Monitor has started.
+        # 2. Monitor can record readings
+        # 3. Card is setup
         while True:
+
+            if stop_checking:
+                time.sleep(self._interval)
+                continue
+
             # There is a possibility that the `monitor` thread is not started yet
             # because it somehow crashed at the very start.
+            if not self._monitor_started and _check_time > self.max_check_timeout:
+                current.card["gpu_profile"].clear()
+                current.card["gpu_profile"].append(
+                    Markdown("## GPU profile failed: %s" % self.error)
+                )
+                current.card["gpu_profile"].refresh()
+                stop_checking = True
+
+            # Try restarting monitor if it hasn't started yet
             if not self._monitor_started:
                 self._start_monitor()
+                _check_time += self._interval
                 time.sleep(self._interval)
                 continue
 
+            # Ensure that we are getting well formatted readings
             readings = self._make_reading()
 
             if readings is None:
@@ -413,6 +450,7 @@ class GPUProfiler:
                 time.sleep(self._interval)
                 continue
 
+            # ensure that the card is setup
             if not self._card_setup_finished:
                 self._setup_card()
                 time.sleep(self._interval)
@@ -642,41 +680,15 @@ class gpu_profile:
     def __call__(self, f):
         @wraps(f)
         def func(s):
-            prof = GPUProfiler(
-                interval=self.interval, artifact_name=self.artifact_prefix + "data"
-            )
-            if self.include_artifacts:
-                setattr(s, self.artifact_prefix + "num_gpus", len(prof.devices))
+            return f(s)
 
-            current.card["gpu_profile"].append(
-                Markdown("# GPU profile for `%s`" % current.pathspec)
-            )
-            current.card["gpu_profile"].append(
-                Markdown(
-                    "_Started at: %s_"
-                    % datetime.now().astimezone().strftime("%Y-%m-%dT%H:%M:%S %z")
-                )
-            )
-            prof._setup_card()
-            current.card["gpu_profile"].refresh()
-            update_thread = threading.Thread(target=prof._update_card, daemon=True)
-            update_thread.start()
-
-            try:
-                f(s)
-            finally:
-                try:
-                    results = prof.finish()
-                except:
-                    results = {"error": "couldn't read profiler results"}
-                if self.include_artifacts:
-                    setattr(s, self.artifact_prefix + "data", results)
-
-        from metaflow import card
-
-        return card(type="blank", id="gpu_profile", refresh_interval=self.interval)(
-            func
-        )
+        from metaflow import gpu_profile
+
+        return gpu_profile(
+            include_artifacts=self.include_artifacts,
+            artifact_prefix=self.artifact_prefix,
+            interval=self.interval,
+        )(func)
 
 
 def translate_to_vegalite(

{ob_metaflow_extensions-1.1.102.dist-info → ob_metaflow_extensions-1.1.103.dist-info}/METADATA

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

{ob_metaflow_extensions-1.1.102.dist-info → ob_metaflow_extensions-1.1.103.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=7-n1jwyvdNkMo3oU9CMF-YABr6jWRqlI-g8u9LexVR0,12464
+metaflow_extensions/outerbounds/plugins/__init__.py,sha256=lF8tJA_llp-VjJeNUFKnwOx9DNgJPTyNWAYzDIaORU8,12639
 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
@@ -18,6 +18,10 @@ metaflow_extensions/outerbounds/plugins/nvcf/heartbeat_store.py,sha256=pOWwm8LFQ
 metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py,sha256=1ayR5YLiEy3XRGVEnRUpsDCH_UyK7z-7A8L5GxI4qOE,9535
 metaflow_extensions/outerbounds/plugins/nvcf/nvcf_cli.py,sha256=YZb5AvbVgUwUJVxRxQ4JqqP8e1RMJr6dZ9U4KkHE-M8,9134
 metaflow_extensions/outerbounds/plugins/nvcf/nvcf_decorator.py,sha256=TKWN3mG2twC40uChfWYVsNKqf3euQIM_YHLAuNf1pvA,6525
+metaflow_extensions/outerbounds/plugins/profilers/deco_injector.py,sha256=oI_C3c64XBm7n88FILqHwn-Nnc5DeT_68I67lM9rXaI,2434
+metaflow_extensions/outerbounds/plugins/profilers/gpu_profile_decorator.py,sha256=gDHQ2sMIp4NuZSzUspbSd8RGdFAoO5mgZAyFcZ2a51Y,2619
+metaflow_extensions/outerbounds/plugins/secrets/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+metaflow_extensions/outerbounds/plugins/secrets/secrets.py,sha256=OSKm0aMKGDimqPLC7qzsUnXCEMpRfRDpQJN9Z8b4jo8,6796
 metaflow_extensions/outerbounds/plugins/snowpark/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark.py,sha256=0R8aFN9MpgWraqiaI6ZF82YpLdFJ1f-3z_-BPRpZfxM,10674
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_cli.py,sha256=ErsVoCQLa33byiykOQzDEeEkRKk0mgffZme43f3jxn4,8747
@@ -28,13 +32,13 @@ metaflow_extensions/outerbounds/plugins/snowpark/snowpark_job.py,sha256=d_5UhXqZ
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_service_spec.py,sha256=AI_kcm1hZV3JRxJkookcH6twiGnAYjk9Dx-MeoYz60Y,8511
 metaflow_extensions/outerbounds/plugins/tensorboard/__init__.py,sha256=9lUM4Cqi5RjrHBRfG6AQMRz8-R96eZC8Ih0KD2lv22Y,1858
 metaflow_extensions/outerbounds/profilers/__init__.py,sha256=wa_jhnCBr82TBxoS0e8b6_6sLyZX0fdHicuGJZNTqKw,29
-metaflow_extensions/outerbounds/profilers/gpu.py,sha256=7yeGa_ji1NE_JpvWne8yjuKhykAphiMNpCIVBvU8aiw,27066
+metaflow_extensions/outerbounds/profilers/gpu.py,sha256=3Er8uKQzfm_082uadg4yn_D4Y-iSCgzUfFmguYxZsz4,27485
 metaflow_extensions/outerbounds/toplevel/__init__.py,sha256=qWUJSv_r5hXJ7jV_On4nEasKIfUCm6_UjkjXWA_A1Ts,90
 metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py,sha256=Zq3OuL1bOod8KJra-Zk8B3gNhSHoWEGteM9T7g0pp6E,1881
 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.102.dist-info/METADATA,sha256=RJFw6nwnROMpizP-sTqxIQCDjszYk23JIva_ps9-IRw,521
-ob_metaflow_extensions-1.1.102.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.1.102.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.1.102.dist-info/RECORD,,
+ob_metaflow_extensions-1.1.103.dist-info/METADATA,sha256=jydJTGUDvrFg4bhCVxtiF4IL584lR73yr2Iv0LvseMA,521
+ob_metaflow_extensions-1.1.103.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.1.103.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.1.103.dist-info/RECORD,,