"""
Module de supervision et d'optimisation de la memoire unifiee Apple Silicon (UMA).
Permet de controler les allocations Metal MLX et de surveiller l'empreinte RAM/VRAM.
"""
from __future__ import annotations
import ctypes
import gc
import os
import re
import struct
import subprocess
import sys
import threading
import time
import urllib.request
import mlx.core as mx
import psutil
from .logger import getMirzaLogger
logger = getMirzaLogger("mirza.memory")
#: Intervalle minimal entre deux purges automatiques.
PURGE_MIN_INTERVAL_SECONDS = 15.0
[docs]
class MemoryManager:
"""
Gestionnaire centralise de la memoire unifiee (RAM CPU + VRAM Metal GPU) pour Apple Silicon.
Sur les SoC Apple (M1/M2/M3/M4), le CPU et le GPU partagent le meme espace
physique de memoire vive unifiee. Module interroge les compteurs du
framework MLX (Metal C++) et de psutil pour une telemetrie precise, et
declenche la liberation proactive des allocations non utilisees.
"""
#: Port du daemon MTPLX, renseigne a l'initialisation. 0 desactive l'appel.
mtplxPort: int = 0
#: Seuil de RSS au-dessus duquel une purge automatique est justifiee.
_purgeThresholdBytes: int = 0
#: Derniere purge automatique, pour borner la frequence.
_lastPurge: float = 0.0
_purgeLock = threading.Lock()
[docs]
@staticmethod
def initializeEnvironment(
vramLimitBytes: int = 16 * 1024 * 1024 * 1024,
cacheLimit: int = 0,
mtplxPort: int = 0,
purgeThresholdGb: float = 0.0,
) -> None:
"""
Configure les limites materielles Metal MLX pour empecher tout debordement de cache.
Parameters
----------
vramLimitBytes : int, optional
Limite maximale d'allocation Metal en octets. Defaut : 16 Go
(16 * 1024**3).
cacheLimit : int, optional
Limite de mise en cache Metal. Defaut : 0, ce qui evite la
retention prolongee de tenseurs temporaires.
mtplxPort : int, optional
Port du daemon MTPLX. 0 desactive l'appel de vidage de cache.
purgeThresholdGb : float, optional
Seuil de RSS declenchant une purge automatique. Defaut : 80 % de
la limite memoire unifiee.
"""
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
MemoryManager.mtplxPort = mtplxPort
if purgeThresholdGb > 0:
MemoryManager._purgeThresholdBytes = int(purgeThresholdGb * (1024**3))
elif vramLimitBytes > 0:
MemoryManager._purgeThresholdBytes = int(vramLimitBytes * 0.8)
mx.set_memory_limit(vramLimitBytes)
mx.set_cache_limit(cacheLimit)
logger.debug("Environnement memoire MLX initialise (limite: %d Go, cache: %d)", vramLimitBytes // (1024**3), cacheLimit)
@staticmethod
def _getProcessFootprint(pid: int) -> int:
"""
Recupere l'empreinte memoire physique reelle (phys_footprint) d'un processus sous macOS en octets.
Inclut les allocations Metal GPU / IOAccelerator ignorees par le RSS standard.
Parameters
----------
pid : int
Identifiant numerique du processus a inspecter.
Returns
-------
int
Taille de l'empreinte memoire physique (phys_footprint) en octets.
"""
if sys.platform == "darwin":
try:
lib = ctypes.CDLL(None)
buf = ctypes.create_string_buffer(512)
if lib.proc_pid_rusage(pid, 4, buf) == 0:
return int(struct.unpack_from("<Q", buf.raw, 72)[0])
except (OSError, AttributeError) as exc:
logger.debug("Impossible d'interroger proc_pid_rusage pour PID %d: %s", pid, exc)
return 0
[docs]
@staticmethod
def getGpuUtilization() -> int:
"""
Interroge l'IOAccelerator via ioreg pour la charge GPU Metal.
Returns
-------
int
Pourcentage d'utilisation, ou 0 si la lecture echoue.
"""
if sys.platform != "darwin":
return 0
try:
out = subprocess.check_output(
["ioreg", "-r", "-d", "1", "-w", "0", "-c", "IOAccelerator"],
timeout=0.2,
)
text = out.decode("latin1", errors="ignore")
m = re.search(r'"Device Utilization %"=(\d+)', text)
if m:
return int(m.group(1))
except Exception as exc:
logger.debug("Erreur lecture charge GPU Metal: %s", exc)
return 0
[docs]
@classmethod
def getMemorySnapshot(cls) -> dict[str, float]:
"""
Capture un instantane exhaustif de l'etat memoire UMA sur Apple Silicon.
Returns
-------
dict[str, float]
Dictionnaire contenant les metriques memoire cles.
"""
proc = psutil.Process()
totalRssBytes = proc.memory_info().rss
countedPids = {proc.pid}
try:
for child in proc.children(recursive=True):
try:
if child.pid not in countedPids:
totalRssBytes += child.memory_info().rss
countedPids.add(child.pid)
except (psutil.NoSuchProcess, psutil.AccessDenied):
continue
except (psutil.NoSuchProcess, psutil.AccessDenied) as exc:
logger.debug("Inspection des processus enfants interrompue: %s", exc)
# Daemons MTPLX autonomes, y compris re-attaches a init
try:
for p in psutil.process_iter(["pid", "cmdline", "memory_info"]):
try:
if p.info["pid"] not in countedPids:
cmd = " ".join(p.info.get("cmdline") or [])
if "mtplx" in cmd:
totalRssBytes += p.info["memory_info"].rss
countedPids.add(p.info["pid"])
except (psutil.NoSuchProcess, psutil.AccessDenied):
continue
except Exception as exc:
logger.debug("Iteration processus systeme: %s", exc)
# Empreinte physique reelle (phys_footprint) sur Apple Silicon
totalFootprintBytes = sum(cls._getProcessFootprint(pid) for pid in countedPids)
sysMem = psutil.virtual_memory()
swapMem = psutil.swap_memory()
activeMem, peakMem, cacheMem = 0, 0, 0
if hasattr(mx, "get_active_memory"):
activeMem = mx.get_active_memory()
peakMem = mx.get_peak_memory()
cacheMem = mx.get_cache_memory()
elif hasattr(mx, "metal"):
activeMem = mx.metal.get_active_memory()
peakMem = mx.metal.get_peak_memory()
cacheMem = mx.metal.get_cache_memory()
# Memoire reelle totale de Mirza
mirzaTotalBytes = max(totalFootprintBytes, totalRssBytes, activeMem + cacheMem)
return {
"processRssMb": totalRssBytes / (1024 * 1024),
"processRssGb": totalRssBytes / (1024**3),
"vramActiveMb": mirzaTotalBytes / (1024 * 1024),
"vramActiveGb": mirzaTotalBytes / (1024**3),
"vramPeakMb": peakMem / (1024 * 1024),
"vramPeakGb": peakMem / (1024**3),
"systemUsedGb": (sysMem.total - sysMem.available) / (1024**3),
"systemTotalGb": sysMem.total / (1024**3),
"systemFreeGb": sysMem.available / (1024**3),
"swapUsedGb": swapMem.used / (1024**3),
}
[docs]
@classmethod
def displayMemoryProfile(cls, stepName: str = "Actuelle") -> None:
"""
Affiche dans le logger l'empreinte memoire formatee et lisible.
Parameters
----------
stepName : str, optional
Nom de l'etape de cycle de vie. Defaut : "Actuelle".
"""
snap = cls.getMemorySnapshot()
logger.info(
"[Profil Memoire : %s] RSS: %.2f Go (%.1f Mo) | VRAM Metal: %.2f Go (Pic: %.2f Go) | RAM Mac: %.2f / %.2f Go",
stepName,
snap["processRssGb"],
snap["processRssMb"],
snap["vramActiveGb"],
snap["vramPeakGb"],
snap["systemUsedGb"],
snap["systemTotalGb"],
)
[docs]
@classmethod
def shouldPurge(cls) -> bool:
"""
Indique si la memoire est sous une pression qui justifie une purge.
La purge n'est gratuite ni en temps ni en cache : ``mx.clear_cache()``
jette les tampons reutilises par Metal, ce qui oblige a les reconstruire
au tour suivant. On ne purge donc que lorsqu'il y a vraiment lieu, et au
plus une fois par intervalle.
Returns
-------
bool
True s'il faut liberer maintenant.
"""
with cls._purgeLock:
now = time.monotonic()
if (now - cls._lastPurge) < PURGE_MIN_INTERVAL_SECONDS:
return False
try:
process = psutil.Process()
used = process.memory_info().rss
except (psutil.NoSuchProcess, psutil.AccessDenied):
return False
cls._lastPurge = now
return used >= cls._purgeThresholdBytes
[docs]
@classmethod
def purgeMemory(cls, force: bool = False) -> None:
"""
Libere les caches Metal et execute le ramasse-miettes Python.
Parameters
----------
force : bool, optional
Si False (defaut), ne purge que sous pression memoire reelle. Si
True, purge inconditionnellement, comme le demande un arret ou une
action explicite de l'utilisateur.
"""
if not force and not cls.shouldPurge():
return
gc.collect()
mx.clear_cache()
try:
import torch
if torch.backends.mps.is_available():
torch.mps.empty_cache()
except Exception as exc:
# torch est une dependance transitive optionnelle : son absence est normale.
logger.debug("Cache MPS torch non vide (torch absent ou indisponible) : %s", exc)
# Inutile sans daemon en ecoute : ce serait un refus de connexion sur le
# chemin chaud de chaque tour.
if cls.mtplxPort:
try:
req = urllib.request.Request(
f"http://127.0.0.1:{cls.mtplxPort}/admin/cache/clear",
data=b"",
headers={"User-Agent": "Mirza-Client/1.0"},
method="POST",
)
with urllib.request.urlopen(req, timeout=0.5):
pass
except Exception as exc:
logger.debug("Cache clear MTPLX: %s", exc)