Right now all of the configuration options remain backwards-compatible with previous versions. However this seems to now be sufficiently abstract as to be at the cusp of rapid 'expandability' in terms of its features. The next step will be to automatically make as many prometheus Gauge's as needed to measure every specificied metric, i.e. without defining them all manually. setattr will be useful for this I believe.

krinsman · krinsman · commit edafd634a4aa · 2020-04-17T19:56:11.000-07:00
diff --git a/nbresuse/__init__.py b/nbresuse/__init__.py
@@ -3,11 +3,15 @@
 from tornado import ioloop
 from traitlets import Bool
 from traitlets import default
+from traitlets import Dict
 from traitlets import Float
 from traitlets import Int
+from traitlets import List
+from traitlets import TraitType
 from traitlets import Union
 from traitlets.config import Configurable
 
+from nbresuse.metrics import PSUtilMetricsLoader
 from nbresuse.prometheus import PrometheusHandler
 
 try:
@@ -17,6 +21,21 @@
     from .utils import Callable
 
 
+class PSUtilMetric(TraitType):
+    """A trait describing the format to specify a metric from the psutil package"""
+
+    info_text = "A dictionary specifying the function/method name, any keyword arguments, and if a named tuple is returned, which attribute of the named tuple to select"
+
+    def validate(self, obj, value):
+        if isinstance(value, dict):
+            keys = list(value.keys())
+            if "name" in keys:
+                keys.remove("name")
+                if all(key in ["kwargs", "attribute"] for key in keys):
+                    return value
+        self.error(obj, value)
+
+
 def _jupyter_server_extension_paths():
     """
     Set up the server extension for collecting metrics
@@ -43,6 +62,25 @@ class ResourceUseDisplay(Configurable):
     Holds server-side configuration for nbresuse
     """
 
