From a1a20871b2e0aa3cd3d5e9ea92e90414b329f656 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 7 May 2026 08:01:21 +0000 Subject: [PATCH] ci(mlx): use mx.get_peak_memory with mx.metal.get_peak_memory fallback Newer MLX deprecates mx.metal.get_peak_memory in favour of the top-level mx.get_peak_memory. The CI was emitting: mx.metal.get_peak_memory is deprecated and will be removed in a future version. Use mx.get_peak_memory instead. Try the new top-level getter first and fall back to the metal one for compatibility with older MLX versions still in the wild. --- tests/studio/run_real_mlx_smoke.py | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/tests/studio/run_real_mlx_smoke.py b/tests/studio/run_real_mlx_smoke.py index a167959337..168e9ad329 100644 --- a/tests/studio/run_real_mlx_smoke.py +++ b/tests/studio/run_real_mlx_smoke.py @@ -97,12 +97,21 @@ def _seed_everything() -> None: def _peak_gpu_gb() -> float: import mlx.core as mx - if mx.metal.is_available(): - try: - return float(mx.metal.get_peak_memory()) / (1024**3) - except Exception: - return 0.0 - return 0.0 + if not mx.metal.is_available(): + return 0.0 + # Newer MLX deprecates mx.metal.get_peak_memory in favour of the + # top-level mx.get_peak_memory; fall back to the old API for + # compatibility with older MLX versions still present in the + # environment. + getter = getattr(mx, "get_peak_memory", None) or getattr( + mx.metal, "get_peak_memory", None + ) + if getter is None: + return 0.0 + try: + return float(getter()) / (1024**3) + except Exception: + return 0.0 def _peak_rss_gb() -> float: