Repository navigation
feat(server): shard ORT sessions across owner workers - #2132
Merged
Merged
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2132 +/- ##
===========================================
+ Coverage 72.10% 80.38% +8.27%
===========================================
Files 12 430 +418
Lines 5227 213120 +207893
Branches 5227 213120 +207893
===========================================
+ Hits 3769 171313 +167544
- Misses 1325 36013 +34688
- Partials 133 5794 +5661
Flags with carried forward coverage won't be shown. Click here to find out more. 🚀 New features to boost your workflow:
|
Add a bounded, opt-in ORT session worker count and build every structurally non-Send engine on its own OS thread. Route session and stateless work deterministically with exact placement and in-flight accounting while preserving the historical single-worker path by default. Share only immutable ORT resources across worker factories, fail closed for unsupported backends and execution modes, isolate worker failures, join all workers on shutdown, and expose per-worker status. Update the concurrency architecture and cover real-thread concurrency, placement, failure, teardown, and capability gates. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Copilot <223556219+Copilot@users.noreply.github.com>
Preserve ORT session-before-environment teardown during startup rollback and make every worker charge the same device, host, and disk memory authorities. Move session reservation, close lease/accounting, and embedding permits into queued commands so caller cancellation cannot orphan work or release admission early. Coordinate worker health and counters under one lifecycle lock with checked stale releases. Restore the Python binding's sendable, GIL-releasing behavior through a typed owner-thread command facade while keeping RustEngine structurally thread-affine. Add regression coverage for each reviewed failure mode. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Copilot <223556219+Copilot@users.noreply.github.com>
Retain a read-only handle to the shared model-wide device, host, and disk ledgers and refresh every resource snapshot from those live authorities. Worker allocation, failure, release, and teardown now reach /v1/resources and metrics without requiring worker 0 to execute another command. Add W=2 driver and HTTP regressions plus shared-tier accounting coverage while preserving the existing single-worker snapshot baseline and shared budget enforcement. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Copilot <223556219+Copilot@users.noreply.github.com>
Drop the command-owned model-qualified lease before completing the close response so callers cannot observe success while the session remains transiently reserved. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
force-pushed
the
justinchuby/ort-worker-sharding
branch
from
August 25, 2026 21:10
3468ae4 to
de15ab2
Compare
This was referenced Aug 26, 2026
justinchuby
added a commit
that referenced
this pull request
Aug 26, 2026
Closes #2158. `main` has not compiled `onnx-runtime-python` on any platform since #2132 merged (~5h at time of writing). ## What was actually wrong Not a missing bound and not a `Cargo.toml` drift. **A deleted `unsafe impl`.** `onnx_genai_engine::Engine` is `!Send` deliberately. `WorkerRuntimeState` (`runtime_state.rs:240`) carries `Rc`, `RefCell` and raw ORT handles behind `ThreadBound = PhantomData<*const ()>`, and says why: > `ThreadBound` makes the whole group structurally `!Send` and `!Sync`, so the compiler refuses to let a future pool hand it to a second thread. A `#[pyclass]` must be `Send + Sync` — CPython may call any method from any thread. Both facts have been true since `2d380b722` (Aug 24), and they coexisted only because the engine crate *also* carried, one layer up: ```rust // SAFETY: `Engine` owns every ORT or native-runtime handle reachable through // its sessions and decode state. ... // This would stop being sound if an execution provider introduced a // non-migratable handle or a field gained unsynchronized shared mutation. unsafe impl Send for Engine {} ``` `526f3174c` — the first commit of #2132, "shard ORT sessions across owner workers" — deleted it. **Correctly**: sharding ORT sessions across owner workers is precisely the expiry condition that comment named for itself. The removal is right and this crate was collateral. The full compiler chain, for the record: ``` Rc<WorkflowRuntime> -> BTreeMap<IterationPolicy, Rc<WorkflowRuntime>> -> RefCell<..> -> WorkerRuntimeState -> WorkflowRuntime -> Box<WorkflowRuntime> -> onnx_genai_engine::Engine -> Mutex<..> -> genai::Engine -> assert_pyclass_send_sync ``` ## The fix The engine is no longer stored in the pyclass. `ThreadOwned<T>` builds `T` on a thread of its own and keeps it there; callers hand it `Box<dyn FnOnce(&mut T) -> R + Send>` and get `R` back. **`ThreadOwned<T>` is `Send + Sync` for every `T`, including `!Send` and `!Sync` ones, with no `unsafe` anywhere in the file.** That is ordinary auto-trait derivation, not an assertion: the struct holds no `T` and no `PhantomData<T>`, only a channel whose element type is `Send` because it says so, independent of `T`. `T` is built, used and dropped on the worker and never crosses a thread; only the job closure and its return value do, and both carry their own `Send` bounds. An unsafe assertion is replaced by a structural argument the compiler checks. Re-adding `unsafe impl Send` *here* would have been strictly worse than where it was: this crate cannot see ORT worker affinity at all. ### Held constant - `try_lock` non-re-entrancy and the exact `ENGINE_IN_USE` wording. The mutex is retained as an in-use token, so a second caller is still **refused**, not silently queued. (Mutation-checked: making it queue **hangs** the suite.) - The poisoned-engine wording. - GIL release at all three `py.detach` sites. - The panic payload still reaches Python. The `Mutex` version unwound into the caller, so losing it would have been a debuggability regression; it is appended to the lost-engine message. ### Improved - The engine is dropped **on the thread that owns its ORT handles**, which under the old `unsafe impl` it was not. - `Drop` joins the worker, so release is ordered rather than eventual. ### Costs, stated plainly One OS thread per `Engine`. One channel round-trip per call — noise against `generate`, but real against `tokenize`, which also now copies its input `&str` because the job outlives the borrow. ## Evidence Reproduced first, at `origin/main` `bc79b3ff1`: `cargo check --locked -p onnx-runtime-python` → **40 errors**. Then, on this branch: | gate | result | |---|---| | `cargo check --locked -p onnx-runtime-ep-cuda -p onnx-runtime-python --features onnx-runtime-python/cuda` — *the exact failing CI step* | **pass** | | `cargo test --locked -p onnx-runtime-python --lib` | **pass**, 30 tests | | `cargo fmt --all -- --check` | **pass** | Nine tests added or rewritten. Six-arm mutation battery — a mutant that **hangs** is killed too, since a handle that queues never answers: | arm | expect | got | |---|---|---| | 0 healthy (positive control) | PASS | PASS | | 1 in-use refusal removed (queues) | HANG | HANG | | 2 teardown join removed | FAIL | FAIL — `drop_waits_for_the_value_to_be_released` | | 3 panic message dropped | FAIL | FAIL — `a_panicking_job_reports_a_lost_worker_with_its_message` | | 4 thread-identity probe rigged to the caller | NOBUILD | NOBUILD | | 5 restored (negative control) | PASS | PASS | Arm 4 is the interesting one: you *cannot* write the "built on the calling thread" mutant, because the payload is `!Send`. The property is enforced by the type system, and the arm records that rather than pretending to test it. The `!Send` tripwire carries its own controls (`u32` must probe `true`, `Rc<()>` must probe `false`), so a probe that silently always returned `false` would fail rather than pass. ## Independent review, and what it changed An independent adversarial review (Opus) went at the soundness claim, deadlock, teardown, panic path, parity and the tests. It confirmed the central claim and found **one real bug**, which is fixed in `55b2df628`. **The bug.** The worker wrote the panic message to a shared `last_panic` slot *after* `catch_unwind` returned, but the caller was unblocked earlier than that — by the reply sender being dropped during the unwind. Two unordered paths, so the call that *caused* the panic could read the slot first and get `None`, losing exactly the message this design set out to preserve. Now the job catches its own panic and sends the message back on the reply channel it was already going to answer on. One path, no ordering question. `last_panic` remains only for later callers, who observe it through the channel close and are therefore already ordered after the write. Proved with a 50 ms delay injected before the `last_panic` write, to make the race deterministic instead of hoping to catch it: | arm | expect | got | |---|---|---| | A0 fixed, no delay | PASS | PASS | | A1 fixed, +50 ms delay | PASS | PASS | | B pre-review design, +50 ms delay | FAIL | **FAIL** | | C pre-review design, no delay | PASS | PASS | **Arm C is the finding behind the finding.** The old design passes 200 repetitions under eight spinning cores. A statistical stress run could not have found this — and mine did not. Worse, my first attempt at this battery filtered with `--exact` against bare test names, matched nothing, and printed `test result: ok. 0 passed`, which my classifier read as PASS. Five vacuous green arms. The battery above asserts the *number of tests that ran*, not the word `ok`. Two smaller findings, also fixed: - A panicking constructor was reported as `SpawnError::Thread` — naming a subsystem that had not failed, since the thread started fine — and discarded the payload. Now `SpawnError::BuildPanicked`, carrying the message, mapped to the model-loading error. - The worker thread was named `nxrt-genai-engine`: 17 bytes against Linux's 15-plus-NUL `TASK_COMM_LEN`. `pthread_setname_np` fails, std ignores the failure, and the name silently never appeared in `top`/`gdb`/`/proc`. Now `nxrt-genai`. Verified fine by the review, for the record: no `T` or `&T` escapes the worker (the `R: Send` bound on the reply is what closes it); callback re-entrancy cannot deadlock because re-entry arrives on the *worker* thread and `try_lock` is cross-thread; `Drop` cannot race an in-flight job because pyo3 drops a `#[pyclass]` only at refcount zero; blocked callers never hold the GIL, so `Python::attach` on the worker always makes progress; and the `!Send` `Probe<T>` tripwire is correct on stable. ## Two Python-observable changes I had described as unchanged The review was right that these were understated, so they are stated here rather than buried: 1. **A first panic is now `RuntimeError`, not `PanicException`.** Previously a panic inside `generate` unwound through `#[pymethods]` into pyo3's `PanicException`, which subclasses `BaseException` — so `except Exception:` did **not** catch it. It now returns the lost-engine `RuntimeError`, which `except Exception:` **does** catch. Arguably an improvement, but it is a semantic change. 2. **The streaming callback runs on the engine's worker thread**, not the calling thread. Observable via `threading.get_ident()`, thread-locals, and anything bound to the calling thread. ## Limitations 1. **These tests do not run in CI.** `onnx-runtime-python` is on the test denylist — *"PyO3 extension crate requires wheel packaging/runtime DLL staging"* — so it is in no lane: `offline-linux` absent, `ort-backed` absent, `lint` absent (checked with `python3` and a positive control; `python` does not exist on the CI image, and a `python … | grep` gate returns empty, which reads as its own alarm value). Its only CI coverage is the advisory `CUDA compile` lane's `cargo check`. The load-bearing property is compile-enforced anyway: `#[pyclass]` emits `assert_pyclass_send_sync`, which is exactly the assertion that broke. 2. **`CUDA compile` is not a required check**, so this PR being green there is advisory. That is the same gap that let #2132 land; tracked in #1966, not re-litigated here. 3. `clippy -p onnx-runtime-python -D warnings` still reports one `vec_init_then_push` at `lib.rs:68`. **Pre-existing and untouched** — byte-identical on `origin/main`, a false positive caused by a `#[cfg(feature = "cuda")]` push. Left alone to keep this diff to the break. 4. I have not run the Python-level API end to end (no model directory here). The three `py.detach` sites and the streaming callback's `Python::attach` from a non-GIL thread are compile-checked and argued, not executed. ## Note on the earlier attribution I posted a bisect naming #2132 and separately said `2d380b722` was "refuted". Re-measured from CI's own step records, both hold, but the emphasis was wrong: `2d380b722` is the **latent enabler** — it made the type genuinely `!Send`; the `unsafe impl` is the only reason that survived four days. | commit | Linux compile step | Windows compile step | |---|---|---| | `52332c7c2` | success (job failed later at step 8, capture-sync = #2149) | success | | `66ca26a5d` | **failure** (step 6) | **failure** (step 6) | And a correction to myself: I had reported the mechanism as unpinned because "no `Rc`/`PhantomData`/`NonNull` was *added* in the window". That was my instrument, not the tree — I searched `git log -S` for **added** types, and a deleted `unsafe impl` is invisible to that search.⚠️ Flagged for squad review: this touches another crate's Python surface and a thread-ownership model. The mechanism is measured; the design choice is arguable. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 26, 2026
…pilation (#2243) Four workspace members are linted by **nothing**, because a *runtime* exclusion was silently doing duty as a *compile* exclusion. One of them was uncompilable on `main` for 22 commits earlier today and no required check said so. ## The mechanism #2058 already established the principle: `cargo clippy` only ever **checks** — it never links and never runs a test binary — so "needs a GPU", "needs wheel staging", "isn't a unit-test target" have no force over it. That PR restored lint coverage to 21 *tested* crates that had lost it for a reason that did not apply. The argument was never extended to the **denied** crates themselves. The `lint` lane stayed: ```python selected = packages - denied ``` and every one of the five `DENYLIST` reasons is a runtime reason: | crate | DENYLIST reason | kind of reason | |---|---|---| | `onnx-genai-ort-sys` | build script downloads a native ORT distribution | build | | `onnx-genai-bench` | benchmark/perf crate; not a unit-test CI target | **execution** | | `onnx-runtime-ep-cuda` | runtime tests require a CUDA GPU | **execution** | | `onnx-genai-python` | requires wheel packaging/runtime DLL staging | **execution** | | `onnx-runtime-python` | requires wheel packaging/runtime DLL staging | **execution** | Measured on `main` before this PR — 4 of 60 members reached by no clippy invocation anywhere: ``` linted by NOTHING anywhere : ['onnx-genai-bench', 'onnx-genai-ort-sys', 'onnx-genai-python', 'onnx-runtime-python'] ``` ## Why it matters, concretely `onnx-runtime-python` did not compile on `main` from `66ca26a5d` (#2132) to `881c3165c` (#2173) — **22 commits, ~7.5 hours**. The only job that builds that crate is `CUDA compile`, via a hand-written `-p onnx-runtime-python` outside the resolver. It is **not a required check**. On the breaking merge commit it reported exactly what it should have: ``` CUDA compile (Linux x86_64) = failure CUDA compile (Windows x86_64) = failure ``` Required checks on `main` are only `Fast (Linux x86_64)` and `Rust quality`. **The signal existed, was correct, and was discarded.** This PR puts that same breakage in front of a required check. ## The two-arm proof I reverted #2173 to reintroduce the real break, and ran both lane definitions against that identical tree: | lane | packages | rc | `E0277` | |---|---|---|---| | **old** (`packages - DENYLIST`) | 55 | **0 — green** | 0 | | **new** (`packages - LINT_EXEMPT`) | 58 | **101 — red** | 78 | The old lane is green **not because it compiled the crate and found it sound, but because the broken crate was never in its input set.** Both are "green"; only one of them is evidence. ## The change - New `LINT_EXEMPT` set, separate from `DENYLIST`, whose entries must name a reason **clippy itself cannot get past**. The docstring says plainly that "needs a GPU at runtime" / "isn't a unit-test target" / "needs wheel staging" do not qualify — that is the #2058 mistake with new nouns. - `lint` lane is now `packages - LINT_EXEMPT` (55 → **58**). - The lint-coverage guard's domain widens from `tested - linted` to `(packages - LINT_EXEMPT) - linted`. **The old domain could not name a denied crate by construction** — the guard whose entire job is to notice unlinted crates was structurally unable to see these four. - Stale-entry validation for `LINT_EXEMPT`, matching the existing `DENYLIST` check. - Fixes the one clippy error this exposed: a cfg-conditional `vec_init_then_push` in `onnx-runtime-python`, rewritten as two mutually exclusive bindings with identical semantics. ### What stays exempt, and why that is honest - **`onnx-runtime-ep-cuda`** — its build script needs a CUDA toolchain the quality runner does not install. It passed clippy locally, but **this box has `nvcc`**, so a local pass here is not evidence about the runner. Recorded as such. It is already linted by the CUDA job. - **`onnx-genai-bench`** — 28 dead-code findings in `src/bin/compare.rs`; an entire `Direct*` reporting path that nothing constructs. Filed as #2242, including the possibility that matters more than the lint: if that path was *meant* to be wired up, benchmark output someone believes is produced is not being produced. Exempted on explicitly temporary grounds rather than left in `DENYLIST` where the exclusion looks principled. `onnx-genai-ort-sys` costs nothing new: six of its dependents are already in the lane, so its build script already runs there. Confirmed no download occurs. ## Tests Two self-test arms; the first is a real discriminator, not a restatement: - `verify --simulate-unlinted onnx-runtime-python (a denied crate)` — under the old domain this returned **0 and named nobody**, because a denied crate could not enter `tested - linted`. Verified by reverting the one-line domain change and re-running: ``` FAIL verify --simulate-unlinted onnx-runtime-python (a denied crate) expected the report to name: ['onnx-runtime-python'] ``` It fails against the previous code and passes against this one. - `verify --simulate-unlinted onnx-genai-bench (lint-exempt)` pins the other direction, so the widened domain cannot quietly grow to include crates the runner cannot build. It correctly does **not** flip between the two versions. Suite is self-counting (`82/82 arms behaved as stated`), so a miscount fails loudly rather than under-reporting. ## Validation ``` workspace test package coverage ok: 55 tested, 5 denied workspace lint coverage ok: 58 linted by 11 clippy invocation(s), 2 lint-exempt workspace test package self-test: 82/82 arms behaved as stated required-lane coverage ok: 55 package(s) tested by ['Fast (Linux x86_64)', 'Rust quality'] cargo clippy --locked --all-targets <58 packages> -- -D warnings -> rc=0 ``` The clippy run was checked for having done real work rather than trusting `rc=0`: 49 compile/check units, with `onnx-runtime-python`, `onnx-genai-python` and `onnx-genai-ort-sys` each confirmed present by name in the log. ## Limitations - This closes the gap for crates clippy can compile on the Linux quality runner. It does **not** make `CUDA compile` required — that is a policy call for the repo owner, and it remains the only job compiling `onnx-runtime-ep-cuda`'s CUDA paths. - `onnx-genai-bench` remains unlinted until #2242 is resolved. - Lane grows by 3 packages; the incremental clippy cost was ~17s locally. Local validation is necessary and not sufficient — merging on required CI via normal auto-merge, no admin, no bypass. Co-authored-by: gaff <gaff@squad.local> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
--ort-session-workers/ONNX_GENAI_ORT_SESSION_WORKERSconfiguration (default1, range1..=64)Sendengine on its owner OS thread using the existing factory readiness handshakeArc<Session>resources while keeping KV state, bindings, allocators, values, samplers, diagnostics, counters, and graph-capture state worker-localSESSION_CONCURRENCY.mdand add real-thread/barrier coverageBuilds on the owner-thread factory work from #2111 and shared ORT session ownership from #2122.
Review hardening
The follow-up at
445d4cfd6addresses the independent review findings:RustEngineoff its construction threadLive resource accounting follow-up
The incremental fix at
3468ae46fretains a live read-only handle to the shared device/host/disk ledgers./v1/resources, metrics, and VRAM override responses now recompute every tier's usage and headroom. Worker-1 allocation and release/failure/teardown are visible without requiring worker 0 to execute another command. Real W=2 driver and HTTP regressions cover allocation and release.Supported scope
W=1remains the default and follows the historical load/batching path.W>1is supported only for a contracted single-decoder ORT workflow where every relevant ORT session reports concurrentRunsupport.This PR makes no broader concurrency claim. Requested
W>1fails closed with typed actionable load errors for:RunsupportThere is no silent fallback to one worker, cross-worker session migration, cross-worker batching, or automatic worker restart.
Python compatibility
The low-level
RustEngineremains structurally thread-affine. The Python pyclass is sendable because it contains only a typed owner-thread command façade; construction, execution, and teardown remain on one OS thread. Generation/tokenization release the GIL while waiting, and concurrent/reentrant calls retain the existing immediateengine is in useerror.Validation
cargo test -p onnx-genai-server --lib --locked --quiet -- --test-threads=1— 336 passed, 2 ignoredcargo test -p onnx-genai-server --test http --locked --quiet— 40 passed, 1 ignoredcargo test -p onnx-genai-engine --lib --locked --quiet— 499 passed, 1 ignoredcargo test -p onnx-genai-python --lib --locked --quiet— 5 passedW>1fail-closed test — passedcargo fmt --all -- --checkcargo clippy -p onnx-genai-server -p onnx-genai-engine --all-targets --locked -- -D warningscargo clippy -p onnx-genai-server --all-targets --features native-backend --locked -- -D warningsEarlier ORT session/decode, multi-session, workflow-conformance, and Python suites remain green. A workspace-wide check remains blocked by the unpopulated
onnx-runtime-cpuinfo/vendor/cpuinfosubmodule.