From be4fc19603e48aa839f6326d9e71aed9fb7920de Mon Sep 17 00:00:00 2001
From: ausimian <nick@ausimian.net>
Date: Sun, 19 Apr 2026 10:48:44 +1000
Subject: [PATCH] Restore async: true across the test suite
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Stream-per-process (M14) + the async enif_send substrate (Phase 1-3)
make concurrent Emily.Backend ops safe, so the ExUnit modules that were
pinned to async: false for concurrency reasons can now run in parallel.

- Refactor `ConformanceHelper` to set `Nx.default_backend/1` in a
  per-test `setup` (pdict-scoped) rather than `Nx.global_default_backend/1`
  in `setup_all`; conformance modules (qwen3, vit, whisper) and the
  `*_full` variants are now async: true.
- Refactor the training `*_full` tests and `transform_test.exs`
  similarly — three MNIST training modules now overlap on `:training_full`.
- `distilbert_test.exs` keeps its own `setup_all` + `global_default_backend`
  and stays async: false: its `:distilbert_full` test drives
  `Nx.Serving.batched_run` through a supervised serving process, which
  cannot see the caller's pdict.
- Flip the process-scoped modules (compiler, fast/*, stream, async_eval,
  debug_flags, bumblebee/fast_kernels, compiler_axon) to async: true.
- Drop the redundant `on_exit` restores on pdict-setting `setup` blocks —
  the test process exits anyway and on_exit runs in a different process.

Full suite (453 tests incl. all tags): 78.8s, 73.5s async, 5.3s sync.
---
 test/emily/async_eval_test.exs                |  2 +-
 test/emily/bumblebee/fast_kernels_test.exs    |  4 +---
 test/emily/compiler_axon_test.exs             |  4 +---
 test/emily/compiler_test.exs                  |  4 +---
 test/emily/conformance/distilbert_test.exs    | 14 ++++++++++++-
 test/emily/conformance/qwen3_full_test.exs    |  8 +++----
 .../conformance/qwen3_quant_full_test.exs     |  8 +++----
 test/emily/conformance/qwen3_test.exs         |  2 +-
 test/emily/conformance/vit_full_test.exs      |  2 +-
 test/emily/conformance/vit_test.exs           |  2 +-
 test/emily/conformance/whisper_full_test.exs  |  2 +-
 test/emily/conformance/whisper_test.exs       |  2 +-
 test/emily/debug_flags_test.exs               |  2 +-
 test/emily/fast/layer_norm_test.exs           |  4 +---
 test/emily/fast/rms_norm_test.exs             |  4 +---
 test/emily/fast/rope_test.exs                 |  4 +---
 test/emily/fast/sdpa_test.exs                 |  4 +---
 test/emily/quantization/transform_test.exs    |  8 +++----
 test/emily/stream_test.exs                    |  2 +-
 test/emily/training/mnist_bf16_full_test.exs  |  8 +++----
 test/emily/training/mnist_cnn_full_test.exs   |  8 +++----
 test/emily/training/mnist_full_test.exs       |  8 +++----
 test/support/conformance_helper.ex            | 21 ++++++++++++-------
 23 files changed, 60 insertions(+), 67 deletions(-)

diff --git a/test/emily/async_eval_test.exs b/test/emily/async_eval_test.exs
index ce9867b..27765ba 100644
--- a/test/emily/async_eval_test.exs
+++ b/test/emily/async_eval_test.exs
@@ -1,5 +1,5 @@
 defmodule Emily.AsyncEvalTest do
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
 
   # Regression tests for the async `Emily.Native.eval/2` path.
   # `eval/2` internally dispatches to `eval_nif/2` (which returns a
diff --git a/test/emily/bumblebee/fast_kernels_test.exs b/test/emily/bumblebee/fast_kernels_test.exs
index c1aadc1..a5ef59d 100644
--- a/test/emily/bumblebee/fast_kernels_test.exs
+++ b/test/emily/bumblebee/fast_kernels_test.exs
@@ -14,7 +14,7 @@ defmodule Emily.Bumblebee.FastKernelsTest do
   `:fast_kernels_full` tag).
   """
 
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
 
   import Emily.BackendGenerators, only: [assert_close: 3]
 
@@ -24,9 +24,7 @@ defmodule Emily.Bumblebee.FastKernelsTest do
   @f32_tol 1.0e-4
 
   setup do
-    prev = Nx.default_backend()
     Nx.default_backend(Emily.Backend)
-    on_exit(fn -> Nx.default_backend(prev) end)
     :ok
   end
 
diff --git a/test/emily/compiler_axon_test.exs b/test/emily/compiler_axon_test.exs
index e59de61..12bfc05 100644
--- a/test/emily/compiler_axon_test.exs
+++ b/test/emily/compiler_axon_test.exs
@@ -9,14 +9,12 @@ defmodule Emily.CompilerAxonTest do
   ops with the same operands as the reference walker.
   """
 
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
 
   import Emily.BackendGenerators, only: [assert_close: 2]
 
   setup do
-    prev = Nx.default_backend()
     Nx.default_backend(Emily.Backend)
-    on_exit(fn -> Nx.default_backend(prev) end)
     :ok
   end
 
diff --git a/test/emily/compiler_test.exs b/test/emily/compiler_test.exs
index ff39cd4..f3cc865 100644
--- a/test/emily/compiler_test.exs
+++ b/test/emily/compiler_test.exs
@@ -19,15 +19,13 @@ defmodule Emily.CompilerTest do
   and `Nx.Defn.Evaluator` to verify equivalence.
   """
 
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
 
   import Nx.Defn
   import Emily.BackendGenerators, only: [assert_close: 2]
 
   setup do
-    prev = Nx.default_backend()
     Nx.default_backend(Emily.Backend)
-    on_exit(fn -> Nx.default_backend(prev) end)
     :ok
   end
 
diff --git a/test/emily/conformance/distilbert_test.exs b/test/emily/conformance/distilbert_test.exs
index 8fe0ed4..4929472 100644
--- a/test/emily/conformance/distilbert_test.exs
+++ b/test/emily/conformance/distilbert_test.exs
@@ -32,7 +32,8 @@ defmodule Emily.Conformance.DistilbertTest do
   """
 
   use ExUnit.Case, async: false
-  use Emily.ConformanceHelper
+
+  import Emily.ConformanceHelper, only: [assert_all_close: 2, assert_all_close: 3]
 
   alias Emily.Bumblebee.FastKernels
 
@@ -40,6 +41,17 @@ defmodule Emily.Conformance.DistilbertTest do
   @moduletag capture_log: true
   @moduletag timeout: 120_000
 
+  # `batched_run` runs through a supervised serving process, which has
+  # its own process dict. Set the backend globally so the worker sees
+  # `Emily.Backend` as its default; keep the module `async: false` to
+  # avoid racing the global with other suites.
+  setup_all do
+    prev = Nx.default_backend()
+    Nx.global_default_backend(Emily.Backend)
+    on_exit(fn -> Nx.global_default_backend(prev) end)
+    :ok
+  end
+
   test ":base" do
     assert {:ok, %{model: model, params: params, spec: spec}} =
              Bumblebee.load_model({:hf, "hf-internal-testing/tiny-random-DistilBertModel"})
diff --git a/test/emily/conformance/qwen3_full_test.exs b/test/emily/conformance/qwen3_full_test.exs
index ccb681d..fdcd61f 100644
--- a/test/emily/conformance/qwen3_full_test.exs
+++ b/test/emily/conformance/qwen3_full_test.exs
@@ -14,7 +14,7 @@ defmodule Emily.Conformance.Qwen3FullTest do
   checkpoint has been republished — all of which are real signals.
   """
 
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
 
   alias Emily.Bumblebee.FastKernels
 
@@ -25,10 +25,8 @@ defmodule Emily.Conformance.Qwen3FullTest do
   @prompt "The quick brown fox jumps over the lazy dog."
   @reference_text " The quick brown fox is a character in the story. The quick brown fox is a character in the story. The quick brown fox is a character in the story"
 
-  setup_all do
-    prev = Nx.default_backend()
-    Nx.global_default_backend(Emily.Backend)
-    on_exit(fn -> Nx.global_default_backend(prev) end)
+  setup do
+    Nx.default_backend(Emily.Backend)
     :ok
   end
 
diff --git a/test/emily/conformance/qwen3_quant_full_test.exs b/test/emily/conformance/qwen3_quant_full_test.exs
index 8da0ea5..b715a28 100644
--- a/test/emily/conformance/qwen3_quant_full_test.exs
+++ b/test/emily/conformance/qwen3_quant_full_test.exs
@@ -18,7 +18,7 @@ defmodule Emily.Conformance.Qwen3QuantFullTest do
   see PLAN.md M10.5.)
   """
 
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
 
   alias Emily.Bumblebee.FastKernels
   alias Emily.Quantization.Transform
@@ -37,10 +37,8 @@ defmodule Emily.Conformance.Qwen3QuantFullTest do
   # quantization noise across all linear layers.
   @reference_text " Let's see, what is the correct answer for this riddle? The answer is a word that contains the letters B, O, U, and R,"
 
-  setup_all do
-    prev = Nx.default_backend()
-    Nx.global_default_backend(Emily.Backend)
-    on_exit(fn -> Nx.global_default_backend(prev) end)
+  setup do
+    Nx.default_backend(Emily.Backend)
     :ok
   end
 
diff --git a/test/emily/conformance/qwen3_test.exs b/test/emily/conformance/qwen3_test.exs
index 6377b6b..ffa55a0 100644
--- a/test/emily/conformance/qwen3_test.exs
+++ b/test/emily/conformance/qwen3_test.exs
@@ -25,7 +25,7 @@ defmodule Emily.Conformance.Qwen3Test do
       mix test --only conformance
   """
 
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
   use Emily.ConformanceHelper
 
   alias Bumblebee.Text.Generation, as: BBGeneration
diff --git a/test/emily/conformance/vit_full_test.exs b/test/emily/conformance/vit_full_test.exs
index ee7e76f..accb739 100644
--- a/test/emily/conformance/vit_full_test.exs
+++ b/test/emily/conformance/vit_full_test.exs
@@ -21,7 +21,7 @@ defmodule Emily.Conformance.VitFullTest do
   are real signals.
   """
 
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
   use Emily.ConformanceHelper
 
   alias Emily.Bumblebee.FastKernels
diff --git a/test/emily/conformance/vit_test.exs b/test/emily/conformance/vit_test.exs
index 853e6bf..baa39f7 100644
--- a/test/emily/conformance/vit_test.exs
+++ b/test/emily/conformance/vit_test.exs
@@ -24,7 +24,7 @@ defmodule Emily.Conformance.VitTest do
       mix test --only conformance
   """
 
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
   use Emily.ConformanceHelper
 
   @moduletag :conformance
diff --git a/test/emily/conformance/whisper_full_test.exs b/test/emily/conformance/whisper_full_test.exs
index 207852e..6559410 100644
--- a/test/emily/conformance/whisper_full_test.exs
+++ b/test/emily/conformance/whisper_full_test.exs
@@ -22,7 +22,7 @@ defmodule Emily.Conformance.WhisperFullTest do
   which are real signals.
   """
 
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
   use Emily.ConformanceHelper
 
   alias Emily.Bumblebee.FastKernels
diff --git a/test/emily/conformance/whisper_test.exs b/test/emily/conformance/whisper_test.exs
index 84aa723..57f7d15 100644
--- a/test/emily/conformance/whisper_test.exs
+++ b/test/emily/conformance/whisper_test.exs
@@ -24,7 +24,7 @@ defmodule Emily.Conformance.WhisperTest do
       mix test --only conformance
   """
 
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
   use Emily.ConformanceHelper
 
   @moduletag :conformance
diff --git a/test/emily/debug_flags_test.exs b/test/emily/debug_flags_test.exs
index 2500e39..9ea3ab1 100644
--- a/test/emily/debug_flags_test.exs
+++ b/test/emily/debug_flags_test.exs
@@ -12,7 +12,7 @@ defmodule Emily.DebugFlagsTest do
       `DebugHelpers.check_bounds!` / `check_nan_inf!`.
   """
 
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
 
   import Emily.TensorHelpers
 
diff --git a/test/emily/fast/layer_norm_test.exs b/test/emily/fast/layer_norm_test.exs
index 0ed6ea9..e1e3cbe 100644
--- a/test/emily/fast/layer_norm_test.exs
+++ b/test/emily/fast/layer_norm_test.exs
@@ -4,7 +4,7 @@ defmodule Emily.Fast.LayerNormTest do
   defn composability, and fused-kernel equivalence under Emily.
   """
 
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
 
   import Emily.BackendGenerators, only: [assert_close: 3]
 
@@ -41,9 +41,7 @@ defmodule Emily.Fast.LayerNormTest do
 
   describe "emily backend (fused path)" do
     setup do
-      prev = Nx.default_backend()
       Nx.default_backend(Emily.Backend)
-      on_exit(fn -> Nx.default_backend(prev) end)
       :ok
     end
 
diff --git a/test/emily/fast/rms_norm_test.exs b/test/emily/fast/rms_norm_test.exs
index f874762..65ed976 100644
--- a/test/emily/fast/rms_norm_test.exs
+++ b/test/emily/fast/rms_norm_test.exs
@@ -15,7 +15,7 @@ defmodule Emily.Fast.RMSNormTest do
       so bit-match isn't expected).
   """
 
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
 
   import Emily.BackendGenerators, only: [assert_close: 3]
 
@@ -49,9 +49,7 @@ defmodule Emily.Fast.RMSNormTest do
 
   describe "emily backend (fused path)" do
     setup do
-      prev = Nx.default_backend()
       Nx.default_backend(Emily.Backend)
-      on_exit(fn -> Nx.default_backend(prev) end)
       :ok
     end
 
diff --git a/test/emily/fast/rope_test.exs b/test/emily/fast/rope_test.exs
index ffd1fa7..988cc86 100644
--- a/test/emily/fast/rope_test.exs
+++ b/test/emily/fast/rope_test.exs
@@ -8,16 +8,14 @@ defmodule Emily.Fast.RoPETest do
   handles, not something rope itself cares about.
   """
 
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
 
   import Emily.BackendGenerators, only: [assert_close: 3]
 
   @f32_tol 1.0e-4
 
   setup do
-    prev = Nx.default_backend()
     Nx.default_backend(Emily.Backend)
-    on_exit(fn -> Nx.default_backend(prev) end)
     :ok
   end
 
diff --git a/test/emily/fast/sdpa_test.exs b/test/emily/fast/sdpa_test.exs
index 91c1941..27e7860 100644
--- a/test/emily/fast/sdpa_test.exs
+++ b/test/emily/fast/sdpa_test.exs
@@ -7,7 +7,7 @@ defmodule Emily.Fast.SDPATest do
   and Bumblebee in-flight shape.
   """
 
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
 
   import Emily.BackendGenerators, only: [assert_close: 3]
 
@@ -15,9 +15,7 @@ defmodule Emily.Fast.SDPATest do
   @bf16_tol 1.0e-2
 
   setup do
-    prev = Nx.default_backend()
     Nx.default_backend(Emily.Backend)
-    on_exit(fn -> Nx.default_backend(prev) end)
     :ok
   end
 
diff --git a/test/emily/quantization/transform_test.exs b/test/emily/quantization/transform_test.exs
index 16aceb4..f1215eb 100644
--- a/test/emily/quantization/transform_test.exs
+++ b/test/emily/quantization/transform_test.exs
@@ -14,15 +14,13 @@ defmodule Emily.Quantization.TransformTest do
   scheme.
   """
 
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
 
   alias Emily.Quantization.Transform
   alias Emily.QuantizedWeight
 
-  setup_all do
-    prev = Nx.default_backend()
-    Nx.global_default_backend(Emily.Backend)
-    on_exit(fn -> Nx.global_default_backend(prev) end)
+  setup do
+    Nx.default_backend(Emily.Backend)
     :ok
   end
 
diff --git a/test/emily/stream_test.exs b/test/emily/stream_test.exs
index ac3499d..21f39cb 100644
--- a/test/emily/stream_test.exs
+++ b/test/emily/stream_test.exs
@@ -1,5 +1,5 @@
 defmodule Emily.StreamTest do
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
 
   describe "Emily.Stream" do
     test "new/1 creates a stream with a worker reference" do
diff --git a/test/emily/training/mnist_bf16_full_test.exs b/test/emily/training/mnist_bf16_full_test.exs
index a7b63f0..107e751 100644
--- a/test/emily/training/mnist_bf16_full_test.exs
+++ b/test/emily/training/mnist_bf16_full_test.exs
@@ -10,7 +10,7 @@ defmodule Emily.Training.MnistBf16FullTest do
   Opt-in — `mix test --only training_full`.
   """
 
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
 
   alias Emily.MnistHelper
 
@@ -18,10 +18,8 @@ defmodule Emily.Training.MnistBf16FullTest do
   @moduletag capture_log: true
   @moduletag timeout: 600_000
 
-  setup_all do
-    prev = Nx.default_backend()
-    Nx.global_default_backend(Emily.Backend)
-    on_exit(fn -> Nx.global_default_backend(prev) end)
+  setup do
+    Nx.default_backend(Emily.Backend)
     :ok
   end
 
diff --git a/test/emily/training/mnist_cnn_full_test.exs b/test/emily/training/mnist_cnn_full_test.exs
index b60af79..d28654f 100644
--- a/test/emily/training/mnist_cnn_full_test.exs
+++ b/test/emily/training/mnist_cnn_full_test.exs
@@ -17,7 +17,7 @@ defmodule Emily.Training.MnistCnnFullTest do
   epochs on MNIST.
   """
 
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
 
   alias Emily.MnistHelper
 
@@ -25,10 +25,8 @@ defmodule Emily.Training.MnistCnnFullTest do
   @moduletag capture_log: true
   @moduletag timeout: 600_000
 
-  setup_all do
-    prev = Nx.default_backend()
-    Nx.global_default_backend(Emily.Backend)
-    on_exit(fn -> Nx.global_default_backend(prev) end)
+  setup do
+    Nx.default_backend(Emily.Backend)
     :ok
   end
 
diff --git a/test/emily/training/mnist_full_test.exs b/test/emily/training/mnist_full_test.exs
index fbbed2a..69cab9d 100644
--- a/test/emily/training/mnist_full_test.exs
+++ b/test/emily/training/mnist_full_test.exs
@@ -18,7 +18,7 @@ defmodule Emily.Training.MnistFullTest do
   Runs once per invocation, no property sweep.
   """
 
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
 
   alias Emily.MnistHelper
 
@@ -26,10 +26,8 @@ defmodule Emily.Training.MnistFullTest do
   @moduletag capture_log: true
   @moduletag timeout: 600_000
 
-  setup_all do
-    prev = Nx.default_backend()
-    Nx.global_default_backend(Emily.Backend)
-    on_exit(fn -> Nx.global_default_backend(prev) end)
+  setup do
+    Nx.default_backend(Emily.Backend)
     :ok
   end
 
diff --git a/test/support/conformance_helper.ex b/test/support/conformance_helper.ex
index 25dd5d5..0ec9a52 100644
--- a/test/support/conformance_helper.ex
+++ b/test/support/conformance_helper.ex
@@ -4,9 +4,10 @@ defmodule Emily.ConformanceHelper do
 
   `use Emily.ConformanceHelper` installs:
 
-    * a `setup_all` block that swaps the global default backend to
-      `Emily.Backend` for the duration of the module and restores it
-      on exit — every conformance suite does this identically;
+    * a per-test `setup` block that swaps the process-scoped default
+      backend to `Emily.Backend` for the duration of the test and
+      restores it on exit — pdict scope (not application env) so
+      modules can run `async: true`;
     * an import of `assert_all_close/2,3`, the tolerance-aware
       comparison we use against reference slices produced by
       HuggingFace Transformers (PyTorch). Mirrors
@@ -16,16 +17,22 @@ defmodule Emily.ConformanceHelper do
   Each conformance module still declares its own `@moduletag`s
   (`:conformance`, `:qwen3_full`, `:vit_full`, …) — those are not
   shared because they gate test selection.
+
+  ## When to avoid this helper
+
+  Tests that drive `Nx.Serving.batched_run` through a supervised
+  serving process cannot rely on the pdict default, because the
+  serving worker is a separate process that falls back to the
+  application env. Those tests must set `Nx.global_default_backend`
+  directly (and run `async: false`).
   """
 
   defmacro __using__(_opts) do
     quote do
       import Emily.ConformanceHelper, only: [assert_all_close: 2, assert_all_close: 3]
 
-      setup_all do
-        prev = Nx.default_backend()
-        Nx.global_default_backend(Emily.Backend)
-        on_exit(fn -> Nx.global_default_backend(prev) end)
+      setup do
+        Nx.default_backend(Emily.Backend)
         :ok
       end
     end