Merge pull request #35 from ausimian/m18-telemetry

M18: observability & fallback telemetry

Author: ausimian

PLAN.md

@@ -918,6 +918,40 @@ ops rotate on/off `via_binary` — make it observable.
 **Exit:** events documented in `Emily.Telemetry` moduledoc; fallback
 warning behavior covered by tests.
 
+**Shipped.**
+
+- `Emily.Telemetry` (`lib/emily/telemetry.ex`) — moduledoc enumerates
+  every event; `memory_stats/0` samples the MLX allocator and emits
+  `[:emily, :memory, :stats]`. `maybe_warn_fallback/2` and
+  `init_dedup_table/0` are the internal helpers; dedup state lives in
+  a named `:public` ETS table owned by `Emily.Application`.
+- Span events via `:telemetry.span/3`: `[:emily, :eval, *]` on
+  `Emily.eval/1`, `[:emily, :to_binary, *]` on both `Emily.to_binary/1`
+  and `Emily.Backend.to_binary/2` (the Nx.to_binary path) with
+  `:shape`/`:dtype`/`:byte_size` metadata, and
+  `[:emily, :fallback, *]` on every `via_binary` entry with
+  `:op`/`:input_shapes`/`:input_dtypes`.
+- One-shot `Logger.warning` per `{op, input_shapes}` pair. Opt-in via
+  `config :emily, :warn_on_fallback, true`; off by default so library
+  consumers and CI logs stay quiet. The telemetry event fires
+  regardless — the log is a dev-time convenience on top of it.
+- **Scope interpretation.** PLAN's "each Native dispatch" was read as
+  the evaluation boundary (`Native.eval` / `Native.to_binary`) rather
+  than wrapping 300+ graph-construction call sites in
+  `Emily.Backend`. Graph-construction NIFs are <10μs and do no work;
+  the evaluation boundary is where MLX actually runs kernels, and
+  it's the point every lazy tensor funnels through. If per-op
+  graph-construction histograms are ever needed, the right answer is
+  a centralised dispatch helper, not per-callback decoration.
+- **Op-name plumbing.** `via_binary/3` → `via_binary/4`;
+  `via_binary_tuple/3` → `via_binary_tuple/4`; `apply_scatter/7` →
+  `/8`. The op name propagates into both the `:telemetry` metadata
+  and the dedup key.
+- Tests: `test/emily/telemetry_test.exs` covers fallback start/stop
+  events, 100-call dedup via `capture_log`, `warn_on_fallback=false`
+  silence, `to_binary` span metadata, and `memory_stats/0` emission.
+  `async: false` because the dedup table is global.
+
 ### M19 — Error surfacing
 
 C++ exceptions propagate through `fine` and surface as

README.md

@@ -92,6 +92,27 @@ weight copy.
 
 See `Emily.Stream` moduledoc for details.
 
+## Observability
+
+Emily emits `:telemetry` events at the evaluation boundary
+(`[:emily, :eval, *]`, `[:emily, :to_binary, *]`) and at every
+`Nx.BinaryBackend` fallback (`[:emily, :fallback, *]`). Attach a
+handler to graph hotspots or detect silent performance cliffs —
+see `Emily.Telemetry` for the full event catalogue.
+
+When a backend callback has no native MLX path, Emily transparently
+falls back to `Nx.BinaryBackend`. The fallback is ~100× slower; to
+get a one-shot `Logger.warning` per `{op, input_shapes}` pair the
+first time each one fires (recommended during development):
+
+```elixir
+# config/dev.exs
+config :emily, :warn_on_fallback, true
+```
+
+The warning is off by default so library consumers and CI logs stay
+quiet. The telemetry event fires regardless.
+
 ## Milestones shipped
 
 - **M0** — NIF scaffold, MLX prebuilt fetch, tensor round-trip.

RELEASE.md

@@ -8,6 +8,39 @@
 
 ## Added
 
