cubestat: metric reg

Author: okuvshynov

cubestat/cubestat.py

@@ -14,7 +14,7 @@
 
 from cubestat.platforms.factory import get_platform
 
-from cubestat.metrics.registry import get_metrics, metrics_configure_argparse
+from cubestat.metrics_registry import get_metrics, metrics_configure_argparse
 
 
 class ViewMode(DisplayMode):
@@ -134,9 +134,20 @@ def start(stdscr, platform, args):
 def main():
     logging.basicConfig(filename='/tmp/cubestat.log', level=logging.INFO)
     parser = argparse.ArgumentParser("cubestat")
-    parser.add_argument('--refresh_ms', '-i', type=int, default=1000, help='Update frequency, milliseconds')
-    parser.add_argument('--buffer_size', type=int, default=500, help='How many datapoints to store. Having it larger than screen width is a good idea as terminal window can be resized')
-    parser.add_argument('--view', type=ViewMode, default=ViewMode.one, choices=list(ViewMode), help='legend/values/time mode. Can be toggled by pressing v.')
+    parser.add_argument(
+        '--refresh_ms', '-i', type=int, default=1000,
+        help='Update frequency (milliseconds)'
+    )
+
+    parser.add_argument(
+        '--buffer_size', type=int, default=500,
+        help='Number of datapoints to store. Consider setting it larger than the screen width to accommodate window resizing.'
+    )
+
+    parser.add_argument(
+        '--view', type=ViewMode, default=ViewMode.one, choices=list(ViewMode),
+        help='Display mode (legend, values, time). Hotkey: "v".'
+    )
 
     metrics_configure_argparse(parser)
     args = parser.parse_args()

cubestat/data.py

@@ -4,8 +4,13 @@
 
 class DataManager:
     def __init__(self, buffer_size):
-        init_series = lambda: collections.deque(maxlen=buffer_size)
-        init_group = lambda: collections.defaultdict(init_series)
+
+        def init_series():
+            return collections.deque(maxlen=buffer_size)
+
+        def init_group():
+            return collections.defaultdict(init_series)
+
         self.data = collections.defaultdict(init_group)
 
     def get_slice(self, series, indent, h_shift, cols, spacing):

cubestat/input.py

@@ -6,19 +6,7 @@ def __init__(self, horizon):
         self.horizon = horizon
         self.hotkeys = [(m.hotkey(), m) for m in self.horizon.metrics.values() if m.hotkey()]
 
-    def handle_input(self):
-        key = self.horizon.screen.stdscr.getch()
-        if key == ord('q') or key == ord('Q'):
-            exit(0)
-        for k, metric in self.hotkeys:
-            if key == ord(k):
-                with self.horizon.lock:
-                    metric.mode = metric.mode.next()
-                    self.horizon.settings_changed = True
-            if key == ord(k.upper()):
-                with self.horizon.lock:
-                    metric.mode = metric.mode.prev()
-                    self.horizon.settings_changed = True
+    def handle_shifts(self, key):
         if key == curses.KEY_UP:
             with self.horizon.lock:
                 if self.horizon.v_shift > 0:
@@ -38,6 +26,8 @@ def handle_input(self):
                 if self.horizon.h_shift > 0:
                     self.horizon.h_shift -= 1
                     self.horizon.settings_changed = True
+
+    def handle_reset(self, key):
         if key == ord('0'):
             with self.horizon.lock:
                 if self.horizon.v_shift > 0:
@@ -46,6 +36,20 @@ def handle_input(self):
                 if self.horizon.h_shift > 0:
                     self.horizon.h_shift = 0
                     self.horizon.settings_changed = True
+
+    def handle_input(self):
+        key = self.horizon.screen.stdscr.getch()
+        if key == ord('q') or key == ord('Q'):
+            exit(0)
+        for k, metric in self.hotkeys:
+            if key == ord(k):
+                with self.horizon.lock:
+                    metric.mode = metric.mode.next()
+                    self.horizon.settings_changed = True
+            if key == ord(k.upper()):
+                with self.horizon.lock:
+                    metric.mode = metric.mode.prev()
+                    self.horizon.settings_changed = True
         if key == ord('v'):
             with self.horizon.lock:
                 self.horizon.view = self.horizon.view.next()
@@ -62,3 +66,5 @@ def handle_input(self):
             with self.horizon.lock:
                 self.horizon.theme = self.horizon.theme.prev()
                 self.horizon.settings_changed = True
