Merge pull request #805 from datalayer-externals/usage

Add usage_request and usage_reply based on psutil

Author: Carreau

ipykernel/kernelbase.py

@@ -16,6 +16,10 @@
 import time
 import uuid
 import warnings
+try:
+    import psutil
+except ImportError:
+    psutil = None
 
 try:
     # jupyter_client >= 5, use tz-aware now
@@ -204,7 +208,7 @@ def _parent_header(self):
         'apply_request',
     ]
     # add deprecated ipyparallel control messages
-    control_msg_types = msg_types + ['clear_request', 'abort_request', 'debug_request']
+    control_msg_types = msg_types + ['clear_request', 'abort_request', 'debug_request', 'usage_request']
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
@@ -860,6 +864,39 @@ async def debug_request(self, stream, ident, parent):
                                       parent, ident)
         self.log.debug("%s", reply_msg)
 
+    # Taken from https://github.com/jupyter-server/jupyter-resource-usage/blob/e6ec53fa69fdb6de8e878974bcff006310658408/jupyter_resource_usage/metrics.py#L16
+    def get_process_metric_value(self, process, name, attribute=None):
+        try:
+            # psutil.Process methods will either return...
+            metric_value = getattr(process, name)()
+            if attribute is not None:  # ... a named tuple
+                return getattr(metric_value, attribute)
+            else:  # ... or a number
+                return metric_value
+        # Avoid littering logs with stack traces
+        # complaining about dead processes
+        except BaseException:
+            return None
+
+    async def usage_request(self, stream, ident, parent):
+        reply_content = {}
+        if psutil is None:
+            return reply_content
+        current_process = psutil.Process()
+        all_processes = [current_process] + current_process.children(recursive=True)
+        process_metric_value = self.get_process_metric_value
+        reply_content['kernel_cpu'] = sum([process_metric_value(process, 'cpu_percent', None) for process in all_processes])
+        reply_content['kernel_memory'] = sum([process_metric_value(process, 'memory_info', 'rss') for process in all_processes])
+        cpu_percent = psutil.cpu_percent()
+        # https://psutil.readthedocs.io/en/latest/index.html?highlight=cpu#psutil.cpu_percent
+        # The first time cpu_percent is called it will return a meaningless 0.0 value which you are supposed to ignore.
+        if cpu_percent != None and cpu_percent != 0.0:
+            reply_content['host_cpu_percent'] = cpu_percent
+        reply_content['host_virtual_memory'] = dict(psutil.virtual_memory()._asdict())
+        reply_msg = self.session.send(stream, 'usage_reply', reply_content,
+                                      parent, ident)
+        self.log.debug("%s", reply_msg)
+
     async def do_debug_request(self, msg):
         raise NotImplementedError