feat(s032): exla loaded on demand in the release, the smoke probe after the tree, EXLA scoring in the brute store

Package run 35600216451: on Linux the EXLA NIF does not load in Burrito's musl ERTS
(__libc_single_threaded: symbol not found) and exla's application start took the release
down; the release now loads exla without starting it and Bumblebee.exla/0 starts it on
demand, remembering a failure as the tier's reason. On macOS the NIF loaded (EXLA=ok).
The vector check ran before the Repo on every OS; it is now a probe child before the
endpoint. Brute.search scores 10,000 rows as one EXLA product (72 ms p50, was 531 ms in
Elixir); the Elixir path stays for machines without the backend, held to the same numbers
by a test.

Signed-off-by: Ayla Croft <aylacroft@proton.me>

Author: HackTuah

lib/trinity/application.ex

@@ -54,7 +54,10 @@ defmodule Trinity.Application do
           Trinity.Tools.Supervisor,
           # Slice 021: approval requests and their decisions, with pending rows reloaded.
           Trinity.Permissions.Gate,
-          Trinity.Sessions.Supervisor,
+          Trinity.Sessions.Supervisor
+        ] ++
+        Trinity.Smoke.probe(Trinity.Smoke.argv()) ++
+        [
           # Start to serve requests, typically the last entry
           TrinityWeb.Endpoint
         ] ++ Trinity.Smoke.children(Trinity.Smoke.argv())

lib/trinity/memory/embedders/bumblebee.ex

@@ -48,11 +48,41 @@ defmodule Trinity.Memory.Embedders.Bumblebee do
     cond do
       match?({:win32, _}, :os.type()) -> {:off, :no_local_backend}
       not Code.ensure_loaded?(EXLA) -> {:off, {:exla, :not_compiled}}
+      match?({:error, _}, exla()) -> {:off, exla()}
       not model_present?() -> {:off, :model_missing}
       true -> :ok
     end
   end
 
+  @doc """
+  Starts the exla application if it is loaded and not running, once: the release loads it
+  without starting it (mix.exs), because its start loads the NIF and that load fails in
+  Burrito's Linux ERTS. `:ok`, or `{:error, {:exla, reason}}` remembered for the run (a NIF
+  that did not load will not load on the next call either).
+  """
+  @spec exla() :: :ok | {:error, {:exla, term()}}
+  def exla do
+    case :persistent_term.get({__MODULE__, :exla}, nil) do
+      nil ->
+        result =
+          case Application.ensure_all_started(:exla) do
+            {:ok, _} -> :ok
+            {:error, reason} -> {:error, {:exla, summarise(reason)}}
+          end
+
+        :persistent_term.put({__MODULE__, :exla}, result)
+        result
+
+      known ->
+        known
+    end
+  end
+
+  # The loader's message, not the whole start_link failure tree.
+  defp summarise({:exla, {reason, _}}), do: summarise(reason)
+  defp summarise(reason) when is_binary(reason), do: String.slice(reason, 0, 200)
+  defp summarise(reason), do: reason |> inspect() |> String.slice(0, 200)
+
   @doc "True when the weights and tokenizer are in the cache (checked offline, no request made)."
   @spec model_present?() :: boolean()
   def model_present? do

lib/trinity/memory/vector_stores/brute.ex

@@ -4,8 +4,12 @@ defmodule Trinity.Memory.VectorStores.Brute do
   @moduledoc """
   The SQLite vector store (slice 032): vectors as float32 bytes on `memories.embedding`,
   a search loads the filter's rows (a persona's semantic tier, the named scopes, the
-  embedder in force's model) and scores them in Elixir by cosine. Measured at G3 for 10,000
-  rows; docs/02 puts brute force at fine to about 10^5, and hnswlib is the step after.
+  embedder in force's model) and scores them by cosine: as one matrix product on EXLA when
+  the NIF is loaded (the local embedder's own backend), in Elixir over decoded lists when
+  it is not (the hosted embedder on a machine without EXLA). Measured at G3 over 10,000
+  rows (docs/perf.md): loading the rows 93 ms, the Elixir cosine 571 ms, the EXLA product
+  6 ms after its first compile for that row count (102 ms). docs/02's "fine to about 10^5"
+  holds for the product, not for the Elixir path; hnswlib is the step after.
   """
   @behaviour Trinity.Memory.VectorStore
 
