diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index ceeebc23a..34119dc70 100644 --- a/ipykernel/kernelbase.py +++ b/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