From 43c84083d501549328c2202723b5c5b08233b469 Mon Sep 17 00:00:00 2001
From: Tian Gao <gaogaotiantian@hotmail.com>
Date: Wed, 18 Mar 2026 16:46:40 -0700
Subject: [PATCH] Lazy import psutil to improve import speed

---
 python/pyspark/shuffle.py | 21 ++++++++++++---------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/python/pyspark/shuffle.py b/python/pyspark/shuffle.py
index f2d8cf949d171..168e0edfaa1fc 100644
--- a/python/pyspark/shuffle.py
+++ b/python/pyspark/shuffle.py
@@ -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
+
+        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())
@@ -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:"):