Restore async: true across the test suite

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

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
 

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
 

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
 

test/emily/conformance/distilbert_test.exs

@@ -32,14 +32,26 @@ 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
 
   @moduletag :conformance
   @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"})

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
 

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
 

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

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

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

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

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

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
 

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
 

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
 

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
 

test/emily/fast/sdpa_test.exs

@@ -7,17 +7,15 @@ 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]
 
   @f32_tol 1.0e-4
   @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
 

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
 

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

test/emily/training/mnist_bf16_full_test.exs

@@ -10,18 +10,16 @@ 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
 
   @moduletag :training_full
   @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
 

test/emily/training/mnist_cnn_full_test.exs

@@ -17,18 +17,16 @@ defmodule Emily.Training.MnistCnnFullTest do
   epochs on MNIST.
   """
 
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
 
   alias Emily.MnistHelper
 
   @moduletag :training_full
   @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
 

test/emily/training/mnist_full_test.exs

@@ -18,18 +18,16 @@ 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
 
   @moduletag :training_full
   @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
 

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