ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

ci(mlx): use mx.get_peak_memory with mx.metal.get_peak_memory fallback

Browse source 
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.
This commit is contained in:
Daniel Han 2026-05-07 08:01:21 +00:00
commit a1a20871b2
1 changed files with 15 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

21
tests/studio/run_real_mlx_smoke.py
View file

@ -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: