Adding support for jupyterlab statusbar-extension #36

Author: gsaibhar

nbresuse/__init__.py

@@ -1,5 +1,7 @@
+from notebook.utils import url_path_join
 from tornado import ioloop
 
+from nbresuse.api import ApiHandler
 from nbresuse.config import ResourceUseDisplay
 from nbresuse.metrics import PSUtilMetricsLoader
 from nbresuse.prometheus import PrometheusHandler
@@ -32,7 +34,9 @@ def load_jupyter_server_extension(nbapp):
     """
     resuseconfig = ResourceUseDisplay(parent=nbapp)
     nbapp.web_app.settings["nbresuse_display_config"] = resuseconfig
+    route_pattern = url_path_join(nbapp.web_app.settings['base_url'], '/api/nbresuse/v1')
+    nbapp.web_app.add_handlers('.*', [(route_pattern, ApiHandler)])
     callback = ioloop.PeriodicCallback(
         PrometheusHandler(PSUtilMetricsLoader(nbapp)), 1000
     )
-    callback.start()
+    callback.start()

nbresuse/api.py

@@ -0,0 +1,52 @@
+import json
+
+from notebook.base.handlers import IPythonHandler
+from tornado import web
+
+try:
+    # since psutil is an optional dependency
+    import psutil
+except ImportError:
+    psutil = None
+
+try:
+    # Traitlets >= 4.3.3
+    from traitlets import Callable
+except ImportError:
+    from .utils import Callable
+
+
+class ApiHandler(IPythonHandler):
+
+    @web.authenticated
+    async def get(self):
+        """
+        Calculate and return current resource usage metrics
+        """
+        config = self.settings['nbresuse_display_config']
+
+        if psutil:
+            cur_process = psutil.Process()
+            all_processes = [cur_process] + cur_process.children(recursive=True)
+
+            # 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
+
+            limits = {'memory': {
+                'rss': mem_limit
+            }}
+            if config.mem_limit:
+                limits['memory']['warn'] = (mem_limit - rss) < (
+                        mem_limit * config.mem_warning_threshold)
+
+            metrics = {
+                'rss': rss,
+                'limits': limits,
+            }
+
+            self.write(json.dumps(metrics))

nbresuse/static/main.js

@@ -5,13 +5,13 @@ define([
     function setupDOM() {
         $('#maintoolbar-container').append(
             $('<div>').attr('id', 'nbresuse-display')
-                      .addClass('btn-group')
-                      .addClass('pull-right')
-            .append(
-                $('<strong>').text('Memory: ')
-            ).append(
+                .addClass('btn-group')
+                .addClass('pull-right')
+                .append(
+                    $('<strong>').text('Memory: ')
+                ).append(
                 $('<span>').attr('id', 'nbresuse-mem')
-                           .attr('title', 'Actively used Memory (updates every 5s)')
+                    .attr('title', 'Actively used Memory (updates every 5s)')
             )
         );
         // FIXME: Do something cleaner to get styles in here?
@@ -24,49 +24,35 @@ define([
     }
 
     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];
+        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() {
+    var displayMetrics = function () {
         if (document.hidden) {
             // Don't poll when nobody is looking
             return;
         }
-        $.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);
+        $.getJSON({
+            url: utils.get_body_data('baseUrl') + 'api/nbresuse/v1',
+            success: function (data) {
+                totalMemoryUsage = humanFileSize(data['rss']);
+
+                var limits = data['limits'];
+                var display = totalMemoryUsage;
 
-                if (maxMemoryUsage[2] <= 0)
-                    return;
-                totalMemoryUsage = humanFileSize(parseFloat(totalMemoryUsage[2]));
-                maxMemoryUsage = humanFileSize(parseFloat(maxMemoryUsage[2]));
+                if (limits['memory']) {
+                    if (limits['memory']['rss']) {
+                        maxMemoryUsage = humanFileSize(limits['memory']['rss']);
+                        display += "/" + maxMemoryUsage
+                    }
+                    if (limits['memory']['warn']) {
+                        $('#nbresuse-display').addClass('nbresuse-warn');
+                    } else {
+                        $('#nbresuse-display').removeClass('nbresuse-warn');
+                    }
+                }
 
-                var display = totalMemoryUsage + "/" + maxMemoryUsage;
                 $('#nbresuse-mem').text(display);
             }
         });
@@ -78,7 +64,7 @@ define([
         // Update every five seconds, eh?
         setInterval(displayMetrics, 1000 * 5);
 
-        document.addEventListener("visibilitychange", function() {
+        document.addEventListener("visibilitychange", function () {
             // Update instantly when user activates notebook tab
             // FIXME: Turn off update timer completely when tab not in focus
             if (!document.hidden) {
@@ -90,4 +76,4 @@ define([
     return {
         load_ipython_extension: load_ipython_extension,
     };
-});
+});