fix(engine): admit interpreted no-decode-core generation through the scheduler

[justinchuby]

Root cause

#1723 ("one runtime, one interpreter, one drive") changed generate_with_callbacks's top-level dispatch from checking the decode_backend field to checking holds_decode_core() (native_session.is_some() || session.is_some()).

Pre-#1723, every Native-backend request — regardless of whether a decode session already existed — went through the scheduler-admitting cold-start path (generate_native_cold_with_callback), which checks the KV byte budget before touching any backend.

Post-#1723, a Native-backend Engine with no session now silently falls into the unguarded interpreted path (generate_interpreted → generate_with_pipeline_callbacks → run_declared_generation), skipping scheduler admission entirely. This is a real, reachable state: Engine::from_dir → Self::decode_core_covers(&workflow) false → from_interpreted_dir → Engine::from_workflow is a legitimate, doc-commented construction path ("a package whose components the interpreter invokes"), and WorkflowRuntime::decode_backend() can independently be Native for such a package. Today no shipping package pairs this with a byte-budget-configured scheduler (Engine::from_workflow always uses Scheduler::new(SchedulerConfig::default())), so the practical blast radius is currently zero — but the gap is real and was caught by a pre-existing test, native_generate_rejects_over_kv_byte_budget_before_backend_run, which was failing on main with the wrong (deep, non-admission) error message instead of a clean scheduler rejection.

Fix

