Merge pull request #25 from yuvipanda/psutil_optional

Make psutil optional dependency of NBResuse

Author: krinsman

.gitignore

@@ -23,3 +23,4 @@ var/
 *.egg-info/
 .installed.cfg
 *.egg
+*.ipynb

README.md

@@ -25,9 +25,11 @@ main toolbar in the notebook itself, refreshing every 5s.
 You can currently install this package from PyPI.
 
 ```bash
-pip install nbresuse
+pip install nbresuse[resources]
 ```
 
+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/metrics.py

@@ -1,6 +1,9 @@
-from typing import NamedTuple
+from typing import NamedTuple, Optional
 
-import psutil
+try:
+    import psutil
+except ImportError:
+    psutil = None
 
 
 class MemoryMetrics(NamedTuple):
@@ -13,30 +16,38 @@ class CPUMetrics(NamedTuple):
     cpu_usage: float
 
 
-def memory_metrics() -> MemoryMetrics:
-    cur_process = psutil.Process()
-    all_processes = [cur_process] + cur_process.children(recursive=True)
+def memory_metrics() -> Optional[MemoryMetrics]:
+    if psutil:
+        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()
+        rss = sum([p.memory_info().rss for p in all_processes])
+        virtual_memory = psutil.virtual_memory().total
 
-    return MemoryMetrics(rss, virtual_memory.total)
+    else:
+        return None
 
+    return MemoryMetrics(rss, virtual_memory)
 
-def cpu_metrics() -> CPUMetrics:
-    cur_process = psutil.Process()
-    all_processes = [cur_process] + cur_process.children(recursive=True)
 
-    cpu_count = psutil.cpu_count()
+def cpu_metrics() -> Optional[CPUMetrics]:
+    if psutil:
+        cur_process = psutil.Process()
+        all_processes = [cur_process] + cur_process.children(recursive=True)
 
-    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 BaseException:
-            return 0
+        cpu_count = psutil.cpu_count()
 
-    cpu_percent = sum([get_cpu_percent(p) for p in 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 BaseException:
+                return 0
+
+        cpu_percent = sum([get_cpu_percent(p) for p in all_processes])
+
+    else:
+        return None
 
     return CPUMetrics(cpu_count * 100.0, cpu_percent)

nbresuse/prometheus.py

@@ -1,6 +1,7 @@
 from notebook.notebookapp import NotebookApp
 from prometheus_client import Gauge
 from tornado import gen
+from typing import Optional
 
 from nbresuse.metrics import cpu_metrics
 from nbresuse.metrics import CPUMetrics
@@ -31,21 +32,25 @@ def __init__(self, nbapp: NotebookApp):
     @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 metrics is not None:
+            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
+            if metrics is not None:
+                TOTAL_CPU_USAGE.set(metrics.cpu_usage)
+                MAX_CPU_USAGE.set(metrics.cpu_max)
+
+    def apply_memory_limits(self, metrics: Optional[MemoryMetrics]) -> Optional[MemoryMetrics]:
+        if metrics is not None:
+            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
+    def apply_cpu_limits(self, metrics: Optional[CPUMetrics]) -> Optional[CPUMetrics]:
+        if metrics is not None:
+            if self.config.cpu_limit > 0:
+                metrics.cpu_max = self.config.cpu_limit
         return metrics

nbresuse/static/main.js

@@ -61,7 +61,7 @@ define([
                 let totalMemoryUsage = metric("total_memory_usage", data);
                 let maxMemoryUsage = metric("max_memory_usage", data);
 
-                if (!totalMemoryUsage || !maxMemoryUsage)
+                if (maxMemoryUsage[2] <= 0)
                     return;
                 totalMemoryUsage = humanFileSize(parseFloat(totalMemoryUsage[2]));
                 maxMemoryUsage = humanFileSize(parseFloat(maxMemoryUsage[2]));

setup.py

@@ -11,7 +11,7 @@
 
 setuptools.setup(
     name="nbresuse",
-    version="0.3.3",
+    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,9 +23,10 @@
         "Programming Language :: Python :: 3",
     ],
     packages=setuptools.find_packages(),
-    install_requires=["psutil>=5.6.0", "notebook>=5.6.0"],
+    install_requires=["notebook>=5.6.0", "prometheus_client"],
     extras_require={
-        "dev": ["autopep8", "pytest", "flake8", "pytest-cov>=2.6.1", "mock"]
+        "resources": ["psutil>=5.6.0"],
+        "dev": ["autopep8", "pytest", "flake8", "pytest-cov>=2.6.1", "mock"],
     },
     data_files=[
         ("share/jupyter/nbextensions/nbresuse", glob("nbresuse/static/*")),