From a4cef8490858d89b7827ef6e62ba652107434373 Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Thu, 29 Oct 2026 04:12:40 -0700 Subject: [PATCH] Reserve run ownership before installing worker deadlines --- CHANGELOG.md | 3 +- docs/cancellation.md | 12 +++++- docs/integration.md | 9 ++++- lib/ortex/model.ex | 6 +++ lib/ortex/native.ex | 1 + lib/ortex/run_control.ex | 2 +- native/ortex/src/lib.rs | 11 ++++++ native/ortex/src/run_control.rs | 37 +++++++++++++++---- test/managed_runs_test.exs | 65 ++++++++++++++++++++++++++++++++- test/run_control_test.exs | 2 +- 10 files changed, 133 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5163e5d..56ab8c5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,10 +22,11 @@ - 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. +- Reserve run ownership before deadlines so duplicate calls cannot cancel another attempt. - Support checksum-verified precompiled CPU NIFs and source-build overrides. - Default to dynamic runtime loading and API 17, without runtime build headers. - Build and validate all seven artifacts, checksum maps, and a package in CI. - Stage attested draft releases only in RiverFinancial/ortex; no Hex publisher. -Upgrade in a fresh VM; native resource layouts changed (including environment V2). +Upgrade in a fresh VM; native resource representations changed (including environment and run-control V2). diff --git a/docs/cancellation.md b/docs/cancellation.md index 3da3dad..bd063b1 100644 --- a/docs/cancellation.md +++ b/docs/cancellation.md @@ -3,8 +3,8 @@ `Ortex.RunControl.new(model, outputs: [...], tag: "...")` prepares immutable options for one model-bound inference attempt. Pass it with `control: control` to `Ortex.run/3`. A different process can request termination with -`Ortex.RunControl.cancel/1`, and `started?/1` reports whether the attempt has been -claimed, including waiting for the session mutex. A claimed control cannot be +`Ortex.RunControl.cancel/1`, and `started?/1` reports whether the native attempt +has been entered, including waiting for the session mutex. A reserved control cannot be reused even if inference fails. Invalid options are rejected before an attempt. Termination calls ONNX Runtime's run-options flag without acquiring the session @@ -49,6 +49,14 @@ 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. +Ownership is atomically reserved in the caller before a watcher or deadline is +installed. Native entry then consumes that reservation exactly once. This also +rejects duplicates while the original worker has not yet reached its NIF: a +zero-deadline or exiting duplicate cannot cancel the original attempt. The native +control resource is V2; a fresh VM is required when upgrading from V1 controls. +`started?/1` remains false for a reserved attempt until native entry, including +entry that waits for the model mutex. + 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 diff --git a/docs/integration.md b/docs/integration.md index ba0d6ac..73bc376 100644 --- a/docs/integration.md +++ b/docs/integration.md @@ -38,6 +38,7 @@ writes to another repository. The broad original PR is superseded by the stack. | 28 | Managed deadlines and caller-exit cancellation | | 29 | Dirty-scheduler measurement and policy | | 30 | Final lifecycle/managed-run soak and review ledger | +| 31 | Atomic control ownership before worker/deadline creation | Each PR targets the preceding branch. Merge in order with merge commits to preserve ancestry. If using squash/rebase merges, rebase and retarget the remaining @@ -59,7 +60,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 216 checks, with two optional ResNet checks +The assembled stack currently has 218 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. @@ -142,6 +143,12 @@ validation PR records exact integration run links without requiring a release ta Runtime cancellation is cooperative and may finish after the caller deadline. Accelerator runtime execution remains unvalidated without its libraries/hardware. +The final ownership fix reserves a control before creating its watcher or deadline. +Regression tests cover a zero-deadline duplicate before native entry, wrong-model +reuse, and 16 concurrent public reservations with exactly one successful owner. +After the fix, another 100,000 bounded calls per runtime retained checked scores, +bounded sampled memory, and flat process counts. Run-control resources are V2. + ## Deferred scope The stack covers the selected CPU/API 17–24 modernization, session/environment diff --git a/lib/ortex/model.ex b/lib/ortex/model.ex index 69f50ff..870423e 100644 --- a/lib/ortex/model.ex +++ b/lib/ortex/model.ex @@ -93,6 +93,12 @@ defmodule Ortex.Model do ) opts = %{opts | control: control} + + case Ortex.Native.reserve_run(model.reference, control.reference) do + :ok -> :ok + {:error, message} -> raise message + end + Ortex.Runner.run(control, opts.timeout, fn -> run_validated(model, inputs, opts) end) else run_validated(model, inputs, opts) diff --git a/lib/ortex/native.ex b/lib/ortex/native.ex index ef0b2bc..d399cf7 100644 --- a/lib/ortex/native.ex +++ b/lib/ortex/native.ex @@ -40,6 +40,7 @@ defmodule Ortex.Native do def new_run_control(_model, _options), do: :erlang.nif_error(:nif_not_loaded) def cancel_run(_control), do: :erlang.nif_error(:nif_not_loaded) def run_started(_control), do: :erlang.nif_error(:nif_not_loaded) + def reserve_run(_model, _control), do: :erlang.nif_error(:nif_not_loaded) def inference_scheduler, do: :erlang.nif_error(:nif_not_loaded) def run_controlled(_model, _inputs, _control), do: :erlang.nif_error(:nif_not_loaded) def run_named_controlled(_model, _inputs, _control), do: :erlang.nif_error(:nif_not_loaded) diff --git a/lib/ortex/run_control.ex b/lib/ortex/run_control.ex index af7df76..508c92f 100644 --- a/lib/ortex/run_control.ex +++ b/lib/ortex/run_control.ex @@ -47,6 +47,6 @@ defmodule Ortex.RunControl do end end - @doc "Returns whether the single attempt has been claimed, including waiting for the session mutex." + @doc "Returns whether the native attempt has been entered, including waiting for the session mutex." def started?(%__MODULE__{reference: reference}), do: Ortex.Native.run_started(reference) end diff --git a/native/ortex/src/lib.rs b/native/ortex/src/lib.rs index 05dba7c..9812542 100644 --- a/native/ortex/src/lib.rs +++ b/native/ortex/src/lib.rs @@ -278,6 +278,17 @@ fn run_started(control: ResourceArc) -> bool { control.started() } +#[rustler::nif] +fn reserve_run( + model: ResourceArc, + control: ResourceArc, +) -> NifResult { + control + .reserve(&model) + .map_err(|e| rustler::Error::Term(Box::new(e.to_string())))?; + Ok(rustler::types::atom::ok()) +} + #[cfg_attr(feature = "benchmark-dirty-cpu", rustler::nif(schedule = "DirtyCpu"))] #[cfg_attr( not(feature = "benchmark-dirty-cpu"), diff --git a/native/ortex/src/run_control.rs b/native/ortex/src/run_control.rs index 304e31f..d7d7db6 100644 --- a/native/ortex/src/run_control.rs +++ b/native/ortex/src/run_control.rs @@ -3,7 +3,7 @@ use crate::model::{self, InferenceOptions, OrtexModel}; use ort::session::{HasSelectedOutputs, OutputSelector, RunOptions}; use ort::Error; use rustler::{Resource, ResourceArc}; -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::{AtomicU8, Ordering}; pub struct PreparedRun { names: Vec, @@ -58,11 +58,11 @@ impl PreparedRun { pub struct RunControl { model: ResourceArc, prepared: PreparedRun, - claimed: AtomicBool, + state: AtomicU8, } // Encoding only copies the opaque resource handle; it cannot mutate run options. impl std::panic::RefUnwindSafe for RunControl {} -#[rustler::resource_impl(name = "OrtexRunControlV1")] +#[rustler::resource_impl(name = "OrtexRunControlV2")] impl Resource for RunControl {} impl RunControl { @@ -71,15 +71,36 @@ impl RunControl { Ok(Self { model, prepared, - claimed: AtomicBool::new(false), + state: AtomicU8::new(0), }) } - pub fn claim(&self, model: &OrtexModel) -> Result<(), Error> { + fn check_model(&self, model: &OrtexModel) -> Result<(), Error> { if !std::ptr::eq(&*self.model, model) { return Err(Error::new("run control belongs to another model")); } - self.claimed - .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + Ok(()) + } + // Reserve ownership before spawning a worker or installing its deadline. + // A rejected duplicate must never acquire authority to cancel another call. + pub fn reserve(&self, model: &OrtexModel) -> Result<(), Error> { + self.check_model(model)?; + self.state + .compare_exchange(0, 1, Ordering::AcqRel, Ordering::Acquire) + .map_err(|_| Error::new("run control has already been used"))?; + Ok(()) + } + pub fn claim(&self, model: &OrtexModel) -> Result<(), Error> { + self.check_model(model)?; + // Direct internal NIF calls can enter from ready (0); public managed + // calls enter from reserved (1). Neither can execute twice (2). + self.state + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |state| { + if state < 2 { + Some(2) + } else { + None + } + }) .map_err(|_| Error::new("run control has already been used"))?; Ok(()) } @@ -90,6 +111,6 @@ impl RunControl { self.prepared.options.terminate() } pub fn started(&self) -> bool { - self.claimed.load(Ordering::Acquire) + self.state.load(Ordering::Acquire) == 2 } } diff --git a/test/managed_runs_test.exs b/test/managed_runs_test.exs index 43b665c..ea9cf75 100644 --- a/test/managed_runs_test.exs +++ b/test/managed_runs_test.exs @@ -135,7 +135,10 @@ defmodule Ortex.ManagedRunsTest do end refute RunControl.started?(control) - assert_raise RuntimeError, ~r/terminate/, fn -> Ortex.run(model, input, control: control) end + + assert_raise RuntimeError, ~r/already been used/, 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 @@ -145,6 +148,66 @@ defmodule Ortex.ManagedRunsTest do assert {_, _} = Ortex.run(model, input, timeout: :infinity) end + test "a duplicate deadline cannot cancel a reserved attempt before its native entry" do + model = load("named_io") + control = RunControl.new(model, outputs: ["sum"]) + :ok = Ortex.Native.reserve_run(model.reference, control.reference) + refute RunControl.started?(control) + input = {Nx.tensor([7.0, 5.0]), Nx.tensor([2.0, 1.0])} + + assert_raise RuntimeError, ~r/already been used/, fn -> + Ortex.run(model, input, control: control, timeout: 0) + end + + refs = + input + |> Tuple.to_list() + |> Enum.map(fn tensor -> + %Nx.Tensor{data: %Ortex.Backend{ref: ref}} = Nx.backend_transfer(tensor, Ortex.Backend) + ref + end) + + [{:tensor, {ref, [2], :f, 32}}] = + Ortex.Native.run_controlled(model.reference, refs, control.reference) + + assert Ortex.Native.to_binary(ref, 32, 2) |> Nx.from_binary(:f32) |> Nx.to_flat_list() == [ + 9.0, + 6.0 + ] + end + + test "concurrent public reservations permit one successful owner" do + model = load("named_io") + control = RunControl.new(model, outputs: ["sum"]) + + results = + 1..16 + |> Task.async_stream( + fn _ -> + try do + {sum} = + Ortex.run(model, {Nx.tensor([7.0, 5.0]), Nx.tensor([2.0, 1.0])}, + control: control, + timeout: 5000 + ) + + {:success, Nx.to_flat_list(sum)} + rescue + error in RuntimeError -> {:rejected, error.message} + end + end, + max_concurrency: 16 + ) + |> Enum.to_list() + + assert Enum.count(results, &(&1 == {:ok, {:success, [9.0, 6.0]}})) == 1 + + assert Enum.count(results, fn + {:ok, {:rejected, message}} -> message =~ "already been used" + _ -> false + end) == 15 + end + test "timed-out protocol workers terminate and leave no late reply or monitor messages" do model = load("named_io") caller = self() diff --git a/test/run_control_test.exs b/test/run_control_test.exs index 21f54e8..0c25468 100644 --- a/test/run_control_test.exs +++ b/test/run_control_test.exs @@ -69,7 +69,7 @@ defmodule Ortex.RunControlTest do control = RunControl.new(model) assert_raise RuntimeError, ~r/another model/, fn -> - Ortex.run(other, inputs(), control: control) + Ortex.run(other, inputs(), control: control, timeout: 0) end assert_raise ArgumentError, fn ->