Lots of new features + v0.2 bump!

- Configurable memory usage warnings
- More configurability via traitlets
- Document memory limit behavior
- Add screenshot for memory usage warning
- Change 'Mem' prefix in display to less cryptic 'Memory'
- Fix primary screenshot to show memory limit too
- Consolidate code into single file

Author: yuvipanda

README.md

@@ -1,6 +1,6 @@
 # nbresuse
 
-![Screenshot](screenshot.png)
+![Screenshot with memory limit](screenshot.png)
 
 NB Resource Usage (nbresuse) is a small extension for Jupyter Notebooks that
 displays an indication of how much resources your current notebook server and
@@ -9,21 +9,44 @@ main toolbar in the notebook itself, refreshing every 5s.
 
 ## Installation
 
-You can currently install this directly from git:
+You can currently install this package from PyPI.
 
+```bash
+pip install nbresuse
 ```
-pip install git+https://github.com/yuvipanda/nbresuse.git
-jupyter serverextension enable --py nbresuse
-jupyter nbextension install --py nbresuse
-```
-
-To enable this extension for all notebooks:
 
+Once installed, you need to enable it for use in your notebook
 ```
-jupyter nbextension enable --py nbresuse
+jupyter serverextension enable --py nbresuse --sys-prefix
+jupyter nbextension install --py nbresuse --sys-prefix
+jupyter nbextension enable --py nbresuse --sys-prefix
 ```
 
+## Configuration
+
+### Memory Limit
+
+`nbresuse` can display a memory limit (but not enforce it). You can set this
+in serveral ways:
+
+1. `MEM_LIMIT` environment variable. This is set by [JupyterHub](http://github.com/jupyterhub/jupyterhub/)
+   if using a spawner that supports it.
+2. In the commandline when starting `jupyter notebook`, as `--ResourceUseDisplay.mem_limit`.
+3. In your jupyter notebook [traitlets](https://traitlets.readthedocs.io/en/stable/) config file
+
+The limit needs to be set as an integer in Bytes.
+
+### Memory usage warning threshold
+
+![Screenshot with memory warning](warn-screenshot.png)
+
+The background of the resource display can be changed to red when the user is near a memory limit.
+The threshold for this warning can be configured as a fraction of the memory limit.
+
+If you want to flash the warning to the user when they are within 10% of the memory limit, you
+can set the parameter `--ResourceUseDisplay.mem_warning_threshold=0.1`.
+
 ## Resources displayed
 
 Currently it only displays Memory usage (just RSS). Other metrics will be
-added in the future as needed.
+added in the future as needed.

nbresuse/__init__.py

@@ -1,17 +1,97 @@
-from nbresuse.handlers import setup_handlers
-# Jupyter Extension points
+import os
+import json
+import psutil
+from traitlets import Float, Int, default
+from traitlets.config import Configurable
+from notebook.utils import url_path_join
+from notebook.base.handlers import IPythonHandler
+
+
+class MetricsHandler(IPythonHandler):
+    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)
+        rss = sum([p.memory_info().rss for p in all_processes])
+
+        limits = {}
+
+        if config.mem_limit != 0:
+            limits['memory'] = {
+                'rss': config.mem_limit
+            }
+            if config.mem_warning_threshold != 0:
+                limits['memory']['warn'] = (config.mem_limit - rss) < (config.mem_limit * config.mem_warning_threshold)
+        metrics = {
+            'rss': rss,
+            'limits': limits,
+        }
+        self.write(json.dumps(metrics))
+
+
 def _jupyter_server_extension_paths():
+    """
+    Set up the server extension for collecting metrics
+    """
     return [{
         'module': 'nbresuse',
     }]
 
 def _jupyter_nbextension_paths():
+    """
+    Set up the notebook extension for displaying metrics
+    """
     return [{
         "section": "notebook",
         "dest": "nbresuse",
         "src": "static",
         "require": "nbresuse/main"
     }]
 
+class ResourceUseDisplay(Configurable):
+    """
+    Holds server-side configuration for nbresuse
+    """
+
+    mem_warning_threshold = Float(
+        0.1,
+        help="""
+        Warn user with flashing lights when memory usage is within this fraction
+        memory limit.
+
+        For example, if memory limit is 128MB, `mem_warning_threshold` is 0.1,
+        we will start warning the user when they use (128 - (128 * 0.1)) MB.
+
+        Set to 0 to disable warning.
+        """,
+        config=True
+    )
+
+    mem_limit = Int(
+        0,
+        config=True,
+        help="""
+        Memory limit to display to the user, in bytes.
+
+        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.
+        """
+    )
+
+    @default('mem_limit')
+    def _mem_limit_default(self):
+        return int(os.environ.get('MEM_LIMIT', 0))
+
 def load_jupyter_server_extension(nbapp):
-    setup_handlers(nbapp.web_app)
+    """
+    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)])

nbresuse/handlers.py

@@ -1,16 +1,23 @@
 define(['jquery', 'base/js/utils'], function ($, utils) {
-    function createDisplayDiv() {
+    function setupDOM() {
         $('#maintoolbar-container').append(
             $('<div>').attr('id', 'nbresuse-display')
                       .addClass('btn-group')
                       .addClass('pull-right')
             .append(
-                $('<strong>').text('Mem: ')
+                $('<strong>').text('Memory: ')
             ).append(
                 $('<span>').attr('id', 'nbresuse-mem')
                            .attr('title', 'Actively used Memory (updates every 5s)')
             )
         );
+        // FIXME: Do something cleaner to get styles in here?
+        $('head').append(
+            $('<style>').html('.nbresuse-warn { background-color: #FFD2D2; color: #D8000C; }')
+        );
+        $('head').append(
+            $('<style>').html('#nbresuse-display { padding: 2px 8px; }')
+        );
     }
 
     var displayMetrics = function() {
@@ -19,15 +26,25 @@ define(['jquery', 'base/js/utils'], function ($, utils) {
             // after the ., but GB should have 2.
             var display = Math.round(data['rss'] / (1024 * 1024));
 
+            var limits = data['limits'];
+            if ('memory' in limits) {
+                if ('rss' in limits['memory']) {
+                    display += " / " + (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) {
-                display += " / " + (data['limits']['memory'] / (1024 * 1024));
             }
             $('#nbresuse-mem').text(display + ' MB');
         });
     }
 
     var load_ipython_extension = function () {
-        createDisplayDiv();
+        setupDOM();
         displayMetrics();
         // Update every five seconds, eh?
         setInterval(displayMetrics, 1000 * 5);

nbresuse/static/main.js

@@ -2,7 +2,7 @@
 
 setuptools.setup(
     name="nbresuse",
-    version='0.1.0',
+    version='0.2.0',
     url="https://github.com/yuvipanda/nbresuse",
     author="Yuvi Panda",
     description="Simple Jupyter extension to show how much resources (RAM) your notebook is using",