+        self.handle_shifts(key)
+        self.handle_reset(key)

cubestat/metrics/__init__.py

@@ -1,6 +1,6 @@
 import subprocess
 from cubestat.metrics.base_metric import base_metric
-from cubestat.metrics.registry import cubestat_metric
+from cubestat.metrics_registry import cubestat_metric
 
 
 @cubestat_metric('darwin')

cubestat/metrics/accel.py

@@ -2,7 +2,7 @@
 import psutil
 
 from cubestat.metrics.base_metric import base_metric
-from cubestat.metrics.registry import cubestat_metric
+from cubestat.metrics_registry import cubestat_metric
 from cubestat.common import DisplayMode
 
 
@@ -43,7 +43,13 @@ def hotkey(self):
 
     @classmethod
     def configure_argparse(cls, parser):
-        parser.add_argument('--cpu', type=CPUMode, default=auto_cpu_mode(), choices=list(CPUMode), help='CPU mode - showing all cores, only cumulative by cluster or both. Can be toggled by pressing c.')
+        parser.add_argument(
+            '--cpu',
+            type=CPUMode,
+            default=auto_cpu_mode(),
+            choices=list(CPUMode),
+            help='Select CPU mode: all cores, cumulative by cluster, or both. Hotkey: "c".'
+        )
 
 
 @cubestat_metric('linux')

cubestat/metrics/cpu.py

@@ -2,7 +2,7 @@
 
 from cubestat.common import SimpleMode, RateReader, label_bytes_per_sec
 from cubestat.metrics.base_metric import base_metric
-from cubestat.metrics.registry import cubestat_metric
+from cubestat.metrics_registry import cubestat_metric
 
 
 class disk_metric(base_metric):
@@ -23,7 +23,13 @@ def hotkey(self):
 
     @classmethod
     def configure_argparse(cls, parser):
-        parser.add_argument('--disk', type=SimpleMode, default=SimpleMode.show, choices=list(SimpleMode), help="Show disk read/write. Can be toggled by pressing d.")
+        parser.add_argument(
+            '--disk',
+            type=SimpleMode,
+            default=SimpleMode.show,
+            choices=list(SimpleMode),
+            help='Show disk read/write rate. Hotkey: "d"'
+        )
 
     def configure(self, conf):
         self.mode = conf.disk

cubestat/metrics/disk.py

@@ -3,7 +3,7 @@
 
 from cubestat.common import DisplayMode
 from cubestat.metrics.base_metric import base_metric
-from cubestat.metrics.registry import cubestat_metric
+from cubestat.metrics_registry import cubestat_metric
 
 
 class GPUMode(DisplayMode):
@@ -38,7 +38,13 @@ def hotkey(self):
 
     @classmethod
     def configure_argparse(cls, parser):
-        parser.add_argument('--gpu', type=GPUMode, default=GPUMode.load_only, choices=list(GPUMode), help='GPU mode - hidden, showing all GPUs load, or showing load and vram usage. Can be toggled by pressing g.')
+        parser.add_argument(
+            '--gpu',
+            type=GPUMode,
+            default=GPUMode.load_only,
+            choices=list(GPUMode),
+            help='GPU mode - hidden, load, or load and vram usage. Hotkey: "g"'
+        )
 
 
 @cubestat_metric('linux')

cubestat/metrics/gpu.py

@@ -1,7 +1,7 @@
 import psutil
 
 from cubestat.metrics.base_metric import base_metric
-from cubestat.metrics.registry import cubestat_metric
+from cubestat.metrics_registry import cubestat_metric
 from cubestat.common import label_bytes
 from cubestat.common import DisplayMode
 

cubestat/metrics/memory.py

@@ -2,7 +2,7 @@
 
 from cubestat.common import SimpleMode, RateReader, label_bytes_per_sec
 from cubestat.metrics.base_metric import base_metric
-from cubestat.metrics.registry import cubestat_metric
+from cubestat.metrics_registry import cubestat_metric
 
 
 class network_metric(base_metric):
@@ -28,7 +28,13 @@ def configure(self, conf):
 
     @classmethod
     def configure_argparse(cls, parser):
-        parser.add_argument('--network', type=SimpleMode, default=SimpleMode.show, choices=list(SimpleMode), help="Show network io. Can be toggled by pressing n.")
+        parser.add_argument(
+            '--network',
+            type=SimpleMode,
+            default=SimpleMode.show,
+            choices=list(SimpleMode),
+            help='Show network io. Hotkey: "n"'
+        )
 
 
 @cubestat_metric('darwin')

