Make metric definitions/handling more abstract to prepare for more highly extensible configuration capapble of using (almost) all psutil metrics.

Author: krinsman

CONTRIBUTING.md

@@ -32,9 +32,11 @@ We recommend using [pipenv](https://docs.pipenv.org/) to make development easier
 4. Do a dev install of nbresuse and its dependencies
 
    ```bash
-   pip install --editable .
+   pip install --editable .[resources]
    ```
 
+   To test the behavior of NBResuse without `psutil` installed, run `pip install --editable .` instead.
+
 5. Install and enable the nbextension for use with Jupyter Classic Notebook.
 
    ```bash

nbresuse/metrics.py

@@ -1,4 +1,5 @@
-from typing import NamedTuple, Optional
+from typing import NamedTuple
+from typing import Optional
 
 try:
     import psutil
@@ -7,47 +8,79 @@
 
 
 class MemoryMetrics(NamedTuple):
-    current_memory: int
-    max_memory: int
+    rss: int
+    virtual_memory: int
 
 
 class CPUMetrics(NamedTuple):
-    cpu_max: float
-    cpu_usage: float
+    cpu_percent: float
+    cpu_count: int
 
 
-def memory_metrics() -> Optional[MemoryMetrics]:
-    if psutil:
-        cur_process = psutil.Process()
-        all_processes = [cur_process] + cur_process.children(recursive=True)
+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)
 
-        rss = sum([p.memory_info().rss for p in all_processes])
-        virtual_memory = psutil.virtual_memory().total
+        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)
+                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 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
+
+
+def memory_metrics() -> Optional[MemoryMetrics]:
+
+    rss = {"metric_name": "memory_info", "metric_attribute": "rss"}
+    rss_value = per_process_metric(**rss)
+
+    virtual_memory = {"metric_name": "virtual_memory", "metric_attribute": "total"}
+    virtual_memory_value = system_metric(**virtual_memory)
+
+    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(rss, virtual_memory)
+    return MemoryMetrics(**memory_metric_values)
 
 
 def cpu_metrics() -> Optional[CPUMetrics]:
-    if psutil:
-        cur_process = psutil.Process()
-        all_processes = [cur_process] + cur_process.children(recursive=True)
 
-        cpu_count = psutil.cpu_count()
+    cpu_percent = {"metric_name": "cpu_percent", "metric_kwargs": {"interval": 0.05}}
+    cpu_percent_value = per_process_metric(**cpu_percent)
 
-        def get_cpu_percent(p):
-            try:
-                return p.cpu_percent(interval=0.05)
-            # Avoid littering logs with stack traces complaining
-            # about dead processes having no CPU usage
-            except BaseException:
-                return 0
+    cpu_count = {"metric_name": "cpu_count"}
+    cpu_count_value = system_metric(**cpu_count)
 
-        cpu_percent = sum([get_cpu_percent(p) for p in all_processes])
+    cpu_metric_values = {"cpu_percent": cpu_percent_value, "cpu_count": cpu_count_value}
 
-    else:
+    if any(value is None for value in cpu_metric_values.values()):
         return None
 
-    return CPUMetrics(cpu_count * 100.0, cpu_percent)
+    return CPUMetrics(**cpu_metric_values)

nbresuse/prometheus.py

@@ -1,7 +1,8 @@
+from typing import Optional
+
 from notebook.notebookapp import NotebookApp
 from prometheus_client import Gauge
 from tornado import gen
-from typing import Optional
 
 from nbresuse.metrics import cpu_metrics
 from nbresuse.metrics import CPUMetrics
@@ -14,45 +15,59 @@
 except ImportError:
     from .utils import Callable
 
-TOTAL_MEMORY_USAGE = Gauge("total_memory_usage", "counter for total memory usage", [])
-
-MAX_MEMORY_USAGE = Gauge("max_memory_usage", "counter for max memory usage", [])
-
-TOTAL_CPU_USAGE = Gauge("total_cpu_usage", "counter for total cpu usage", [])
-
-MAX_CPU_USAGE = Gauge("max_cpu_usage", "counter for max cpu usage", [])
-
 
 class PrometheusHandler(Callable):
     def __init__(self, nbapp: NotebookApp):
         super().__init__()
         self.config = nbapp.web_app.settings["nbresuse_display_config"]
         self.session_manager = nbapp.session_manager
 
+        self.TOTAL_MEMORY_USAGE = Gauge(
+            "total_memory_usage", "counter for total memory usage", []
+        )
+        self.MAX_MEMORY_USAGE = Gauge(
+            "max_memory_usage", "counter for max memory usage", []
+        )
+
+        self.TOTAL_CPU_USAGE = Gauge(
+            "total_cpu_usage", "counter for total cpu usage", []
+        )
+        self.MAX_CPU_USAGE = Gauge("max_cpu_usage", "counter for max cpu usage", [])
+
     @gen.coroutine
     def __call__(self, *args, **kwargs):
-        metrics = self.apply_memory_limits(memory_metrics())
-        if metrics is not None:
-            TOTAL_MEMORY_USAGE.set(metrics.current_memory)
-            MAX_MEMORY_USAGE.set(metrics.max_memory)
+        memory_metric_values = memory_metrics()
+        if memory_metric_values is not None:
+            self.TOTAL_MEMORY_USAGE.set(memory_metric_values.rss)
+            self.MAX_MEMORY_USAGE.set(self.apply_memory_limit(memory_metric_values))
         if self.config.track_cpu_percent:
-            metrics = self.apply_cpu_limits(cpu_metrics())
-            if metrics is not None:
-                TOTAL_CPU_USAGE.set(metrics.cpu_usage)
-                MAX_CPU_USAGE.set(metrics.cpu_max)
+            cpu_metric_values = cpu_metrics()
+            if cpu_metric_values is not None:
+                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_limits(self, metrics: Optional[MemoryMetrics]) -> Optional[MemoryMetrics]:
-        if metrics is not None:
+    def apply_memory_limit(
+        self, memory_metric_values: Optional[MemoryMetrics]
+    ) -> Optional[int]:
+        if memory_metric_values is None:
+            return None
+        else:
             if callable(self.config.mem_limit):
-                metrics.max_memory = self.config.mem_limit(rss=metrics.current_memory)
+                return self.config.mem_limit(rss=memory_metric_values.rss)
             elif self.config.mem_limit > 0:  # mem_limit is an Int
-                metrics.max_memory = self.config.mem_limit
-        return metrics
+                return self.config.mem_limit
+            else:
+                return memory_metric_values.virtual_memory
 
-    def apply_cpu_limits(self, metrics: Optional[CPUMetrics]) -> Optional[CPUMetrics]:
-        if metrics is not None:
+    def apply_cpu_limit(
+        self, cpu_metric_values: Optional[CPUMetrics]
+    ) -> Optional[float]:
+        if cpu_metric_values is None:
+            return None
+        else:
             if callable(self.config.cpu_limit):
-                metrics.cpu_max = self.config.cpu_limit(cpu_percent=metrics.cpu_usage)
+                return self.config.cpu_limit(cpu_percent=cpu_metric_values.cpu_percent)
             elif self.config.cpu_limit > 0.0:  # cpu_limit is a Float
-                metrics.cpu_max = self.config.cpu_limit
-        return metrics
+                return self.config.cpu_limit
+            else:
+                return 100.0 * cpu_metric_values.cpu_count