Adding support for jupyterlab statusbar-extension #36

CONTRIBUTING.md

@@ -17,7 +17,7 @@ We recommend using [pipenv](https://docs.pipenv.org/) to make development easier
    ```
 
 2. Create an environment that will hold our dependencies.
-   
+
    ```bash
    cd nbresuse
    pipenv --python 3.6
@@ -32,11 +32,9 @@ 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 .[resources]
+   pip install --editable .[dev]
    ```
 
-   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
@@ -73,7 +71,7 @@ the pre-commit hook should take care of how it should look. Here is how to set u
     ```bash
     pre-commit run
     ```
-	
+
 which should run any autoformatting on your code
 and tell you about any errors it couldn't fix automatically.
 You may also install [black integration](https://github.com/ambv/black#editor-integration)

README.md

@@ -22,11 +22,9 @@ main toolbar in the notebook itself, refreshing every 5s.
 You can currently install this package from PyPI.
 
 ```bash
-pip install nbresuse[resources]
+pip install nbresuse
 ```
 
-The above command will install NBResuse along with `psutil` Python package (which is used for getting hardware usage information from the system). If you would like to install NBResuse _without_ `psutil` (in which case NBResuse does essentially nothing), run `pip install nbresuse` instead.
-
 **If your notebook version is < 5.3**, you need to enable the extension manually.
 
 ```

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,6 +34,10 @@ def load_jupyter_server_extension(nbapp):
     """
     resuseconfig = ResourceUseDisplay(parent=nbapp)
     nbapp.web_app.settings["nbresuse_display_config"] = resuseconfig
+    base_url = nbapp.web_app.settings["base_url"]
+    nbapp.web_app.add_handlers(
+        ".*", [(url_path_join(base_url, "/metrics"), ApiHandler)]
+    )
     callback = ioloop.PeriodicCallback(
         PrometheusHandler(PSUtilMetricsLoader(nbapp)), 1000
     )

nbresuse/api.py

@@ -0,0 +1,72 @@
+import json
+from concurrent.futures import ThreadPoolExecutor
+
+import psutil
+from notebook.base.handlers import IPythonHandler
+from tornado import web
+from tornado.concurrent import run_on_executor
+
+try:
+    # Traitlets >= 4.3.3
+    from traitlets import Callable
+except ImportError:
+    from .utils import Callable
+
+
+class ApiHandler(IPythonHandler):
+
+    executor = ThreadPoolExecutor(max_workers=5)
+
+    @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)
+
+        # 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}
+
+        # Optionally get CPU information
+        if config.track_cpu_percent:
+            cpu_count = psutil.cpu_count()
+            cpu_percent = await self._get_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.update(cpu_percent=cpu_percent, cpu_count=cpu_count)
+
+        self.write(json.dumps(metrics))
+
+    @run_on_executor
+    def _get_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])

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({
+        $.getJSON({
             url: utils.get_body_data('baseUrl') + 'metrics',
-            success: function(data) {
-                let totalMemoryUsage = metric("total_memory_usage", data);
-                let maxMemoryUsage = metric("max_memory_usage", data);
+            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) {

nbresuse/tests/test_basic.py

@@ -25,7 +25,7 @@ def test_import_serverextension(self):
 
         # mock a notebook app
         nbapp_mock = MagicMock()
-        nbapp_mock.web_app.settings = {}
+        nbapp_mock.web_app.settings = {"base_url": ""}
 
         # mock these out for unit test
         with patch("tornado.ioloop.PeriodicCallback") as periodic_callback_mock, patch(

setup.py

@@ -11,7 +11,7 @@
 
 setuptools.setup(
     name="nbresuse",
-    version="0.4.0",
+    version="0.3.4",
     url="https://github.com/yuvipanda/nbresuse",
     author="Yuvi Panda",
     description="Simple Jupyter extension to show how much resources (RAM) your notebook is using",
@@ -23,10 +23,9 @@
         "Programming Language :: Python :: 3",
     ],
     packages=setuptools.find_packages(),
-    install_requires=["notebook>=5.6.0", "prometheus_client"],
+    install_requires=["notebook>=5.6.0", "prometheus_client", "psutil>=5.6.0"],
     extras_require={
-        "resources": ["psutil>=5.6.0"],
-        "dev": ["autopep8", "black", "pytest", "flake8", "pytest-cov>=2.6.1", "mock"],
+        "dev": ["autopep8", "black", "pytest", "flake8", "pytest-cov>=2.6.1", "mock"]
     },
     data_files=[
         ("share/jupyter/nbextensions/nbresuse", glob("nbresuse/static/*")),