ob-metaflow-extensions 1.1.42rc0__tar.gz → 1.1.43__tar.gz

{ob-metaflow-extensions-1.1.42rc0 → ob-metaflow-extensions-1.1.43}/PKG-INFO

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

ob-metaflow-extensions-1.1.43/metaflow_extensions/outerbounds/__init__.py

@@ -0,0 +1,9 @@
+import metaflow.metaflow_config_funcs
+
+from metaflow_extensions.outerbounds.remote_config import init_config
+
+# we want to overide OSS Metaflow's initialization behavior with our own to support remote configs
+# we're reassigning the METAFLOW_CONFIG variable because all downstream settings rely on it and
+# users still have the power to overriden them with environment variables
+metaflow.metaflow_config_funcs.METAFLOW_CONFIG = init_config()
+metaflow.metaflow_config_funcs.init_config = init_config

{ob-metaflow-extensions-1.1.42rc0 → ob-metaflow-extensions-1.1.43}/metaflow_extensions/outerbounds/config/__init__.py

@@ -1,3 +1,5 @@
 DEFAULT_AWS_CLIENT_PROVIDER = "obp"
 
 DEFAULT_AZURE_CLIENT_PROVIDER = "obp"
+
+DEFAULT_GCP_CLIENT_PROVIDER = "obp"

{ob-metaflow-extensions-1.1.42rc0 → ob-metaflow-extensions-1.1.43}/metaflow_extensions/outerbounds/plugins/__init__.py

@@ -1,4 +1,5 @@
 import os
+import json
 import tempfile
 from contextlib import contextmanager
 
@@ -161,3 +162,81 @@ class ObpAzureAuthProvider(object):
 
 
 AZURE_CLIENT_PROVIDERS_DESC = [("obp", ".ObpAzureAuthProvider")]
+
+import threading
+import time
+
+_gcp_client_cache = dict()
+
+
+def _get_cache_key():
+    return os.getpid(), threading.get_ident()
+
+
+class ObpGcpAuthProvider(object):
+    name = "obp"
+
+    @staticmethod
+    def get_gs_storage_client(*args, **kwargs):
+
+        import sys
+        from metaflow_extensions.outerbounds.plugins.auth_server import get_token
+
+        cache_key = _get_cache_key()
+        if _gcp_client_cache.get(cache_key):
+            # Don't cache the client for more than 5 minutes as it may have
+            # expired.
+            if _gcp_client_cache[cache_key]._created_at < time.time() - 300:
+                del _gcp_client_cache[cache_key]
+            else:
+                return _gcp_client_cache[cache_key]
+
+        from hashlib import sha256
+        from metaflow.util import get_username
+
+        user = get_username()
+
+        token_info = get_token("/generate/gcp")
+        token_file = "/tmp/obp_token." + sha256(user.encode("utf-8")).hexdigest()[:16]
+        credentials_file = (
+            "/tmp/obp_credentials." + sha256(user.encode("utf-8")).hexdigest()[:16]
+        )
+
+        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)
+
+        credentials_json = {
+            "type": "external_account",
+            "audience": f"//iam.googleapis.com/projects/{token_info['gcpProjectNumber']}/locations/global/workloadIdentityPools/{token_info['gcpWorkloadIdentityPool']}/providers/{token_info['gcpWorkloadIdentityPoolProvider']}",
+            "subject_token_type": "urn:ietf:params:oauth:token-type:jwt",
+            "token_url": "https://sts.googleapis.com/v1/token",
+            "service_account_impersonation_url": f"https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/{token_info['gcpServiceAccountEmail']}:generateAccessToken",
+            "credential_source": {
+                "file": token_file,
+                "format": {"type": "text"},
+            },
+        }
+
+        with tempfile.NamedTemporaryFile("w", delete=False) as f:
+            f.write(json.dumps(credentials_json))
+            tmp_credentials_file = f.name
+        os.rename(tmp_credentials_file, credentials_file)
+
+        os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = credentials_file
+        from google.cloud import storage
+
+        storage_client = storage.Client(project=token_info["gcpProjectId"])
+        storage_client._created_at = time.time()
+        _gcp_client_cache[cache_key] = storage_client
+        return storage_client
+
+    @staticmethod
+    def get_credentials(scopes, *args, **kwargs):
+        import google.auth
+
+        return google.auth.default(scopes=scopes)
+
+
+GCP_CLIENT_PROVIDERS_DESC = [("obp", ".ObpGcpAuthProvider")]

