Bound managed inference and cancel work on caller exit

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.
 

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.

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.

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

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)

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

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])),

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

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)
     }
 

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

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<model::OrtexModel>,
     options: model::InferenceOptions,