Bound the MLX buffer cache instead of letting it track the machine

src/ezmsg/sigproc/asarray.py

@@ -80,17 +80,93 @@ def _detect_backend(data) -> str:
     raise TypeError(f"Unrecognized array type: {type(data)} (module={module})")
 
 
+_MLX_CACHE_LIMIT_APPLIED: float | None = None
+"""Value this process has already handed to ``mx.set_cache_limit``.
+
+Module-level because the limit is a property of the process, not of a node: two
+``AsArray`` nodes converting to MLX in one process share one allocator.
+"""
+
+
+def _apply_mlx_cache_limit(limit_mb: float) -> None:
+    """Bound the MLX buffer cache for this process. Idempotent.
+
+    MLX caches every freed buffer in a multimap keyed by *exact* byte size, and
+    only reuses one within ``min(2 * size, size + 2 * page_size)`` of the
+    request -- above ~32 KiB that is effectively an exact match. A graph whose
+    message length varies therefore mints a permanent new size class per length,
+    in a cache whose default limit is the whole machine (23 GiB on a 24 GiB
+    host). Measured on a 30 kHz feature chain that sees 40 distinct message
+    lengths over an hour: 6015 MiB of physical footprint unbounded, 966 MiB at
+    512 MiB, with *higher* throughput at the limit because there is less memory
+    pressure.
+
+    Eviction is LRU, which is what makes a limit the right tool rather than a
+    blunt one: the steady-state shape is re-touched every message and stays at
+    the head, while one-off shapes from a stall fall to the tail and are freed
+    first. Measured, the hot allocation is unaffected (0.97x) after 39 rare
+    shapes have been evicted past a 128 MiB limit.
+    """
+    global _MLX_CACHE_LIMIT_APPLIED
+    if _MLX_CACHE_LIMIT_APPLIED == limit_mb:
+        return
+    import mlx.core as mx
+
+    if _MLX_CACHE_LIMIT_APPLIED is not None:
+        ez.logger.warning(
+            f"MLX cache limit already set to {_MLX_CACHE_LIMIT_APPLIED} MiB in this process; "
+            f"overriding with {limit_mb} MiB. The limit is process-global, so the last AsArray "
+            "node to convert wins -- give every MLX-targeting AsArray in a process the same "
+            "mlx_cache_limit_mb."
+        )
+    mx.set_cache_limit(int(limit_mb * 1024 * 1024))
+    _MLX_CACHE_LIMIT_APPLIED = limit_mb
+
+
 class AsArraySettings(ez.Settings):
     backend: ArrayBackend = ArrayBackend.numpy
     """Target array backend."""
 
     dtype: str | None = None
     """Target dtype as a string (e.g. "float32", "float64"). None keeps the original dtype."""
 
+    mlx_cache_limit_mb: float | None = 512.0
+    """Cap the MLX buffer cache (MiB) for the process that runs this node.
+
+    Applied only when :attr:`backend` is MLX, once per process, on the first
+    message. ``None`` leaves MLX's default, which is the size of the machine.
+
+    Sizing: one *distinct message shape* costs roughly **50x the message
+    payload** in cached buffers -- about 20 intermediates across a typical chain,
+    each keeping its own size class. Measured cache for a steady-state chain,
+    against ``samples x channels x 4`` bytes per message: 46x at 256 ch x 1200
+    samples, 67x at 256 ch x 300. So::
+
+        limit_MiB ~= 50 * message_MiB * (distinct shapes to keep hot)
+
+    A 256-channel, 300-sample float32 message is 0.29 MiB, so ~15 MiB per shape
+    and the 512 MiB default holds ~26 distinct shapes. Steady state needs only
+    one; the rest is headroom for the varying-length messages a stall produces.
+    Raise it if the graph legitimately cycles through many shapes; the floor is
+    one working set (~1x the ``50 *`` term), and 0 disables caching entirely at
+    a measured 40% throughput cost.
+
+    This is a process-global MLX setting, so it is shared with anything else
+    using MLX in the same process. The default suits streaming graphs; large
+    offline batch work in the same process may want it raised or set to
+    ``None``."""
+
 
 class AsArrayTransformer(BaseTransformer[AsArraySettings, AxisArray, AxisArray]):
     def _process(self, message: AxisArray) -> AxisArray:
         target_backend = str(self.settings.backend)
