(tvm) branch main updated: [Fix] Lazy import of "psutil" in disco process pool (#16752)

[tq...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 7683bc23b1 [Fix] Lazy import of "psutil" in disco process pool (#16752)
7683bc23b1 is described below

commit 7683bc23b1b0152710231ea3b4b5fd7669c70799
Author: Ruihang Lai <ru...@cs.cmu.edu>
AuthorDate: Wed Mar 20 08:50:59 2024 -0400

    [Fix] Lazy import of "psutil" in disco process pool (#16752)
    
    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: <function DiscoPopenWorker.__del__ at 0x7f2c922bfe20>
    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 e91d855953..1ad8659d60 100644
--- a/python/tvm/runtime/disco/process_pool.py
+++ b/python/tvm/runtime/disco/process_pool.py
@@ -20,8 +20,6 @@ import os
 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)