Defer the psutil import until it is actually needed (#1551)

psutil is optional and its import is not cheap, yet it was paid for on
every kernel startup even though most sessions never ask for usage
information or need to enumerate child processes.

It is now imported through `_get_psutil()`, which caches the (possibly
None) result on first use.

Author: minrk

ipykernel/kernelbase.py

@@ -61,18 +61,29 @@
 from .iostream import OutStream
 from .utils import LazyDict, _async_in_context
 
-psutil: t.Any | None
-try:
-    import psutil as _psutil
-except ImportError:
-    psutil = None
-else:
-    psutil = _psutil
+psutil: t.Any | None = None
+_NO_SUCH_PROCESS: tuple[type[BaseException], ...] = ()
+_psutil_import_attempted = False
+
+
+def _get_psutil() -> t.Any | None:
+    """Import psutil on first use, caching the (possibly None) result.
+
+    psutil is optional and its import is not cheap, so we avoid paying for
+    it unless something actually needs process/resource-usage information.
+    """
+    global psutil, _NO_SUCH_PROCESS, _psutil_import_attempted  # noqa: PLW0603
+    if not _psutil_import_attempted:
+        _psutil_import_attempted = True
+        try:
+            import psutil as _psutil
+        except ImportError:
+            pass
+        else:
+            psutil = _psutil
+            _NO_SUCH_PROCESS = (psutil.NoSuchProcess,)
+    return psutil
 
-if psutil is None:
-    _NO_SUCH_PROCESS: tuple[type[BaseException], ...] = ()
-else:
-    _NO_SUCH_PROCESS = (psutil.NoSuchProcess,)
 
 _AWAITABLE_MESSAGE: str = (
     "For consistency across implementations, it is recommended that `{func_name}`"
@@ -1174,6 +1185,7 @@ async def usage_request(self, stream, ident, parent):
         if not self.session:
             return
         reply_content = {"hostname": socket.gethostname(), "pid": os.getpid()}
+        psutil = _get_psutil()
         if psutil is None:
             reply_content["cpu_count"] = os.cpu_count()
             reply_msg = self.session.send(stream, "usage_reply", reply_content, parent, ident)
@@ -1505,6 +1517,7 @@ def _process_children(self):
         - including parents and self with killpg
         - including all children that may have forked-off a new group
         """
+        psutil = _get_psutil()
         if psutil is None:
             return []
 

tests/test_kernel_direct.py

@@ -157,6 +157,7 @@ async def test_usage_request_without_psutil(kernel, monkeypatch):
     import ipykernel.kernelbase as kernelbase
 
     monkeypatch.setattr(kernelbase, "psutil", None)
+    monkeypatch.setattr(kernelbase, "_psutil_import_attempted", True)
     reply = await kernel.test_control_message("usage_request", {})
     content = reply["content"]
 
@@ -173,6 +174,7 @@ async def test_child_process_fallbacks_without_psutil(kernel, monkeypatch):
     import ipykernel.kernelbase as kernelbase
 
     monkeypatch.setattr(kernelbase, "psutil", None)
+    monkeypatch.setattr(kernelbase, "_psutil_import_attempted", True)
 
     assert kernel._process_children() == []
     kernel._signal_children(signal.SIGTERM)