From 7e913aad46c3301c43414d2c3a51cdc2edc2890d Mon Sep 17 00:00:00 2001 From: Peter Shoukry Date: Sat, 19 Sep 2026 12:12:36 +0300 Subject: [PATCH 1/2] Add an intra_threads option to Ortex.load ONNX Runtime sizes a session's intra-op thread pool at one thread per physical core unless told otherwise, and keeps every thread busy while the session runs. One model stream can then occupy most of a machine, which leaves nothing for a second model or for the rest of the program. ort exposes SessionBuilder::with_intra_threads; this passes it through as an optional fourth argument to Ortex.load/4: Ortex.load(path, [:cpu], 3, intra_threads: 4) Unset, the session is built exactly as before, so ONNX Runtime's own default and its ORT_INTRA_OP_NUM_THREADS environment fallback still apply. A value that is not a positive integer, or an unknown option, raises ArgumentError before the NIF is called. --- lib/ortex.ex | 11 ++++++++++- lib/ortex/model.ex | 16 ++++++++++++++-- lib/ortex/native.ex | 2 +- native/ortex/src/lib.rs | 3 ++- native/ortex/src/model.rs | 16 +++++++++++++--- test/ortex_test.exs | 30 ++++++++++++++++++++++++++++++ 6 files changed, 70 insertions(+), 8 deletions(-) diff --git a/lib/ortex.ex b/lib/ortex.ex index 51b0e45..0d30c66 100644 --- a/lib/ortex.ex +++ b/lib/ortex.ex @@ -18,14 +18,23 @@ defmodule Ortex do `config.exs` where `EXECUTION_PROVIDERS` is a list of strings of which execution providers to enable. + ## Options + + * `:intra_threads` - the number of threads ONNX Runtime spreads a single + inference across. By default ONNX Runtime uses one thread per physical + core, so one busy model can occupy most of the machine; a smaller pool + trades per-inference latency for cores left free for other work, such + as a second model. Must be a positive integer. + ## Examples iex> Ortex.load("./models/tinymodel.onnx") iex> Ortex.load("./models/tinymodel.onnx", [:cuda, :cpu]) iex> Ortex.load("./models/tinymodel.onnx", [:cpu], 0) + iex> Ortex.load("./models/tinymodel.onnx", [:cpu], 3, intra_threads: 2) """ - defdelegate load(path, eps \\ [:cpu], opt \\ 3), to: Ortex.Model + defdelegate load(path, eps \\ [:cpu], opt \\ 3, options \\ []), to: Ortex.Model @doc """ Run a forward pass through a model. diff --git a/lib/ortex/model.ex b/lib/ortex/model.ex index d8e550f..1e59aa1 100644 --- a/lib/ortex/model.ex +++ b/lib/ortex/model.ex @@ -22,8 +22,10 @@ defmodule Ortex.Model do defstruct [:reference] @doc false - def load(path, eps \\ [:cpu], opt \\ 3) do - case Ortex.Native.init(path, eps, opt) do + def load(path, eps \\ [:cpu], opt \\ 3, options \\ []) do + intra_threads = options |> Keyword.validate!(intra_threads: nil) |> intra_threads!() + + case Ortex.Native.init(path, eps, opt, intra_threads) do {:error, msg} -> raise msg @@ -32,6 +34,16 @@ defmodule Ortex.Model do end end + defp intra_threads!(intra_threads: nil), do: nil + + defp intra_threads!(intra_threads: threads) when is_integer(threads) and threads > 0, + do: threads + + defp intra_threads!(intra_threads: threads) do + raise ArgumentError, + "expected :intra_threads to be a positive integer, got: #{inspect(threads)}" + end + @doc false def run(%Ortex.Model{} = model, tensor) when not is_tuple(tensor) do run(model, {tensor}) diff --git a/lib/ortex/native.ex b/lib/ortex/native.ex index ee16c07..aca8672 100644 --- a/lib/ortex/native.ex +++ b/lib/ortex/native.ex @@ -25,7 +25,7 @@ defmodule Ortex.Native do # When loading a NIF module, dummy clauses for all NIF function are required. # NIF dummies usually just error out when called when the NIF is not loaded, as that should never normally happen. - def init(_model_path, _execution_providers, _optimization_level), + def init(_model_path, _execution_providers, _optimization_level, _intra_threads), do: :erlang.nif_error(:nif_not_loaded) def run(_model, _inputs), do: :erlang.nif_error(:nif_not_loaded) diff --git a/native/ortex/src/lib.rs b/native/ortex/src/lib.rs index 0523105..f7ba674 100644 --- a/native/ortex/src/lib.rs +++ b/native/ortex/src/lib.rs @@ -22,9 +22,10 @@ fn init( model_path: String, eps: Vec, opt: i32, + intra_threads: Option, ) -> NifResult> { let eps = utils::map_eps(env, eps); - let model = model::init(model_path, eps, opt) + let model = model::init(model_path, eps, opt, intra_threads) .map_err(|e| rustler::Error::Term(Box::new(e.to_string())))?; Ok(ResourceArc::new(model)) } diff --git a/native/ortex/src/model.rs b/native/ortex/src/model.rs index f5e04b3..221a0d5 100644 --- a/native/ortex/src/model.rs +++ b/native/ortex/src/model.rs @@ -4,7 +4,7 @@ //! # Examples //! //! ``` -//! let model = init("./models/resnet50.onnx", vec![])?; +//! let model = init("./models/resnet50.onnx", vec![], 3, None)?; //! let (inputs, outputs) = show(model)?; //! ``` @@ -31,16 +31,26 @@ unsafe impl Sync for OrtexModel {} /// Creates a model given the path to the model and vector of execution providers. /// The execution providers are Atoms from Erlang/Elixir. +/// +/// `intra_threads` caps the thread pool one inference is spread across. `None` +/// leaves the choice to ONNX Runtime, which defaults to one thread per physical +/// core and honours `ORT_INTRA_OP_NUM_THREADS` when set in the OS environment. pub fn init( model_path: String, eps: Vec, opt: i32, + intra_threads: Option, ) -> Result { // TODO: send tracing logs to erlang/elixir _somehow_ // tracing_subscriber::fmt::init(); - let session = Session::builder()? - .with_optimization_level(map_opt_level(opt))? + let mut builder = Session::builder()?.with_optimization_level(map_opt_level(opt))?; + + if let Some(threads) = intra_threads { + builder = builder.with_intra_threads(threads)?; + } + + let session = builder .with_execution_providers(eps)? .commit_from_file(model_path)?; diff --git a/test/ortex_test.exs b/test/ortex_test.exs index bb41f8b..bc19d2f 100644 --- a/test/ortex_test.exs +++ b/test/ortex_test.exs @@ -23,6 +23,36 @@ defmodule OrtexTest do assert result |> Nx.backend_transfer() |> Nx.argmax(axis: 1) == Nx.tensor([499]) end + describe "intra_threads" do + test "a session capped at one intra-op thread still runs the model" do + model = Ortex.load("./models/tinymodel.onnx", [:cpu], 3, intra_threads: 1) + + {%Nx.Tensor{shape: {1, 10}}, %Nx.Tensor{shape: {1, 10}}, %Nx.Tensor{shape: {1, 10}}} = + Ortex.run(model, { + Nx.broadcast(0, {1, 100}) |> Nx.as_type(:s32), + Nx.broadcast(0.0, {1, 100}) |> Nx.as_type(:f32) + }) + end + + test "is not required" do + assert %Ortex.Model{} = Ortex.load("./models/tinymodel.onnx", [:cpu], 3, []) + end + + test "must be a positive integer" do + for bad <- [0, -1, 2.0, "4"] do + assert_raise ArgumentError, ~r/intra_threads/, fn -> + Ortex.load("./models/tinymodel.onnx", [:cpu], 3, intra_threads: bad) + end + end + end + + test "an unknown option is refused rather than ignored" do + assert_raise ArgumentError, ~r/intra_thread/, fn -> + Ortex.load("./models/tinymodel.onnx", [:cpu], 3, intra_thread: 2) + end + end + end + test "Nx.Serving with tinymodel" do model = Ortex.load("./models/tinymodel.onnx") From 89721cccdb0564636a406c08c874f1b6c6b8f2ce Mon Sep 17 00:00:00 2001 From: Peter Shoukry Date: Sat, 19 Sep 2026 13:41:39 +0300 Subject: [PATCH 2/2] Validate intra_threads by value rather than by the option list Matching the whole validated keyword list meant a second option would break every clause. Read the key out of the validated options instead. --- lib/ortex/model.ex | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/lib/ortex/model.ex b/lib/ortex/model.ex index 1e59aa1..14ade2a 100644 --- a/lib/ortex/model.ex +++ b/lib/ortex/model.ex @@ -23,7 +23,8 @@ defmodule Ortex.Model do @doc false def load(path, eps \\ [:cpu], opt \\ 3, options \\ []) do - intra_threads = options |> Keyword.validate!(intra_threads: nil) |> intra_threads!() + options = Keyword.validate!(options, intra_threads: nil) + intra_threads = validate_intra_threads!(options[:intra_threads]) case Ortex.Native.init(path, eps, opt, intra_threads) do {:error, msg} -> @@ -34,12 +35,12 @@ defmodule Ortex.Model do end end - defp intra_threads!(intra_threads: nil), do: nil + defp validate_intra_threads!(nil), do: nil - defp intra_threads!(intra_threads: threads) when is_integer(threads) and threads > 0, + defp validate_intra_threads!(threads) when is_integer(threads) and threads > 0, do: threads - defp intra_threads!(intra_threads: threads) do + defp validate_intra_threads!(threads) do raise ArgumentError, "expected :intra_threads to be a positive integer, got: #{inspect(threads)}" end