From 570aae9e81c69bb2a6f1a4f92c326a0ab1e5e8aa Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Tue, 19 Mar 2024 23:38:15 -0400 Subject: [PATCH] [Fix] Lazy import of "psutil" in disco process pool Prior to this PR, module "psutil" is imported at the top level of the disco process pool. The pool will try to kill all the processes at the time of destruction (when `__del__` is implicitly invoked). The `__del__` function eventually calls into a function that uses `pstuil`. But it is possible that the top-level `psutil` has already been released by Python, which leads to a KeyError as follows: ``` Exception ignored in: Traceback (most recent call last): File "/home/ruihangl/Workspace/tvm/python/tvm/runtime/disco/process_pool.py", line 67, in __del__ File "/home/ruihangl/Workspace/tvm/python/tvm/runtime/disco/process_pool.py", line 81, in kill File "/home/ruihangl/Workspace/tvm/python/tvm/runtime/disco/process_pool.py", line 162, in _kill_child_processes File "/home/ruihangl/Workspace/miniconda3/envs/python311/lib/python3.11/site-packages/psutil/__init__.py", line 323, in __init__ File "/home/ruihangl/Workspace/miniconda3/envs/python311/lib/python3.11/site-packages/psutil/__init__.py", line 353, in _init File "/home/ruihangl/Workspace/miniconda3/envs/python311/lib/python3.11/site-packages/psutil/_pslinux.py", line 1738, in __init__ File "/home/ruihangl/Workspace/miniconda3/envs/python311/lib/python3.11/site-packages/psutil/_common.py", line 864, in get_procfs_path KeyError: 'psutil' ``` This PR fixes the issue by lazily importing `psutil` when needed. --- python/tvm/runtime/disco/process_pool.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tvm/runtime/disco/process_pool.py b/python/tvm/runtime/disco/process_pool.py index e91d855953b4..1ad8659d6088 100644 --- a/python/tvm/runtime/disco/process_pool.py +++ b/python/tvm/runtime/disco/process_pool.py @@ -20,8 +20,6 @@ import subprocess import sys -import psutil - from tvm._ffi import register_func from tvm.runtime import ShapeTuple @@ -158,6 +156,8 @@ def _kill_child_processes(pid): pid : int The given parameter id. """ + import psutil # pylint: disable=import-outside-toplevel + try: parent = psutil.Process(pid) children = parent.children(recursive=True)