cubestat/metrics/network.py

@@ -1,6 +1,6 @@
 from cubestat.common import DisplayMode, label10
 from cubestat.metrics.base_metric import base_metric
-from cubestat.metrics.registry import cubestat_metric
+from cubestat.metrics_registry import cubestat_metric
 
 
 class PowerMode(DisplayMode):
@@ -44,4 +44,10 @@ def hotkey(self):
 
     @classmethod
     def configure_argparse(cls, parser):
-        parser.add_argument('--power', type=PowerMode, default=PowerMode.combined, choices=list(PowerMode), help='Power mode - off, showing breakdown CPU/GPU/ANE load, or showing combined usage. Can be toggled by pressing p.')
+        parser.add_argument(
+            '--power',
+            type=PowerMode,
+            default=PowerMode.combined,
+            choices=list(PowerMode),
+            help='Power: hidden, CPU/GPU/ANE breakdown, or combined usage. Hotkey: "p"'
+        )

cubestat/metrics/power.py

@@ -3,7 +3,7 @@
 import subprocess
 
 from cubestat.metrics.base_metric import base_metric
-from cubestat.metrics.registry import cubestat_metric
+from cubestat.metrics_registry import cubestat_metric
 from cubestat.common import SimpleMode, label_bytes
 
 
@@ -22,7 +22,13 @@ def key(cls):
 
     @classmethod
     def configure_argparse(cls, parser):
-        parser.add_argument('--swap', type=SimpleMode, default=SimpleMode.show, choices=list(SimpleMode), help="Show swap . Can be toggled by pressing s.")
+        parser.add_argument(
+            '--swap',
+            type=SimpleMode,
+            default=SimpleMode.show,
+            choices=list(SimpleMode),
+            help='swap show/hide. Hotkey: "s"'
+        )
 
     def configure(self, conf):
         self.mode = conf.swap
@@ -54,10 +60,18 @@ def read(self, _context):
         res = {}
         try:
             swap_stats = subprocess.run(["sysctl", "vm.swapusage"], capture_output=True, text=True)
+            swap_stats.check_returncode()  # Raise CalledProcessError if exit status is non-zero
             tokens = swap_stats.stdout.strip().split(' ')
+            if len(tokens) < 8:
+                raise IndexError("Invalid sysctl output")
             res['swap used'] = self._parse_memstr(tokens[7])
-        except:
-            logging.error("unable to get swap stats.")
+        except subprocess.CalledProcessError as e:
+            logging.error(f"sysctl command failed: {e}")
+        except IndexError as e:
+            logging.error(f"Invalid sysctl output: {e}")
+        except Exception as e:
+            logging.error(f"Unexpected error: {e}")
+
         return res
 
 

cubestat/metrics/swap.py

@@ -1,3 +1,5 @@
+import importlib
+import pkgutil
 import sys
 
 _metrics = []
@@ -24,5 +26,15 @@ def get_metrics(args):
     }
 
 
-# to run the annotations and register metrics
-from cubestat.metrics import cpu, gpu, memory, accel, swap, network, disk, power
+# Dynamically discover and import metrics
+def import_submodules(package_name):
+    package = importlib.import_module(package_name)
+    for _, module_name, is_pkg in pkgutil.iter_modules(package.__path__):
+        full_module_name = f"{package_name}.{module_name}"
+        importlib.import_module(full_module_name)
+        if is_pkg:
+            import_submodules(full_module_name)
+
+
+# Import all submodules of cubestat.metrics
+import_submodules("cubestat.metrics")

cubestat/metrics/registry.py renamed to cubestat/metrics_registry.py

@@ -4,7 +4,13 @@
 
 class MacOSPlatform:
     def __init__(self, interval_ms) -> None:
-        cmd = ['sudo', 'powermetrics', '-f', 'plist', '-i', str(interval_ms), '-s', 'cpu_power,gpu_power,ane_power,network,disk']
+        cmd = [
+            'sudo',
+            'powermetrics',
+            '-f', 'plist',
+            '-i', str(interval_ms),
+            '-s', 'cpu_power,gpu_power,ane_power,network,disk'
+        ]
         self.powermetrics = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
         # getting first line here to allow user to enter sudo credentials before
         # curses initialization.