Skip to content
20 changes: 15 additions & 5 deletions ipykernel/kernelbase.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,8 @@ class Kernel(SingletonConfigurable):
# attribute to override with a GUI
eventloop = Any(None)

processes: t.Dict[str, psutil.Process] = {}

@observe("eventloop")
def _update_eventloop(self, change):
"""schedule call to eventloop from IOLoop"""
Expand Down Expand Up @@ -953,10 +955,8 @@ async def debug_request(self, stream, ident, parent):
reply_msg = self.session.send(stream, "debug_reply", reply_content, 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)
Expand All @@ -971,12 +971,22 @@ async def usage_request(self, stream, ident, parent):
reply_content = {"hostname": socket.gethostname(), "pid": os.getpid()}
current_process = psutil.Process()
all_processes = [current_process] + current_process.children(recursive=True)
process_metric_value = self.get_process_metric_value
# Ensure 1) self.processes is updated to only current subprocesses
# and 2) we reuse processes when possible (needed for accurate CPU)
self.processes = {
process.pid: self.processes.get(process.pid, process) for process in all_processes
}
reply_content["kernel_cpu"] = sum(
[process_metric_value(process, "cpu_percent", None) for process in all_processes]
[
self.get_process_metric_value(process, "cpu_percent", None)
for process in self.processes.values()
]
)
reply_content["kernel_memory"] = sum(
[process_metric_value(process, "memory_info", "rss") for process in all_processes]
[
self.get_process_metric_value(process, "memory_info", "rss")
for process in self.processes.values()
]
)
cpu_percent = psutil.cpu_percent()
# https://psutil.readthedocs.io/en/latest/index.html?highlight=cpu#psutil.cpu_percent
Expand Down