ob-metaflow-extensions 1.1.45rc3__py2.py3-none-any.whl → 1.5.1__py2.py3-none-any.whl

metaflow_extensions/outerbounds/profilers/gpu.py

@@ -13,6 +13,7 @@ import json
 import sys
 from tempfile import TemporaryDirectory
 from subprocess import check_output, Popen
+import subprocess
 from datetime import datetime, timedelta
 from functools import wraps
 from collections import namedtuple
@@ -186,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:
@@ -315,28 +319,48 @@ 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):
+    def __init__(
+        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": {},
+            "charts": {},
+            "reading_duration": {},
+            "error_component": None,
+        }
+        self._monitor_started = False
+        self._start_monitor()
+
+    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=interval, duration=monitor_batch_duration
+                interval=self._interval, duration=self._monitor_batch_duration
             )
             self._monitor_thread = threading.Thread(
                 target=self._monitor._monitor_update_thread, daemon=True
             )
+            self.error = None
             self._monitor_thread.start()
-            self._interval = interval
-
-        self._card_comps = {"max_utilization": {}, "charts": {}, "reading_duration": {}}
-        self._card_created = False
+            self._monitor_started = True
 
     def finish(self):
         ret = {
@@ -376,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)
             )
@@ -383,22 +413,84 @@ 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:
                 print("GPU Profiler readings are none", file=sys.stderr)
                 time.sleep(self._interval)
                 continue
+
+            # ensure that the card is setup
+            if not self._card_setup_finished:
+                self._setup_card()
+                time.sleep(self._interval)
+                continue
+
             _update_utilization(readings, self._card_comps["max_utilization"])
             _update_charts(readings, self._card_comps["charts"])
             current.card["gpu_profile"].refresh()
             time.sleep(self._interval)
 
-    def _setup_card(self, artifact_name):
+    def _setup_card(self):
         from metaflow import current
 
         results = self._make_reading()
+        if "profile" not in results:
+            if self._card_comps["error_component"] is None:
+                self._card_comps["error_component"] = Markdown(
+                    "## GPU profile failed: %s" % results["error"]
+                )
+                current.card["gpu_profile"].append(self._card_comps["error_component"])
+            else:
+                self._card_comps["error_component"].update(
+                    Markdown("## GPU profile failed: %s" % results["error"])
+                )
+                current.card["gpu_profile"].refresh()
+            return
+
         els = current.card["gpu_profile"]
+        self._card_comps["error_component"] = None
+        els.clear()
+
+        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")
+            )
+        )
 
         def _drivers():
             els.append(Markdown("## Drivers"))
@@ -444,7 +536,7 @@ class GPUProfiler:
                 Table(data=_rows, headers=["Device ID", "Max GPU %", "Max memory"])
             )
             els.append(
-                Markdown(f"Detailed data saved in an artifact `{artifact_name}`")
+                Markdown(f"Detailed data saved in an artifact `{self.artifact_name}`")
             )
             return rows
 
@@ -479,20 +571,36 @@ class GPUProfiler:
         _interconnect()
         self._card_comps["max_utilization"] = _utilization()
         self._card_comps["charts"] = _plots()
+        self._card_setup_finished = True
 
     def _read_versions(self):
         def parse(r, s):
             return r.search(s).group(1).strip().decode("utf-8")
 
         try:
-            out = check_output(["nvidia-smi"])
+            result = subprocess.run(
+                ["nvidia-smi"],
+                check=True,  # This will raise a CalledProcessError if the command fails
+                stdout=subprocess.PIPE,  # Capture stdout
+                stderr=subprocess.PIPE,  # Capture stderr
+            )
+            # Access the standard output
+            out = result.stdout
             return parse(DRIVER_VER, out), parse(CUDA_VER, out), None
         except FileNotFoundError:
             return None, None, "nvidia-smi not found"
         except AttributeError:
             return None, None, "nvidia-smi output is unexpected"
