diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index efdf877..ada4f16 100644 --- a/CONTRIBUTING.md +++ b/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) diff --git a/README.md b/README.md index e67ee7f..dba1e4c 100644 --- a/README.md +++ b/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. ``` diff --git a/nbresuse/__init__.py b/nbresuse/__init__.py index e4d6a4f..f25bc29 100644 --- a/nbresuse/__init__.py +++ b/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 ) diff --git a/nbresuse/api.py b/nbresuse/api.py new file mode 100644 index 0000000..b2d7f23 --- /dev/null +++ b/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]) diff --git a/nbresuse/static/main.js b/nbresuse/static/main.js index a1facca..581a775 100644 --- a/nbresuse/static/main.js +++ b/nbresuse/static/main.js @@ -5,13 +5,13 @@ define([ function setupDOM() { $('#maintoolbar-container').append( $('
').attr('id', 'nbresuse-display') - .addClass('btn-group') - .addClass('pull-right') - .append( - $('').text('Memory: ') - ).append( + .addClass('btn-group') + .addClass('pull-right') + .append( + $('').text('Memory: ') + ).append( $('').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) { diff --git a/nbresuse/tests/test_basic.py b/nbresuse/tests/test_basic.py index 5d1b6cd..a761ba9 100644 --- a/nbresuse/tests/test_basic.py +++ b/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( diff --git a/setup.py b/setup.py index 92cbfa2..90c04bd 100644 --- a/setup.py +++ b/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/*")),