{ob-metaflow-extensions-1.1.42rc0 → ob-metaflow-extensions-1.1.43}/metaflow_extensions/outerbounds/plugins/auth_server.py

@@ -1,37 +1,22 @@
-import os
-import json
-import requests
 from urllib.parse import urlparse
 
+import requests
+
 
 def read_mf_config():
-    # Read configuration from $METAFLOW_HOME/config_<profile>.json.
-    home = os.environ.get("METAFLOW_HOME", "~/.metaflowconfig")
-    profile = os.environ.get("METAFLOW_PROFILE")
-    path_to_config = os.path.join(home, "config.json")
-    if profile:
-        path_to_config = os.path.join(home, "config_%s.json" % profile)
-    path_to_config = os.path.expanduser(path_to_config)
-    config = {}
-    if os.path.exists(path_to_config):
-        with open(path_to_config, encoding="utf-8") as f:
-            return json.load(f)
-    elif profile:
-        from metaflow.exception import MetaflowException
-
-        raise MetaflowException(
-            "Unable to locate METAFLOW_PROFILE '%s' in '%s')" % (profile, home)
-        )
-    return config
+    # this should be overridden with the resolved remote config here:
+    # obp-python-packages/ob-metaflow-extensions/metaflow_extensions/outerbounds/__init__.py
+    from metaflow.metaflow_config_funcs import METAFLOW_CONFIG
 
+    return METAFLOW_CONFIG
 
-def get_token(url_path):
+
+def get_token_url_and_headers(url_path):
     from metaflow.metaflow_config import (
         SERVICE_HEADERS,
-        from_conf,
         SERVICE_URL,
     )
-    from metaflow.exception import MetaflowException
+    from metaflow.metaflow_config import SERVICE_HEADERS, SERVICE_URL
 
     # Infer auth host from metadata service URL, unless it has been
     # specified explicitly. Take the MDS host and replace first part of
@@ -44,6 +29,13 @@ def get_token(url_path):
     assert url_path.startswith("/")
     url = "https://" + authServer + url_path
     headers = SERVICE_HEADERS
+    return url, headers
+
+
+def get_token(url_path):
+    from metaflow.exception import MetaflowException
+
+    url, headers = get_token_url_and_headers(url_path)
     try:
         r = requests.get(url, headers=headers)
         r.raise_for_status()

ob-metaflow-extensions-1.1.43/metaflow_extensions/outerbounds/profilers/__init__.py

@@ -0,0 +1 @@
+from .gpu import gpu_profile

ob-metaflow-extensions-1.1.43/metaflow_extensions/outerbounds/profilers/gpu.py

