fix memory_info() of psutil

avulanov · Jul 24, 2014 · dcf03a9 · dcf03a9
1 parent 67e6eba
commit dcf03a9
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/python/pyspark/shuffle.py b/python/pyspark/shuffle.py
@@ -28,9 +28,12 @@
 
     def get_used_memory():
         """ Return the used memory in MB """
-        self = psutil.Process(os.getpid())
-        return self.memory_info().rss >> 20
-
+        process = psutil.Process(os.getpid())
+        if hasattr(process, "memory_info"):
+            info = process.memory_info()
+        else:
+            info = process.get_memory_info()
+        return info.rss >> 20
 except ImportError:
 
     def get_used_memory():