From dfccd17ab99093bc2ecb9b6639f4747cceb62dc9 Mon Sep 17 00:00:00 2001
From: Cheng <zcbenz@gmail.com>
Date: Sat, 14 Dec 2024 12:50:13 +0900
Subject: [PATCH] Use psutil to get memory info on Windows (#1700)

---
 python/tests/test_array.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/python/tests/test_array.py b/python/tests/test_array.py
index 1e1ef8bdb7..ef3c6dd2ee 100644
--- a/python/tests/test_array.py
+++ b/python/tests/test_array.py
@@ -2,14 +2,20 @@
 
 import gc
 import operator
+import os
 import pickle
-import resource
+import platform
 import sys
 import unittest
 import weakref
 from copy import copy, deepcopy
 from itertools import permutations
 
+if platform.system() == "Windows":
+    import psutil
+else:
+    import resource
+
 import mlx.core as mx
 import mlx_tests
 import numpy as np
@@ -1932,7 +1938,11 @@ def test_deep_graphs(self):
 
     def test_siblings_without_eval(self):
         def get_mem():
-            return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+            if platform.system() == "Windows":
+                process = psutil.Process(os.getpid())
+                return process.memory_info().peak_wset
+            else:
+                return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
 
         key = mx.array([1, 2])