+- M18 — Observability & fallback telemetry. Makes silent
+  `via_binary` round-trips and long-running memory drift observable
+  without changing any op semantics.
+  - **`Emily.Telemetry`** (`lib/emily/telemetry.ex`) — moduledoc
+    enumerates every event, `memory_stats/0` samples the MLX
+    allocator and emits `[:emily, :memory, :stats]`.
+  - **Span events** via `:telemetry.span/3`:
+    `[:emily, :eval, *]` wraps `Emily.eval/1`;
+    `[:emily, :to_binary, *]` wraps both `Emily.to_binary/1` and
+    `Emily.Backend.to_binary/2` (the Nx.to_binary path) with
+    `:shape`, `:dtype`, `:byte_size` metadata;
+    `[:emily, :fallback, *]` wraps every `via_binary` / `via_binary_tuple`
+    entry with `:op`, `:input_shapes`, `:input_dtypes` metadata.
+  - **One-shot fallback warning** per `{op, input_shapes}` pair via
+    a `:public, :named_table` ETS dedup set owned by
+    `Emily.Application`. **Off by default** — library consumers and
+    CI logs stay quiet; opt in with
+    `config :emily, :warn_on_fallback, true` (typically in
+    `config/dev.exs`) when chasing the Whisper-before-M8 class of
+    bug where forward-pass time silently lands on BinaryBackend.
+    The telemetry event fires regardless of the config.
+  - **Op-name plumbing**: `via_binary/3` → `via_binary/4` and
+    `via_binary_tuple/3` → `via_binary_tuple/4` take a leading
+    op-name atom; `apply_scatter/7` → `/8` threads it through for
+    `indexed_add` / `indexed_put`. All call sites updated.
+  - **Tests** (`test/emily/telemetry_test.exs`): fallback start/stop
+    event with op+shape metadata (via `Nx.reduce`), 100-call dedup
+    capture_log assertion, `warn_on_fallback=false` silence path,
+    `to_binary` span metadata with byte_size, and `memory_stats/0`
+    emission. `async: false` because the ETS dedup is global.
+  - **No `mix.exs` change**: `:telemetry` is already a transitive dep
+    via Nx.
+
 - M17 — Conv-pool training (native window ops). Lifted `window_sum`,
   `window_max`, `window_min`, `window_product`, `window_scatter_max`,
   and `window_scatter_min` off the `via_binary` fallback onto native

lib/emily.ex

@@ -46,7 +46,12 @@ defmodule Emily do
   """
   @spec to_binary(t()) :: binary()
   def to_binary(tensor) do
-    Native.to_binary(Emily.MlxStream.default_worker(), tensor)
+    metadata = %{shape: Native.shape(tensor), dtype: Native.dtype(tensor)}
+
+    :telemetry.span([:emily, :to_binary], metadata, fn ->
+      bytes = Native.to_binary(Emily.MlxStream.default_worker(), tensor)
+      {bytes, Map.put(metadata, :byte_size, byte_size(bytes))}
+    end)
   end
 
   @doc "Return the tensor's shape as a list of non-negative ints."
@@ -65,6 +70,8 @@ defmodule Emily do
   """
   @spec eval(t()) :: :ok
   def eval(tensor) do
-    Native.eval(Emily.MlxStream.default_worker(), tensor)
+    :telemetry.span([:emily, :eval], %{}, fn ->
+      {Native.eval(Emily.MlxStream.default_worker(), tensor), %{}}
+    end)
   end
 end

lib/emily/application.ex

@@ -4,6 +4,8 @@ defmodule Emily.Application do
 
   @impl true
   def start(_type, _args) do
+    Emily.Telemetry.init_dedup_table()
+
     children = [
       {Emily.MlxStream, name: Emily.MlxStream.Default}
     ]

lib/emily/backend.ex

@@ -122,8 +122,16 @@ defmodule Emily.Backend do
   defp ensure_binary(bs) when is_bitstring(bs), do: :erlang.list_to_bitstring([bs])
 
   @impl true
-  def to_binary(%T{data: %B{ref: r}, type: {_, bits}} = tensor, limit) do
-    bin = Native.to_binary(worker(), r)
+  def to_binary(%T{data: %B{ref: r}, shape: shape, type: type} = tensor, limit) do
+    metadata = %{shape: shape, dtype: type}
+
+    bin =
+      :telemetry.span([:emily, :to_binary], metadata, fn ->
+        bytes = Native.to_binary(worker(), r)
+        {bytes, Map.put(metadata, :byte_size, byte_size(bytes))}
+      end)
+
+    {_, bits} = type
     elem_bits = effective_elem_bits(bits)
     size = Nx.size(tensor)
 
