Add back the optional CPU tracking for the endpoint

Author: jtpio

nbresuse/api.py

@@ -1,8 +1,10 @@
 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
@@ -12,6 +14,9 @@
 
 
 class ApiHandler(IPythonHandler):
+
+    executor = ThreadPoolExecutor(max_workers=5)
+
     @web.authenticated
     async def get(self):
         """
@@ -38,4 +43,30 @@ async def get(self):
 
         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])