From 672ced831365ad06c8b6881ffcd8491613f04a04 Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Thu, 29 Oct 2026 03:15:57 -0700 Subject: [PATCH] Bound managed inference and cancel work on caller exit --- CHANGELOG.md | 1 + docs/cancellation.md | 29 ++++++ docs/integration.md | 4 +- lib/ortex.ex | 9 ++ lib/ortex/model.ex | 19 +++- lib/ortex/run_control.ex | 6 +- lib/ortex/run_options.ex | 8 +- lib/ortex/runner.ex | 111 +++++++++++++++++++++ lib/ortex/telemetry.ex | 5 +- lib/ortex/timeout_error.ex | 6 ++ native/ortex/src/lib.rs | 3 +- test/managed_runs_test.exs | 197 +++++++++++++++++++++++++++++++++++++ test/telemetry_test.exs | 16 +++ 13 files changed, 406 insertions(+), 8 deletions(-) create mode 100644 lib/ortex/runner.ex create mode 100644 lib/ortex/timeout_error.ex create mode 100644 test/managed_runs_test.exs diff --git a/CHANGELOG.md b/CHANGELOG.md index b1b1e8c..5163e5d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,7 @@ - Retain numeric storage/Serving support and reject unsupported batching value types early. - Summarize ONNX protobuf graphs without a runtime, including versions, operators, initializer estimates, and ZipMap. - Cancel one model-bound inference attempt with immutable `Ortex.RunControl` options. +- Bound caller wait with per-run timeouts and cancel managed inference when its caller exits. - Support checksum-verified precompiled CPU NIFs and source-build overrides. - Default to dynamic runtime loading and API 17, without runtime build headers. diff --git a/docs/cancellation.md b/docs/cancellation.md index 70cc362..3da3dad 100644 --- a/docs/cancellation.md +++ b/docs/cancellation.md @@ -26,3 +26,32 @@ from concurrent processes. They also cancel a queued attempt while another holds the model mutex, check post-cancellation scores, prevent reuse/cross-model use, and exercise controlled string and ZipMap outputs. Both runtime 1.17.3 and 1.24.4 are covered. Long-running provider kernels still need provider-specific validation. + +## Managed deadlines and ownership + +`Ortex.run(model, inputs, timeout: milliseconds)` creates a control automatically +and raises `Ortex.TimeoutError` if caller wait reaches its deadline. An explicit +control can also be paired with a timeout. Zero expires before starting a worker; +`:infinity` is the default. Validation/control preparation precedes the deadline; +it then covers input transfer, mutex waiting, inference, and output conversion. + +Finite deadlines and explicit controls use a separate native worker plus an owner +watcher. The watcher monitors the caller and worker. Caller death or a deadline +abort requests runtime termination and kills the BEAM worker; any executing NIF +retains its resources until it returns. Caller wait is bounded, while native +termination remains cooperative. An uninterruptible kernel can keep the session +occupied until it finishes. Ordinary unbounded calls without a control keep the +synchronous path and its allocation profile. + +The caller owns an Erlang reply alias. On return, exception, or timeout the alias +is removed, pending reply/monitor messages are flushed, and late replies are +ignored. Worker exceptions are re-raised with their original stack. Completed +attempts do not trigger watcher cancellation, so rejecting a reused or wrong-model +control cannot cancel its legitimate in-flight attempt. + +Tests exercise actual loop deadlines, caller death, deadlines behind an occupied +mutex, post-cancel inference, worker exception stacks, resource ownership, worker +termination, response/deadline races, and lack of late messages. Telemetry emits +the deadline and cancellation mode in the caller's span, with TimeoutError as an +exception event. Control preparation uses dirty CPU scheduling independently of +the dirty I/O pool used by inference mutex waits. diff --git a/docs/integration.md b/docs/integration.md index 656cefd..1a4efcf 100644 --- a/docs/integration.md +++ b/docs/integration.md @@ -35,6 +35,7 @@ writes to another repository. The broad original PR is superseded by the stack. | 25 | Numeric storage and Serving support contract | | 26 | Runtime-independent ONNX graph summaries | | 27 | One-attempt native cancellation controls | +| 28 | Managed deadlines and caller-exit cancellation | Each PR targets the preceding branch. Merge or retarget in order; the integration branch is for validation rather than a replacement broad PR. CI on each branch @@ -54,7 +55,7 @@ ORTEX_BUILD=true ORT_DYLIB_PATH=/does-not-exist/libonnxruntime.so MIX_ENV=test \ mix run --no-start scripts/check_binary_without_runtime.exs ``` -The assembled stack currently has 206 checks, with two optional ResNet checks +The assembled stack currently has 216 checks, with two optional ResNet checks excluded when that external checkpoint is absent. Its fixture contains only synthetic data. Memory loading checks file/bytes score and descriptor parity, source-owner process exit, malformed-model recovery, and concurrent loading. @@ -127,7 +128,6 @@ alone does not prove those platform or constructor combinations passed. The initial runtime/build foundation is not the whole capability backlog. Continue with separate focused PRs for: -- Managed timeout and caller-exit handling around native cancellation. - Per-run profiling toggles require API 25 and remain deferred beyond the API 17/24 runtime matrix; session profiling is available now. - Dirty-scheduler classification under representative concurrent inference. diff --git a/lib/ortex.ex b/lib/ortex.ex index 429033a..6a147eb 100644 --- a/lib/ortex.ex +++ b/lib/ortex.ex @@ -179,6 +179,15 @@ defmodule Ortex do different process to cancel the inference attempt. Its immutable output/tag options replace per-call `:outputs`/`:tag`, which cannot be passed alongside it. + `:timeout` accepts milliseconds (0 through 4_294_967_295) or `:infinity` + (default). Finite deadlines and explicit controls use a monitored worker; + caller exit or timeout requests native termination. Deadlines include input + transfer, session-mutex waiting, inference, and result conversion, after option + and control preparation. Timeout raises `Ortex.TimeoutError` and bounds the + caller's wait; native termination is cooperative, so a running kernel may + finish later and hold the session until it returns. Ordinary unbounded calls + without a control keep their synchronous execution path. + Ortex.run(model, %{"features" => features}, outputs: ["scores"], tag: "risk-score") """ defdelegate run(model, tensors, opts), to: Ortex.Model diff --git a/lib/ortex/model.ex b/lib/ortex/model.ex index b52be66..69f50ff 100644 --- a/lib/ortex/model.ex +++ b/lib/ortex/model.ex @@ -80,7 +80,24 @@ defmodule Ortex.Model do @doc false def run(%__MODULE__{} = model, inputs, opts) do opts = Ortex.RunOptions.validate!(opts) - Ortex.Telemetry.run(model, inputs, opts, fn -> run_validated(model, inputs, opts) end) + + Ortex.Telemetry.run(model, inputs, opts, fn -> + if opts.control || opts.timeout != :infinity do + control = + opts.control || + Ortex.RunControl.new( + model, + Enum.reject([outputs: opts.outputs, tag: opts.tag], fn {_, value} -> + is_nil(value) + end) + ) + + opts = %{opts | control: control} + Ortex.Runner.run(control, opts.timeout, fn -> run_validated(model, inputs, opts) end) + else + run_validated(model, inputs, opts) + end + end) end defp run_validated(%__MODULE__{reference: reference}, inputs, opts) diff --git a/lib/ortex/run_control.ex b/lib/ortex/run_control.ex index a017c00..af7df76 100644 --- a/lib/ortex/run_control.ex +++ b/lib/ortex/run_control.ex @@ -18,6 +18,10 @@ defmodule Ortex.RunControl do cancellation is harmless and never cancels later runs using another control. A control retains its model and cannot be used with a different model. Selected output buffers are not preallocated or shared between runs. + + Controlled calls use a monitored worker and request cancellation if their + caller exits. Add `timeout: milliseconds` to `Ortex.run/3` to bound caller wait; + the constructor accepts only output selection and tag options. """ @enforce_keys [:reference] defstruct [:reference, :outputs, :tag] @@ -26,7 +30,7 @@ defmodule Ortex.RunControl do def new(%Ortex.Model{} = model, opts \\ []) do options = Ortex.RunOptions.validate!(opts) - if options.control, + if Keyword.has_key?(opts, :control) or Keyword.has_key?(opts, :timeout), do: raise(ArgumentError, "run control options must contain only outputs and tag") case Ortex.Native.new_run_control(model.reference, options) do diff --git a/lib/ortex/run_options.ex b/lib/ortex/run_options.ex index 12bd6e0..5090480 100644 --- a/lib/ortex/run_options.ex +++ b/lib/ortex/run_options.ex @@ -6,7 +6,7 @@ defmodule Ortex.RunOptions do do: raise(ArgumentError, "expected a keyword list of run options") keys = Keyword.keys(opts) - unknown = keys -- [:outputs, :tag, :control] + unknown = keys -- [:outputs, :tag, :control, :timeout] if unknown != [], do: raise(ArgumentError, "unknown Ortex run options: #{inspect(unknown)}") if length(keys) != length(Enum.uniq(keys)), @@ -24,13 +24,17 @@ defmodule Ortex.RunOptions do :control -> is_struct(value, Ortex.RunControl) + + :timeout -> + value == :infinity or (is_integer(value) and value >= 0 and value <= 4_294_967_295) end unless valid, do: raise(ArgumentError, "invalid Ortex run option #{inspect(key)}: #{inspect(value)}") end - options = Map.merge(%{outputs: nil, tag: nil, control: nil}, Map.new(opts)) + options = + Map.merge(%{outputs: nil, tag: nil, control: nil, timeout: :infinity}, Map.new(opts)) if options.control do if Enum.any?(keys, &(&1 in [:outputs, :tag])), diff --git a/lib/ortex/runner.ex b/lib/ortex/runner.ex new file mode 100644 index 0000000..d4f9f55 --- /dev/null +++ b/lib/ortex/runner.ex @@ -0,0 +1,111 @@ +defmodule Ortex.Runner do + @moduledoc false + + # The caller owns a reply alias; late responses are dropped when it expires. + # A separate watcher owns the native worker and cancels it on caller death. + def run(control, timeout, fun) do + deadline = if timeout == :infinity, do: :infinity, else: now() + timeout + if timeout == 0, do: timed_out(control, timeout) + owner = self() + reply = :erlang.alias() + {watcher, monitor} = spawn_monitor(fn -> watch(owner, reply, control, fun) end) + + try do + receive do + {^reply, outcome} -> + if expired?(deadline), do: abort(watcher, reply, control, timeout) + + case outcome do + {:return, value} -> value + {:raise, kind, reason, stack} -> :erlang.raise(kind, reason, stack) + end + + {:DOWN, ^monitor, :process, ^watcher, reason} -> + Ortex.RunControl.cancel(control) + exit({:ortex_run_watcher, reason}) + after + remaining(deadline) -> abort(watcher, reply, control, timeout) + end + after + :erlang.unalias(reply) + Process.demonitor(monitor, [:flush]) + + receive do + {^reply, _} -> :ok + after + 0 -> :ok + end + end + end + + defp watch(owner, reply, control, fun) do + owner_monitor = Process.monitor(owner) + watcher = self() + + {worker, worker_monitor} = + spawn_monitor(fn -> + outcome = + try do + {:return, fun.()} + catch + kind, reason -> {:raise, kind, reason, __STACKTRACE__} + end + + send(watcher, {reply, self(), outcome}) + end) + + try do + action = + receive do + {^reply, ^worker, outcome} -> + send(reply, {reply, outcome}) + :complete + + {:DOWN, ^owner_monitor, :process, ^owner, _reason} -> + :cancel + + {:DOWN, ^worker_monitor, :process, ^worker, reason} -> + send(reply, {reply, {:raise, :exit, {:ortex_run_worker, reason}, []}}) + :cancel + + {:abort, ^reply} -> + :cancel + end + + if action == :cancel, do: cancel_worker(control, worker) + catch + kind, reason -> + stack = __STACKTRACE__ + cancel_worker(control, worker) + :erlang.raise(kind, reason, stack) + after + Process.demonitor(owner_monitor, [:flush]) + Process.demonitor(worker_monitor, [:flush]) + end + end + + defp cancel_worker(control, worker) do + # An executing NIF retains native resources even after its BEAM worker dies. + try do + Ortex.RunControl.cancel(control) + after + Process.exit(worker, :kill) + end + end + + defp abort(watcher, reply, control, timeout) do + send(watcher, {:abort, reply}) + timed_out(control, timeout) + end + + defp timed_out(control, timeout) do + Ortex.RunControl.cancel(control) + raise Ortex.TimeoutError, timeout: timeout + end + + defp now, do: System.monotonic_time(:millisecond) + defp expired?(:infinity), do: false + defp expired?(deadline), do: now() >= deadline + defp remaining(:infinity), do: :infinity + defp remaining(deadline), do: max(deadline - now(), 0) +end diff --git a/lib/ortex/telemetry.ex b/lib/ortex/telemetry.ex index df03fc7..29aee8b 100644 --- a/lib/ortex/telemetry.ex +++ b/lib/ortex/telemetry.ex @@ -11,7 +11,8 @@ defmodule Ortex.Telemetry do `:providers`, and `:optimization_level`. A `nil` provider option means that the session inherits environment providers. Successful loads add the opaque `:model` reference. Run metadata includes the model reference, `:tag`, - `:selected_outputs`, and `:inputs` summaries; successful runs add `:outputs`. + `:selected_outputs`, `:timeout`, `:cancellable`, and `:inputs` summaries; + successful runs add `:outputs`. Summaries include tensor shapes/dtypes or container kind/size, without tensor values or model bytes. Named I/O summaries use binary keys; positional ones use lists in argument/result order. @@ -39,6 +40,8 @@ defmodule Ortex.Telemetry do model: model.reference, tag: options.tag, selected_outputs: options.outputs, + timeout: options.timeout, + cancellable: options.control != nil or options.timeout != :infinity, inputs: summarize(inputs) } diff --git a/lib/ortex/timeout_error.ex b/lib/ortex/timeout_error.ex new file mode 100644 index 0000000..4b622c2 --- /dev/null +++ b/lib/ortex/timeout_error.ex @@ -0,0 +1,6 @@ +defmodule Ortex.TimeoutError do + @moduledoc "Raised when a bounded inference call exceeds its caller deadline." + defexception [:timeout] + @impl true + def message(%{timeout: timeout}), do: "Ortex inference timed out after #{timeout}ms" +end diff --git a/native/ortex/src/lib.rs b/native/ortex/src/lib.rs index 8abc279..43cbd2f 100644 --- a/native/ortex/src/lib.rs +++ b/native/ortex/src/lib.rs @@ -228,7 +228,8 @@ fn take_runtime_logs( rustler::init!("Elixir.Ortex.Native"); -#[rustler::nif(schedule = "DirtyIo")] +// Option preparation must not queue behind inference's dirty-I/O mutex waits. +#[rustler::nif(schedule = "DirtyCpu")] fn new_run_control( model: ResourceArc, options: model::InferenceOptions, diff --git a/test/managed_runs_test.exs b/test/managed_runs_test.exs new file mode 100644 index 0000000..43b665c --- /dev/null +++ b/test/managed_runs_test.exs @@ -0,0 +1,197 @@ +defmodule Ortex.ManagedRunsTest do + use ExUnit.Case, async: true + alias Ortex.RunControl + + defp load(name), do: Ortex.load(Path.expand("fixtures/#{name}.onnx", __DIR__), intra_threads: 1) + defp loop_inputs(n), do: {Nx.tensor(n, type: :s64), Nx.tensor(0, type: :s64)} + + defp wait_started(control, deadline) do + unless RunControl.started?(control) do + if System.monotonic_time(:millisecond) >= deadline, do: flunk("native run was not claimed") + Process.sleep(1) + wait_started(control, deadline) + end + end + + test "bounded named runs preserve output selection and tensor ownership" do + result = + Task.async(fn -> + model = load("named_io") + + Ortex.run(model, %{"z_left" => Nx.tensor([7.0, 5.0]), "a_right" => Nx.tensor([2.0, 1.0])}, + outputs: ["Δ"], + tag: "bounded", + timeout: 5000 + ) + end) + |> Task.await() + + :erlang.garbage_collect() + assert Map.keys(result) == ["Δ"] + assert Nx.to_flat_list(result["Δ"]) == [5.0, 4.0] + end + + test "a real loop times out and releases the session for subsequent inference" do + model = load("cancellable_loop") + + error = + assert_raise Ortex.TimeoutError, fn -> + Ortex.run(model, loop_inputs(1_000_000_000_000), timeout: 5) + end + + assert error.timeout == 5 + {result} = Ortex.run(model, loop_inputs(10), timeout: 5000) + assert Nx.to_number(result) == 10 + end + + test "caller death cancels its native worker and releases the model mutex" do + model = load("cancellable_loop") + control = RunControl.new(model) + on_exit(fn -> RunControl.cancel(control) end) + + {owner, monitor} = + spawn_monitor(fn -> Ortex.run(model, loop_inputs(1_000_000_000_000), control: control) end) + + wait_started(control, System.monotonic_time(:millisecond) + 5000) + Process.exit(owner, :kill) + assert_receive {:DOWN, ^monitor, :process, ^owner, :killed} + {result} = Ortex.run(model, loop_inputs(10), timeout: 5000) + assert Nx.to_number(result) == 10 + end + + test "a queued deadline returns while another attempt still owns the session" do + model = load("cancellable_loop") + control = RunControl.new(model) + on_exit(fn -> RunControl.cancel(control) end) + + owner = + Task.async(fn -> + try do + Ortex.run(model, loop_inputs(1_000_000_000_000), control: control) + rescue + error -> {:error, error} + end + end) + + wait_started(control, System.monotonic_time(:millisecond) + 5000) + assert_raise Ortex.TimeoutError, fn -> Ortex.run(model, loop_inputs(10), timeout: 5) end + assert Process.alive?(owner.pid) + :ok = RunControl.cancel(control) + assert {:error, %RuntimeError{}} = Task.await(owner, 5000) + {result} = Ortex.run(model, loop_inputs(10), timeout: 5000) + assert Nx.to_number(result) == 10 + end + + test "worker exceptions retain their error and inference stack" do + model = load("named_io") + + try do + Ortex.run(model, %{"wrong" => Nx.tensor([1.0, 2.0]), "a_right" => Nx.tensor([3.0, 4.0])}, + timeout: 5000 + ) + + flunk("expected invalid input") + rescue + error in RuntimeError -> + assert error.message =~ "Unknown input" + assert Enum.any?(__STACKTRACE__, fn {module, _, _, _} -> module == Ortex.Model end) + end + end + + test "rejecting a reused control does not cancel its original in-flight attempt" do + model = load("cancellable_loop") + control = RunControl.new(model) + on_exit(fn -> RunControl.cancel(control) end) + + owner = + Task.async(fn -> + try do + Ortex.run(model, loop_inputs(1_000_000_000_000), control: control) + rescue + error -> {:error, error} + end + end) + + wait_started(control, System.monotonic_time(:millisecond) + 5000) + + assert_raise RuntimeError, ~r/already been used/, fn -> + Ortex.run(model, loop_inputs(10), control: control, timeout: 5000) + end + + reference = owner.ref + refute_receive {^reference, _}, 20 + assert Process.alive?(owner.pid) + :ok = RunControl.cancel(control) + assert {:error, %RuntimeError{}} = Task.await(owner, 5000) + end + + test "zero and invalid deadlines are explicit and control construction remains separate" do + model = load("named_io") + input = {Nx.tensor([7.0, 5.0]), Nx.tensor([2.0, 1.0])} + control = RunControl.new(model) + + assert_raise Ortex.TimeoutError, fn -> + Ortex.run(model, input, control: control, timeout: 0) + end + + refute RunControl.started?(control) + assert_raise RuntimeError, ~r/terminate/, fn -> Ortex.run(model, input, control: control) end + + for timeout <- [-1, 4_294_967_296, 1.5, nil, :later] do + assert_raise ArgumentError, fn -> Ortex.run(model, input, timeout: timeout) end + end + + assert_raise ArgumentError, fn -> RunControl.new(model, timeout: 5) end + assert {_, _} = Ortex.run(model, input, timeout: :infinity) + end + + test "timed-out protocol workers terminate and leave no late reply or monitor messages" do + model = load("named_io") + caller = self() + control = RunControl.new(model) + + assert_raise Ortex.TimeoutError, fn -> + Ortex.Runner.run(control, 50, fn -> + send(caller, {:worker, self()}) + Process.sleep(:infinity) + end) + end + + assert_receive {:worker, worker} + monitor = Process.monitor(worker) + assert_receive {:DOWN, ^monitor, :process, ^worker, _}, 5000 + assert_no_late_replies() + end + + test "successful and deadline-racing replies do not accumulate in the caller mailbox" do + model = load("named_io") + + for _ <- 1..25 do + control = RunControl.new(model) + + try do + assert :completed = + Ortex.Runner.run(control, 1, fn -> + Process.sleep(1) + :completed + end) + rescue + Ortex.TimeoutError -> :ok + end + end + + assert_no_late_replies() + end + + defp assert_no_late_replies do + receive do + {reference, _} = message when is_reference(reference) -> + flunk("late reply: #{inspect(message)}") + + {:DOWN, _, :process, _, _} = message -> + flunk("late monitor: #{inspect(message)}") + after + 20 -> :ok + end + end +end diff --git a/test/telemetry_test.exs b/test/telemetry_test.exs index b8a2d95..c1b51e7 100644 --- a/test/telemetry_test.exs +++ b/test/telemetry_test.exs @@ -103,4 +103,20 @@ defmodule Ortex.TelemetryTest do Ortex.run(sequence, {}) assert_receive {[:ortex, :run, :stop], _, %{outputs: [%{kind: :sequence, size: 0}]}} end + + test "deadlines emit a timeout exception in the caller span" do + model = load() + + assert_raise Ortex.TimeoutError, fn -> + Ortex.run(model, {Nx.tensor([1.0, 2.0]), Nx.tensor([3.0, 4.0])}, timeout: 0, tag: "bounded") + end + + assert_receive {[:ortex, :run, :start], _, %{timeout: 0, cancellable: true, tag: "bounded"}} + + assert_receive {[:ortex, :run, :exception], %{duration: d}, + %{reason: %Ortex.TimeoutError{timeout: 0}}} + + assert d >= 0 + refute_receive {[:ortex, :run, :stop], _, _} + end end