ob-metaflow-extensions 1.1.99__tar.gz → 1.1.100__tar.gz

{ob-metaflow-extensions-1.1.99 → ob-metaflow-extensions-1.1.100}/PKG-INFO

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

{ob-metaflow-extensions-1.1.99 → ob-metaflow-extensions-1.1.100}/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
@@ -315,7 +316,23 @@ 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"
+    ):
+        self._interval = interval
+        self._monitor_batch_duration = monitor_batch_duration
+        self.artifact_name = artifact_name
+        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,
@@ -327,16 +344,14 @@ class GPUProfiler:
         else:
             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 = {
@@ -384,21 +399,60 @@ class GPUProfiler:
             return
 
         while True:
+            # 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:
+                self._start_monitor()
+                time.sleep(self._interval)
+                continue
+
             readings = self._make_reading()
+
             if readings is None:
                 print("GPU Profiler readings are none", file=sys.stderr)
                 time.sleep(self._interval)
                 continue
+
+            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 +498,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 +533,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,7 +642,9 @@ class gpu_profile:
     def __call__(self, f):
         @wraps(f)
         def func(s):
-            prof = GPUProfiler(interval=self.interval)
+            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))
 
@@ -585,7 +657,7 @@ class gpu_profile:
                     % datetime.now().astimezone().strftime("%Y-%m-%dT%H:%M:%S %z")
                 )
             )
-            prof._setup_card(self.artifact_prefix + "data")
+            prof._setup_card()
             current.card["gpu_profile"].refresh()
             update_thread = threading.Thread(target=prof._update_card, daemon=True)
             update_thread.start()

{ob-metaflow-extensions-1.1.99 → ob-metaflow-extensions-1.1.100}/ob_metaflow_extensions.egg-info/PKG-INFO

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

{ob-metaflow-extensions-1.1.99 → ob-metaflow-extensions-1.1.100}/setup.py

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