Add item-count and TTL bounds to the compile-closure cache by Chapaman · Pull Request #140 · elixir-nx/emlx

Chapaman · 2026-09-02T22:32:58Z

closes #139

#137's :emlx_compile_closures table grew without a cap. This adds compile_cache_max_items and compile_cache_ttl (both :infinity by default) so a long-running serving process can opt in without changing today's unbounded behaviour.

…#139) elixir-nx#137's table grew without a cap. compile_cache_max_items and compile_cache_ttl both default to :infinity so serving can opt in without changing the unbounded default. Overflow is FIFO; TTL is from insert time. No byte cap — Erlang term size wouldn't track MLX native memory.

polvalente

Looking good, but I think we can improve performance!

polvalente · 2026-09-02T23:23:09Z

+  Evicting a closure does **not** free the compiled MLX program held in
+  `:emlx_native_dispatch_cache`.


I think it should indeed evict it

Couldn't drop the native dispatch row here — it's shared across compile-closure keys. Any suggestions?

polvalente · 2026-09-02T23:26:03Z

+  @doc false
+  def fetch(table, key, opts) do
+    ttl = bound!(opts, :ttl)
+    # Validate both knobs on fetch and put so a bad config fails at first use.
+    _max_items = bound!(opts, :max_items)
+
+    case :ets.lookup(table, key) do
+      [{^key, fun, inserted_ms}] ->
+        if expired?(inserted_ms, ttl) do
+          :ets.delete(table, key)
+          :miss
+        else
+          {:ok, fun}
+        end
+
+      [] ->
+        :miss
+    end
+  end


I think this module should become a process, such that the TTL expiration logic lives inside a loop of Process.send_after(self(), :expire_entries, div(ttl, 2)) + its corresponding handle_info

Cache reads would only do the lookup with {:ok, result} vs {:error, :cache_miss}

CompileCache is now a GenServer under EMLX.Application. Reads are a plain ETS lookup ({:ok, fun} / {:error, :cache_miss}). TTL lives in handle_info(:expire_entries) on send_after(div(ttl, 2)), skipped when ttl is :infinity at start (a running sweeper re-reads app env each tick).

polvalente · 2026-09-02T23:31:49Z

+  defp evict_overflow(table, max_items) do
+    overflow = :ets.info(table, :size) - max_items
+
+    if overflow > 0 do
+      oldest_keys =
+        :ets.foldl(fn {key, _fun, ms}, acc -> [{ms, key} | acc] end, [], table)
+        |> Enum.sort_by(&elem(&1, 0))
+        |> Enum.take(overflow)
+        |> Enum.map(&elem(&1, 1))
+
+      Enum.each(oldest_keys, &:ets.delete(table, &1))
+    else
+      :ok
+    end


I think :ets.select_count could be used to find the amount of items we need to evict. Also if we keep a counter that we insert in the table, we can easily select for the items that should be deleted by matching on index <= current_counter - max_items to prune

Switched to an insert index via :ets.update_counter and :ets.select_delete of index <= current - max_items.

polvalente · 2026-09-02T23:35:34Z

+  other compilers, and it works identically here. EMLX additionally keeps a
+  persistent, structural (shape/op-based, not object-identity) dispatch-key
+  cache across calls (see `dispatch_key/3` in the source), which is what makes
+  `Nx.Defn.Graph.run/3`'s per-call re-tracing and structurally-identical-but-
+  distinct call sites (e.g. many copies of the same layer in a model) cheap
+  too — but a caller-held `Nx.Defn.compile/3` closure is always cheaper still.
+
+  ### Compile-closure cache bounds
+
+  `:emlx_compile_closures` is unbounded by default. Long-running serving can
+  cap it at runtime (not `compile_env` — no recompile needed):
+
+      config :emlx, compile_cache_max_items: 1024
+      config :emlx, compile_cache_ttl: :timer.minutes(30)
+
+  Both accept `:infinity` (the default) or a positive integer. TTL is in
+  milliseconds, counted from insert time — a hit does not refresh it.
+  Overflow evicts the oldest inserts (FIFO). Two concurrent misses on the
+  same key both compile; the later insert wins.
+
+  Evicting a closure does **not** free the compiled MLX program in
+  `:emlx_native_dispatch_cache`. See `EMLX.CompileCache`.


