Skip to content

[SPARK-56067][PYTHON] Lazy import psutil to improve import speed #54897

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
gaogaotiantian wants to merge 1 commit into
base: master
Choose a base branch
from
+12 −9
Open

[SPARK-56067][PYTHON] Lazy import psutil to improve import speed #54897

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 12 additions & 9 deletions python/pyspark/shuffle.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,13 +35,19 @@
)
from pyspark.util import fail_on_stopiteration

try:
import psutil
process = None

process = None

def get_used_memory():
"""Return the used memory in MiB"""
def get_used_memory():
"""Return the used memory in MiB"""
try:
import psutil
Copy link
Contributor

@devin-petersohn devin-petersohn Mar 20, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe this will result in the try...except block getting executed multiple times (the import runs each time the function is called).

We could use a strategy that lazily imports, but caches the import in a global variable instead of having Python's import system cache it.

_psutil = None
_psutil_checked = False

def get_used_memory():
    global _psutil, _psutil_checked
    if not _psutil_checked:
        try:
            import psutil
            _psutil = psutil
        except ImportError:
            pass
        _psutil_checked = True
    if _psutil is not None:
        # psutil path
    else:
        # fallback path

Copy link
Contributor Author

@gaogaotiantian gaogaotiantian Mar 20, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The import cache check is basically a dict check in sys.modules, which is super fast compared to the rest of the function. I don't think it's worth it to make the code more complicated. Getting memory usage is expensive (involves IO normally) so optimizing the rest of tis function does not give us much. I prefer to keep code simple unless we have proof that the performance difference is observable.


has_psutil = True
except ImportError:
has_psutil = False

if has_psutil:
global process
if process is None or process._pid != os.getpid():
process = psutil.Process(os.getpid())
Expand All @@ -51,10 +57,7 @@ def get_used_memory():
info = process.get_memory_info()
return info.rss >> 20

except ImportError:

def get_used_memory():
"""Return the used memory in MiB"""
else:
if platform.system() == "Linux":
for line in open("/proc/self/status"):
if line.startswith("VmRSS:"):
Expand Down