# SPDX-License-Identifier: Apache-2.0 """ Process-level memory enforcer for oMLX. The memory guard tier says how much memory oMLX leaves for everything else on the Mac: safe keeps ~20% of RAM (6-16 GB) free for heavy apps next to oMLX balanced keeps ~8% of RAM (3-8 GB) free for light apps aggressive keeps 2% of RAM (1.5-4 GB) for the OS and may push half of other apps' active memory into the compressor custom uses a user-pinned ceiling (2 GB static reserve) Ceiling = min(static_ceiling, dynamic_ceiling, metal_cap_effective): static_ceiling = total_ram - tier_reserve dynamic_ceiling = omlx_phys + free + inactive + active file cache + reclaim_ratio * other_apps_active - tier_reserve (custom: the user-pinned ceiling) metal_cap_effective = metal_cap + omlx_cpu_footprint The dynamic ceiling moves with other apps every poll, so the reserve stays free when they grow. The Metal cap limits GPU (wired) memory only, so oMLX's CPU-side footprint is added on top of it before comparing it with total process usage. """ from __future__ import annotations import asyncio import inspect import logging import os import subprocess import time from contextlib import suppress from typing import TYPE_CHECKING, Any import mlx.core as mx from . import settings as _settings from .engine.base import BaseNonStreamingEngine from .utils import psutil_compat from .utils.image import clear_image_decode_cache from .utils.metal_sync import unreleased_graphics_bytes from .utils.proc_memory import ( get_graphics_footprint, get_phys_footprint, release_free_malloc_memory, ) if TYPE_CHECKING: from .engine_pool import EnginePool from .model_settings import ModelSettingsManager from .settings import GlobalSettings logger = logging.getLogger(__name__) # Memory each tier leaves for the OS and other apps: (share of RAM, min, max). _TIER_RESERVE: dict[str, tuple[float, int, int]] = { "safe": (0.20, 6 * 1024**3, 16 * 1024**3), "balanced": (0.08, 3 * 1024**3, 8 * 1024**3), "aggressive": (0.02, int(1.5 * 1024**3), 4 * 1024**3), } # `custom` keeps a small static reserve so a typed ceiling above RAM stays # panic safe. _CUSTOM_STATIC_RESERVE = 2 * 1024**3 _VALID_TIERS = frozenset((*_TIER_RESERVE, "custom")) # Share of other apps' active memory a tier may push into the compressor. # macOS compresses 2-3x, so half of active is reclaimable without swap. _OTHER_APP_RECLAIM_RATIO: dict[str, float] = { "safe": 0.0, "balanced": 0.0, "aggressive": 0.5, } def tier_reserve_bytes(tier: str, total_ram: int) -> int: """Memory ``tier`` leaves for the OS and other apps on a ``total_ram`` Mac.""" if tier == "custom": return _CUSTOM_STATIC_RESERVE share, low, high = _TIER_RESERVE.get(tier, _TIER_RESERVE["balanced"]) return int(min(high, max(low, total_ram * share))) # Default soft watermark per tier. A saved 0.85 from older configs is treated # as the legacy default so every tier can use its own default. _LEGACY_SOFT_THRESHOLD = 0.85 _SOFT_THRESHOLD_BY_TIER: dict[str, float] = { "safe": 0.85, "balanced": 0.90, "aggressive": 0.95, "custom": 0.90, } # Default hard watermark per tier; a saved 0.95 is the legacy default. _LEGACY_HARD_THRESHOLD = 0.95 _HARD_THRESHOLD_BY_TIER: dict[str, float] = { "safe": 0.95, "balanced": 0.95, "aggressive": 0.98, "custom": 0.97, } # Fraction of the hard ceiling used by the adaptive prefill chunk sizer. _PREFILL_HEADROOM_SAFETY: dict[str, float] = { "safe": 0.90, "balanced": 0.92, "aggressive": 0.97, "custom": 0.95, } # Fraction of the effective physical cap used by the pre-chunk prediction # guard. MLX backpressure holds a chunk at the target when the prediction is # short, so aggressive and custom can run close to the ceiling. _PREFILL_ABORT_MARGIN: dict[str, float] = { "safe": 0.90, "balanced": 0.93, "aggressive": 0.97, "custom": 0.95, } # Last-resort active-request brake. Normal hard pressure starts at the # hard watermark (usually 95% of ceiling); pinned workloads are only aborted # after the process crosses the actual ceiling by a small margin, or stays # over the ceiling for consecutive polls. _EMERGENCY_OVER_CEILING_MARGIN_BYTES = 2 * 1024**3 _EMERGENCY_OVER_CEILING_POLLS = 2 _HOT_CACHE_RESERVATION_SLACK_BYTES = 512 * 1024**2 # Longest a prefill waits for released dirty hot-cache blocks to reach SSD. _HOT_CACHE_WRITE_WAIT_S = 5.0 def _format_gb(b: int) -> str: """Format bytes as GB string.""" return f"{b / 1024**3:.1f}GB" def get_macos_vm_stats() -> dict[str, int] | None: """Snapshot of mach `vm_statistics64` in bytes. Returns None on non-macOS or when the host call fails. ~0.8 us per call so this is safe inside the enforcer poll loop and inside per-chunk memcheck. The dict always exposes the first four page counters used by the dynamic ceiling math. Those counters are stable at the front of `vm_statistics64`; speculative and compressed counters are included when the kernel fills the relevant tail fields. A max-sized `host_info64_t` buffer avoids pinning oMLX to an SDK-specific tail layout. """ return psutil_compat.get_macos_vm_stats() def get_iogpu_wired_limit_bytes() -> int: """Read the kernel's `iogpu.wired_limit_mb` sysctl in bytes. Returns 0 when the value is unset (`0` in sysctl means "use the system default", typically ~75% of RAM) or when the read fails. Callers should treat 0 as "limit unknown / not enforced" and fall back to a different source (e.g. mx.device_info()'s working set size). """ try: out = subprocess.run( ["/usr/sbin/sysctl", "-n", "iogpu.wired_limit_mb"], capture_output=True, text=True, timeout=2, ) value_mb = int(out.stdout.strip()) if value_mb <= 0: return 0 return value_mb * 1024**2 except (subprocess.SubprocessError, ValueError, OSError): return 0 def _get_max_metal_working_set_bytes() -> int: """Apple's default Metal cap as reported by MLX (~75% of RAM). `mx.set_wired_limit` refuses any value above this when the kernel iogpu.wired_limit_mb is unset (= 0). """ try: info = mx.device_info() size = int(info.get("max_recommended_working_set_size", 0) or 0) return max(0, size) except Exception: # noqa: BLE001 return 0 def get_effective_metal_cap_bytes() -> int: """Effective per-process Metal allocation cap. Uses the kernel iogpu.wired_limit_mb when explicitly set (> 0). Otherwise falls back to Apple's max_recommended_working_set_size. This is the value above which `mx.set_wired_limit` will reject the request, so callers should clamp against it before calling MLX. """ sysctl_cap = get_iogpu_wired_limit_bytes() if sysctl_cap > 0: return sysctl_cap return _get_max_metal_working_set_bytes() def _wired_limit_suggestion_bytes(desired_bytes: int) -> int: """Clamp and align a wired-limit recommendation to leave 5% of RAM. Wiring within a few GiB of physical RAM invites jetsam during a large-model load burst, and a jetsammed/hard-killed process strands its wired allocation at kernel level until reboot (#2184). 5% matches the field-stable margin from that report (488 GiB stable on a 512 GiB box, 510 GiB crash-looped). On small-memory machines 5% is below the tier static reserve, so their recommendation is unchanged. Only shapes the suggested sysctl value in logs and the admin banner; the enforcement path still honors whatever the kernel sysctl allows, including user-set values above this recommendation. ``iogpu.wired_limit_mb`` accepts whole MiB. Floor the recommendation to that unit so the displayed command cannot cross the byte-level safety threshold and the UI compares against a realizable value. Resolved through the settings module at call time so tests that patch omlx.settings.get_system_memory control this the same way they control the static ceiling. """ try: total = int(_settings.get_system_memory()) except Exception: # noqa: BLE001 return desired_bytes if total <= 0: return desired_bytes suggestion = max(0, min(desired_bytes, total - total // 20)) mib = 1024**2 return suggestion // mib * mib def _apply_metal_wired_limit(desired_bytes: int) -> tuple[int, int | None]: """Try to raise Metal wired limit for this process to `desired_bytes`. Returns (applied_bytes, previous_bytes). `applied_bytes` is what we actually told MLX (clamped to the kernel iogpu.wired_limit_mb if it is lower); `previous_bytes` is what MLX reports the prior limit was, or None on failure / older macOS where the call is unavailable. Emits a WARNING when the kernel sysctl caps us below the recommended wired limit so the user sees the hint in logs in addition to the admin UI red banner. The recommended value is clamped below physical RAM (_wired_limit_suggestion_bytes) so following it cannot push the OS into jetsam during a large-model load (#2184). When iogpu.wired_limit_mb is unset (0), leave Apple's default Metal cap active instead of calling mx.set_wired_limit with the same default cap. The scheduler still clamps against get_effective_metal_cap_bytes(); this only avoids changing MLX allocator state unless the user explicitly raised the kernel cap. """ if desired_bytes <= 0: return 0, None suggestion = _wired_limit_suggestion_bytes(desired_bytes) sysctl_cap = get_iogpu_wired_limit_bytes() if sysctl_cap <= 0: effective_cap = get_effective_metal_cap_bytes() if effective_cap > 0 and effective_cap < suggestion: logger.warning( "Metal cap (%s, Apple max_recommended_working_set_size) is " "below the oMLX static ceiling (%s); leaving Apple's default " "Metal cap active because iogpu.wired_limit_mb is unset. " "Raise it with: sudo sysctl iogpu.wired_limit_mb=%d", _format_gb(effective_cap), _format_gb(desired_bytes), suggestion // (1024**2), ) else: logger.debug( "Skipping mx.set_wired_limit because iogpu.wired_limit_mb is " "unset (target=%s, Apple cap=%s)", _format_gb(desired_bytes), _format_gb(effective_cap), ) return 0, None try: _total = int(_settings.get_system_memory()) except Exception: # noqa: BLE001 _total = 0 if _total > 0: _reserve = _total // 20 if sysctl_cap > _total - _reserve: # The kernel cap leaves the OS less headroom than the # recommendation floor. A big-model load burst can jetsam the # server at this setting, and the stranded wired memory then # needs a reboot to reclaim (#2184). logger.warning( "iogpu.wired_limit_mb (%s) leaves the OS less than %s of " "physical RAM (%s); a large model load can trigger jetsam " "at this setting and strand the wired memory until reboot " "(#2184). Recommended maximum: %d MB.", _format_gb(sysctl_cap), _format_gb(_reserve), _format_gb(_total), (_total - _reserve) // (1024**2), ) effective_cap = sysctl_cap capped = effective_cap > 0 and effective_cap < desired_bytes applied = effective_cap if capped else desired_bytes try: previous = mx.set_wired_limit(applied) if capped and effective_cap < suggestion: logger.warning( "Metal cap (%s, %s) is below the oMLX static ceiling (%s); " "Metal will clamp allocations to the cap and panic if a " "request exceeds it. Raise it with: sudo sysctl " "iogpu.wired_limit_mb=%d", _format_gb(effective_cap), "kernel iogpu.wired_limit_mb", _format_gb(desired_bytes), suggestion // (1024**2), ) return applied, int(previous) except Exception as exc: # noqa: BLE001 # Older macOS (<15) or the API just isn't available. Log + skip. logger.warning( "mx.set_wired_limit(%s) failed; Metal will use its default cap (%s)", _format_gb(applied), exc, ) return 0, None def preview_tier_ceilings() -> dict[str, dict[str, Any]]: """Ceiling each tier would set right now, for the settings previews. ``custom`` reports only its static and Metal limits because its ceiling is the user's draft value. """ metal_cap_raw = get_effective_metal_cap_bytes() previews: dict[str, dict[str, Any]] = {} for tier in ("safe", "balanced", "aggressive", "custom"): enforcer = ProcessMemoryEnforcer(engine_pool=None, memory_guard_tier=tier) enforcer._effective_metal_cap_bytes = metal_cap_raw terms = ( {"free": 0, "inactive": 0, "other_apps": 0, "dynamic": 0} if tier == "custom" else enforcer._dynamic_ceiling_terms() ) static = enforcer._get_static_ceiling() metal_cap = ( metal_cap_raw + enforcer._cpu_footprint_bytes() if metal_cap_raw > 0 else 0 ) components = { name: value for name, value in ( ("static", static), ("dynamic", terms["dynamic"]), ("metal_cap", metal_cap), ) if value > 0 } ceiling = min(components.values()) if components else 0 previews[tier] = { "reserve_bytes": enforcer._tier_reserve_bytes(), "free_bytes": terms["free"], "inactive_bytes": terms["inactive"], "other_apps_bytes": terms["other_apps"], "static_bytes": static, "dynamic_bytes": terms["dynamic"], "metal_cap_bytes": metal_cap, "ceiling_bytes": ceiling, "binding": next( (name for name, value in components.items() if value == ceiling), "", ), } return previews class ProcessMemoryEnforcer: """ Background task that enforces process-level memory limits. Polls usage every poll_interval seconds. On every tick it recomputes the dynamic ceiling from system_available, so other-app pressure is reflected immediately without restarting the enforcer. """ def __init__( self, engine_pool: EnginePool, memory_guard_tier: str = "balanced", memory_guard_custom_ceiling_gb: float = 0.0, poll_interval: float = 1.0, settings_manager: ModelSettingsManager | None = None, prefill_memory_guard: bool = True, global_settings: GlobalSettings | None = None, soft_threshold: float | None = None, hard_threshold: float | None = None, prefill_safe_zone_ratio: float = 0.89, prefill_min_chunk_tokens: int = 256, ): """ Initialize the process memory enforcer. Args: engine_pool: The engine pool to evict models from. memory_guard_tier: One of "safe", "balanced", "aggressive", "custom". Picks the memory reserved for other apps and how much of their active memory may be compressed. "custom" uses memory_guard_custom_ceiling_gb directly for the dynamic ceiling instead of computing from vm_stat. memory_guard_custom_ceiling_gb: Custom ceiling in GB. Only used when tier == "custom". Clamped by static_ceiling and metal_cap so a too-large value is panic-safe. poll_interval: Seconds between memory checks. settings_manager: Optional settings manager for TTL checks. prefill_memory_guard: When False, returns a ceiling of 0 so callers treat the limit as disabled. global_settings: Optional global settings for idle timeout. soft_threshold: Optional explicit fraction of ceiling that triggers soft action. None, or the legacy 0.85 default, uses the tier default instead. hard_threshold: Optional fraction of ceiling that triggers hard action (LRU/non-pinned aborts, loading aborts, and idle reclaim). None, or the legacy 0.95 default, uses the tier default instead. prefill_safe_zone_ratio: Fraction of hard cap below which prefill runs at full chunk size; above triggers adaptive shrink. prefill_min_chunk_tokens: Floor for adaptive shrink. """ self._engine_pool = engine_pool self._memory_guard_tier = self._normalize_tier(memory_guard_tier) self._memory_guard_custom_ceiling_bytes = max( 0, int(memory_guard_custom_ceiling_gb * 1024**3) ) self._active_poll_interval = poll_interval self._loaded_idle_poll_interval = 10.0 self._unloaded_idle_poll_interval = 30.0 self._current_poll_interval = poll_interval self._settings_manager = settings_manager self._prefill_memory_guard = prefill_memory_guard self._global_settings = global_settings self._soft_threshold_override = self._normalize_threshold_override( soft_threshold, _LEGACY_SOFT_THRESHOLD ) self._soft_threshold = self._get_soft_threshold() self._hard_threshold_override = self._normalize_threshold_override( hard_threshold, _LEGACY_HARD_THRESHOLD ) self._hard_threshold = self._get_hard_threshold() self._prefill_headroom_safety = self._get_prefill_headroom_safety() self._prefill_safe_zone_ratio = prefill_safe_zone_ratio self._prefill_min_chunk_tokens = prefill_min_chunk_tokens self._task: asyncio.Task | None = None self._loop: asyncio.AbstractEventLoop | None = None self._wake_event: asyncio.Event | None = None self._running = False self._activity_hint_until = 0.0 # Most recently observed pressure level, consumed by scheduler / # admission control. Updated on every poll iteration. self._pressure_level: str = "ok" self._over_ceiling_polls: int = 0 # Consecutive hard-pressure polls we have deferred destructive # enforcement while a fat Metal buffer pool drains at the next # prefill chunk / step boundary. See the grace check in # _check_and_enforce. self._pressure_reclaim_grace_polls: int = 0 # Last value passed to mx.set_wired_limit (0 if not yet applied # or the call failed). Used by the admin dashboard to surface a # warning when the kernel iogpu.wired_limit_mb is below this. self._metal_wired_limit_request: int = 0 # Cached Metal cap used by the background poll loop. Reading the Apple # default cap falls back to mx.device_info(), so keep it out of active # decode ticks once the enforcer is running. self._effective_metal_cap_bytes: int | None = None # Engine types we've already complained about in # ``_propagate_memory_limit``'s "scheduler unreachable" path. # Prevents the per-poll warning from spamming logs while keeping # the first occurrence loud enough to alert CI / oncall. self._scheduler_resolve_warned: set[str] = set() @staticmethod def _normalize_tier(tier: str) -> str: t = (tier or "").strip().lower() if t not in _VALID_TIERS: return "balanced" return t @staticmethod def _normalize_threshold_override( value: float | None, legacy_default: float ) -> float | None: if value is None: return None threshold = float(value) if threshold <= 0: return None if abs(threshold - legacy_default) < 1e-9: return None return threshold def _refresh_tier_thresholds(self) -> None: self._soft_threshold = self._get_soft_threshold() self._hard_threshold = self._get_hard_threshold() self._prefill_headroom_safety = self._get_prefill_headroom_safety() def _get_soft_threshold(self) -> float: if self._soft_threshold_override is not None: return self._soft_threshold_override return _SOFT_THRESHOLD_BY_TIER[self._memory_guard_tier] def _get_hard_threshold(self) -> float: if self._hard_threshold_override is not None: return self._hard_threshold_override return _HARD_THRESHOLD_BY_TIER[self._memory_guard_tier] def _get_prefill_headroom_safety(self) -> float: return _PREFILL_HEADROOM_SAFETY[self._memory_guard_tier] @property def memory_guard_tier(self) -> str: return self._memory_guard_tier @memory_guard_tier.setter def memory_guard_tier(self, value: str) -> None: new_tier = self._normalize_tier(value) if new_tier == self._memory_guard_tier: return old = self._memory_guard_tier self._memory_guard_tier = new_tier self._refresh_tier_thresholds() if self._running: if self._prefill_memory_guard: self._refresh_effective_metal_cap_bytes() # The kernel-cap banner compares against this; the MLX wired # limit itself is re-armed on the next start. self._metal_wired_limit_request = _wired_limit_suggestion_bytes( self._get_static_ceiling() ) self._propagate_memory_limit() logger.info(f"Memory guard tier changed: {old} -> {new_tier}") @property def memory_guard_custom_ceiling_bytes(self) -> int: return self._memory_guard_custom_ceiling_bytes @memory_guard_custom_ceiling_bytes.setter def memory_guard_custom_ceiling_bytes(self, value: int) -> None: new_value = max(0, int(value)) if new_value == self._memory_guard_custom_ceiling_bytes: return old = self._memory_guard_custom_ceiling_bytes self._memory_guard_custom_ceiling_bytes = new_value if self._running: if self._prefill_memory_guard: self._refresh_effective_metal_cap_bytes() self._propagate_memory_limit() logger.info( "Memory guard custom ceiling changed: %s -> %s", _format_gb(old), _format_gb(new_value), ) @property def is_running(self) -> bool: """Whether the enforcement loop is active.""" return self._running def start(self) -> None: """Start the background enforcement loop. Also mirrors the static ceiling into MLX's wired-memory limit when the user explicitly raised iogpu.wired_limit_mb. When the kernel sysctl is unset, the scheduler still clamps against Apple's default Metal cap, but oMLX leaves MLX allocator state untouched. """ if self._running: return self._pressure_reclaim_grace_polls = 0 if self._prefill_memory_guard: self._refresh_effective_metal_cap_bytes() self._running = True self._propagate_memory_limit() ceiling = self._get_hard_limit_bytes() if self._prefill_memory_guard: static_ceiling = self._get_static_ceiling() applied, previous = _apply_metal_wired_limit(static_ceiling) # Store the *recommended* limit (static ceiling clamped below # physical RAM, see _wired_limit_suggestion_bytes) rather than # the post-clamp applied value. The admin UI compares this # against the live iogpu.wired_limit_mb so a kernel cap below # the recommendation triggers the red sysctl-command banner; # clamping keeps the banner from telling users to wire nearly # all of physical RAM (#2184). self._metal_wired_limit_request = _wired_limit_suggestion_bytes( static_ceiling ) if applied > 0: logger.info( "Metal wired limit raised: %s -> %s " "(target=%s, iogpu sysctl cap=%s)", _format_gb(previous or 0), _format_gb(applied), _format_gb(static_ceiling), _format_gb(get_iogpu_wired_limit_bytes()), ) self._task = asyncio.create_task(self._enforcement_loop()) logger.info( f"Process memory enforcer started " f"(tier={self._memory_guard_tier}, " f"ceiling={_format_gb(ceiling)}, " f"interval={self._active_poll_interval}s)" ) def wake(self, *, active: bool = False) -> None: """Wake the polling loop before its current sleep timeout expires. ``active=True`` keeps the loop on the fast interval briefly. This covers request/model-load entry points before an engine's active-request collectors are visible to ``_select_poll_interval``. """ if active: self._activity_hint_until = max( self._activity_hint_until, time.monotonic() + max(2.0, self._active_poll_interval * 2), ) event = self._wake_event loop = self._loop if event is None or loop is None or loop.is_closed(): return try: running_loop = asyncio.get_running_loop() except RuntimeError: running_loop = None if running_loop is loop: event.set() else: loop.call_soon_threadsafe(event.set) def _get_static_ceiling(self) -> int: """Total RAM minus the tier reserve.""" from .settings import get_system_memory system_bytes = get_system_memory() return max(0, system_bytes - self._tier_reserve_bytes(system_bytes)) def _tier_reserve_bytes(self, system_bytes: int | None = None) -> int: """Memory the current tier leaves for the OS and other apps.""" if system_bytes is None: from .settings import get_system_memory system_bytes = get_system_memory() return tier_reserve_bytes(self._memory_guard_tier, int(system_bytes)) def _get_dynamic_ceiling(self) -> int: """Tier-aware reclaimable-memory ceiling. custom: Returns the user-supplied ceiling verbatim (clamped >= 0). min() with static / metal_cap still applies in `_get_hard_limit_bytes` so out-of-range input is panic safe. safe / balanced / aggressive: omlx_phys + free + inactive + active file cache + other_apps_active * ratio - tier_reserve free / inactive / active come from `host_statistics64` (recomputed every call — never cached). Subtracting the tier reserve keeps that much memory free for other apps, so the ceiling shrinks as they grow. File-backed pages are dropped without compression, so every tier counts them. Only aggressive counts part of other apps' anonymous active memory, which macOS must compress to hand over. oMLX's own CPU pages are removed from that term; its Metal buffers are wired and never appear there. Speculative and purgeable pages are subsets of free / inactive, so we deliberately do not add them (would double count). VM stats failure: falls back to psutil_compat.virtual_memory().available (= roughly free + inactive on macOS, similar elsewhere). On macOS that compat layer avoids psutil's HOST_VM_INFO64 adapter and uses a cached vm_stat fallback only if the fast host_statistics64 path is unavailable. If all telemetry is unavailable, fall back to the static ceiling so server health endpoints and the enforcer keep running. """ if self._memory_guard_tier == "custom": return max(0, self._memory_guard_custom_ceiling_bytes) return self._dynamic_ceiling_terms()["dynamic"] def _dynamic_ceiling_terms(self) -> dict[str, int]: """Terms of the reserve-tier dynamic ceiling, for the guard and previews.""" omlx_usage = get_phys_footprint() reserve = self._tier_reserve_bytes() terms = { "omlx": omlx_usage, "free": 0, "inactive": 0, "other_apps": 0, "reserve": reserve, } stats = get_macos_vm_stats() if stats is None: try: available = int(psutil_compat.virtual_memory().available) except Exception as exc: # noqa: BLE001 logger.warning( "Memory guard could not read available memory; " "using static ceiling fallback: %s", exc, ) return {**terms, "dynamic": self._get_static_ceiling()} terms["free"] = available else: terms["free"] = stats["free"] # File-backed pages still in the active queue (a model file just # read, say) are dropped without compression. Inactive may hold # some of them already, so only the excess over it is added. file_active = max(0, stats.get("external", 0) - stats["inactive"]) terms["inactive"] = stats["inactive"] + file_active ratio = _OTHER_APP_RECLAIM_RATIO[self._memory_guard_tier] if ratio > 0: own_cpu = max(0, omlx_usage - get_graphics_footprint()) anon_active = stats["active"] - file_active - own_cpu terms["other_apps"] = int(max(0, anon_active) * ratio) reclaimable = terms["free"] + terms["inactive"] + terms["other_apps"] # Other apps may already hold the reserve. Stay at 1 byte rather than # 0, which callers read as "guard disabled". terms["dynamic"] = max(1, omlx_usage + reclaimable - reserve) return terms def _get_hard_limit_bytes(self) -> int: """Final hard ceiling = min(static, dynamic, metal_cap). Thin wrapper over ``_get_ceiling_breakdown`` that discards the component breakdown. Hot callers that don't need to know which ceiling is binding should keep using this helper. `metal_cap` is the effective Metal allocation cap (kernel iogpu.wired_limit_mb when set, otherwise Apple's max_recommended_working_set_size). Including it here means oMLX never plans allocations above what Metal will actually accept, so users who have not raised iogpu.wired_limit_mb still get a safe (smaller) ceiling rather than a panic. For `custom` tier the dynamic ceiling (vm_stat-based) is skipped; the user-specified value is capped only by static (total - 2 GB) and metal_cap. Returns 0 if the memory guard is disabled (callers treat 0 as "no limit"). """ return self._get_ceiling_breakdown()["hard_limit"] def _get_ceiling_breakdown(self) -> dict[str, int]: """Compute the hard limit AND the three component ceilings. Returns a dict with keys ``static``, ``dynamic``, ``metal_cap``, ``hard_limit`` (= min of the three non-zero values, or 0 when the guard is disabled). Used by ``_propagate_memory_limit`` to push the breakdown to schedulers so the prefill-rejection error message can identify which constraint is binding and suggest the right remedy. Single computation so the subprocess to ``sysctl`` (inside ``get_effective_metal_cap_bytes``) only fires once per call. ``metal_cap`` is in total-usage units: the kernel cap limits GPU (wired) memory, so oMLX's CPU footprint is added to it. """ if not self._prefill_memory_guard: return {"static": 0, "dynamic": 0, "metal_cap": 0, "hard_limit": 0} static_ceiling = self._get_static_ceiling() if self._memory_guard_tier == "custom": dynamic_ceiling = max(0, self._memory_guard_custom_ceiling_bytes) else: dynamic_ceiling = self._get_dynamic_ceiling() metal_cap_raw = self._get_effective_metal_cap_bytes() metal_cap = ( metal_cap_raw + self._cpu_footprint_bytes() if metal_cap_raw > 0 else 0 ) # An unset custom ceiling (0) must not read as "guard disabled". candidates = [static_ceiling] if dynamic_ceiling > 0: candidates.append(dynamic_ceiling) if metal_cap > 0: candidates.append(metal_cap) return { "static": static_ceiling, "dynamic": dynamic_ceiling, "metal_cap": metal_cap, "hard_limit": min(candidates), } @staticmethod def _cpu_footprint_bytes() -> int: """oMLX footprint outside Metal buffers (0 without the graphics ledger).""" graphics = get_graphics_footprint() if graphics <= 0: return 0 return max(0, get_phys_footprint() - graphics) def get_final_ceiling(self) -> int: """Public accessor used by engine_pool pre-load admission.""" return self._get_hard_limit_bytes() def get_ceiling_breakdown(self) -> dict[str, int]: """Public accessor for the component ceilings. The engine pool reads this when a load is refused so the error can name the binding constraint, the same way the scheduler's prefill-rejection message does from its propagated copies. """ return self._get_ceiling_breakdown() def get_admission_ceiling(self) -> int: """Best-effort pre-load ceiling that survives disabling the guard. ``get_final_ceiling()`` returns 0 when the prefill memory guard is off, which used to disable engine-pool pre-load eviction entirely: loading a second model would overcommit physical memory and thrash the machine instead of evicting the LRU model (#2290). This accessor keeps load-time eviction working independently of the guard toggle by falling back to the static ceiling (total RAM minus the tier reserve) when the guard is disabled. The Metal cap and the dynamic (vm_stat) ceiling are deliberately excluded from the fallback: with the guard off nothing is wired, allocations beyond Apple's recommended working set stay pageable, and workloads legitimately run above that cap. The engine pool treats this fallback as best-effort — it evicts idle LRU models to fit under it but never refuses a load, preserving the unguarded "no hard limits" semantics. """ if self._prefill_memory_guard: return self._get_hard_limit_bytes() return self._get_static_ceiling() def get_residency_ceiling(self) -> int: """Stable ceiling for *where weights live*, not for admitting a load. Admission must respect memory free right now — overcommitting thrashes the machine. Choosing whether a model's PLE table stays resident or falls back to mmap is a throughput call, and the instantaneous ceiling is actively wrong for it: engine_pool asks right after the previous model unloaded, before the OS has returned its pages, so ``dynamic`` dips and a table that fits gets pushed to SSD for the rest of that engine's life. Measured on a model swap here: 35.3 -> 14.1 tok/s. Uses only the two components that don't move with instantaneous pressure. Returns 0 when the guard is off, like the other accessors. """ if not self._prefill_memory_guard: return 0 candidates = [ value for value in ( self._get_static_ceiling(), self._get_effective_metal_cap_bytes(), ) if value > 0 ] return min(candidates) if candidates else 0 def get_admission_soft_target(self) -> int: """Soft watermark that pre-load admission evicts down to (#2319). The admission check used to evict idle models only when the projected total exceeded the final ceiling, while every other pressure mechanism (pressure levels, prefill-headroom eviction) targets the soft watermark. A second model admitted into the soft..ceiling band kept both models resident through the load, pushing the process into hard-pressure swap until the first request's prefill guard finally evicted the old one. Exposing the soft watermark lets the engine pool evict idle LRU models down to the same target *before* the new weights start allocating; load refusal stays governed by the admission ceiling. Returns 0 when no ceiling is available (callers fall back to ceiling-only admission). """ ceiling = self.get_admission_ceiling() if ceiling <= 0: return 0 return int(ceiling * self._soft_threshold) def _get_abort_limit_bytes(self) -> int: """Stable physical cap used to ABORT an in-flight prefill. Deliberately excludes the dynamic ceiling: that value jitters every poll with other-app pressure, and a transient dip must not kill a near-complete prefill whose usage actually fits the physical envelope. We use ``min(static_ceiling, metal_cap)`` — exactly the limit ``start()`` arms via ``mx.set_wired_limit`` — so allocating up to it cannot trigger a Metal clamp/panic. The dynamic ceiling still governs chunk-size throttling and admission elsewhere; this is only the last-resort kill threshold. Returns 0 when the guard is disabled (callers treat 0 as "no limit" and fall back to the dynamic hard limit). """ if not self._prefill_memory_guard: return 0 static_ceiling = self._get_static_ceiling() metal_cap = self._get_effective_metal_cap_bytes() if metal_cap > 0: return min(static_ceiling, metal_cap + self._cpu_footprint_bytes()) return static_ceiling def _get_prefill_abort_margin(self) -> float: """Tier-specific prediction margin for pre-chunk safety checks.""" return _PREFILL_ABORT_MARGIN[self._memory_guard_tier] def _soft_bytes(self) -> int: """Soft watermark: ceiling * soft_threshold.""" ceiling = self._get_hard_limit_bytes() if ceiling <= 0: return 0 return int(ceiling * self._soft_threshold) def _hard_bytes(self) -> int: """Hard watermark: ceiling * hard_threshold.""" ceiling = self._get_hard_limit_bytes() if ceiling <= 0: return 0 return int(ceiling * self._hard_threshold) def _current_usage_bytes(self) -> int: """Process memory usage as seen by macOS jetsam. During active requests this must not call MLX/Metal APIs from the background enforcer thread. The scheduler records the last mx.get_active_memory() sample on the MLX executor thread; the enforcer combines that cached value with the kernel phys_footprint ledger. When no request is active we keep the legacy direct MLX telemetry path so idle/status accounting remains as precise as before. """ phys = get_phys_footprint() if self._has_active_requests(): active, mlx_bytes = self._cached_executor_mlx_memory_bytes() fresh = False else: active = mx.get_active_memory() mlx_bytes = active + mx.get_cache_memory() fresh = True return max(active, phys - unreleased_graphics_bytes(mlx_bytes, fresh=fresh)) def _is_emergency_pressure(self, current: int, ceiling: int) -> bool: """Return True only for pressure beyond the configured ceiling. This deliberately does not fire at the hard watermark. A pinned model means "do not unload"; active requests remain protected until the process is at or above the real ceiling and needs a last-resort brake. """ if ceiling <= 0 or current < ceiling: self._over_ceiling_polls = 0 return False self._over_ceiling_polls += 1 if current >= ceiling + _EMERGENCY_OVER_CEILING_MARGIN_BYTES: return True return self._over_ceiling_polls >= _EMERGENCY_OVER_CEILING_POLLS def _refresh_effective_metal_cap_bytes(self) -> int: """Refresh the cached effective Metal cap outside the poll hot path.""" self._effective_metal_cap_bytes = get_effective_metal_cap_bytes() return self._effective_metal_cap_bytes def _get_effective_metal_cap_bytes(self) -> int: """Return the cached Metal cap, populating it on first use.""" if self._effective_metal_cap_bytes is None: return self._refresh_effective_metal_cap_bytes() return self._effective_metal_cap_bytes def _has_active_requests(self) -> bool: """Best-effort active-request detection without touching MLX.""" for entry in self._engine_pool._entries.values(): engine = getattr(entry, "engine", None) if engine is None: continue has_active_requests = getattr(engine, "has_active_requests", None) if not callable(has_active_requests): continue try: if has_active_requests() is True: return True except Exception: return True return False def _cached_executor_active_memory_bytes(self) -> int: """Max MLX active-memory sample recorded by scheduler executor threads.""" return self._cached_executor_mlx_memory_bytes()[0] def _cached_executor_mlx_memory_bytes(self) -> tuple[int, int]: """Max (active, active + pool) MLX samples from scheduler executors.""" active = 0 total = 0 for entry in self._engine_pool._entries.values(): scheduler = self._resolve_scheduler(entry) if scheduler is None: continue getter = getattr(scheduler, "get_cached_mlx_active_memory_bytes", None) total_getter = getattr(scheduler, "get_cached_mlx_memory_bytes", None) try: value = ( getter() if callable(getter) else getattr(scheduler, "_last_mlx_active_memory_bytes", 0) ) value_total = total_getter() if callable(total_getter) else value except Exception: continue if isinstance(value, (int, float)): active = max(active, int(value)) if isinstance(value_total, (int, float)): total = max(total, int(value_total)) return active, max(active, total) @staticmethod def _nonnegative_bytes(value: Any) -> int | None: if isinstance(value, bool) or not isinstance(value, (int, float)): return None return max(0, int(value)) @classmethod def _nonnegative_byte_attr(cls, obj: Any, name: str) -> int | None: try: value = getattr(obj, name) except Exception: return None return cls._nonnegative_bytes(value) def _hot_cache_budget(self) -> Any | None: config = getattr(self._engine_pool, "_scheduler_config", None) if config is None: return None budget = getattr(config, "hot_cache_budget", None) if budget is None: return None max_bytes = self._nonnegative_byte_attr(budget, "max_bytes") total_bytes = self._nonnegative_byte_attr(budget, "total_bytes") if max_bytes is None or max_bytes <= 0 or total_bytes is None: return None return budget def _hot_cache_max_bytes(self) -> int: budget = self._hot_cache_budget() if budget is not None: return self._nonnegative_byte_attr(budget, "max_bytes") or 0 config = getattr(self._engine_pool, "_scheduler_config", None) if config is None: return 0 return self._nonnegative_byte_attr(config, "hot_cache_max_size") or 0 @staticmethod def _manager_hot_cache_bytes(manager: Any) -> int: try: stats = manager.get_stats() size = ProcessMemoryEnforcer._nonnegative_byte_attr( stats, "hot_cache_size_bytes" ) if size is not None: return size except Exception: pass return ( ProcessMemoryEnforcer._nonnegative_byte_attr( manager, "_hot_cache_total_bytes" ) or 0 ) def _hot_cache_used_bytes(self) -> int: budget = self._hot_cache_budget() if budget is not None: return self._nonnegative_byte_attr(budget, "total_bytes") or 0 total = 0 seen_managers: set[int] = set() for entry in self._engine_pool._entries.values(): scheduler = self._resolve_scheduler(entry) if scheduler is None: continue manager = getattr(scheduler, "paged_ssd_cache_manager", None) if manager is None or id(manager) in seen_managers: continue seen_managers.add(id(manager)) total += self._manager_hot_cache_bytes(manager) return total def _hot_cache_reserved_bytes(self) -> int: max_bytes = self._hot_cache_max_bytes() if max_bytes <= 0: return 0 used = self._hot_cache_used_bytes() return min(max_bytes, used + _HOT_CACHE_RESERVATION_SLACK_BYTES) def _scheduler_limit_bytes( self, process_limit: int, *, reserved: int | None = None ) -> int: if process_limit <= 0: return 0 if reserved is None: reserved = self._hot_cache_reserved_bytes() if reserved <= 0: return process_limit return max(1, process_limit - reserved) def _active_hot_cache_block_hashes(self) -> set[bytes]: hashes: set[bytes] = set() for entry in self._engine_pool._entries.values(): scheduler = self._resolve_scheduler(entry) if scheduler is None: continue getter = getattr(scheduler, "get_active_hot_cache_block_hashes", None) if not callable(getter): continue try: hashes.update(bytes(h) for h in getter()) except Exception: logger.debug("Failed to read active hot-cache block hashes") return hashes def _shrink_hot_cache_for_pressure(self, current: int, target: int) -> int: """Try to shrink hot cache enough to move process usage toward target.""" if current <= target: return 0 freed = self._shrink_hot_cache_by(current - target) if freed > 0: release_free_malloc_memory() return freed async def release_hot_cache_for_prefill(self, reserved_bytes: int) -> int: """Shrink the hot cache until its reservation drops by ``reserved_bytes``. Blocks of admitted or queued requests stay. Nothing is released when the other blocks cannot cover the drop. Returns the actual drop. """ before = self._hot_cache_reserved_bytes() if before <= 0 or reserved_bytes <= 0: return 0 # The reservation is used + slack, capped at max_bytes. target_used = max( 0, before - reserved_bytes - _HOT_CACHE_RESERVATION_SLACK_BYTES ) excess = self._hot_cache_used_bytes() - target_used if excess <= 0: return 0 protected = self._active_hot_cache_block_hashes() releasable = getattr(self._hot_cache_budget(), "releasable_bytes", None) if callable(releasable) and releasable(protected) < excess: return 0 await asyncio.to_thread(self._release_hot_cache_blocks, excess, protected) released = max(0, before - self._hot_cache_reserved_bytes()) if released > 0: self._propagate_memory_limit() return released def _release_hot_cache_blocks( self, bytes_to_free: int, protected_hashes: set[bytes] ) -> None: if self._shrink_hot_cache_by(bytes_to_free, protected_hashes) <= 0: return # Dirty blocks stay in RAM until the SSD writer has written them. deadline = time.monotonic() + _HOT_CACHE_WRITE_WAIT_S seen_managers: set[int] = set() for entry in self._engine_pool._entries.values(): scheduler = self._resolve_scheduler(entry) manager = getattr(scheduler, "paged_ssd_cache_manager", None) wait = getattr(manager, "wait_for_pending_writes", None) if callable(wait) and id(manager) not in seen_managers: seen_managers.add(id(manager)) wait(max(0.0, deadline - time.monotonic())) release_free_malloc_memory() def _shrink_hot_cache_by( self, bytes_to_free: int, protected_hashes: set[bytes] | None = None ) -> int: """Free up to ``bytes_to_free`` of unprotected hot-cache blocks by LRU.""" hot_used = self._hot_cache_used_bytes() if hot_used <= 0 or bytes_to_free <= 0: return 0 target_hot_bytes = max(0, hot_used - bytes_to_free) if protected_hashes is None: protected_hashes = self._active_hot_cache_block_hashes() budget = self._hot_cache_budget() if budget is not None: shrink = getattr(budget, "shrink_to", None) if callable(shrink): freed = int( shrink(target_hot_bytes, protected_hashes=protected_hashes) or 0 ) if freed > 0: logger.warning( "Shrank shared hot cache under memory pressure: " "freed=%s target_hot=%s protected=%d", _format_gb(freed), _format_gb(target_hot_bytes), len(protected_hashes), ) return freed freed_total = 0 remaining_to_free = bytes_to_free seen_managers: set[int] = set() for entry in self._engine_pool._entries.values(): if remaining_to_free <= 0: break scheduler = self._resolve_scheduler(entry) if scheduler is None: continue manager = getattr(scheduler, "paged_ssd_cache_manager", None) if manager is None or id(manager) in seen_managers: continue seen_managers.add(id(manager)) manager_used = self._manager_hot_cache_bytes(manager) if manager_used <= 0: continue target_manager_bytes = max(0, manager_used - remaining_to_free) shrink = getattr(manager, "shrink_hot_cache_to", None) if not callable(shrink): continue freed = int( shrink(target_manager_bytes, protected_hashes=protected_hashes) or 0 ) freed_total += freed remaining_to_free = max(0, remaining_to_free - freed) if freed_total > 0: logger.warning( "Shrank hot cache under memory pressure: freed=%s protected=%d", _format_gb(freed_total), len(protected_hashes), ) return freed_total # Pool bytes below which a pressure-triggered clear is not worth the # per-clear IOGPUFamily refcount cost (matches the scheduler's own # _periodic_clear_threshold_bytes floor). Above it, a clear meaningfully # returns memory to the OS. _POOL_RECLAIM_FLOOR = 2 * 1024**3 # Max consecutive hard-pressure polls destructive enforcement is deferred # while a reclaimable pool drains. Prefill chunks run ~5s at full step # size, so 5 one-second polls cover at least one chunk boundary. _PRESSURE_RECLAIM_GRACE_POLLS_MAX = 5 def _pool_bytes(self) -> int: """MLX buffer-pool size, 0 when unreadable (mocked mx in tests).""" try: return int(mx.get_cache_memory()) except (TypeError, ValueError): return 0 def _request_scheduler_cache_reclaim(self, freed_hot: int) -> int: """Ask each scheduler to run its step-boundary Metal cache clear. Routes the reclaim through the scheduler's shipped, lock-protected, engine-stream-synchronized ``_sync_and_clear_cache`` path (the same one the periodic clear uses) instead of touching Metal from the enforcer thread — the enforcer must never touch Metal directly. Setting the per-scheduler flag is GIL-atomic; the actual clear fires on the inference thread at the next step boundary, even under load — unlike ``request_idle_reclaim``, which never fires while requests are running and so cannot recover a busy server from hard pressure. Fires before destructive hard-pressure enforcement when the MLX buffer pool alone holds a meaningful amount, or after a hot-cache shrink just freed references so those newly pooled buffers can return to the OS. """ # A non-numeric reading (e.g. a wholesale-mocked mx in unit tests) # cannot justify a clear; _pool_bytes treats it as an empty pool. pool_bytes = self._pool_bytes() if freed_hot <= 0 and pool_bytes <= self._POOL_RECLAIM_FLOOR: return 0 requested = 0 for entry in self._engine_pool._entries.values(): scheduler = self._resolve_scheduler(entry) request = getattr(scheduler, "request_pressure_reclaim", None) if callable(request): try: request() except Exception as exc: # noqa: BLE001 logger.warning("Pressure cache reclaim request failed: %s", exc) continue requested += 1 if requested: logger.info( "Requested pressure cache reclaim on %d scheduler(s) " "(freed_hot=%s, pool=%s)", requested, _format_gb(freed_hot), _format_gb(pool_bytes), ) return requested def get_pressure_level(self) -> str: """Return cached pressure level: 'ok', 'soft', or 'hard'. Consumed by scheduler `_schedule_waiting` and HTTP admission control. Updated on every enforcer poll iteration. """ return self._pressure_level if self._running else "ok" @property def prefill_memory_guard(self) -> bool: """Whether prefill memory guard is enabled.""" return self._prefill_memory_guard @prefill_memory_guard.setter def prefill_memory_guard(self, value: bool) -> None: self._prefill_memory_guard = value if self._running: self._propagate_memory_limit() logger.info(f"Prefill memory guard: {'enabled' if value else 'disabled'}") @staticmethod def _resolve_scheduler(entry: Any) -> Any | None: """Resolve the watermark target (Scheduler or DFlash guard) from an EnginePool entry. Most engines (BatchedEngine, VLMBatchedEngine) wrap the scheduler as ``entry.engine._engine.engine.scheduler`` (AsyncEngineCore → EngineCore → Scheduler). Some non-streaming engines may expose ``entry.engine.scheduler`` directly — including ``DFlashEngine`` in fallback mode, whose ``scheduler`` property resolves the fallback engine's real scheduler. DFlash's *primary* speculative path has no scheduler at all, so it exposes a lightweight ``_prefill_guard`` that carries the same watermark attrs; tried last so standard engines resolve unchanged. Returns None if nothing resolves. """ eng = entry.engine if eng is None: return None sched = getattr(eng, "scheduler", None) if sched is not None: return sched inner = getattr(eng, "_engine", None) if inner is not None: inner_engine = getattr(inner, "engine", None) if inner_engine is not None: sched = getattr(inner_engine, "scheduler", None) if sched is not None: return sched # DFlash primary mode bypasses the scheduler entirely; the enforcer # still needs to push the ceiling somewhere, so it lands on the # engine's ``_prefill_guard`` (None for every non-DFlash engine). return getattr(eng, "_prefill_guard", None) def _propagate_memory_limit(self) -> None: """Propagate ceiling-derived watermarks to all schedulers. Called on every enforcer tick so the dynamic ceiling reaches the schedulers as fast as the poll interval allows. """ breakdown = self._get_ceiling_breakdown() ceiling = breakdown["hard_limit"] metal_cap_raw = ( self._get_effective_metal_cap_bytes() if self._prefill_memory_guard else 0 ) abort_limit = self._get_abort_limit_bytes() hot_cache_reserved = ( self._hot_cache_reserved_bytes() if ceiling > 0 or abort_limit > 0 else 0 ) # Clamp to the reservation: the usage-side exclusion must never exceed # what the ceiling actually gave up, or a transient hot-cache overshoot # past max_bytes would net-weaken the guard exactly under pressure. hot_cache_used = ( min(self._hot_cache_used_bytes(), hot_cache_reserved) if hot_cache_reserved > 0 else 0 ) scheduler_ceiling = self._scheduler_limit_bytes( ceiling, reserved=hot_cache_reserved ) soft_limit = ( int(scheduler_ceiling * self._soft_threshold) if scheduler_ceiling > 0 else 0 ) hard_watermark = ( int(scheduler_ceiling * self._hard_threshold) if scheduler_ceiling > 0 else 0 ) scheduler_abort_limit = self._scheduler_limit_bytes( abort_limit, reserved=hot_cache_reserved ) admission_paused = self._pressure_level != "ok" for entry in self._engine_pool._entries.values(): scheduler = self._resolve_scheduler(entry) if scheduler is None: engine = getattr(entry, "engine", None) if engine is None: # Discovered-but-not-loaded entry. There is no # scheduler to propagate to yet and that is normal, # not a wrapper break, so skip silently. Warning here # would fire on a routine startup before any model is # loaded and turn the signal into noise. continue if getattr(engine, "_loaded", True) is False: # EnginePool clears entry.engine after stop(), but an # enforcer tick can observe the engine in the small # teardown window after its scheduler has been released. # Treat that like an unloaded entry, not a wrapper break. continue if getattr( engine, "_prefill_memory_guard_managed_externally", False ): # DistributedBatchedEngine is a coordinator-side proxy. # Its rank processes own the schedulers and receive their # guard budgets from the signed deployment, so there is no # coordinator scheduler for this enforcer to update. continue if ( type(engine).__name__ == "DFlashEngine" and getattr(engine, "_fallback_engine", None) is None ): continue if getattr(engine, "is_diffusion_model", False): continue if isinstance(engine, BaseNonStreamingEngine): # TTS/STT/STS/Embedding/Reranker engines run on the MLX # executor without a Scheduler, so an unresolvable # scheduler is their normal shape, not a wrapper break. # Warning here reads as a guard regression (#2312). engine.set_memory_soft_limit(soft_limit) continue # Silent no-op was the failure mode that originally hid # the dead memory guard: a wrapper-chain change made # ``_resolve_scheduler()`` return None on a loaded engine # and the loop kept iterating without complaining. Surface # it now — once per engine type per enforcer lifetime so # the regression is loud in CI / oncall but a misconfigured # engine polled every second doesn't spam. engine_type = type(engine).__name__ if engine_type not in self._scheduler_resolve_warned: self._scheduler_resolve_warned.add(engine_type) logger.warning( "ProcessMemoryEnforcer: could not resolve " "scheduler for engine type %s — prefill memory " "guard will not propagate to this engine. " "Verify the wrapper chain " "(engine._engine.engine.scheduler) still holds.", engine_type, ) continue scheduler._memory_limit_bytes = soft_limit scheduler._memory_hard_limit_bytes = scheduler_ceiling scheduler._memory_hard_watermark_bytes = hard_watermark scheduler._memory_abort_limit_bytes = scheduler_abort_limit scheduler._prefill_abort_margin = self._get_prefill_abort_margin() # Propagate the component ceilings too so the rejection # message in ``_preflight_memory_check`` can name the binding # constraint and steer the user toward the right remedy # (close apps for dynamic, raise sysctl for metal_cap, raise # tier or reduce context for static). scheduler._memory_static_ceiling_bytes = breakdown["static"] scheduler._memory_dynamic_ceiling_bytes = breakdown["dynamic"] scheduler._memory_metal_cap_bytes = breakdown["metal_cap"] scheduler._memory_metal_cap_raw_bytes = metal_cap_raw scheduler._memory_hot_cache_reserved_bytes = hot_cache_reserved # Usage-side counterpart of the reservation above: targets whose # usage read is raw phys_footprint (the DFlash primary guard) # subtract this so serialized hot-cache CPU bytes are not charged # both here and in the reserved ceiling. The Scheduler reads its # own live counter instead and ignores this attr. scheduler._memory_hot_cache_used_bytes = hot_cache_used # Tier name disambiguates dynamic = computed reclaimable # (safe/balanced/aggressive) from dynamic = user-pinned # custom_ceiling_bytes (custom). The advice ladder needs # the distinction to point at the right knob. scheduler._memory_guard_tier = self._memory_guard_tier scheduler._prefill_memory_guard = self._prefill_memory_guard scheduler._admission_paused = admission_paused scheduler._prefill_headroom_safety = self._prefill_headroom_safety scheduler._prefill_safe_zone_ratio = self._prefill_safe_zone_ratio scheduler._prefill_min_chunk_tokens = self._prefill_min_chunk_tokens bg = getattr(scheduler, "batch_generator", None) if bg is not None and hasattr(bg, "_memory_limit_bytes"): bg._memory_limit_bytes = soft_limit bg._memory_hard_limit_bytes = scheduler_ceiling def _walk_store_cache_caps(self) -> None: """Walk each scheduler's store-cache gate one step per poll (#1383). Driven on every enforcement tick, not just on pressure transitions, so the cap converges ±1 per poll toward its pressure-driven target (ok -> max_num_seqs, soft/hard -> 1). Decoupled from `_propagate_memory_limit` to avoid double-stepping the cap when a transition fires. """ for entry in self._engine_pool._entries.values(): scheduler = self._resolve_scheduler(entry) if scheduler is None: continue adjust = getattr(scheduler, "adjust_store_cache_cap", None) if adjust is not None: adjust(self._pressure_level) async def _abort_loaded_requests_for_memory_emergency(self) -> int: """Abort active requests on loaded models without unloading them.""" aborted_total = 0 for entry in self._engine_pool._entries.values(): engine = getattr(entry, "engine", None) if engine is None: continue abort_all = getattr(engine, "abort_all_requests", None) if not callable(abort_all): continue try: result = abort_all() if inspect.isawaitable(result): result = await result except Exception as exc: # noqa: BLE001 logger.warning( "Emergency memory abort failed for '%s': %s", getattr(entry, "model_id", ""), exc, ) continue if isinstance(result, (int, float)): aborted_total += max(0, int(result)) return aborted_total def _find_lru_busy_non_pinned_victim_locked(self) -> str | None: """Find a non-pinned loaded model that is busy but abortable. Caller must hold the engine-pool lock. This is used only at hard pressure, after idle victims have already been considered. """ candidates: list[tuple[float, str]] = [] for mid, entry in self._engine_pool._entries.items(): if ( getattr(entry, "engine", None) is None or getattr(entry, "is_pinned", False) or getattr(entry, "is_loading", False) ): continue busy = getattr(entry, "in_use", 0) > 0 if not busy: engine = getattr(entry, "engine", None) has_active = getattr(engine, "has_active_requests", None) if callable(has_active): try: busy = has_active() is True except Exception: # noqa: BLE001 busy = True if busy: candidates.append((getattr(entry, "last_access", 0.0), mid)) if not candidates: return None candidates.sort() return candidates[0][1] async def stop(self) -> None: """Stop the background enforcement loop.""" self._running = False self._pressure_reclaim_grace_polls = 0 if self._wake_event is not None: self._wake_event.set() if self._task: self._task.cancel() with suppress(asyncio.CancelledError): await self._task self._task = None self._wake_event = None self._loop = None logger.info("Process memory enforcer stopped") async def _enforcement_loop(self) -> None: """Main polling loop.""" self._loop = asyncio.get_running_loop() self._wake_event = asyncio.Event() while self._running: if self._wake_event is not None: self._wake_event.clear() try: await self._check_and_enforce() await self._check_ttl() except asyncio.CancelledError: break except Exception as e: logger.error(f"Process memory enforcer error: {e}") interval = self._select_poll_interval() self._current_poll_interval = interval if self._wake_event is None: await asyncio.sleep(interval) continue with suppress(TimeoutError): await asyncio.wait_for(self._wake_event.wait(), timeout=interval) def _select_poll_interval(self) -> float: """Choose the next polling interval from current engine activity.""" if self._pressure_level != "ok": return self._active_poll_interval if time.monotonic() < self._activity_hint_until: return self._active_poll_interval has_loaded = False for entry in self._engine_pool._entries.values(): if getattr(entry, "is_loading", False): return self._active_poll_interval engine = getattr(entry, "engine", None) if engine is None: continue has_loaded = True has_active_requests = getattr(engine, "has_active_requests", None) if not callable(has_active_requests): return self._active_poll_interval try: if has_active_requests(): return self._active_poll_interval except Exception: # If activity detection itself fails, bias toward safety. return self._active_poll_interval if has_loaded: return self._loaded_idle_poll_interval return self._unloaded_idle_poll_interval async def _check_ttl(self) -> None: """Check and unload models that exceeded their TTL.""" if self._settings_manager is None: return await self._engine_pool.check_ttl_expirations( self._settings_manager, global_idle_timeout_seconds=( self._global_settings.idle_timeout.idle_timeout_seconds if self._global_settings else None ), ) async def _check_and_enforce(self) -> None: """Check current memory and enforce 2-watermark policy. The ceiling is recomputed on every tick (dynamic ceiling moves with system_available), so watermarks shift as other apps take or release memory. Pressure levels: - ok (current < soft): no action, ensure admission unpaused. - soft (soft <= current < hard): LRU non-pinned eviction + signal schedulers to pause new admissions (in-flight requests proceed). - hard (current >= hard): LRU evict, abort a sole non-pinned victim's in-flight requests, abort in-progress model loads, and request idle reclaim when no victim exists. Pinned active requests are only aborted under emergency pressure beyond the real ceiling. """ # Always propagate so the scheduler sees the latest ceiling / # admission_paused, even when usage stays below the soft mark. self._propagate_memory_limit() ceiling = self._get_hard_limit_bytes() if ceiling <= 0: self._pressure_level = "ok" self._over_ceiling_polls = 0 self._pressure_reclaim_grace_polls = 0 return current = self._current_usage_bytes() soft = int(ceiling * self._soft_threshold) hard = int(ceiling * self._hard_threshold) # Reclaim decoded CPU images before pausing admission or evicting # models. Request-owned references may remain, so remeasure usage. if current >= soft: dropped_images = await asyncio.to_thread(clear_image_decode_cache) if dropped_images: current = self._current_usage_bytes() prev_level = self._pressure_level emergency = self._is_emergency_pressure(current, ceiling) if current < soft: new_level = "ok" elif current < hard: new_level = "soft" else: new_level = "hard" if new_level == "ok": # The pressure episode ended (drain worked or the load finished): # the next one gets a fresh reclaim-grace budget. Must happen # before the ok-level early return below. self._pressure_reclaim_grace_polls = 0 if new_level != prev_level: self._pressure_level = new_level self._propagate_memory_limit() logger.info( f"Memory pressure level: {prev_level} -> {new_level} " f"(current={_format_gb(current)}, " f"soft={_format_gb(soft)}, hard={_format_gb(hard)}, " f"ceiling={_format_gb(ceiling)})" ) # When pooled Metal buffers can be returned at the next inference # boundary, let that non-destructive reclaim run before evicting idle # models (soft), shrinking shared hot cache or aborting work (hard). A # failed / unavailable request deliberately falls through to the # established enforcement path. if ( new_level != "ok" and not emergency and os.environ.get("OMLX_DISABLE_PRESSURE_RECLAIM") != "1" and self._pressure_reclaim_grace_polls < self._PRESSURE_RECLAIM_GRACE_POLLS_MAX ): requested = self._request_scheduler_cache_reclaim(0) if requested: self._pressure_reclaim_grace_polls += 1 logger.info( "%s memory pressure: deferring enforcement (poll %d/%d) " "while pooled Metal buffers drain on %d scheduler(s)", new_level.capitalize(), self._pressure_reclaim_grace_polls, self._PRESSURE_RECLAIM_GRACE_POLLS_MAX, requested, ) # This gate is non-destructive and is documented to walk # once per enforcement tick, including while pressure actions # are deferred. self._walk_store_cache_caps() return if new_level == "hard": freed_hot = await asyncio.to_thread( self._shrink_hot_cache_for_pressure, current, soft, ) # Return reclaimable Metal memory to the OS. The shrink above only # drops mx.array references into MLX's buffer-cache pool, which is # pinned by set_cache_limit(total) (the #300 panic guard), so the # freed bytes never leave the process and phys_footprint does not # drop — the enforcer then wrongly concludes "no evictable models" # and livelocks until restart. Env gate # OMLX_DISABLE_PRESSURE_RECLAIM=1 restores stock behavior. if os.environ.get("OMLX_DISABLE_PRESSURE_RECLAIM") != "1": self._request_scheduler_cache_reclaim(freed_hot) if freed_hot > 0: current = self._current_usage_bytes() emergency = self._is_emergency_pressure(current, ceiling) if current < soft: recovered_level = "ok" elif current < hard: recovered_level = "soft" else: recovered_level = "hard" if recovered_level != new_level: logger.info( "Memory pressure after hot-cache shrink: %s -> %s " "(current=%s, freed_hot=%s)", new_level, recovered_level, _format_gb(current), _format_gb(freed_hot), ) new_level = recovered_level self._pressure_level = new_level self._propagate_memory_limit() if new_level == "ok": # Still walk the store-cache cap so it can recover toward # max_num_seqs while pressure stays low (#1383). self._walk_store_cache_caps() return # Recover below soft regardless of level — prevents oscillation # at the boundary. target = soft async with self._engine_pool._lock: while self._current_usage_bytes() > target: pending = self._engine_pool._find_pending_unload_ready_locked() if pending is not None: await self._engine_pool._unload_pending_if_idle_locked(pending) continue victim = self._engine_pool._find_lru_victim() if victim is not None: loaded_non_pinned = [ mid for mid, e in self._engine_pool._entries.items() if e.engine is not None and not e.is_pinned ] if new_level == "hard" or len(loaded_non_pinned) > 1: # Evict idle LRU victims cleanly. At hard pressure even # the last idle non-pinned model should be unloaded; it # has no active KV to preserve and keeping it resident # does not reduce pressure. entry = self._engine_pool._entries.get(victim) if ( entry and entry.engine is not None and hasattr(entry.engine, "abort_all_requests") ): result = await entry.engine.abort_all_requests() aborted = max(0, int(result or 0)) if aborted > 0: logger.warning( f"Aborted {aborted} requests on " f"'{victim}' before eviction" ) logger.warning( f"Evicting model '{victim}' (pressure={new_level})" ) await self._engine_pool._unload_engine(victim) continue # soft: leave in-flight alone — admission pause already # signaled, eviction can't help further without aborts. break # No non-pinned victim. Loaded models are pinned or no loaded # engines exist at all. if new_level == "hard": busy_victim = self._find_lru_busy_non_pinned_victim_locked() if busy_victim is not None: entry = self._engine_pool._entries.get(busy_victim) aborted = 0 if ( entry and entry.engine is not None and hasattr(entry.engine, "abort_all_requests") ): aborted = await entry.engine.abort_all_requests() if not emergency: logger.warning( "Hard memory pressure: aborted %d request(s) on " "'%s' and kept model loaded", aborted, busy_victim, ) break if entry is not None: self._engine_pool._mark_pending_unload_locked( busy_victim, "hard memory pressure", abort_requested=True, ) unloaded = ( await self._engine_pool._unload_pending_if_idle_locked( busy_victim ) ) if not unloaded: # Not drained yet: schedule the poller, mirroring # request_unload's own fallback, or the latch never clears. self._engine_pool._schedule_pending_unload_locked( busy_victim ) logger.warning( "Hard memory pressure: requested abort/unload for " "'%s' (aborted=%d)", busy_victim, aborted, ) break # Hard only: abort any in-progress model loads. aborted_any = False for entry in self._engine_pool._entries.values(): if entry.is_loading and not entry.abort_loading: logger.warning( f"Aborting in-progress load of " f"'{entry.model_id}' (hard memory pressure)" ) entry.abort_loading = True aborted_any = True if not aborted_any: has_loaded = any( e.engine is not None for e in self._engine_pool._entries.values() ) if has_loaded: if emergency: emergency_current = self._current_usage_bytes() else: emergency_current = 0 if emergency and emergency_current >= ceiling: aborted = await ( self._abort_loaded_requests_for_memory_emergency() ) if aborted > 0: logger.warning( "Emergency memory pressure: aborted " "%d in-flight request(s) " "(current=%s, ceiling=%s); models " "kept loaded.", aborted, _format_gb(emergency_current), _format_gb(ceiling), ) break # Nothing to evict (all pinned) and no load to # abort — but the resident footprint may still hold # reclaimable Metal transients from a finished turn. # Ask each loaded scheduler to trim them between # turns. This only sets a flag; the actual reclaim # runs on the inference thread when it is idle, so # we never touch Metal from the enforcer thread. requested = 0 for entry in self._engine_pool._entries.values(): sched = self._resolve_scheduler(entry) if sched is not None and hasattr( sched, "request_idle_reclaim" ): sched.request_idle_reclaim() requested += 1 logger.warning( "Hard memory pressure, no evictable models " "and no loads in progress: requested " "idle reclaim on %d scheduler(s).", requested, ) else: logger.warning("Hard memory pressure but no models loaded.") # soft + all pinned: nothing to do beyond admission pause. break # Re-evaluate level after eviction completes so admission state # reflects post-eviction reality on the next propagate. post_current = self._current_usage_bytes() # Recompute ceiling again — eviction may free phys, shifting the # dynamic ceiling. post_ceiling = self._get_hard_limit_bytes() post_soft = int(post_ceiling * self._soft_threshold) if post_ceiling > 0 else 0 post_hard = int(post_ceiling * self._hard_threshold) if post_ceiling > 0 else 0 if post_ceiling <= 0 or post_current < post_soft: post_level = "ok" elif post_current < post_hard: post_level = "soft" else: post_level = "hard" if post_ceiling <= 0 or post_current < post_ceiling: self._over_ceiling_polls = 0 if post_level == "ok": # Same reset as the pre-action check: eviction inside this tick # may already have ended the pressure episode. self._pressure_reclaim_grace_polls = 0 if post_level != self._pressure_level: self._pressure_level = post_level self._propagate_memory_limit() logger.info( f"Memory pressure post-eviction: {new_level} -> {post_level} " f"(current={_format_gb(post_current)})" ) # Walk each scheduler's store-cache gate ±1 toward its # pressure-driven target every poll (#1383). self._walk_store_cache_caps() def get_status(self) -> dict: """Get enforcer status for monitoring endpoints. Reports the same `max(active, phys_footprint)` value the enforcer uses internally so admin UI / /health utilization matches the watermark the enforcer is actually comparing against. """ breakdown = self._get_ceiling_breakdown() if self._running else {} ceiling = breakdown.get("hard_limit", 0) reserve = self._tier_reserve_bytes() if self._running else 0 static_ceiling = self._get_static_ceiling() if self._running else 0 dynamic_ceiling = self._get_dynamic_ceiling() if self._running else 0 current = self._current_usage_bytes() if self._running else 0 hot_reserved = self._hot_cache_reserved_bytes() if self._running else 0 scheduler_ceiling = ( self._scheduler_limit_bytes(ceiling, reserved=hot_reserved) if self._running else 0 ) scheduler_abort = ( self._scheduler_limit_bytes( self._get_abort_limit_bytes(), reserved=hot_reserved ) if self._running else 0 ) soft = int(ceiling * self._soft_threshold) if ceiling > 0 else 0 hard = int(ceiling * self._hard_threshold) if ceiling > 0 else 0 return { "enabled": self._running, "memory_guard_tier": self._memory_guard_tier, "memory_guard_custom_ceiling_bytes": self._memory_guard_custom_ceiling_bytes, "ceiling_bytes": ceiling, "ceiling_formatted": _format_gb(ceiling), "static_ceiling_bytes": static_ceiling, "static_ceiling_formatted": _format_gb(static_ceiling), "dynamic_ceiling_bytes": dynamic_ceiling, "dynamic_ceiling_formatted": _format_gb(dynamic_ceiling), "metal_cap_bytes": breakdown.get("metal_cap", 0), "metal_cap_raw_bytes": ( self._get_effective_metal_cap_bytes() if self._running else 0 ), "reserve_bytes": reserve, "reserve_formatted": _format_gb(reserve), "other_app_reclaim_ratio": _OTHER_APP_RECLAIM_RATIO.get( self._memory_guard_tier, 0.0 ), "hot_cache_reserved_bytes": hot_reserved, "hot_cache_reserved_formatted": _format_gb(hot_reserved), "scheduler_ceiling_bytes": scheduler_ceiling, "scheduler_ceiling_formatted": _format_gb(scheduler_ceiling), "scheduler_abort_limit_bytes": scheduler_abort, "scheduler_abort_limit_formatted": _format_gb(scheduler_abort), "soft_threshold": self._soft_threshold, "hard_threshold": self._hard_threshold, "soft_bytes": soft, "soft_formatted": _format_gb(soft), "hard_bytes": hard, "hard_formatted": _format_gb(hard), "current_bytes": current, "current_formatted": _format_gb(current), "pressure_level": self._pressure_level if self._running else "ok", "utilization": (current / ceiling if ceiling > 0 else 0.0), "poll_interval_seconds": ( self._current_poll_interval if self._running else 0.0 ), }