report memory and cpu metrics to prometheus

witzatom · witzatom · commit 0e8dc5fa103f · 2020-02-05T12:13:04.000+01:00
diff --git a/nbresuse/__init__.py b/nbresuse/__init__.py
@@ -1,100 +1,17 @@
 import os
-import json
-import psutil
+
+from tornado import ioloop
 from traitlets import Bool, Float, Int, Union, default
 from traitlets.config import Configurable
-from notebook.utils import url_path_join
-from notebook.base.handlers import IPythonHandler
-from tornado import web
+
+from nbresuse.prometheus import PrometheusHandler
 
 try:
     # Traitlets >= 4.3.3
     from traitlets import Callable
 except ImportError:
     from .utils import Callable
 
-from concurrent.futures import ThreadPoolExecutor
-from tornado.concurrent import run_on_executor
-
-class MetricsHandler(IPythonHandler):
-    def initialize(self):
-        super().initialize()
-        self.cpu_percent = 0
-        
-        # https://www.tornadoweb.org/en/stable/concurrent.html#tornado.concurrent.run_on_executor
-        self.executor = ThreadPoolExecutor(max_workers=10)
-
-        self.cpu_count = psutil.cpu_count()
-
-    @run_on_executor
-    def update_cpu_percent(self, all_processes):
-
-        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:
-                return 0
-
-        return sum([get_cpu_percent(p) for p in all_processes])
-
-    @web.authenticated
-    async def get(self):
-        """
-        Calculate and return current resource usage metrics
-        """
-        config = self.settings['nbresuse_display_config']
-        cur_process = psutil.Process()
-        all_processes = [cur_process] + cur_process.children(recursive=True)     
-        limits = {}
-
-        # Get memory information
-        rss = sum([p.memory_info().rss for p in all_processes])
-
-        if callable(config.mem_limit):
-            mem_limit = config.mem_limit(rss=rss)
-        else: # mem_limit is an Int
-            mem_limit = config.mem_limit
-
-        # A better approach would use cpu_affinity to account for the
-        # fact that the number of logical CPUs in the system is not
-        # necessarily the same as the number of CPUs the process
-        # can actually use. But cpu_affinity isn't available for OS X.
-        cpu_count = psutil.cpu_count()
-
-        if config.track_cpu_percent:
-            self.cpu_percent = await self.update_cpu_percent(all_processes)
-
-        if config.mem_limit != 0:
-            limits['memory'] = {
-                'rss': mem_limit
-            }
-            if config.mem_warning_threshold != 0:
-                limits['memory']['warn'] = (mem_limit - rss) < (mem_limit * config.mem_warning_threshold)
-
-        # Optionally get CPU information
-        if config.track_cpu_percent:
-            self.cpu_percent = await self.update_cpu_percent(all_processes)
-
-            if config.cpu_limit != 0:
-                limits['cpu'] = {
-                    'cpu': config.cpu_limit
-                }
-                if config.cpu_warning_threshold != 0:
-                    limits['cpu']['warn'] = (config.cpu_limit - self.cpu_percent) < (config.cpu_limit * config.cpu_warning_threshold)
-
-        metrics = {
-            'rss': rss,
-            'limits': limits,
-        }
-        if config.track_cpu_percent:
-            metrics.update(cpu_percent=self.cpu_percent,
-                               cpu_count=self.cpu_count)
-
-        self.log.debug("NBResuse metrics: %s", metrics)
-        self.write(json.dumps(metrics))
-
 
 def _jupyter_server_extension_paths():
     """
@@ -104,6 +21,7 @@ def _jupyter_server_extension_paths():
         'module': 'nbresuse',
     }]
 
+
 def _jupyter_nbextension_paths():
     """
     Set up the notebook extension for displaying metrics
@@ -115,6 +33,7 @@ def _jupyter_nbextension_paths():
         "require": "nbresuse/main"
     }]
 
+
 class ResourceUseDisplay(Configurable):
     """
     Holds server-side configuration for nbresuse
@@ -142,7 +61,7 @@ class ResourceUseDisplay(Configurable):
         Note that this does not actually limit the user's memory usage!
 
         Defaults to reading from the `MEM_LIMIT` environment variable. If
-        set to 0, no memory limit is displayed.
+        set to 0, the max memory available is displayed.
         """
     ).tag(config=True)
 