@@ -37,13 +41,55 @@ defmodule Trinity.Memory.VectorStores.Brute do
           e.embedding_model == ^model and not is_nil(e.embedding)
     )
     |> Repo.all()
-    |> Enum.map(fn e ->
-      %{id: e.id, score: Embedder.cosine(query, Embedder.from_binary(e.embedding)), entry: e}
-    end)
+    |> score(query)
     |> Enum.sort_by(& &1.score, :desc)
     |> Enum.take(k)
   end
 
+  @doc "Cosine of the query against every row's vector, as hits; the EXLA product when the NIF is loaded, Elixir otherwise."
+  @spec score([Entry.t()], Embedder.vector()) :: [Trinity.Memory.VectorStore.hit()]
+  def score([], _query), do: []
+
+  def score(rows, query) do
+    scores = if exla?(), do: exla_scores(rows, query), else: elixir_scores(rows, query)
+    Enum.zip_with(rows, scores, fn e, s -> %{id: e.id, score: s, entry: e} end)
+  end
+
+  @doc "Cosines in Elixir over decoded lists (the path without EXLA); public so the suite can hold the two paths to the same numbers."
+  @spec elixir_scores([Entry.t()], Embedder.vector()) :: [float()]
+  def elixir_scores(rows, query),
+    do: Enum.map(rows, &Embedder.cosine(query, Embedder.from_binary(&1.embedding)))
+
+  @doc "Cosines as one (rows x dim) . dim product on the EXLA host client; a row of another width scores 0.0, as `cosine/2` refuses it."
+  @spec exla_scores([Entry.t()], Embedder.vector()) :: [float()]
+  def exla_scores(rows, query) do
+    dim = length(query)
+    fit = Enum.filter(rows, &(byte_size(&1.embedding) == dim * 4))
+    scores = if fit == [], do: %{}, else: product(fit, query, dim)
+    Enum.map(rows, &Map.get(scores, &1.id, 0.0))
+  end
+
+  defp product(fit, query, dim) do
+    m =
+      fit
+      |> Enum.map(& &1.embedding)
+      |> IO.iodata_to_binary()
+      |> Nx.from_binary(:f32, backend: EXLA.Backend)
+      |> Nx.reshape({length(fit), dim})
+
+    q = Nx.tensor(query, type: :f32, backend: EXLA.Backend)
+    norms = Nx.multiply(Nx.LinAlg.norm(m, axes: [1]), Nx.LinAlg.norm(q))
+
+    m
+    |> Nx.dot(q)
+    |> Nx.divide(Nx.select(Nx.equal(norms, 0.0), 1.0, norms))
+    |> Nx.to_flat_list()
+    |> then(&Map.new(Enum.zip(Enum.map(fit, fn e -> e.id end), &1)))
+  end
+
+  defp exla?,
+    do: Code.ensure_loaded?(EXLA.Backend) and Trinity.Memory.Embedders.Bumblebee.exla() == :ok
+
   @impl true
   def delete(id) do
     Repo.update_all(from(e in Entry, where: e.id == ^id),

lib/trinity/smoke.ex

@@ -82,7 +82,7 @@ defmodule Trinity.Smoke do
   """
   @spec exla_line() :: String.t()
   def exla_line do
-    if Code.ensure_loaded?(EXLA) do
+    if Code.ensure_loaded?(EXLA) and Trinity.Memory.Embedders.Bumblebee.exla() == :ok do
       try do
         t = Nx.tensor([1.0, 2.0], backend: EXLA.Backend)
         [3.0] = Nx.to_flat_list(Nx.sum(t))
@@ -94,7 +94,7 @@ defmodule Trinity.Smoke do
           "TRINITY_SMOKE_EXLA=failed:#{inspect({kind, reason}) |> String.slice(0, 200)}"
       end
     else
-      "TRINITY_SMOKE_EXLA=failed:not_compiled"
+      "TRINITY_SMOKE_EXLA=failed:#{inspect(Trinity.Memory.Embedders.Bumblebee.exla()) |> String.slice(0, 200)}"
     end
   end
 
@@ -168,13 +168,46 @@ defmodule Trinity.Smoke do
   def children(args) do
     if requested?(args) do
       markdown = markdown_line()
-      # The 032 lines need the Repo and the memory supervisor, which have started by now;
-      # they are cheap (one tiny tensor, three rows rolled back) and computed here for the
-      # same race the markdown line lost once.
-      rest = [exla_line(), vec_line(), semantic_line()]
-      [{Task, fn -> run(&IO.puts/1, &System.halt/1, markdown, rest) end}]
+      # The 032 lines were computed by `probe/1`, a child placed after the Repo, the memory
+      # supervisor and the sessions (package run 35600216451: computed here, before the
+      # tree, the vector check found no Repo); the Task only reads them.
+      [{Task, fn -> run(&IO.puts/1, &System.halt/1, markdown, probed()) end}]
     else
       []
     end
   end
+
+  @doc """
+  The child that computes the 032 lines (slice 032): placed in `Trinity.Application` just
+  before the endpoint, so the Repo, the memory supervisor and the sessions are up. Its
+  `start_link/1` does the work synchronously and answers `:ignore`, so the supervisor waits
+  for it and starts no process; `children/1`'s Task reads the result.
+  """
+  @spec probe([String.t()]) :: [Supervisor.child_spec()]
+  def probe(args) do
+    if requested?(args), do: [Trinity.Smoke.Probe], else: []
+  end
+
+  @doc false
+  @spec probed() :: [String.t()]
+  def probed do
+    :persistent_term.get({__MODULE__, :probed}, [
+      "TRINITY_SMOKE_EXLA=failed:not_probed",
+      "TRINITY_SMOKE_VEC=failed:not_probed",
+      "TRINITY_SMOKE_SEMANTIC=off:not_probed"
+    ])
+  end
+
+  defmodule Probe do
+    @moduledoc false
+    use Boundary, top_level?: true, deps: [Trinity, Trinity.Smoke]
+
+    def child_spec(_), do: %{id: __MODULE__, start: {__MODULE__, :start_link, []}}
+
+    def start_link do
+      lines = [Trinity.Smoke.exla_line(), Trinity.Smoke.vec_line(), Trinity.Smoke.semantic_line()]
+      :persistent_term.put({Trinity.Smoke, :probed}, lines)
+      :ignore
+    end
+  end
 end

lib/trinity_web/live/memory_live.ex

@@ -319,7 +319,7 @@ defmodule TrinityWeb.MemoryLive do
               @status != :on && "bg-warning/10"
             ]}
           >
-            {Semantic.describe(@status)}
+            {Semantic.describe(@status)}.
             <span :if={@status != :on}>{gettext("Full-text search still works.")}</span>
           </p>
           <div

mix.exs

@@ -47,6 +47,13 @@ defmodule Trinity.MixProject do
     [
       desktop: [
         steps: [:assemble, &Burrito.wrap/1],
+        # Slice 032: exla is loaded, not started, in the release. Its application start
+        # loads the NIF, and in Burrito's musl ERTS on Linux that load fails
+        # (`__libc_single_threaded: symbol not found`, package run 35600216451), which took
+        # the whole release down at boot. `Trinity.Memory.Embedders.Bumblebee.exla/0`
+        # starts it on demand and a failure is the tier off with that reason, not a crash.
+        # On a Windows host exla is not declared at all (`exla_deps/0`).
+        applications: exla_release_applications(),
         burrito: [
           targets: [
             linux_x86_64: [os: :linux, cpu: :x86_64],
@@ -96,6 +103,13 @@ defmodule Trinity.MixProject do
   # `{:off, :no_local_backend}`, and the semantic tier is off until a Windows backend exists
   # (NOTES decision 3). The lock file carries exla either way; `mix deps.get` leaves an
   # undeclared lock entry alone.
+  defp exla_release_applications do
+    case :os.type() do
+      {:win32, _} -> []
+      _ -> [exla: :load]
+    end
+  end
+
   defp exla_deps do
     case :os.type() do
       {:win32, _} -> []