This docstring seems a bit too verbose, especially since it repeats a lot of what's described in the CompileCache moduledoc

Dropped the bounds subsection from the EMLX moduledoc; it now points at EMLX.CompileCache.

Co-authored-by: Paulo Valente <16843419+polvalente@users.noreply.github.com>

TTL now lives in handle_info(:expire_entries) so fetches are a plain ETS lookup. Overflow uses an insert counter and select_delete instead of folding the table. The dispatch-program table stays unbounded — its keys are structural and shared across compile-closure entries.

polvalente

Almost there! I think CompileCache should lean more into it's GenServer nature

polvalente · 2026-09-03T00:57:36Z

  @doc false
  def init do
-    case :ets.whereis(@compile_closure_table) do
-      :undefined ->
-        :ets.new(@compile_closure_table, [
-          :named_table,
-          :public,
-          :set,
-          read_concurrency: true,
-          write_concurrency: true
-        ])
-
-      _ ->
-        :ok
-    end
+    EMLX.CompileCache.ensure_table()
  end


We don't need to ensure_table here, we can just trust the default CompileCache init.

In fact I think we don't even need this init function

polvalente · 2026-09-03T00:58:31Z

    eval_fn =
-      case :ets.lookup(@compile_closure_table, cache_key) do
-        [{^cache_key, cached}] ->
+      case EMLX.CompileCache.fetch(EMLX.CompileCache.table(), cache_key) do


fetch shouldn't be table-aware, but process aware, defaulting to the global named process

polvalente · 2026-09-03T00:59:07Z

+          EMLX.CompileCache.put(
+            EMLX.CompileCache.table(),
+            cache_key,
+            built,
+            EMLX.CompileCache.opts()
+          )


Same for here, we should just be passing cache_key and built

polvalente · 2026-09-03T01:00:19Z

+EMLX caches jit/compile closures in an ETS table so repeated calls of the same
+function and shapes skip retrace. The table is unbounded by default. For a
+long-running serving process you can cap it:
+
+```elixir
+config :emlx, compile_cache_max_items: 1024          # or :infinity
+config :emlx, compile_cache_ttl: :timer.minutes(30)  # milliseconds, or :infinity
+```
+
+Overflow drops the oldest inserts. TTL is counted from insert time, not last
+access. See `EMLX.CompileCache`.
+


Let's remove this from here as it's too much detail for the readme. We can have the moduledoc for EMLX reference EMLX.CompileCache, and the moduledoc from it has the nitty-gritty details

polvalente · 2026-09-03T01:00:59Z

+  defp env_key(:max_items), do: :compile_cache_max_items
+  defp env_key(:ttl), do: :compile_cache_ttl


how about we use config :emlx, EMLX.CompileCache, [....] as the config path?

polvalente · 2026-09-03T01:04:53Z

+        :ets.new(table, [
+          :named_table,


The table should not be named, but saved in the process state, and all interactions with it pass through the genserver interface. This serializes access to it, but the process is very lightweight so this isn't really a penalty

polvalente reviewed Sep 2, 2026

View reviewed changes

Chapaman and others added 2 commits September 2, 2026 20:43

Apply batched suggestions from code review

187dd74

Co-authored-by: Paulo Valente <16843419+polvalente@users.noreply.github.com>

polvalente reviewed Sep 3, 2026

View reviewed changes

		Evicting a closure does not free the compiled MLX program held in
		`:emlx_native_dispatch_cache`.

		defp env_key(:max_items), do: :compile_cache_max_items
		defp env_key(:ttl), do: :compile_cache_ttl

Uh oh!

Conversation

Chapaman commented Sep 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

polvalente left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Chapaman Sep 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

polvalente left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Chapaman commented Sep 2, 2026 •

edited

Loading

Chapaman Sep 2, 2026 •

edited

Loading