@@ -0,0 +1,701 @@
+from metaflow.cards import Markdown, Table, VegaChart
+from functools import wraps
+import threading
+from datetime import datetime
+from metaflow import current
+from metaflow.cards import Table, Markdown, VegaChart, Image
+import time
+from typing import List, Dict, Any, Union
+import re
+import os
+import uuid
+import json
+import sys
+from tempfile import TemporaryDirectory
+from subprocess import check_output, Popen
+from datetime import datetime, timedelta
+from functools import wraps
+from collections import namedtuple
+
+# Card plot styles
+MEM_COLOR = "#0c64d6"
+GPU_COLOR = "#ff69b4"
+
+NVIDIA_TS_FORMAT = "%Y/%m/%d %H:%M:%S"
+
+
+DRIVER_VER = re.compile(b"Driver Version: (.+?) ")
+CUDA_VER = re.compile(b"CUDA Version:(.*) ")
+
+MONITOR_FIELDS = [
+    "timestamp",
+    "gpu_utilization",
+    "memory_used",
+    "memory_total",
+]
+
+MONITOR = """nvidia-smi --query-gpu=pci.bus_id,timestamp,utilization.gpu,memory.used,memory.total --format=csv,noheader,nounits -l {interval};"""
+ProcessUUID = namedtuple("ProcessUUID", ["uuid", "start_time", "end_time"])
+
+
+def _get_uuid(time_duration=600):
+    frmt_str = "%Y-%m-%d-%H-%M-%S"
+    # Create a datetime range between the timerange values using current date as start date and time_duration as end date
+    start_date = datetime.now()
+    end_date = start_date + timedelta(seconds=time_duration)
+    datetime_range = start_date = (
+        datetime.now().strftime(frmt_str) + "_" + end_date.strftime(frmt_str)
+    )
+    uuid_str = uuid.uuid4().hex.replace("-", "") + "_" + datetime_range
+    return ProcessUUID(uuid_str, start_date, end_date)
+
+
+class AsyncProcessManager:
+    """
+    This class is responsible for managing the nvidia SMI subprocesses
+    """
+
+    processes: Dict[str, Dict] = {
+        # "procid": {
+        #     "proc": subprocess.Popen,
+        #     "started": time.time()
+        # }
+    }
+
+    @classmethod
+    def _register_process(cls, procid, proc):
+        cls.processes[procid] = {
+            "proc": proc,
+            "started": time.time(),
+        }
+
+    @classmethod
+    def get(cls, procid):
+        proc_dict = cls.processes.get(procid, None)
+        if proc_dict is not None:
+            return proc_dict["proc"], proc_dict["started"]
+        return None, None
+
+    @classmethod
+    def spawn(cls, procid, cmd, file):
+        proc = Popen(cmd, stdout=file)
+        cls._register_process(procid, proc)
+
+    @classmethod
+    def remove(cls, procid, delete_item=True):
+        if procid in cls.processes:
+            if cls.processes[procid]["proc"].stdout is not None:
+                cls.processes[procid]["proc"].stdout.close()
+            cls.processes[procid]["proc"].terminate()
+            cls.processes[procid]["proc"].wait()
+            if delete_item:
+                del cls.processes[procid]
+
+    @classmethod
+    def cleanup(cls):
+        for procid in cls.processes:
+            cls.remove(procid, delete_item=False)
+        cls.processes.clear()
+
+    @classmethod
+    def is_running(cls, procid):
+        if procid not in cls.processes:
+            return False
+        return cls.processes[procid]["proc"].poll() is None
+
+
+def _parse_timestamp(timestamp):
+    try:
+        ts = timestamp.split(".")[0]
+        return datetime.strptime(ts, NVIDIA_TS_FORMAT)
+    except ValueError:
+        return None
+
+
+class GPUMonitor:
+    """
+    The `GPUMonitor` class is designed to monitor GPU usage.
+
+    When an instance of `GPUMonitor` is created, it initializes with a specified `interval` and `duration`.
+    The `duration` is the timeperiod it will run the NVIDIA SMI command for and the `interval` is the timeperiod between each reading.
+    The class exposes a `_monitor_update_thread` method which runs as a background thread that continuously updates the GPU usage readings.
+    It will keep running unitl the `_finished` flag is set to `True`.
+
+    The class will statefully manage the the spawned NVIDI-SMI processes.
+    It will start a new NVIDI-SMI process after the current one has ran for the specified `duration`.
+    At a time this class will only maintain readings for the `_current_process` and will have all the aggregated
+    readings for the past processes stored in the `_past_readings` dictionary.
+    When a process finishes completion, the readings are appended to the `_past_readings` dictionary and a new process is started.
+
+    If the caller of this class wishes to read the GPU usage, they can call the `read` method which will return the readings in a dictionary format.
+    The `read` method will aggregate the readings from the `_current_readings` and `_past_readings`.
+    """
+
+    _started_processes: List[ProcessUUID] = []
+
+    _current_process: Union[ProcessUUID, None] = None
+
+    _current_readings: Dict[str, Any] = {}
+
+    _past_readings: Dict[str, Any] = {}
+
+    def __init__(self, interval=1, duration=300) -> None:
+        self._tempdir = TemporaryDirectory(prefix="gpu_card_monitor", dir="./")
+        self._interval = interval
+        self._duration = duration
+        self._finished = False
+
+    @property
+    def _current_file(self):
+        if self._current_process is None:
+            return None
+        return os.path.join(self._tempdir.name, self._current_process.uuid + ".csv")
+
+    def get_file_name(self, uuid):
+        return os.path.join(self._tempdir.name, uuid + ".csv")
+
+    def create_new_monitor(self):
+        uuid = _get_uuid(self._duration)
+        file = open(self.get_file_name(uuid.uuid), "w")
+        cmd = MONITOR.format(interval=self._interval, time_duration=self._duration)
+        AsyncProcessManager.spawn(uuid.uuid, ["bash", "-c", cmd], file)
+        self._started_processes.append(uuid)
+        self._current_process = uuid
+        return uuid
+
+    def clear_current_monitor(self):
+        if self._current_process is None:
+            return
+        AsyncProcessManager.remove(self._current_process.uuid)
+        self._current_process = None
+
+    def current_process_has_ended(self):
+        if self._current_process is None:
+            return True
+        return datetime.now() > self._current_process.end_time
+
+    def current_process_is_running(self):
+        if self._current_process is None:
+            return False
+        return AsyncProcessManager.is_running(self._current_process.uuid)
+
+    def _read_monitor(self):
+        """
+        Reads the monitor file and returns the readings in a dictionary format
+        """
+        all_readings = []
+        if self._current_file is None:
+            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:
+            for line in _monitor_out.readlines():
+                data = {}
+                fields = [f.strip() for f in line.split(",")]
+                if len(fields) == len(all_fields):
+                    # strip subsecond resolution from timestamps that doesn't align across devices
+                    for idx, _f in enumerate(all_fields):
+                        data[_f] = fields[idx]
+                    all_readings.append(data)
+                else:
+                    # expect that the last line may be truncated
+                    break
+
+        # Convert to dictionary format
+        devdata = {}
+        for reading in all_readings:
+            gpu_id = reading["gpu_id"]
+            if "timestamp" not in reading:
+                continue
+            if _parse_timestamp(reading["timestamp"]) is None:
+                continue
+            reading["timestamp"] = reading["timestamp"].split(".")[0]
+            if gpu_id not in devdata:
+                devdata[gpu_id] = {}
+
+            for i, field in enumerate(MONITOR_FIELDS):
+                if field not in devdata[gpu_id]:
+                    devdata[gpu_id][field] = []
+                devdata[gpu_id][field].append(reading[field])
+        return devdata
+
+    def _update_readings(self):
+        """
+        Core update function that checks if the current process has ended and if so, it will create a new monitor
+        otherwise sets the current readings to the readings from the monitor file
+        """
+        if self.current_process_has_ended() or not self.current_process_is_running():
+            self._update_past_readings()
+            self.clear_current_monitor()
+            self.create_new_monitor()
+            # Sleep for 1 seconds to allow the new process to start and we can make a reading
+            time.sleep(1)
+
+        readings = self._read_monitor()
+        if readings is None:
+            return
+        self._current_readings = readings
+
+    @staticmethod
+    def _make_full_reading(current, past):
+        if current is None:
+            return past
+        for gpu_id in current:
+            if gpu_id not in past:
+                past[gpu_id] = {}
+            for field in MONITOR_FIELDS:
+                if field not in past[gpu_id]:
+                    past[gpu_id][field] = []
+                past[gpu_id][field].extend(current[gpu_id][field])
+        return past
+
+    def read(self):
+        return self._make_full_reading(
+            self._current_readings, json.loads(json.dumps(self._past_readings))
+        )
+
+    def _update_past_readings(self):
+        if self._current_readings is None:
+            return
+        self._past_readings = self._make_full_reading(
+            self._current_readings, json.loads(json.dumps(self._past_readings))
+        )
+        self._current_readings = None
+
+    def cleanup(self):
+        self._finished = True
+        AsyncProcessManager.cleanup()
+        self._tempdir.cleanup()
+
+    def _monitor_update_thread(self):
+        while not self._finished:
+            self._update_readings()
+            time.sleep(self._interval)
+
+
+def _get_ts_range(_range):
+    if _range == "":
+        return "*No readings available*"
+    return "*Time range of charts: %s*" % _range
+
+
+def _update_utilization(results, md_dict):
+    for device, data in results["profile"].items():
+        if device not in md_dict:
+            print(
+                "Device %s not found in the GPU card layout. Skipping..." % device,
+                file=sys.stderr,
+            )
+            continue
+        md_dict[device]["gpu"].update(
+            "%2.1f%%" % max(map(float, data["gpu_utilization"]))
+        )
+        md_dict[device]["memory"].update("%dMB" % max(map(float, data["memory_used"])))
+
+
+def _update_charts(results, md_dict):
+    for device, data in results["profile"].items():
+        try:
+            if device not in md_dict:
+                continue
+            gpu_plot, mem_plot, ts_range = profile_plots(
+                device,
+                data["timestamp"],
+                data["gpu_utilization"],
+                data["memory_used"],
+                data["memory_total"],
+            )
+            md_dict[device]["gpu"].update(gpu_plot)
+            md_dict[device]["memory"].update(mem_plot)
+            md_dict[device]["reading_duration"].update(_get_ts_range(ts_range))
+        except ValueError as e:
+            # This is thrown when the date is unparsable. We can just safely ignore this.
+            print("ValueError: Could not parse date \n%s" % str(e), file=sys.stderr)
+
+
+# This code is adapted from: https://github.com/outerbounds/monitorbench
+class GPUProfiler:
+    def __init__(self, interval=1, monitor_batch_duration=200):
+        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.devices = self._read_devices()
+            self._monitor = GPUMonitor(
+                interval=interval, duration=monitor_batch_duration
+            )
+            self._monitor_thread = threading.Thread(
+                target=self._monitor._monitor_update_thread, daemon=True
+            )
+            self._monitor_thread.start()
+            self._interval = interval
+
+        self._card_comps = {"max_utilization": {}, "charts": {}, "reading_duration": {}}
+        self._card_created = False
+
+    def finish(self):
+        ret = {
+            "error": self.error,
+            "cuda_version": self.cuda_ver,
+            "driver_version": self.driver_ver,
+        }
+        if self.error:
+            return ret
+        else:
+            ret["devices"] = self.devices
+            ret["profile"] = self._monitor.read()
+            ret["interconnect"] = {
+                "data": self.interconnect_data,
+                "legend": self.interconnect_legend,
+            }
+            self._monitor.cleanup()
+            return ret
+
+    def _make_reading(self):
+        ret = {
+            "error": self.error,
+            "cuda_version": self.cuda_ver,
+            "driver_version": self.driver_ver,
+        }
+        if self.error:
+            return ret
+        else:
+            ret["devices"] = self.devices
+            ret["profile"] = self._monitor.read()
+            ret["interconnect"] = {
+                "data": self.interconnect_data,
+                "legend": self.interconnect_legend,
+            }
+            return ret
+
+    def _update_card(self):
+        if len(self.devices) == 0:
+            current.card["gpu_profile"].clear()
+            current.card["gpu_profile"].append(
+                Markdown("## GPU profile failed: %s" % self.error)
+            )
+            current.card["gpu_profile"].refresh()
+
+            return
+
+        while True:
+            readings = self._make_reading()
+            if readings is None:
+                print("GPU Profiler readings are none", file=sys.stderr)
+                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):
+        from metaflow import current
+
+        results = self._make_reading()
+        els = current.card["gpu_profile"]
+
+        def _drivers():
+            els.append(Markdown("## Drivers"))
+            els.append(
+                Table(
+                    [[results["cuda_version"], results["driver_version"]]],
+                    headers=["NVidia driver version", "CUDA version"],
+                )
+            )
+
+        def _devices():
+            els.append(Markdown("## Devices"))
+            rows = [
+                [d["device_id"], d["name"], d["memory"]] for d in results["devices"]
+            ]
+            els.append(Table(rows, headers=["Device ID", "Device type", "GPU memory"]))
+
+        def _interconnect():
+            if results["interconnect"]["data"] and results["interconnect"]["legend"]:
+                els.append(Markdown("## Interconnect"))
+                interconnect_data = results["interconnect"]["data"]
+                rows = list(interconnect_data.values())
+                rows = [list(transpose_row) for transpose_row in list(zip(*rows))]
+                els.append(Table(rows, headers=list(interconnect_data.keys())))
+                els.append(Markdown("#### Legend"))
+                els.append(
+                    Table(
+                        [list(results["interconnect"]["legend"].values())],
+                        headers=list(results["interconnect"]["legend"].keys()),
+                    )
+                )
+
+        def _utilization():
+            els.append(Markdown("## Maximum utilization"))
+            rows = {}
+            for d in results["devices"]:
+                rows[d["device_id"]] = {
+                    "gpu": Markdown("0%"),
+                    "memory": Markdown("0MB"),
+                }
+            _rows = [[Markdown(k)] + list(v.values()) for k, v in rows.items()]
+            els.append(
+                Table(data=_rows, headers=["Device ID", "Max GPU %", "Max memory"])
+            )
+            els.append(
+                Markdown(f"Detailed data saved in an artifact `{artifact_name}`")
+            )
+            return rows
+
+        def _plots():
+            els.append(Markdown("## GPU utilization and memory usage over time"))
+
+            rows = {}
+            for d in results["devices"]:
+                gpu_plot, mem_plot, ts_range = profile_plots(
+                    d["device_id"], [], [], [], []
+                )
+                rows[d["device_id"]] = {
+                    "gpu": VegaChart(gpu_plot),
+                    "memory": VegaChart(mem_plot),
+                    "reading_duration": Markdown(_get_ts_range(ts_range)),
+                }
+            for k, v in rows.items():
+                els.append(Markdown("### GPU Utilization for device : %s" % k))
+                els.append(v["reading_duration"])
+                els.append(
+                    Table(
+                        data=[
+                            [Markdown("GPU Utilization"), v["gpu"]],
+                            [Markdown("Memory usage"), v["memory"]],
+                        ]
+                    )
+                )
+            return rows
+
+        _drivers()
+        _devices()
+        _interconnect()
+        self._card_comps["max_utilization"] = _utilization()
+        self._card_comps["charts"] = _plots()
+
+    def _read_versions(self):
+        def parse(r, s):
+            return r.search(s).group(1).strip().decode("utf-8")
+
+        try:
+            out = check_output(["nvidia-smi"])
+            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"
+
+    def _read_devices(self):
+        out = check_output(
+            [
+                "nvidia-smi",
+                "--query-gpu=name,pci.bus_id,memory.total",
+                "--format=csv,noheader",
+            ]
+        )
+        return [
+            dict(
+                zip(("name", "device_id", "memory"), (x.strip() for x in l.split(",")))
+            )
+            for l in out.decode("utf-8").splitlines()
+        ]
+
+    def _read_multi_gpu_interconnect(self):
+        """
+        parse output of `nvidia-smi tomo -m`, such as this sample:
+
+            GPU0    GPU1    CPU Affinity    NUMA Affinity
+            GPU0     X      NV2     0-23            N/A
+            GPU1    NV2      X      0-23            N/A
+
+        returns two dictionaries describing multi-GPU topology:
+            data: {index: [GPU0, GPU1, ...], GPU0: [X, NV2, ...], GPU1: [NV2, X, ...], ...}
+            legend_items: {X: 'Same PCI', NV2: 'NVLink 2', ...}
+        """
+        try:
+            import re
+
+            ansi_escape = re.compile(r"(\x9B|\x1B\[)[0-?]*[ -\/]*[@-~]")
+
+            out = check_output(["nvidia-smi", "topo", "-m"])
+            rows = out.decode("utf-8").split("\n")
+
+            header = ansi_escape.sub("", rows[0]).split("\t")[1:]
+            data = {}
+            data["index"] = []
+            data |= {k: [] for k in header}
+
+            for i, row in enumerate(rows[1:]):
+                row = ansi_escape.sub("", row).split()
+                if len(row) == 0:
+                    continue
+                if row[0].startswith("GPU"):
+                    data["index"].append(row[0])
+                    for key, val in zip(header, row[1:]):
+                        data[key].append(val)
+                elif row[0].startswith("Legend"):
+                    break
+
+            legend_items = {}
+            for legend_row in rows[i:]:
+                if legend_row == "" or legend_row.startswith("Legend"):
+                    continue
+                res = legend_row.strip().split(" = ")
+                legend_items[res[0].strip()] = res[1].strip()
+
+            return data, legend_items
+
+        except:
+            return None, None
+
+
+class gpu_profile:
+    def __init__(
+        self,
+        include_artifacts=True,
+        artifact_prefix="gpu_profile_",
+        interval=1,
+    ):
+        self.include_artifacts = include_artifacts
+        self.artifact_prefix = artifact_prefix
+        self.interval = interval
+
+    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))
+
+            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
+        )
+
+
+def translate_to_vegalite(
+    tstamps,
+    vals,
+    description,
+    y_label,
+    legend,
+    line_color=None,
+    percentage_format=False,
+):
+    # Preprocessing for Vega-Lite
+    # Assuming tstamps is a list of datetime objects and vals is a list of values
+    data = [{"tstamps": str(t), "vals": v} for t, v in zip(tstamps, vals)]
+
+    # Base Vega-Lite spec
+    vega_lite_spec = {
+        "$schema": "https://vega.github.io/schema/vega-lite/v5.json",
+        "description": description,
+        "data": {"values": data},
+        "width": 600,
+        "height": 400,
+        "encoding": {
+            "x": {"field": "tstamps", "type": "temporal", "axis": {"title": "Time"}},
+            "y": {
+                "field": "vals",
+                "type": "quantitative",
+                "axis": {
+                    "title": y_label,
+                    **({"format": "%"} if percentage_format else {}),
+                },
+            },
+        },
+        "layer": [
+            {
+                "mark": {
+                    "type": "line",
+                    "color": line_color if line_color else "blue",
+                    "tooltip": True,
+                    "description": legend,  # Adding legend as description
+                },
+                "encoding": {"tooltip": [{"field": "tstamps"}, {"field": "vals"}]},
+            }
+        ],
+    }
+
+    return vega_lite_spec
+
+
+def profile_plots(device_id, ts, gpu, mem_used, mem_total):
+    tstamps = [datetime.strptime(t, NVIDIA_TS_FORMAT) for t in ts]
+    gpu = [i / 100 for i in list(map(float, gpu))]
+    mem = [float(used) / float(total) for used, total in zip(mem_used, mem_total)]
+    time_stamp_range = ""
+    if len(tstamps) > 1:
+        max_time = max(tstamps).strftime(NVIDIA_TS_FORMAT)
+        min_time = min(tstamps).strftime(NVIDIA_TS_FORMAT)
+        time_stamp_range = "%s to %s" % (min_time, max_time)
+
+    gpu_plot = translate_to_vegalite(
+        tstamps,
+        gpu,
+        "GPU utilization",
+        "GPU utilization",
+        "device: %s" % device_id,
+        line_color=GPU_COLOR,
+        percentage_format=True,
+    )
+    mem_plot = translate_to_vegalite(
+        tstamps,
+        mem,
+        "Percentage Memory utilization",
+        "Percentage Memory utilization",
+        "device: %s" % device_id,
+        line_color=MEM_COLOR,
+        percentage_format=True,
+    )
+    return gpu_plot, mem_plot, time_stamp_range
+
+
+if __name__ == "__main__":
+    prof = GPUProfiler(monitor_batch_duration=10)
+
+    def _write_json_file(data, filename):
+        with open(filename, "w") as f:
+            json.dump(data, f, indent=4)
+
+    import time
+
+    for i in range(15):
+        time.sleep(1)
+        _write_json_file(prof._monitor.read(), "gpu_profile.json")
+
+    print(json.dumps(prof.finish()))