@@ -668,7 +676,7 @@ defmodule Emily.Backend do
         Native.reshape(w, r, Tuple.to_list(out.shape)) |> wrap(out, w)
 
       true ->
-        via_binary(out, [input, indices], &Nx.gather(&1, &2, opts))
+        via_binary(:gather, out, [input, indices], &Nx.gather(&1, &2, opts))
     end
   end
 
@@ -753,7 +761,7 @@ defmodule Emily.Backend do
         Native.unquote(native_name)(w, ref(t), axis, reverse, true) |> wrap(out, w)
       else
         nx_fun = unquote(nx_name)
-        via_binary(out, [t], &apply(Nx, nx_fun, [&1, opts]))
+        via_binary(nx_fun, out, [t], &apply(Nx, nx_fun, [&1, opts]))
       end
     end
   end
@@ -779,7 +787,7 @@ defmodule Emily.Backend do
     if float_like?(type) do
       batched_matmul(out, a, contract_a, batch_a, b, contract_b, batch_b)
     else
-      via_binary(out, [a, b], &Nx.dot(&1, contract_a, batch_a, &2, contract_b, batch_b))
+      via_binary(:dot, out, [a, b], &Nx.dot(&1, contract_a, batch_a, &2, contract_b, batch_b))
     end
   end
 
@@ -990,10 +998,10 @@ defmodule Emily.Backend do
   def conv(out, input, kernel, opts) do
     cond do
       opts[:batch_group_size] > 1 ->
-        via_binary(out, [input, kernel], &Nx.conv(&1, &2, opts))
+        via_binary(:conv, out, [input, kernel], &Nx.conv(&1, &2, opts))
 
       match?({:c, _}, out.type) ->
-        via_binary(out, [input, kernel], &Nx.conv(&1, &2, opts))
+        via_binary(:conv, out, [input, kernel], &Nx.conv(&1, &2, opts))
 
       true ->
         w = worker()
@@ -1062,40 +1070,68 @@ defmodule Emily.Backend do
   # those scalars land on the current global default — which is
   # `Emily.Backend` during conformance tests — and the resulting
   # mixed-backend operand list crashes inside BinaryBackend's op.
-  defp via_binary(%T{} = out, tensors, fun) when is_list(tensors) do
-    result =
-      Nx.with_default_backend(Nx.BinaryBackend, fn ->
-        tensors |> transfer_all() |> then(&apply(fun, &1))
-      end)
+  defp via_binary(op, %T{} = out, tensors, fun) when is_atom(op) and is_list(tensors) do
+    metadata = fallback_metadata(op, tensors)
+    Emily.Telemetry.maybe_warn_fallback(op, metadata.input_shapes)
 
-    from_binary(out, Nx.to_binary(result), [])
+    :telemetry.span([:emily, :fallback], metadata, fn ->
+      result =
+        Nx.with_default_backend(Nx.BinaryBackend, fn ->
+          tensors |> transfer_all() |> then(&apply(fun, &1))
+        end)
+
+      {from_binary(out, Nx.to_binary(result), []), metadata}
+    end)
   end
 
   # Same pattern, but the op returns a tuple of tensors. `outs` is a
   # tuple of output templates matching arity; positions are zipped.
-  defp via_binary_tuple(outs, tensors, fun) when is_tuple(outs) and is_list(tensors) do
-    result_tuple =
-      Nx.with_default_backend(Nx.BinaryBackend, fn ->
-        tensors |> transfer_all() |> then(&apply(fun, &1))
-      end)
+  defp via_binary_tuple(op, outs, tensors, fun)
+       when is_atom(op) and is_tuple(outs) and is_list(tensors) do
+    metadata = fallback_metadata(op, tensors)
+    Emily.Telemetry.maybe_warn_fallback(op, metadata.input_shapes)
+
+    :telemetry.span([:emily, :fallback], metadata, fn ->
+      result_tuple =
+        Nx.with_default_backend(Nx.BinaryBackend, fn ->
+          tensors |> transfer_all() |> then(&apply(fun, &1))
+        end)
+
+      result =
+        outs
+        |> Tuple.to_list()
+        |> Enum.zip(Tuple.to_list(result_tuple))
+        |> Enum.map(fn {out, r} -> from_binary(out, Nx.to_binary(r), []) end)
+        |> List.to_tuple()
+
+      {result, metadata}
+    end)
+  end
 