-        except:
-            return None, None, "nvidia-smi error"
+        except subprocess.CalledProcessError as e:
+            _error_message = "nvidia-smi error (CalledProcessError calling nvidia-smi)"
+            if e.stderr is not None:
+                _error_message = (
+                    "nvidia-smi error (CalledProcessError stderr) \n %s \n %s"
+                    % (e.stderr.decode("utf-8"), e.stdout.decode("utf-8"))
+                )
+            return None, None, _error_message
+        except Exception as e:
+            return None, None, "nvidia-smi error (unknown error) \n%s" % str(e)
 
     def _read_devices(self):
         out = check_output(
@@ -572,39 +680,15 @@ class gpu_profile:
     def __call__(self, f):
         @wraps(f)
         def func(s):
-            prof = GPUProfiler(interval=self.interval)
-            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(self.artifact_prefix + "data")
-            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(

metaflow_extensions/outerbounds/remote_config.py

@@ -8,10 +8,14 @@ import requests
 from metaflow.exception import MetaflowException
 from requests.models import HTTPError
 from metaflow_extensions.outerbounds.plugins.perimeters import (
-    CURRENT_PERIMETER_KEY,
-    CURRENT_PERIMETER_URL,
+    get_perimeter_config_url_if_set_in_ob_config,
 )
 
+
+class OuterboundsConfigException(MetaflowException):
+    _OB_CONFIG_EXCEPTION = True
+
+
 OBP_REMOTE_CONFIG_KEY = "OBP_METAFLOW_CONFIG_URL"
 HOSTNAME_KEY = "OBP_API_SERVER"
 AUTH_KEY = "METAFLOW_SERVICE_AUTH_KEY"
@@ -32,7 +36,7 @@ def read_config_from_local() -> Optional[Path]:
 
     # we should error because the user wants a specific config
     if profile:
-        raise MetaflowException(
+        raise OuterboundsConfigException(
             f"Unable to locate METAFLOW_PROFILE {profile} in {config_path}"
         )
 
@@ -55,11 +59,24 @@ def resolve_config_from_remote(remote_url: str, auth_token: str) -> Dict[str, st
         response.raise_for_status()
         data = response.json()
         return data["config"]
-    except HTTPError:
-        raise MetaflowException(
-            "Error fetching resolving configuration. Make sure you have run \
-                                `outerbounds configure` with the correct value"
-        )
+    except HTTPError as e:
+        if e.response.status_code >= 500:
+            raise OuterboundsConfigException(
+                "Error resolving outerbounds configuration [status:%s]. Please reach out to  "
+                "the outerbounds support team to help resolve this issue."
+                % e.response.status_code
+            )
+        elif e.response.status_code == 403:
+            raise OuterboundsConfigException(
+                "Outerbounds token validity expired [status:%s]. Please re-run the `outerbounds configure` "
+                "command with a new magic string from the UI." % e.response.status_code
+            )
+        else:
+            raise OuterboundsConfigException(
+                "Failed to fetch the outerbounds configuration string [status:%s]. Please reach out to  "
+                "the outerbounds support team to help resolve this issue."
+                % e.response.status_code
+            )
 
 
 def init_config() -> Dict[str, str]:
@@ -70,8 +87,8 @@ def init_config() -> Dict[str, str]:
     because steps are executed in subprocesses (local) or environments which expect environment variables to be set.
     """
     _init_debug("starting initialization")
-
-    if config_json := os.environ.get(CONFIG_READ_ONCE_KEY):
+    config_json = os.environ.get(CONFIG_READ_ONCE_KEY)
+    if config_json:
         _init_debug("reading config from environment")
         return json.loads(config_json)
 
@@ -82,18 +99,19 @@ def init_config() -> Dict[str, str]:
     try:
         remote_config = json.loads(config_path.read_text())
     except ValueError:
-        raise MetaflowException(
-            "Error decoding your metaflow config. Please run the `outerbounds configure` \
-                                 command with the string provided in the Outerbounds dashboard"
+        raise OuterboundsConfigException(
+            "Error decoding your metaflow config. Please re-run the `outerbounds configure` "
+            "command with a new magic string from the UI."
         )
 
+    perimeter_config_url = get_perimeter_config_url_if_set_in_ob_config()
+    if perimeter_config_url:
+        remote_config[OBP_REMOTE_CONFIG_KEY] = perimeter_config_url
+
     # users still have a legacy format and that's ok.
     if OBP_REMOTE_CONFIG_KEY not in remote_config:
         return remote_config
 
-    if CURRENT_PERIMETER_KEY in os.environ and CURRENT_PERIMETER_URL in os.environ:
-        remote_config[OBP_REMOTE_CONFIG_KEY] = os.environ[CURRENT_PERIMETER_URL]
-
     metaflow_config = resolve_config_from_remote(
         remote_url=remote_config[OBP_REMOTE_CONFIG_KEY],
         auth_token=remote_config[AUTH_KEY],
@@ -112,3 +130,22 @@ def _init_debug(*args, **kwargs):
         init_str = "ob_extension_init:"
         kwargs["file"] = sys.stderr
         print(init_str, *args, **kwargs)
+
+
+def reload_config():
+    """
+    This function is used to reload the config. Currently its a best effort implementation
+    that will only reload auth token.
+    """
+    _init_debug("reloading config")
+    if CONFIG_READ_ONCE_KEY in os.environ:
+        del os.environ[CONFIG_READ_ONCE_KEY]
+    config = init_config()
+    import metaflow.metaflow_config
+
+    metaflow.metaflow_config.SERVICE_AUTH_KEY = config.get(AUTH_KEY)
+    metaflow.metaflow_config.SERVICE_HEADERS["x-api-key"] = config.get(AUTH_KEY)
+    if config:
+        _init_debug("reloaded config")
+    else:
+        _init_debug("no config to reload")

metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py

@@ -5,6 +5,142 @@
 __version__ = "v1"
 __mf_extensions__ = "ob"
 
-# To support "from metaflow import get_aws_client"
-from metaflow.plugins.aws.aws_client import get_aws_client
+from metaflow_extensions.outerbounds.toplevel.s3_proxy import (
+    get_aws_client_with_s3_proxy,
+    get_S3_with_s3_proxy,
+)
+
+_S3_PROXY_CONFIG = None
+
+
+def set_s3_proxy_config(config):
+    global _S3_PROXY_CONFIG
+    _S3_PROXY_CONFIG = config
+
+
+def clear_s3_proxy_config():
+    global _S3_PROXY_CONFIG
+    _S3_PROXY_CONFIG = None
+
+
+def get_s3_proxy_config():
+    global _S3_PROXY_CONFIG
+    if _S3_PROXY_CONFIG is None:
+        set_s3_proxy_config(get_s3_proxy_config_from_env())
+    return _S3_PROXY_CONFIG
+
+
+# TODO: Refactor out the _S3_PROXY_CONFIG global variable and instead use the function that
+# extracts it from the environment variables.
+
+import os
+import json
+
+
+def get_s3_proxy_config_from_env():
+    env_conf = os.environ.get("METAFLOW_S3_PROXY_USER_CODE_CONFIG")
+    if env_conf:
+        return json.loads(env_conf)
+    return None
+
+
+# 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
+    import os
+
+    from metaflow_extensions.outerbounds.plugins import USE_CSPR_ROLE_ARN_IF_SET
+    from metaflow_extensions.outerbounds.plugins.aws.assume_role import (
+        OBP_ASSUME_ROLE_ARN_ENV_VAR,
+    )
+
+    # Check if the assume_role decorator has set a role ARN via environment variable
+    # This takes precedence over CSPR but not over explicitly passed role_arn
+    if role_arn is None:
+        decorator_role_arn = os.environ.get(OBP_ASSUME_ROLE_ARN_ENV_VAR)
+        if decorator_role_arn:
+            role_arn = decorator_role_arn
+
+    if module == "s3" and get_s3_proxy_config() is not None:
+        return get_aws_client_with_s3_proxy(
+            module,
+            with_error,
+            role_arn,
+            session_vars,
+            client_params,
+            get_s3_proxy_config(),
+        )
+
+    client = 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,
+    )
+
+    return client
+
+
+# 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
+    import os
+    from metaflow_extensions.outerbounds.plugins import USE_CSPR_ROLE_ARN_IF_SET
+    from metaflow_extensions.outerbounds.plugins.aws.assume_role import (
+        OBP_ASSUME_ROLE_ARN_ENV_VAR,
+    )
+
+    # Check if the assume_role decorator has set a role ARN via environment variable
+    # This takes precedence over CSPR but not over explicitly passed role
+    if "role" not in kwargs or kwargs["role"] is None:
+        decorator_role_arn = os.environ.get(OBP_ASSUME_ROLE_ARN_ENV_VAR)
+        if decorator_role_arn:
+            kwargs["role"] = decorator_role_arn
+        else:
+            kwargs["role"] = USE_CSPR_ROLE_ARN_IF_SET
+
+    # Check if S3 proxy is active using module variable (like CSPR)
+    if get_s3_proxy_config() is not None:
+        return get_S3_with_s3_proxy(get_s3_proxy_config(), *args, **kwargs)
+
+    return metaflow.plugins.datatools.s3.S3(*args, **kwargs)
+
+
+# Setting the S3 client docstring in order to ensure that
+# stubs get generated properly.
+import metaflow.plugins.datatools.s3
+
+S3.__doc__ = metaflow.plugins.datatools.s3.S3.__doc__
+
 from .. import profilers
+from ..plugins.snowflake import Snowflake
+from ..plugins.checkpoint_datastores import nebius_checkpoints, coreweave_checkpoints
+from ..plugins.aws import assume_role
+from . import ob_internal
+from ..plugins.apps.core import (
+    AppDeployer,
+    DeployedApp,
+    bake_image,
+    BakedImage,
+    package_code,
+    PackagedCode,
+)

metaflow_extensions/outerbounds/toplevel/ob_internal.py

@@ -0,0 +1,4 @@
+from ..plugins.kubernetes.pod_killer import PodKiller
+from ..plugins.fast_bakery.baker import bake_image as internal_bake_image
+from ..plugins.apps import core as app_core
+from ..plugins.apps.core import AppDeployer

metaflow_extensions/outerbounds/toplevel/plugins/ollama/__init__.py

@@ -0,0 +1 @@
+__mf_promote_submodules__ = ["plugins.ollama"]

metaflow_extensions/outerbounds/toplevel/plugins/optuna/__init__.py

@@ -0,0 +1 @@
+__mf_promote_submodules__ = ["plugins.optuna"]

metaflow_extensions/outerbounds/toplevel/plugins/snowflake/__init__.py

@@ -0,0 +1 @@
+__mf_promote_submodules__ = ["plugins.snowflake"]

metaflow_extensions/outerbounds/toplevel/plugins/torchtune/__init__.py

@@ -0,0 +1 @@
+__mf_promote_submodules__ = ["plugins.torchtune"]

metaflow_extensions/outerbounds/toplevel/plugins/vllm/__init__.py

@@ -0,0 +1 @@
+__mf_promote_submodules__ = ["plugins.vllm"]

metaflow_extensions/outerbounds/toplevel/s3_proxy.py

@@ -0,0 +1,88 @@
+from metaflow_extensions.outerbounds.plugins import USE_CSPR_ROLE_ARN_IF_SET
+from metaflow.metaflow_config import AWS_SECRETS_MANAGER_DEFAULT_REGION
+from metaflow_extensions.outerbounds.plugins.s3_proxy.constants import (
+    DEFAULT_PROXY_HOST,
+    DEFAULT_PROXY_PORT,
+)
+
+
+def get_aws_client_with_s3_proxy(
+    module,
+    with_error=False,
+    role_arn=None,
+    session_vars=None,
+    client_params=None,
+    s3_config=None,
+):
+    if not client_params:
+        client_params = {}
+
+    client_params["region_name"] = client_params.get(
+        "region_name", s3_config.get("region")
+    )
+    client_params["endpoint_url"] = s3_config.get(
+        "endpoint_url", f"http://{DEFAULT_PROXY_HOST}:{DEFAULT_PROXY_PORT}"
+    )
+
+    import metaflow.plugins.aws.aws_client
+
+    client = 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,
+    )
+
+    def override_s3_proxy_host_header(request, **kwargs):
+        region = kwargs["region_name"]
+        request.headers["Host"] = f"s3.{region}.amazonaws.com"
+        if "x-ob-write-to" not in request.headers and "write_mode" in s3_config:
+            request.headers["x-ob-write-to"] = s3_config.get("write_mode")
+
+    client.meta.events.register("before-sign", override_s3_proxy_host_header)
+
+    return client
+
+
+def get_S3_with_s3_proxy(s3_config, *args, **kwargs):
+    if "region_name" not in kwargs:
+        kwargs["region_name"] = s3_config.get(
+            "region", AWS_SECRETS_MANAGER_DEFAULT_REGION
+        )
+
+    kwargs["endpoint_url"] = s3_config.get(
+        "endpoint_url", f"http://{DEFAULT_PROXY_HOST}:{DEFAULT_PROXY_PORT}"
+    )
+
+    import metaflow.plugins.datatools.s3
+
+    mf_s3 = metaflow.plugins.datatools.s3.S3(*args, **kwargs)
+
+    # Override reset_client to ensure proxy endpoint is preserved
+    original_reset_client = mf_s3._s3_client.reset_client
+
+    def proxy_reset_client():
+        original_reset_client()
+        import boto3
+
+        proxy_client = boto3.client(
+            "s3",
+            region_name=kwargs.get("region_name", s3_config.get("region")),
+            endpoint_url=s3_config.get("endpoint_url"),
+        )
+        mf_s3._s3_client._s3_client = proxy_client
+
+    mf_s3._s3_client.reset_client = proxy_reset_client
+    mf_s3._s3_client.reset_client()
+
+    def override_s3_proxy_host_header(request, **kwargs):
+        region = kwargs["region_name"]
+        request.headers["Host"] = f"s3.{region}.amazonaws.com"
+        if "x-ob-write-to" not in request.headers and "write_mode" in s3_config:
+            request.headers["x-ob-write-to"] = s3_config.get("write_mode")
+
+    mf_s3._s3_client._s3_client.meta.events.register(
+        "before-sign", override_s3_proxy_host_header
+    )
+    return mf_s3

{ob_metaflow_extensions-1.1.45rc3.dist-info → ob_metaflow_extensions-1.5.1.dist-info}/METADATA

@@ -1,13 +1,13 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.1.45rc3
+Version: 1.5.1
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial
 Description-Content-Type: text/markdown
 Requires-Dist: boto3
 Requires-Dist: kubernetes
-Requires-Dist: ob-metaflow (==2.11.0.4)
+Requires-Dist: ob-metaflow (==2.19.15.3)
 
 # Outerbounds platform package