Repository navigation
fix(memory): reclaim MLX pool before hot-cache shrink - #2635
Conversation
|
Thanks for the clean fix. I reviewed it against current main and I like that the grace generalizes the existing abort deferral instead of adding a new mechanism, and every bail-out path (emergency, disabled reclaim, small pool, missing scheduler, failed request) still falls through to the old enforcement. One thing that settled the safety question for me: since set_cache_limit pins the pool, the old immediate shrink never actually lowered phys_footprint until the step-boundary clear ran anyway, so requesting the clear first is strictly better ordering rather than a loosening of the guard. The idle-engine case is covered too because has_work stays true while the clear is pending. Live hot-cache retention numbers can come from the reporter's setup in #2581. Merging. |
Summary
Root cause
ProcessMemoryEnforcercompares a process footprint that includes MLX's reclaimable buffer pool against its watermarks. At hard pressure it previously sized and performed hot-cache shrink before asking the scheduler to clear that much larger pool. When pooled buffers dominated the footprint, the computed hot-cache target repeatedly fell to zero even though a synchronized pool clear could recover the pressure without destroying reusable prompt state.The reclaim remains scheduler-owned: this change only sets the existing pressure-reclaim flag, and the actual Metal clear still runs at the synchronized inference-step boundary.
Validation
uv run pytest -q tests/test_process_memory_enforcer.py: 163 passeduv run pytest -q tests/test_process_memory_enforcer.py tests/test_scheduler.py tests/test_hot_cache.py -m 'not slow and not integration': 418 passedSIM117finding excludedgit diff --checkpassedThis validates ordering and bounded fallback in CPU tests. It does not yet claim live M5 physical-memory recovery, hot-cache retention, or concurrent multi-engine performance; those remain follow-up workload measurements.
Closes #2581