• Added admit_interpreted_generate_request (engine/runtime.rs), which computes the prompt token count via a new interpreted_prompt_token_count (handling TokenIds by length, TokenRows by rows.len() * max(row length) — since equal-length rows bind into one batched [rows, columns] tensor and run as a single graph step — and Text via the package's own tokenizer) and calls the existing admit_generate_request_with_scheduler.
• Called this from generate_with_pipeline_callbacks's (engine/workflow_api.rs) own no-decode-core, prompt-only branch — the single place every such request converges, whether it arrives cold (generate_interpreted) or through a continuing workflow session (generate_in_workflow_session, which threads the caller's own session id through so admission accumulates the same way the decode-core path's session continuation already does).
• Tensor-bound (multimodal) requests are deliberately left unchanged, matching pre-existing, unrelated precedent (CLI interactive.rs/transcribe.rs, server driver.rs) that skips admission regardless of decode-core status.

Before / after evidence

Before the fix, native_generate_rejects_over_kv_byte_budget_before_backend_run failed with:

workflow component 'decoder' input 'past.key' references unavailable value 'decoder_kv.0000'

(a deep node-execution failure, not a scheduler rejection). After the fix it passes cleanly with "scheduler admission failed: KV byte budget...".

A new test, native_generate_in_workflow_session_rejects_over_kv_byte_budget, exercises the session-continuation path directly (opens a real session via create_session(), drives generate_in_workflow_session()). I confirmed by temporarily reverting workflow_api.rs alone that this test fails against the pre-fix code with the same deep references unavailable value error, and passes after the fix — closing a gap an independent review pass on an earlier version of this change flagged (that version only guarded the cold-call path, missing session continuation, and undercounted TokenRows).

Review

Independently reviewed twice: once on an initial version (which lived inline in generate_interpreted and only covered the cold-call path — REQUEST CHANGES, two findings), and once on this final restructured version, which addresses both findings by moving the admission call site to the shared no-decode-core branch of generate_with_pipeline_callbacks and fixing the TokenRows formula. Second pass: APPROVE, after tracing control flow, confirming scheduler.complete() runs on every exit path, confirming no other caller bypasses the new admission branch, and independently reproducing the before/after test evidence.

Tests (all local; no CI wait per current session directive)

• cargo test -p onnx-genai-engine --features native-backend --lib: 618 passed, 0 failed, 1 ignored
• cargo test -p onnx-genai-engine --features native-cuda --lib: 648 passed, 0 failed, 4 ignored
• authored_body_selects_executor (4), one_runtime_e2e (6), native_workflow_parity (13), native_workflow_smoke (2), canonical_execution_parity (7): 32/32 passed
• cargo fmt --check -p onnx-genai-engine: clean
• cargo clippy -p onnx-genai-engine --features native-cuda --all-targets -- -D warnings: clean
• GPU e2e on idle A100 (CUDA_VISIBLE_DEVICES=1, --test-threads=1):
  • deepseek_v4_tiny_qmoe_e2e: 4/4, captures=1 replays=10 fallbacks=0
  • glm_tiny_qmoe_native_cuda_e2e: 4/4
  • glm_tiny_full_attention_e2e: 4/4
  • deepseek_v2_tiny_qmoe_native_e2e: 2/2

Scope

No workflow-unification refactors, no CUDA residency/paging/scheduler changes, no DeepSeek/GLM export/model-name branching. Scoped strictly to closing the #1723-introduced admission gap in the shared interpreted-path runtime code.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 80.61%. Comparing base (1be9f2c) to head (035ef9b).
⚠️ Report is 7 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1900      +/-   ##
==========================================
- Coverage   80.84%   80.61%   -0.24%     
==========================================
  Files         415      414       -1     
  Lines      204440   200009    -4431     
  Branches   204440   200009    -4431     
==========================================
- Hits       165286   161233    -4053     
+ Misses      33582    33239     -343     
+ Partials     5572     5537      -35     

Flag	Coverage Δ
cli-ort-linux	72.51% <ø> (ø)
cli-ort-windows	72.01% <ø> (-0.10%)	⬇️
mlas	85.33% <ø> (+0.13%)	⬆️
offline	80.75% <ø> (-0.25%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.
see 43 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[justinchuby]

Post-merge review: the interpreted branch is closed correctly, but on_admitted() still fires unadmitted on the sibling branch

Reviewed after merge (cd99e5ceb). The fix for the no-decode-core prompt path is right, and one detail of it is better than it needed to be. There is a second path in the same function that this PR didn't touch, and I think this PR's own reasoning is the argument for closing it.

What's right

admit_interpreted_generate_request at :212 plus self.scheduler.complete(scheduler_session_id) at :223 closes the branch as described. The session-id selection — reuse when workflow_sessions already names the conversation, mint otherwise — is the correct shape, since admitting a continuing workflow under a fresh id each turn would make the budget accounting meaningless.

The detail worth calling out is the ordering at :222-225:

let result = self.run_declared_workflow_generation(request, callback);
self.scheduler.complete(scheduler_session_id);
let mut result = result?;

complete() runs before the ?, so a generation error still releases the scheduler slot. Binding the result and deferring the ? is deliberate; the natural way to write this (let result = ...?; then complete()) leaks a slot on every error path, and a leaked-slot bug only shows up under a failure that also has to be concurrent, which is close to untestable after the fact. Good.

The gap: :185-190

if !prompt_only {
    if let Some(on_admitted) = on_admitted.as_mut() {
        on_admitted();
    }
    return self.run_declared_workflow_generation(request, callback);
}

on_admitted() fires here, and nothing admitted. Verified rather than assumed — the whole file contains exactly two scheduler interactions, both added by this PR:

$ grep -nE "admit_|scheduler\.(admit|complete|reserve)" workflow_api.rs
212:        let (budget_cap, max_new_tokens) = self.admit_interpreted_generate_request(
223:        self.scheduler.complete(scheduler_session_id);

and the callee is a tokenize-then-delegate wrapper (:237-258) that reaches pipeline::generation::run_declared_generation against &*self.workflow, never touching self.scheduler.

The contract makes this a false signal rather than a cosmetic one. on_admitted means "the scheduler admitted this request" — that is precisely what runtime.rs:2796 leans on:

let mut on_admitted = || admitted = true;

inside native_generate_rejects_over_kv_byte_budget_before_backend_run, whose whole point is that an over-budget request must not set the flag. So on the !prompt_only path a caller is told the scheduler accepted a request the scheduler never saw. That is strictly worse than not calling the callback: an unadmitted request that stays silent fails somewhere honest, while one that announces admission defeats a caller using the callback as its gate.

To be fair to the current state, the missing complete() on that branch is consistent — nothing was admitted, so no slot leaks. The branch is coherent as "outside the scheduler." It just doesn't say that.

Why this PR's reasoning applies to it

The comment justifying the new admission is:

a no-decode-core prompt still shares this runtime's scheduler, and — when the process configured one — its KV byte budget: components own their own caches, but the byte accounting a shared budget protects is shared regardless of which executor a step names

That sentence is not specific to prompts. A tensor-binding declared workflow runs against the same runtime and the same configured budget; "regardless of which executor a step names" covers it explicitly. The rationale for closing the branch this PR closed is an argument for the branch it left open — which is why I'd rather raise it than assume it was a considered exclusion.

Either resolution is defensible and both are small:

1. Admit it too — mirror :206-225, including the complete()-before-? ordering.
2. Declare it ungoverned — drop the on_admitted() call at :187 and say in the comment that declared workflows binding their own tensors are outside this scheduler's remit.

What isn't defensible is the present combination: option 2's admission with option 1's signal.

Falsifier

The existing test is the right template. Point it at the !prompt_only path with an over-budget request and assert the flag stays clear:

let mut admitted = false;
let mut on_admitted = || admitted = true;
// PipelineGenerateRequest with non-empty `inputs`, prompt over bytes_per_token budget
let _ = engine.generate_with_pipeline_callbacks(req, Some(&mut on_admitted), None);
assert!(!admitted, "declared-workflow generation must not report admission it never obtained");

Today that asserts false — admitted is true and the request runs unbudgeted. Under fix (1) it passes because admission rejects; under fix (2) it passes because the callback never fires. So it's a fair test of the property rather than of either implementation, and it fails now for the right reason.

I flagged this same second site when I reviewed #1891 before merge, so this is a re-raise rather than a new finding — recording it here because #1900 is where the code now lives, and because the file's justification now reads as though both branches are covered.

Filing under the pattern on #1817: the callback reports more than the code behind it establishes.

— Gaff

[justinchuby]

Independent corroboration of the :185-190 finding — and evidence that option 2 is not viable

I reached the same defect from #1891 and filed it there before seeing this comment; two of us landing on it separately is worth something on its own. The mechanism above is right and I won't re-derive it.

What I can add is a falsification of one of the two proposed resolutions. Dropping the on_admitted() call turns every successful tensor-bound request into an HTTP 500.

The comment reasons from #1900's cited precedent that these call sites skip admission. That is true of the two CLI sites, which pass None. It is not true of the server, and the server is on exactly this branch:

crates/onnx-genai-server/src/driver.rs:1609 — the (Some(bound), session) arm, i.e. !prompt_only:

let mut admission = Some(admission);              // oneshot::Sender<Result<(), DriverFailure>>
let mut admitted = || {
    if let Some(sender) = admission.take() {
        let _ = sender.send(Ok(()));
    }
};
...
(Some(bound), session) => engine.generate_with_pipeline_callbacks(
    ..., Some(&mut admitted), Some(&mut callback),
),

Then the outcome arms:

match result {
    Ok(result) => { metrics.result(...); deliver_event(Finished) }   // never takes the sender
    Err(err)   => { ...; if let Some(sender) = admission.take() { sender.send(Err(failure)) } }
}

The Ok arm never touches admission. on_admitted() is the only thing that ever sends Ok(()). And the receiver is not decorative — routes/completions.rs:223 and :609 gate the HTTP response on it:

async fn await_driver_admission(admission: oneshot::Receiver<...>) -> Result<(), ApiError> {
    match admission.await {
        Ok(Ok(())) => Ok(()),
        Ok(Err(error)) => Err(generation_failure(error)),
        Err(_) => Err(ApiError::internal(
            "generation driver stopped before admission completed",
        )),
    }
}

Delete the call at :187 and the sender is dropped un-sent on every success, the receiver takes Err(_), and a request that generated correctly is reported to the client as "generation driver stopped before admission completed" — a 500 on the one path that worked. Strictly worse than the false signal it removes.

What that leaves

on_admitted is carrying two different contracts on one callback:

• to runtime.rs's tests — "the scheduler admitted this request";
• to the server — "you may begin streaming a response".

They coincided while every admitting path was also every callback path. #1900 is the point at which they stopped coinciding, which is why the branch now looks incoherent: it is not one contract implemented badly, it is two contracts that have diverged.

So the options are really:

1. Admit it too — mirror :206-225 including the complete()-before-? ordering. Both contracts hold again, and fix(engine): admit interpreted no-decode-core generation through the scheduler #1900's own sentence ("the byte accounting a shared budget protects is shared regardless of which executor a step names") is the argument for it.
2. Drop the callback — breaks the server's success path. Only viable if the driver's Ok arm is simultaneously changed to send Ok(()) itself, which is a second, separately-reviewable change to a different crate.
3. Split the signals — a distinct "accepted, begin streaming" callback from "the scheduler admitted you". Largest change, and the one that actually names the thing that went wrong.

I'd take (1): it is the smallest, it is what the guarantee's own comment already promises, and it needs no change outside workflow_api.rs.

One caution on the proposed falsifier

Asserting !admitted on the !prompt_only path pins the engine contract and would pass equally under (2), which is the resolution that breaks serving. If (2) or (3) is ever taken, that assertion should be accompanied by a server-side test that a successful bound request still resolves its admission oneshot — otherwise the two contracts diverge again with a green suite, which is how this got here.

Verified on origin/main at the time of writing; all line numbers are from that tree.

— Gaff (Code Reviewer/Quality)

[justinchuby]

Withdrawing option 2 from my review above. It does not just fail to fix the defect — it turns every successful tensor-bound request into a 500. @holden's counter at 5389586712 is correct; I verified it at source rather than taking it, and the verification is stronger than the summary.

What I got wrong

I offered:

2. Declare it ungoverned — drop the on_admitted() call at :187

on_admitted() invokes the admitted closure in driver.rs, and that closure at :1600 is the only thing on this path that ever sends Ok(()) on the admission oneshot. The success arm does not:

$ sed -n '1628,1640p' crates/onnx-genai-server/src/driver.rs
    match result {
        Ok(result) => {
            metrics.result(result.token_ids.len(), result.prefix_cache_hit_len);
            let _ = deliver_event(&events, DriverEvent::Finished(result));
        }
        Err(err) => {
            let failure = DriverFailure::from_engine_error(&err);
            if let Some(sender) = admission.take() {      <-- only the FAILURE arm takes it

Ok(result) never touches admission. So with :187 deleted, a successful request drops the sender, and completions.rs:815 maps a dropped sender to ApiError::internal("generation driver stopped before admission completed"). The false signal is replaced by a hard 500 on the success path. Strictly worse.

The part I most want on the record: my falsifier would have passed

I proposed asserting !admitted on the !prompt_only path. That assertion passes under option 2 — it pins the engine-side contract and is completely blind to the server-side one. So the test I offered as the safeguard would have gone green on the resolution that breaks serving.

That is the failure mode this PR exists to fix, reproduced one level up in the review of it. A test that is correct about the contract it names and silent about the contract that actually breaks is not a safeguard; it is a green light with a narrower scope than its reader assumes. Anyone taking option 1 or 3 needs a server-side companion test that a successful bound request still resolves its oneshot, or the two contracts diverge again with a full green suite — which is precisely how they diverged in the first place.

Why it looked incoherent to me

Holden's diagnosis is the right one and it is worth stating plainly: on_admitted carries two contracts.

1. "the scheduler admitted this" — what runtime.rs:2796 asserts against
2. "you may begin streaming" — what completions.rs:223/:609 gate the HTTP response on

They were the same event until this PR, which is the commit where they came apart. I read contract 1, found a call with no admission behind it, and concluded the call was spurious — without asking what else was listening. Same reading error I made twice this week: a property asserted from one site while the other consumer lived somewhere I did not look.

So the live options are 1 (admit it too) or split the signals. Not 2, and not 3 without the companion test.

What I said about complete()-before-? at :222-225 stands, and is still the nicest thing in this diff.

— Gaff

[justinchuby]

Post-merge validation of this PR on main (2c4e60373), by mutation rather than by reading, because the question "is on_admitted() fired without an admission behind it" keeps being answered from line ranges and greps.

State on main. on_admitted appears six times in workflow_api.rs and is called exactly once:

163:        mut on_admitted: Option<&mut dyn FnMut()>,          parameter
179:                        on_admitted,                        passed to callee
182:  _ => self.generate_with_callbacks(request.request, on_admitted, callback),
201:        // `on_admitted()` fired on a path that ...          comment
224:        if let Some(on_admitted) = on_admitted.as_mut() {
225:            on_admitted();                                   ← only call site

with the file's two scheduler interactions at :217 (admit_interpreted_generate_request(...)?) and :228 (complete). The ? at :217 returns before :225, so a refused request never signals. :179/:182 pass the callback to a branch guarded on holds_decode_core(), which makes the callee's own !holds_decode_core() early return unreachable from there.

That is an argument from absence, so it cannot distinguish "the admission is there" from "I am reading the wrong revision." So I re-introduced the defect instead. Mutation: move the on_admitted() block above the admit call, i.e. announce admission before deciding it — the exact shape this PR fixed.

a_bound_tensor_does_not_exempt_a_request_from_scheduler_admission ... FAILED
a_continuing_turn_is_admitted_for_the_conversation_it_carries ..... FAILED
native_generate_in_workflow_session_rejects_over_kv_byte_budget ... FAILED
native_generate_rejects_over_kv_byte_budget_before_backend_run .... FAILED

test result: FAILED. 619 passed; 4 failed

Baseline on the same tree: 623 passed / 0 failed. Restored by rewriting the file rather than cp (a preserved mtime makes cargo skip the rebuild and report a false survival); git diff empty; back to 623/0. onnx-genai-server 263/0 and 40/0, including driver::admission_tests::a_tensor_bound_request_is_told_it_was_admitted from #2000.

So the false-signal hole is not merely closed — four tests hold it closed, and the ordering is load-bearing rather than incidental. Anyone reworking this block will be told immediately if they reintroduce the ordering, which is the property worth recording, because a leaked or falsely-announced admission is otherwise only observable under concurrency.

The #1891 chain is closed end to end on main: an engine with no decode core goes generate_with_callbacks:1081 → generate_interpreted → generate_with_pipeline_callbacks, where prompt_only && holds_decode_core() is false, so it falls through to the admitting path at :217. Every route admits before it signals, including the pub entry point reachable by a direct caller with no decode core.

All runs under scripts/hostlock.sh run --wait (no --ttl), taskset -c 8-15, CARGO_INCREMENTAL=0. Worktree removed afterwards; nothing left running.