Skip to content

fix(memory): reclaim MLX pool before hot-cache shrink - #2635

Merged
jundot merged 1 commit into
jundot:mainfrom
PhilipJohnBasile:agent/fix-omlx-2581-memory-reclaim
Aug 15, 2026
Merged

jundot merged 1 commit into
jundot:mainfrom
PhilipJohnBasile:agent/fix-omlx-2581-memory-reclaim

Conversation

@PhilipJohnBasile

Copy link
Copy Markdown
Contributor

Summary

  • request a synchronized MLX buffer-pool reclaim before destructive hard-pressure enforcement
  • defer hot-cache shrink, eviction, and request abort for at most five polls, and only when a scheduler accepted the reclaim request
  • preserve immediate enforcement for emergency pressure, disabled reclaim, small or unreadable pools, missing schedulers, reclaim-request failures, and exhausted grace
  • keep the existing per-tick store-cache concurrency adjustment active during the grace

Root cause

ProcessMemoryEnforcer compares 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 passed
  • uv 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 passed
  • targeted Ruff check passed with the pre-existing unrelated SIM117 finding excluded
  • git diff --check passed
  • independent frozen-diff review: approved with no blocking findings

This 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

@PhilipJohnBasile
PhilipJohnBasile marked this pull request as ready for review August 13, 2026 18:03
@jundot

jundot commented Aug 15, 2026

Copy link
Copy Markdown
Owner

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

2 participants