+        if target_backend == "mlx" and self.settings.mlx_cache_limit_mb is not None:
+            # Here rather than in the Unit's initialize(): this transformer is
+            # also used bare (offline chains, benchmarks), and the limit has to
+            # be set in whichever process actually converts, which for a
+            # multi-process graph is not the one that built it -- set_cache_limit
+            # does not survive the spawn.
+            _apply_mlx_cache_limit(self.settings.mlx_cache_limit_mb)
         dtype_str = self.settings.dtype
         data = message.data
 

src/ezmsg/sigproc/util/mlx_metal_common.py

@@ -72,22 +72,48 @@ def chunked_scan(x_flat, n_samples, chunk_sizes, state, launch_fn):
     ``launch_fn``. Kernels must emit their state at ``valid_length - 1`` rather
     than after the padding.
 
+    Chunk sizes bound the set of Metal kernel *specializations*; they do not by
+    themselves bound the set of *buffer sizes*. MLX caches freed buffers in a
+    multimap keyed by exact byte size and only reuses one within
+    ``min(2 * size, size + 2 * page_size)`` of the request -- effectively an
+    exact match above ~32 KiB -- so every distinct intermediate length becomes a
+    permanent new size class in a cache whose default limit is the whole
+    machine. Concatenating the *padded* chunks and trimming once therefore
+    allocates on the multiple-of-chunk_size grid rather than at ``n_samples``,
+    which measured 31% less cached memory over 40 distinct input lengths.
+
+    Trimming once at the end is only correct because padding can occur on the
+    final chunk alone: every earlier iteration has ``remaining > chunk_size``,
+    so ``valid == chunk_size`` and the chunk is emitted whole. Were an interior
+    chunk padded, its padding would sit between two runs of valid samples and
+    the single trim would return garbage, so the loop asserts that invariant
+    rather than trusting whoever next edits the size-selection rule.
+
     Returns ``(y_combined, final_state)``.
     """
     y_chunks = []
     start = 0
     max_chunk_size = chunk_sizes[-1]
+    padded_len = 0
     while start < n_samples:
         remaining = n_samples - start
         chunk_size = next((size for size in chunk_sizes if size >= remaining), max_chunk_size)
         valid = min(remaining, chunk_size)
         end = start + valid
         x_chunk = x_flat[:, start:end]
         if valid < chunk_size:
+            if end != n_samples:
+                raise AssertionError(
+                    f"chunked_scan padded an interior chunk (valid={valid}, chunk_size={chunk_size}, "
+                    f"end={end}, n_samples={n_samples}); the single trim below would return padding as data."
+                )
             x_chunk = mx.pad(x_chunk, [(0, 0), (0, chunk_size - valid)])
         valid_length = mx.array([valid], dtype=mx.uint32)
         y_chunk, state = launch_fn(x_chunk, state, chunk_size, valid_length)
-        y_chunks.append(y_chunk[:, :valid])
+        y_chunks.append(y_chunk)
+        padded_len += chunk_size
         start = end
-    y_combined = y_chunks[0] if len(y_chunks) == 1 else mx.concatenate(y_chunks, axis=-1)
+    y_padded = y_chunks[0] if len(y_chunks) == 1 else mx.concatenate(y_chunks, axis=-1)
+    # Exact length is part of the contract: callers reshape to n_samples.
+    y_combined = y_padded if padded_len == n_samples else y_padded[:, :n_samples]
     return y_combined, state

tests/unit/test_asarray.py

@@ -13,6 +13,7 @@
     _get_backend_module,
 )
 from tests.helpers.empty_time import check_empty_result, make_empty_msg, make_msg
+from tests.helpers.util import requires_mlx
 
 # -- Helpers ------------------------------------------------------------------
 
@@ -187,3 +188,93 @@ def test_empty_time_cross_backend(backend):
     # The time dimension should still be 0.
     time_idx = result.dims.index("time")
     assert result.data.shape[time_idx] == 0
+
+
+# -- MLX buffer cache limit ----------------------------------------------------
+
+
+def test_mlx_cache_limit_default_is_set():
+    """A default, not None: the growth it prevents is invisible in RSS.
+
+    MLX buffers are IOKit allocations, so a graph whose message length varies
+    can put gigabytes into the allocator's per-size cache while RSS stays flat.
+    Users do not go looking for a knob they have no symptom for.
+    """
+    assert AsArraySettings().mlx_cache_limit_mb == 512.0
+
+
+def test_mlx_cache_limit_not_applied_for_numpy_target(monkeypatch):
+    """Converting TO numpy must not touch a process-global MLX setting."""
+    import ezmsg.sigproc.asarray as asarray_module
+
+    calls = []
+    monkeypatch.setattr(asarray_module, "_apply_mlx_cache_limit", lambda mb: calls.append(mb))
+    proc = AsArrayTransformer(AsArraySettings(backend=ArrayBackend.numpy, mlx_cache_limit_mb=128.0))
+    proc(make_msg())
+    assert calls == []
+
+
+@requires_mlx
+def test_mlx_cache_limit_applied_once_for_mlx_target(monkeypatch):
+    import ezmsg.sigproc.asarray as asarray_module
+
+    calls = []
+    monkeypatch.setattr(asarray_module, "_apply_mlx_cache_limit", lambda mb: calls.append(mb))
+    proc = AsArrayTransformer(AsArraySettings(backend=ArrayBackend.mlx, mlx_cache_limit_mb=128.0))
+    for _ in range(3):
+        proc(make_msg())
+    # Called per message, but the applier itself is what dedupes -- see below.
+    assert calls == [128.0, 128.0, 128.0]
+
+
+@requires_mlx
+def test_apply_mlx_cache_limit_is_idempotent_and_warns_on_conflict(monkeypatch):
+    """The limit is process-global, so a second, different value is a conflict.
+
+    Silently letting the last node win would make the effective limit depend on
+    which unit happened to convert first.
+    """
+    import mlx.core as mx
+
+    import ezmsg.sigproc.asarray as asarray_module
+
+    sets = []
+    monkeypatch.setattr(mx, "set_cache_limit", lambda n: sets.append(n))
+    monkeypatch.setattr(asarray_module, "_MLX_CACHE_LIMIT_APPLIED", None)
+    warnings = []
+    monkeypatch.setattr(asarray_module.ez.logger, "warning", lambda msg, *a: warnings.append(msg))
+
+    asarray_module._apply_mlx_cache_limit(128.0)
+    asarray_module._apply_mlx_cache_limit(128.0)  # same value: no-op
+    assert sets == [128 * 1024 * 1024]
+    assert warnings == []
+
+    asarray_module._apply_mlx_cache_limit(256.0)  # different value: warn, override
+    assert sets == [128 * 1024 * 1024, 256 * 1024 * 1024]
+    assert len(warnings) == 1 and "process-global" in warnings[0]
+
+
+@requires_mlx
+def test_mlx_cache_limit_actually_bounds_the_cache():
+    """End to end: the setting reaches the allocator and caps it."""
+    import mlx.core as mx
+
+    import ezmsg.sigproc.asarray as asarray_module
+
+    previous = mx.set_cache_limit(2**40)
+    applied = asarray_module._MLX_CACHE_LIMIT_APPLIED
+    try:
+        asarray_module._MLX_CACHE_LIMIT_APPLIED = None
+        proc = AsArrayTransformer(AsArraySettings(backend=ArrayBackend.mlx, mlx_cache_limit_mb=32.0))
+        proc(make_msg())
+        mx.clear_cache()
+        # Churn many distinct sizes; without a limit this cache grows unbounded.
+        for n in range(1, 60):
+            a = mx.zeros((256, 300 * n))
+            mx.eval(a)
+            del a
+        assert mx.get_cache_memory() <= 32 * 1024 * 1024
+    finally:
+        mx.clear_cache()
+        mx.set_cache_limit(previous)
+        asarray_module._MLX_CACHE_LIMIT_APPLIED = applied

tests/unit/test_mlx_metal_common.py

@@ -0,0 +1,115 @@
+import numpy as np
+
+from tests.helpers.util import requires_mlx
+
+
+@requires_mlx
+def test_chunked_scan_trims_to_exact_length():
+    """The padded-concat optimization must not change the output contract.
+
+    ``chunked_scan`` concatenates the kernels' *padded* chunks so the allocation
+    lands on a multiple-of-chunk_size grid rather than at ``n_samples`` (fewer
+    MLX buffer size classes). The single trim at the end is what keeps the
+    result exact, and callers reshape to ``n_samples`` immediately afterwards --
+    so an off-by-one here surfaces as a reshape error, not as wrong data.
+    """
+    import mlx.core as mx
+
+    from ezmsg.sigproc.util.mlx_metal_common import chunked_scan
+
+    def launch(x_chunk, state, cs, valid_length):
+        # A stand-in kernel with the real ones' contract: output is the FULL
+        # chunk width, state carries forward.
+        assert x_chunk.shape[-1] == cs
+        return x_chunk * 2.0, state + 1
+
+    n_channels = 3
+    for chunk_sizes in ((32,), (32, 128), (16, 64, 256)):
+        for n_samples in (1, 15, 16, 17, 31, 32, 33, 100, 128, 129, 257, 1000):
+            x = mx.array(np.arange(n_channels * n_samples, dtype=np.float32).reshape(n_channels, n_samples))
+            y, state = chunked_scan(x, n_samples, chunk_sizes, mx.array(0), launch)
+            assert y.shape == (n_channels, n_samples), f"chunk_sizes={chunk_sizes} n_samples={n_samples} gave {y.shape}"
+            # Padding must never leak into the result.
+            np.testing.assert_allclose(np.asarray(y), np.asarray(x) * 2.0)
+
+
+@requires_mlx
+def test_chunked_scan_pads_only_the_final_chunk():
+    """Trimming once is sound only while padding is confined to the last chunk.
+
+    Padding an interior chunk would put padding *between* two runs of valid
+    samples and the single trim would return it as data. The invariant holds
+    structurally -- padding needs ``remaining < chunk_size``, which makes that
+    chunk the tail -- so this pins the property across size sets and lengths
+    rather than trying to construct a violation (there isn't one to construct;
+    the assertion in ``chunked_scan`` guards future edits to the selection rule).
+    """
+    import mlx.core as mx
+
+    from ezmsg.sigproc.util.mlx_metal_common import chunked_scan
+
+    for chunk_sizes in ((32,), (32, 128), (16, 64, 256), (128, 32)):
+        for n_samples in (1, 15, 31, 32, 33, 100, 128, 129, 257, 1000):
+            seen = []
+
+            def launch(x_chunk, state, cs, valid_length, _seen=seen):
+                _seen.append((int(np.asarray(valid_length)[0]), cs))
+                return x_chunk, state
+
+            chunked_scan(mx.zeros((2, n_samples)), n_samples, chunk_sizes, mx.array(0), launch)
+            padded = [i for i, (valid, cs) in enumerate(seen) if valid < cs]
+            assert padded in ([], [len(seen) - 1]), (
+                f"chunk_sizes={chunk_sizes} n_samples={n_samples}: padded chunks at {padded} of {len(seen)}"
+            )
+
+
+@requires_mlx
+def test_chunked_scan_allocates_fewer_size_classes_than_per_chunk_trim():
+    """The point of the change, measured against the implementation it replaced.
+
+    Compares like with like: identical inputs, held live so their own size
+    classes are not what is being counted, so the delta is the scan's
+    intermediates alone. A MiB figure would be machine-dependent; the ratio is
+    the property worth pinning.
+    """
+    import mlx.core as mx
+
+    from ezmsg.sigproc.util.mlx_metal_common import chunked_scan
+
+    chunk_sizes = (32, 1024)
+
+    def launch(x_chunk, state, cs, valid_length):
+        return x_chunk * 2.0, state
+
+    def per_chunk_trim(x_flat, n_samples, sizes, state, launch_fn):
+        """The pre-change behavior: trim every chunk, concatenate at n_samples."""
+        y_chunks, start = [], 0
+        while start < n_samples:
+            remaining = n_samples - start
+            cs = next((s for s in sizes if s >= remaining), sizes[-1])
+            valid = min(remaining, cs)
+            x_chunk = x_flat[:, start : start + valid]
+            if valid < cs:
+                x_chunk = mx.pad(x_chunk, [(0, 0), (0, cs - valid)])
+            y, state = launch_fn(x_chunk, state, cs, mx.array([valid], dtype=mx.uint32))
+            y_chunks.append(y[:, :valid])
+            start += valid
+        y = y_chunks[0] if len(y_chunks) == 1 else mx.concatenate(y_chunks, axis=-1)
+        return y, state
+
+    lengths = [300 * k for k in range(1, 21)] * 3
+    inputs = [mx.zeros((256, n)) for n in lengths]  # held live throughout
+    mx.eval(*inputs)
+
+    def cache_for(scan):
+        mx.clear_cache()
+        for x, n in zip(inputs, lengths):
+            y, _ = scan(x, n, chunk_sizes, mx.array(0), launch)
+            mx.eval(y)
+        used = mx.get_cache_memory()
+        mx.clear_cache()
+        return used
+
+    before = cache_for(per_chunk_trim)
+    after = cache_for(chunked_scan)
+    assert after < before, f"padded concat cached more, not less: {after} vs {before} bytes"