+    process_memory_metrics = List(
+        trait=PSUtilMetric(),
+        default_value=[{"name": "memory_info", "attribute": "rss"}],
+    )
+
+    system_memory_metrics = List(
+        trait=PSUtilMetric(),
+        default_value=[{"name": "virtual_memory", "attribute": "total"}],
+    )
+
+    process_cpu_metrics = List(
+        trait=PSUtilMetric(),
+        default_value=[{"name": "cpu_percent", "kwargs": {"interval": 0.05}}],
+    )
+
+    system_cpu_metrics = List(
+        trait=PSUtilMetric(), default_value=[{"name": "cpu_count"}]
+    )
+
     mem_warning_threshold = Float(
         default_value=0.1,
         help="""
@@ -117,5 +155,7 @@ def load_jupyter_server_extension(nbapp):
     """
     resuseconfig = ResourceUseDisplay(parent=nbapp)
     nbapp.web_app.settings["nbresuse_display_config"] = resuseconfig
-    callback = ioloop.PeriodicCallback(PrometheusHandler(nbapp), 1000)
+    callback = ioloop.PeriodicCallback(
+        PrometheusHandler(PSUtilMetricsLoader(nbapp)), 1000
+    )
     callback.start()
diff --git a/nbresuse/metrics.py b/nbresuse/metrics.py
@@ -1,86 +1,81 @@
-from typing import NamedTuple
-from typing import Optional
-
 try:
     import psutil
 except ImportError:
     psutil = None
 
-
-class MemoryMetrics(NamedTuple):
-    rss: int
-    virtual_memory: int
-
-
-class CPUMetrics(NamedTuple):
-    cpu_percent: float
-    cpu_count: int
-
-
-def per_process_metric(metric_name, metric_kwargs={}, metric_attribute=None):
-    if psutil is None:
-        return None
-    else:
-        current_process = psutil.Process()
-        all_processes = [current_process] + current_process.children(recursive=True)
-
-        def get_per_process_metric(
-            process, metric_name, metric_kwargs, metric_attribute=None
-        ):
-            try:
-                metric_value = getattr(process, metric_name)(**metric_kwargs)
-                if metric_attribute is not None:
-                    return getattr(metric_value, metric_attribute)
+from notebook.notebookapp import NotebookApp
+
+
+class PSUtilMetricsLoader:
+    def __init__(self, nbapp: NotebookApp):
+        self.config = nbapp.web_app.settings["nbresuse_display_config"]
+        self.nbapp = nbapp
+
+    def process_metric(self, name, kwargs={}, attribute=None):
+        if psutil is None:
+            return None
+        else:
+            current_process = psutil.Process()
+            all_processes = [current_process] + current_process.children(recursive=True)
+
+            def get_process_metric(process, name, kwargs, attribute=None):
+                try:
+                    # psutil.Process methods will either return...
+                    metric_value = getattr(process, name)(**kwargs)
+                    if attribute is not None:  # ... a named tuple
+                        return getattr(metric_value, attribute)
+                    else:  # ... or a number
+                        return metric_value
+                # Avoid littering logs with stack traces
+                # complaining about dead processes
+                except BaseException:
+                    return 0
+
+            process_metric_value = lambda process: get_process_metric(
+                process, name, kwargs, attribute
+            )
+
+            return sum([process_metric_value(process) for process in all_processes])
+
+    def system_metric(self, name, kwargs={}, attribute=None):
+        if psutil is None:
+            return None
+        else:
+            # psutil functions will either return...
+            metric_value = getattr(psutil, name)(**kwargs)
+            if attribute is not None:  # ... a named tuple
+                return getattr(metric_value, attribute)
+            else:  # ... or a number
                 return metric_value
-            # Avoid littering logs with stack traces
-            # complaining about dead processes
-            except BaseException:
-                return 0
-
-        per_process_metric_value = lambda process: get_per_process_metric(
-            process, metric_name, metric_kwargs, metric_attribute
-        )
-
-        return sum([per_process_metric_value(process) for process in all_processes])
 
+    def get_metric_values(self, metrics, metric_type):
+        metric_types = {"process": self.process_metric, "system": self.system_metric}
+        metric_value = metric_types[metric_type]  # Switch statement
 
-def system_metric(metric_name, metric_kwargs={}, metric_attribute=None):
-    if psutil is None:
-        return None
-    else:
-        metric_value = getattr(psutil, metric_name)(**metric_kwargs)
-        if metric_attribute is not None:
-            return getattr(metric_value, metric_attribute)
-        return metric_attribute
+        metric_values = {}
+        for metric in metrics:
+            name = metric["name"]
+            if metric.get("attribute", False):
+                name += "_" + metric.get("attribute")
+            metric_values.update({name: metric_value(**metric)})
+        return metric_values
 
+    def metrics(self, process_metrics, system_metrics):
 
-def memory_metrics() -> Optional[MemoryMetrics]:
+        metric_values = self.get_metric_values(process_metrics, "process")
+        metric_values.update(self.get_metric_values(system_metrics, "system"))
 
-    rss = {"metric_name": "memory_info", "metric_attribute": "rss"}
-    rss_value = per_process_metric(**rss)
+        if any(value is None for value in metric_values.values()):
+            return None
 
-    virtual_memory = {"metric_name": "virtual_memory", "metric_attribute": "total"}
-    virtual_memory_value = system_metric(**virtual_memory)
+        return metric_values
 
-    memory_metric_values = {"rss": rss_value, "virtual_memory": virtual_memory_value}
-
-    if any(value is None for value in memory_metric_values.values()):
-        return None
-
-    return MemoryMetrics(**memory_metric_values)
-
-
-def cpu_metrics() -> Optional[CPUMetrics]:
-
-    cpu_percent = {"metric_name": "cpu_percent", "metric_kwargs": {"interval": 0.05}}
-    cpu_percent_value = per_process_metric(**cpu_percent)
-
-    cpu_count = {"metric_name": "cpu_count"}
-    cpu_count_value = system_metric(**cpu_count)
-
-    cpu_metric_values = {"cpu_percent": cpu_percent_value, "cpu_count": cpu_count_value}
-
-    if any(value is None for value in cpu_metric_values.values()):
-        return None
+    def memory_metrics(self):
+        return self.metrics(
+            self.config.process_memory_metrics, self.config.system_memory_metrics
+        )
 
-    return CPUMetrics(**cpu_metric_values)
+    def cpu_metrics(self):
+        return self.metrics(
+            self.config.process_cpu_metrics, self.config.system_cpu_metrics
+        )
diff --git a/nbresuse/prometheus.py b/nbresuse/prometheus.py
@@ -2,12 +2,8 @@
 
 from notebook.notebookapp import NotebookApp
 from prometheus_client import Gauge
-from tornado import gen
 
-from nbresuse.metrics import cpu_metrics
-from nbresuse.metrics import CPUMetrics
-from nbresuse.metrics import memory_metrics
-from nbresuse.metrics import MemoryMetrics
+from nbresuse.metrics import PSUtilMetricsLoader
 
 try:
     # Traitlets >= 4.3.3
@@ -17,10 +13,11 @@
 
 
 class PrometheusHandler(Callable):
-    def __init__(self, nbapp: NotebookApp):
+    def __init__(self, metricsloader: PSUtilMetricsLoader):
         super().__init__()
-        self.config = nbapp.web_app.settings["nbresuse_display_config"]
-        self.session_manager = nbapp.session_manager
+        self.metricsloader = metricsloader
+        self.config = metricsloader.config
+        self.session_manager = metricsloader.nbapp.session_manager
 
         self.TOTAL_MEMORY_USAGE = Gauge(
             "total_memory_usage", "counter for total memory usage", []
@@ -34,40 +31,39 @@ def __init__(self, nbapp: NotebookApp):
         )
         self.MAX_CPU_USAGE = Gauge("max_cpu_usage", "counter for max cpu usage", [])
 
-    @gen.coroutine
-    def __call__(self, *args, **kwargs):
-        memory_metric_values = memory_metrics()
+    async def __call__(self, *args, **kwargs):
+        memory_metric_values = self.metricsloader.memory_metrics()
         if memory_metric_values is not None:
-            self.TOTAL_MEMORY_USAGE.set(memory_metric_values.rss)
+            self.TOTAL_MEMORY_USAGE.set(memory_metric_values["memory_info_rss"])
             self.MAX_MEMORY_USAGE.set(self.apply_memory_limit(memory_metric_values))
         if self.config.track_cpu_percent:
-            cpu_metric_values = cpu_metrics()
+            cpu_metric_values = self.metricsloader.cpu_metrics()
             if cpu_metric_values is not None:
-                self.TOTAL_CPU_USAGE.set(cpu_metric_values.cpu_percent)
+                self.TOTAL_CPU_USAGE.set(cpu_metric_values["cpu_percent"])
                 self.MAX_CPU_USAGE.set(self.apply_cpu_limit(cpu_metric_values))
 
-    def apply_memory_limit(
-        self, memory_metric_values: Optional[MemoryMetrics]
-    ) -> Optional[int]:
+    def apply_memory_limit(self, memory_metric_values) -> Optional[int]:
         if memory_metric_values is None:
             return None
         else:
             if callable(self.config.mem_limit):
-                return self.config.mem_limit(rss=memory_metric_values.rss)
+                return self.config.mem_limit(
+                    rss=memory_metric_values["memory_info_rss"]
+                )
             elif self.config.mem_limit > 0:  # mem_limit is an Int
                 return self.config.mem_limit
             else:
-                return memory_metric_values.virtual_memory
+                return memory_metric_values["virtual_memory_total"]
 
-    def apply_cpu_limit(
-        self, cpu_metric_values: Optional[CPUMetrics]
-    ) -> Optional[float]:
+    def apply_cpu_limit(self, cpu_metric_values) -> Optional[float]:
         if cpu_metric_values is None:
             return None
         else:
             if callable(self.config.cpu_limit):
-                return self.config.cpu_limit(cpu_percent=cpu_metric_values.cpu_percent)
+                return self.config.cpu_limit(
+                    cpu_percent=cpu_metric_values["cpu_percent"]
+                )
             elif self.config.cpu_limit > 0.0:  # cpu_limit is a Float
                 return self.config.cpu_limit
             else:
-                return 100.0 * cpu_metric_values.cpu_count
+                return 100.0 * cpu_metric_values["cpu_count"]
diff --git a/setup.py b/setup.py
@@ -11,7 +11,7 @@
 
 setuptools.setup(
     name="nbresuse",
-    version="0.3.4",
+    version="0.4.0",
     url="https://github.com/yuvipanda/nbresuse",
     author="Yuvi Panda",
     description="Simple Jupyter extension to show how much resources (RAM) your notebook is using",
