Make NBResuse more flexible.

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/__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="""
@@ -93,7 +131,8 @@ def _mem_limit_default(self):
         """,
     ).tag(config=True)
 
-    cpu_limit = Float(
+    cpu_limit = Union(
+        trait_types=[Float(), Callable()],
         default_value=0,
         help="""
         CPU usage limit to display to the user.
@@ -116,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()

nbresuse/metrics.py

@@ -1,53 +1,81 @@
-from typing import NamedTuple, Optional
-
 try:
     import psutil
 except ImportError:
     psutil = None
 
-
-class MemoryMetrics(NamedTuple):
-    current_memory: int
-    max_memory: int
-
-
-class CPUMetrics(NamedTuple):
-    cpu_max: float
-    cpu_usage: float
-
-
-def memory_metrics() -> Optional[MemoryMetrics]:
-    if psutil:
-        cur_process = psutil.Process()
-        all_processes = [cur_process] + cur_process.children(recursive=True)
-
-        rss = sum([p.memory_info().rss for p in all_processes])
-        virtual_memory = psutil.virtual_memory().total
-
-    else:
-        return None
-
-    return MemoryMetrics(rss, virtual_memory)
-
-
-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()
-
-        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_percent = sum([get_cpu_percent(p) for p in all_processes])
-
-    else:
-        return None
-
-    return CPUMetrics(cpu_count * 100.0, cpu_percent)
+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
+
+    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
+
+        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):
+
+        metric_values = self.get_metric_values(process_metrics, "process")
+        metric_values.update(self.get_metric_values(system_metrics, "system"))
+
+        if any(value is None for value in metric_values.values()):
+            return None
+
+        return metric_values
+
+    def memory_metrics(self):
+        return self.metrics(
+            self.config.process_memory_metrics, self.config.system_memory_metrics
+        )
+
+    def cpu_metrics(self):
+        return self.metrics(
+            self.config.process_cpu_metrics, self.config.system_cpu_metrics
+        )

nbresuse/prometheus.py

@@ -1,56 +1,69 @@
+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
-from nbresuse.metrics import memory_metrics
-from nbresuse.metrics import MemoryMetrics
+from nbresuse.metrics import PSUtilMetricsLoader
 
 try:
     # Traitlets >= 4.3.3
     from traitlets import Callable
 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", [])
+class PrometheusHandler(Callable):
+    def __init__(self, metricsloader: PSUtilMetricsLoader):
+        super().__init__()
+        self.metricsloader = metricsloader
+        self.config = metricsloader.config
+        self.session_manager = metricsloader.nbapp.session_manager
 
-MAX_CPU_USAGE = Gauge("max_cpu_usage", "counter for max cpu usage", [])
+        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", [])
 
-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
-
-    @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)
+    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["memory_info_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 = self.metricsloader.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[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.max_memory)
+                return self.config.mem_limit(
+                    rss=memory_metric_values["memory_info_rss"]
+                )
             elif self.config.mem_limit > 0:  # mem_limit is an Int
-                metrics.max_memory = self.config.mem_limit
-        return metrics
-
-    def apply_cpu_limits(self, metrics: Optional[CPUMetrics]) -> Optional[CPUMetrics]:
-        if metrics is not None:
-            if self.config.cpu_limit > 0:
-                metrics.cpu_max = self.config.cpu_limit
-        return metrics
+                return self.config.mem_limit
+            else:
+                return memory_metric_values["virtual_memory_total"]
+
+    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"]
+                )
+            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"]

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",