@@ -151,7 +70,7 @@ def _mem_limit_default(self):
         return int(os.environ.get('MEM_LIMIT', 0))
 
     track_cpu_percent = Bool(
-        default_value=False,
+        default_value=True,
         help="""
         Set to True in order to enable reporting of CPU usage statistics.
         """
@@ -178,19 +97,20 @@ def _mem_limit_default(self):
         Note that this does not actually limit the user's CPU usage!
 
         Defaults to reading from the `CPU_LIMIT` environment variable. If
-        set to 0, no CPU usage limit is displayed.
+        set to 0, the total CPU count available is displayed.
         """
     ).tag(config=True)
 
     @default('cpu_limit')
     def _cpu_limit_default(self):
         return float(os.environ.get('CPU_LIMIT', 0))
 
+
 def load_jupyter_server_extension(nbapp):
     """
     Called during notebook start
     """
     resuseconfig = ResourceUseDisplay(parent=nbapp)
     nbapp.web_app.settings['nbresuse_display_config'] = resuseconfig
-    route_pattern = url_path_join(nbapp.web_app.settings['base_url'], '/metrics')
-    nbapp.web_app.add_handlers('.*', [(route_pattern, MetricsHandler)])
+    callback = ioloop.PeriodicCallback(PrometheusHandler(nbapp), 1000)
+    callback.start()
diff --git a/nbresuse/metrics.py b/nbresuse/metrics.py
@@ -0,0 +1,47 @@
+from typing import NamedTuple
+
+import psutil
+
+
+class MemoryMetrics(NamedTuple):
+    current_memory: int
+    max_memory: int
+
+
+class CPUMetrics(NamedTuple):
+    cpu_max: float
+    cpu_usage: float
+
+
+def memory_metrics() -> MemoryMetrics:
+    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()
+
+    return MemoryMetrics(
+        rss,
+        virtual_memory.total
+    )
+
+
+def cpu_metrics() -> CPUMetrics:
+    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:
+            return 0
+    cpu_percent = sum([get_cpu_percent(p) for p in all_processes])
+
+    return CPUMetrics(
+        cpu_count * 100.0,
+        cpu_percent
+    )
diff --git a/nbresuse/prometheus.py b/nbresuse/prometheus.py
@@ -0,0 +1,64 @@
+from notebook.notebookapp import NotebookApp
+from prometheus_client import Gauge
+from tornado import gen
+
+from nbresuse.metrics import CPUMetrics, MemoryMetrics, cpu_metrics, memory_metrics
+
+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',
+    []
+)
+
+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())
+        TOTAL_MEMORY_USAGE.set(metrics.current_memory)
+        MAX_MEMORY_USAGE.set(metrics.max_memory)
+        if self.config.track_cpu_percent:
+            metrics = self.apply_cpu_limits(cpu_metrics())
+            TOTAL_CPU_USAGE.set(metrics.cpu_usage)
+            MAX_CPU_USAGE.set(metrics.cpu_max)
+
+    def apply_memory_limits(self, metrics: MemoryMetrics) -> MemoryMetrics:
+        if callable(self.config.mem_limit):
+            metrics.max_memory = self.config.mem_limit(rss=metrics.max_memory)
+        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: CPUMetrics) -> CPUMetrics:
+        if self.config.cpu_limit > 0:
+            metrics.cpu_max = self.config.cpu_limit
+        return metrics
diff --git a/nbresuse/static/main.js b/nbresuse/static/main.js
@@ -1,4 +1,7 @@
-define(['jquery', 'base/js/utils'], function ($, utils) {
+define([
+    'jquery',
+    'base/js/utils'
+], function ($, utils) {
     function setupDOM() {
         $('#maintoolbar-container').append(
             $('<div>').attr('id', 'nbresuse-display')
@@ -20,32 +23,54 @@ define(['jquery', 'base/js/utils'], function ($, utils) {
         );
     }
 
+    function humanFileSize(size) {
+        var i = Math.floor( Math.log(size) / Math.log(1024) );
+        return ( size / Math.pow(1024, i) ).toFixed(1) * 1 + ' ' + ['B', 'kB', 'MB', 'GB', 'TB'][i];
+    }
+
+
+    function metric(metric_name, text, multiple=false) {
+        var regex = new RegExp("^" + metric_name + "\{?([^ \}]*)\}? (.*)$", "gm");
+        var matches = [];
+        var match;
+
+        do{
+            match = regex.exec(text);
+            if (match){
+                matches.push(match)
+            }
+        }
+        while (match);
+
+        if (!multiple) {
+            if (matches.length > 0)
+                return matches[0];
+            return null;
+        }else
+            return matches;
+    }
+
     var displayMetrics = function() {
         if (document.hidden) {
             // Don't poll when nobody is looking
             return;
         }
-        $.getJSON(utils.get_body_data('baseUrl') + 'metrics', function(data) {
-            // FIXME: Proper setups for MB and GB. MB should have 0 things
-            // after the ., but GB should have 2.
-            var display = Math.round(data['rss'] / (1024 * 1024));
+        $.ajax({
+            url: utils.get_body_data('baseUrl') + 'metrics',
+            success: function(data) {
+                let totalMemoryUsage = metric("total_memory_usage", data);
+                let maxMemoryUsage = metric("max_memory_usage", data);
 
-            var limits = data['limits'];
-            if ('memory' in limits) {
-                if ('rss' in limits['memory']) {
-                    display += " / " + Math.round(limits['memory']['rss'] / (1024 * 1024));
-                }
-                if (limits['memory']['warn']) {
-                    $('#nbresuse-display').addClass('nbresuse-warn');
-                } else {
-                    $('#nbresuse-display').removeClass('nbresuse-warn');
-                }
-            }
-            if (data['limits']['memory'] !== null) {
+                if (!totalMemoryUsage || !maxMemoryUsage)
+                    return;
+                totalMemoryUsage = humanFileSize(parseFloat(totalMemoryUsage[2]));
+                maxMemoryUsage = humanFileSize(parseFloat(maxMemoryUsage[2]));
+
+                var display = totalMemoryUsage + "/" + maxMemoryUsage;
+                $('#nbresuse-mem').text(display);
             }
-            $('#nbresuse-mem').text(display + ' MB');
         });
-    }
+    };
 
     var load_ipython_extension = function () {
         setupDOM();
