From a6e0270bf6a5ba0c384376199b0ad5cb6504ef07 Mon Sep 17 00:00:00 2001 From: Ihar Hrachyshka Date: Sun, 1 Feb 2026 17:30:01 -0500 Subject: [PATCH] tests: harden memory leak check in test_siblings_without_eval When running the test suite on darwin as part of nixpkgs package build, we often hit the following error: ``` > mx.synchronize() > t() > gc.collect() > expected = get_mem() > for _ in range(100): > t() > used = get_mem() > > self.assertEqual(expected, used) > E AssertionError: 273678336 != 273694720 ``` You can find an example of the full failure log here: https://hydra.nixos.org/build/320731525 The difference between actual and expected is 16k (one page size on aarch64-darwin). The test seems brittle as it depends on *peak* memory, not real memory at the moment of measurement, so anything during the run of the process that could push the allocator to allocate another page can break the test. This patch switches the test case to measure real memory using psutil on all platforms. It also makes sure that macos CI installs psutil (via .[dev] dependency group). Signed-off-by: Ihar Hrachyshka --- .github/actions/build-macos/action.yml | 4 ++-- python/tests/test_array.py | 13 +++---------- 2 files changed, 5 insertions(+), 12 deletions(-) diff --git a/.github/actions/build-macos/action.yml b/.github/actions/build-macos/action.yml index f04ed3d63f..5da840b819 100644 --- a/.github/actions/build-macos/action.yml +++ b/.github/actions/build-macos/action.yml @@ -12,12 +12,12 @@ runs: run: | pip install --upgrade pip pip install cmake setuptools typing_extensions - pip install -e . -v + pip install -e ".[dev]" -v - name: Install tests dependencies shell: bash -l {0} run: | - pip install numpy torch tensorflow + pip install tensorflow - name: Run Python tests shell: bash -l {0} diff --git a/python/tests/test_array.py b/python/tests/test_array.py index 4efed9dac9..4d673d5b4a 100644 --- a/python/tests/test_array.py +++ b/python/tests/test_array.py @@ -11,14 +11,10 @@ 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 +import psutil try: import tensorflow as tf @@ -2087,11 +2083,8 @@ def test_deep_graphs(self): def test_siblings_without_eval(self): def get_mem(): - if platform.system() == "Windows": - process = psutil.Process(os.getpid()) - return process.memory_info().peak_wset - else: - return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + process = psutil.Process(os.getpid()) + return process.memory_info().rss key = mx.array([1, 2])