-    outs
-    |> Tuple.to_list()
-    |> Enum.zip(Tuple.to_list(result_tuple))
-    |> Enum.map(fn {out, r} -> from_binary(out, Nx.to_binary(r), []) end)
-    |> List.to_tuple()
+  defp fallback_metadata(op, tensors) do
+    %{
+      op: op,
+      input_shapes: Enum.map(tensors, & &1.shape),
+      input_dtypes: Enum.map(tensors, & &1.type)
+    }
   end
 
   defp transfer_all(tensors),
     do: Enum.map(tensors, &Nx.backend_transfer(&1, Nx.BinaryBackend))
 
   @impl true
   def reduce(out, t, acc, opts, fun),
-    do: via_binary(out, [t, acc], &Nx.reduce(&1, &2, opts, fun))
+    do: via_binary(:reduce, out, [t, acc], &Nx.reduce(&1, &2, opts, fun))
 
   @impl true
   def window_reduce(out, t, acc, window_shape, opts, fun),
-    do: via_binary(out, [t, acc], &Nx.window_reduce(&1, &2, window_shape, opts, fun))
+    do:
+      via_binary(
+        :window_reduce,
+        out,
+        [t, acc],
+        &Nx.window_reduce(&1, &2, window_shape, opts, fun)
+      )
 
   # M17: window reductions lifted off via_binary. MLX has no native
   # window_* primitive — each op is composed as pad → as_strided
@@ -1234,15 +1270,33 @@ defmodule Emily.Backend do
   # indices; correctness on duplicates with indexed_put is best-effort.
   @impl true
   def indexed_add(out, t, indices, updates, opts) do
-    apply_scatter(out, t, indices, updates, opts, :scatter_add, &Nx.indexed_add(&1, &2, &3, opts))
+    apply_scatter(
+      :indexed_add,
+      out,
+      t,
+      indices,
+      updates,
+      opts,
+      :scatter_add,
+      &Nx.indexed_add(&1, &2, &3, opts)
+    )
   end
 
   @impl true
   def indexed_put(out, t, indices, updates, opts) do
-    apply_scatter(out, t, indices, updates, opts, :scatter, &Nx.indexed_put(&1, &2, &3, opts))
+    apply_scatter(
+      :indexed_put,
+      out,
+      t,
+      indices,
+      updates,
+      opts,
+      :scatter,
+      &Nx.indexed_put(&1, &2, &3, opts)
+    )
   end
 
-  defp apply_scatter(out, t, indices, updates, opts, native_fun, fallback) do
+  defp apply_scatter(op, out, t, indices, updates, opts, native_fun, fallback) do
     axes = opts[:axes] || Enum.to_list(0..(tuple_size(t.shape) - 1))
     indices_shape = Tuple.to_list(indices.shape)
 
@@ -1255,7 +1309,7 @@ defmodule Emily.Backend do
       apply(Native, native_fun, [w, ref(t), idx_refs, updates_ref, axes])
       |> wrap(out, w)
     else
-      via_binary(out, [t, indices, updates], fallback)
+      via_binary(op, out, [t, indices, updates], fallback)
     end
   end
 
@@ -1388,7 +1442,7 @@ defmodule Emily.Backend do
         {wrap(q_ref, q_out, w), wrap(r_ref, r_out, w)}
 
       :complete ->
-        via_binary_tuple({q_out, r_out}, [t], &Nx.LinAlg.qr(&1, opts))
+        via_binary_tuple(:qr, {q_out, r_out}, [t], &Nx.LinAlg.qr(&1, opts))
     end
   end