ob-metaflow-extensions-1.1.43/metaflow_extensions/outerbounds/remote_config.py

@@ -0,0 +1,108 @@
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Optional
+
+import requests
+from metaflow.exception import MetaflowException
+from requests.models import HTTPError
+
+OBP_REMOTE_CONFIG_KEY = "OBP_METAFLOW_CONFIG_URL"
+HOSTNAME_KEY = "OBP_API_SERVER"
+AUTH_KEY = "METAFLOW_SERVICE_AUTH_KEY"
+PERIMETER_KEY = "OBP_PERIMETER"
+CONFIG_READ_ONCE_KEY = "__REMOTE_CONFIG_HAS_BEEN_RESOLVED__"
+
+
+def read_config_from_local() -> Optional[Path]:
+    default_path = Path.home() / ".metaflowconfig"
+    home = Path(os.environ.get("METAFLOW_HOME", default_path))
+
+    profile = os.environ.get("METAFLOW_PROFILE")
+    config_path = home / f"config_{profile}.json" if profile else home / "config.json"
+
+    if config_path.exists() and config_path.is_file():
+        _init_debug(f"using config from {config_path}")
+        return config_path
+
+    # we should error because the user wants a specific config
+    if profile:
+        raise MetaflowException(
+            f"Unable to locate METAFLOW_PROFILE {profile} in {config_path}"
+        )
+
+    # there's no config and that's ok. Metaflow uses environment variables as its primary way to set values
+    # and will fallback to local settings if no config is present
+    _init_debug(f"no config present at path {config_path}")
+    return None
+
+
+def resolve_config_from_remote(remote_url: str, auth_token: str) -> dict[str, str]:
+    _init_debug(f"retrieving config from {remote_url}")
+
+    headers = {"x-api-key": auth_token}
+    try:
+        response = requests.get(remote_url, headers=headers)
+        _init_debug(
+            f"response\nstatus code: {response.status_code}\nbody: {response.text}"
+        )
+
+        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"
+        )
+
+
+def init_config() -> dict[str, str]:
+    """
+    OSS Metaflow reads the config file on every step initialization. This is because OSS assumes config files change
+    relatively infrequently. We want to avoid config values changing between flow steps. Our solution to prevent this
+    is to read a config once and cache it on an environment variable. Environment variables carry over between steps
+    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):
+        _init_debug("reading config from environment")
+        return json.loads(config_json)
+
+    config_path = read_config_from_local()
+    if not config_path:
+        return {}
+
+    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"
+        )
+
+    # users still have a legacy format and that's ok.
+    if OBP_REMOTE_CONFIG_KEY not in remote_config:
+        return remote_config
+
+    metaflow_config = resolve_config_from_remote(
+        remote_url=remote_config[OBP_REMOTE_CONFIG_KEY],
+        auth_token=remote_config[AUTH_KEY],
+    )
+    metaflow_config[AUTH_KEY] = remote_config[AUTH_KEY]
+
+    # set cache
+    os.environ[CONFIG_READ_ONCE_KEY] = json.dumps(metaflow_config)
+    return metaflow_config
+
+
+DEBUG_CONFIG = os.environ.get("METAFLOW_DEBUG_CONFIG")
+
+
+def _init_debug(*args, **kwargs):
+    if DEBUG_CONFIG:
+        init_str = "ob_extension_init:"
+        kwargs["file"] = sys.stderr
+        print(init_str, *args, **kwargs)

ob-metaflow-extensions-1.1.43/metaflow_extensions/outerbounds/toplevel/plugins/gcp/__init__.py

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

{ob-metaflow-extensions-1.1.42rc0 → ob-metaflow-extensions-1.1.43}/ob_metaflow_extensions.egg-info/PKG-INFO

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

{ob-metaflow-extensions-1.1.42rc0 → ob-metaflow-extensions-1.1.43}/ob_metaflow_extensions.egg-info/SOURCES.txt

@@ -1,15 +1,18 @@
 README.md
 setup.py
 metaflow_extensions/outerbounds/__init__.py
+metaflow_extensions/outerbounds/remote_config.py
 metaflow_extensions/outerbounds/config/__init__.py
 metaflow_extensions/outerbounds/plugins/__init__.py
 metaflow_extensions/outerbounds/plugins/auth_server.py
-metaflow_extensions/outerbounds/plugins/perimeters.py
 metaflow_extensions/outerbounds/plugins/kubernetes/__init__.py
 metaflow_extensions/outerbounds/plugins/kubernetes/kubernetes_client.py
+metaflow_extensions/outerbounds/profilers/__init__.py
+metaflow_extensions/outerbounds/profilers/gpu.py
 metaflow_extensions/outerbounds/toplevel/__init__.py
 metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py
 metaflow_extensions/outerbounds/toplevel/plugins/azure/__init__.py
+metaflow_extensions/outerbounds/toplevel/plugins/gcp/__init__.py
 metaflow_extensions/outerbounds/toplevel/plugins/kubernetes/__init__.py
 ob_metaflow_extensions.egg-info/PKG-INFO
 ob_metaflow_extensions.egg-info/SOURCES.txt

ob-metaflow-extensions-1.1.43/ob_metaflow_extensions.egg-info/requires.txt

@@ -0,0 +1,3 @@
+boto3
+kubernetes
+ob-metaflow==2.11.0.4

{ob-metaflow-extensions-1.1.42rc0 → ob-metaflow-extensions-1.1.43}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_namespace_packages
 from pathlib import Path
 
 
-version = "1.1.42rc0"
+version = "1.1.43"
 this_directory = Path(__file__).parent
 long_description = (this_directory / "README.md").read_text()
 
@@ -15,5 +15,5 @@ setup(
     packages=find_namespace_packages(include=["metaflow_extensions.*"]),
     long_description=long_description,
     long_description_content_type="text/markdown",
-    install_requires=["boto3", "kubernetes", "ob-metaflow == 2.11.0.3"],
+    install_requires=["boto3", "kubernetes", "ob-metaflow == 2.11.0.4"],
 )

ob-metaflow-extensions-1.1.42rc0/metaflow_extensions/outerbounds/__init__.py

@@ -1,5 +0,0 @@
-from metaflow_extensions.outerbounds.plugins.perimeters import (
-    override_metaflow_profile_with_perimeter,
-)
-
-override_metaflow_profile_with_perimeter()

ob-metaflow-extensions-1.1.42rc0/metaflow_extensions/outerbounds/plugins/perimeters.py

@@ -1,27 +0,0 @@
-import os
-import fcntl
-from os import path
-import json
-
-
-def override_metaflow_profile_with_perimeter():
-    # If OBP_CONFIG_DIR is set, use that, otherwise use METAFLOW_HOME
-    # If neither are set, use ~/.metaflowconfig
-    obp_config_dir = path.expanduser(
-        os.environ.get(
-            "OBP_CONFIG_DIR", os.environ.get("METAFLOW_HOME", "~/.metaflowconfig")
-        )
-    )
-
-    file_path = os.path.join(obp_config_dir, "ob_config.json")
-
-    if os.path.exists(file_path):
-        # Acquire a shared read lock on the file
-        fd = os.open(file_path, os.O_RDONLY)
-        fcntl.flock(fd, fcntl.LOCK_SH)
-
-        with open(file_path, "r") as f:
-            ob_config = json.loads(f.read())
-
-        if "OB_CURRENT_PERIMETER" in ob_config:
-            os.environ["METAFLOW_PROFILE"] = ob_config["OB_CURRENT_PERIMETER"]

ob-metaflow-extensions-1.1.42rc0/ob_metaflow_extensions.egg-info/requires.txt

@@ -1,3 +0,0 @@
-boto3
-kubernetes
-ob-metaflow==2.11.0.3