ob-metaflow-extensions 1.1.99__py2.py3-none-any.whl → 1.1.100__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
@@ -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.dist-info → ob_metaflow_extensions-1.1.100.dist-info}/METADATA

@@ -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.dist-info → ob_metaflow_extensions-1.1.100.dist-info}/RECORD

@@ -28,13 +28,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=a5YZAepujuP0uDqG9UpXBlZS3wjUt4Yv8CjybXqeT2c,24342
+metaflow_extensions/outerbounds/profilers/gpu.py,sha256=7yeGa_ji1NE_JpvWne8yjuKhykAphiMNpCIVBvU8aiw,27066
 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.99.dist-info/METADATA,sha256=oBe6bSzZL9iBxorXqhVdGW_yKtv4JupzrO9gXsF21lw,520
-ob_metaflow_extensions-1.1.99.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.1.99.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.1.99.dist-info/RECORD,,
+ob_metaflow_extensions-1.1.100.dist-info/METADATA,sha256=JHOEhTwZ33-QppRgWmYPlwFz-yJC8EsjM8bZsai0hA0,521
+ob_metaflow_extensions-1.1.100.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.1.100.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.1.100.dist-info/RECORD,,