fix(engine): the interpreted path admits through the scheduler (#1891)

crates/onnx-genai-engine/src/engine/runtime.rs

@@ -276,7 +276,7 @@ fn native_workspace_query_rows(
 }
 
 impl Engine {
-    fn admit_generate_request_with_scheduler(
+    pub(crate) fn admit_generate_request_with_scheduler(
         &mut self,
         session_id: SessionId,
         prompt_tokens: usize,
@@ -2730,9 +2730,14 @@ mod tests {
         assert_eq!(native_workspace_query_rows(3, None, 8, None), 3);
     }
 
+    /// An engine whose package the interpreter drives.
+    ///
+    /// No fused decode core, so `holds_decode_core()` is false and generation
+    /// takes the interpreted path — the one #1723 introduced and left without
+    /// scheduler admission. `budget_bytes` is the scheduler's whole KV byte
+    /// budget, which is what decides whether a request is admitted at all.
     #[cfg(feature = "native-backend")]
-    #[test]
-    fn native_generate_rejects_over_kv_byte_budget_before_backend_run() -> anyhow::Result<()> {
+    fn interpreted_engine_with_byte_budget(budget_bytes: u64) -> anyhow::Result<Engine> {
         let tokenizer_path = Path::new(env!("CARGO_MANIFEST_DIR"))
             .join("../../tests/fixtures/tiny-llm/tokenizer.json")
             .canonicalize()?;
@@ -2748,7 +2753,7 @@ mod tests {
             ModelKvConfig::known(10, 1),
             0,
         )?;
-        let mut engine = Engine {
+        Ok(Engine {
             workflow: Box::new(crate::pipeline::generation::test_decoder_runtime()?),
             workflow_sessions: HashMap::new(),
             workflow_session_counter: 0,
@@ -2762,7 +2767,7 @@ mod tests {
             decode_path: ModelDecodePath::Generic,
             scheduler: Scheduler::with_byte_budget(
                 scheduler_config,
-                onnx_genai_scheduler::ByteBudget::new(10),
+                onnx_genai_scheduler::ByteBudget::new(budget_bytes),
             ),
             governor,
             sessions: HashMap::new(),
@@ -2787,7 +2792,13 @@ mod tests {
             last_speculative_stats: SpeculativeStats::default(),
             connector: ConnectorBridge::null(),
             _environment: None,
-        };
+        })
+    }
+
+    #[cfg(feature = "native-backend")]
+    #[test]
+    fn native_generate_rejects_over_kv_byte_budget_before_backend_run() -> anyhow::Result<()> {
+        let mut engine = interpreted_engine_with_byte_budget(10)?;
         let mut request = GenerateRequest::new(GeneratePrompt::TokenIds(vec![1]));
         request.options.max_new_tokens = 1;
         request.options.stop_on_eos = false;
@@ -2811,6 +2822,53 @@ mod tests {
         Ok(())
     }
 
+    /// The positive half of #1891, which the rejection test above cannot reach.
+    ///
+    /// A rejection test can pass on an engine that never admits anything --
+    /// including the pre-fix engine, whose interpreted path had no admission
+    /// wired in at all and simply fired `on_admitted()` unconditionally. This
+    /// asserts the two facts that distinguish a real admission from no
+    /// admission: the callback fires only *after* the scheduler accepted the
+    /// request, and the reservation it took is handed back afterwards.
+    ///
+    /// The fixture cannot finish a generation (the interpreted decoder wants a
+    /// KV value no component produces), so the observable is scheduler state,
+    /// not success. `complete()` runs on the error path too, which is the point
+    /// -- a request that admits and then fails must not leak its bytes or its
+    /// running-batch slot.
+    #[cfg(feature = "native-backend")]
+    #[test]
+    fn admitted_interpreted_request_signals_admission_and_releases_its_reservation()
+    -> anyhow::Result<()> {
+        let mut engine = interpreted_engine_with_byte_budget(1 << 20)?;
+        let mut request = GenerateRequest::new(GeneratePrompt::TokenIds(vec![1]));
+        request.options.max_new_tokens = 1;
+        request.options.stop_on_eos = false;
+
+        let mut admitted = false;
+        let mut on_admitted = || admitted = true;
+        let outcome = engine.generate_with_callbacks(request, Some(&mut on_admitted), None);
+
+        if let Err(error) = &outcome {
+            let error = error.to_string();
+            assert!(
+                !error.contains("scheduler admission failed"),
+                "a generous budget must admit; instead: {error}"
+            );
+        }
+        assert!(
+            admitted,
+            "admission callback must fire once the scheduler accepts the request"
+        );
+        assert_eq!(
+            engine.scheduler.running_count(),
+            0,
+            "the admitted request's scheduler slot must be released once it finishes, \
+             success or failure"
+        );
+        Ok(())
+    }
+
     #[cfg(feature = "native-backend")]
     #[test]
     fn native_backend_batching_capability_reports_single_sequence() -> anyhow::Result<()> {

crates/onnx-genai-engine/src/engine/workflow_api.rs

@@ -184,10 +184,6 @@ impl Engine {
                 _ => self.generate_with_callbacks(request.request, on_admitted, callback),
             };
         }
-        if let Some(on_admitted) = on_admitted.as_mut() {
-            on_admitted();
-        }
-        let runtime = &*self.workflow;
         // A prompt is text; a workflow input declaring the `prompt_tokens` role
         // wants ids. Encoding it here, with the package's own tokenizer, is
         // what lets `generate("hello")` work on a package whose components the
@@ -196,19 +192,63 @@ impl Engine {
         // the one thing it has the tokenizer for.
         let mut request = request;
         if let crate::config::GeneratePrompt::Text(text) = &request.request.prompt {
-            let tokenizer = runtime.package_tokenizer().context(
+            let tokenizer = self.workflow.package_tokenizer().context(
                 "this package declares a prompt_tokens input but ships no tokenizer, so a text                  prompt cannot be encoded for it; supply token ids instead",
             )?;
             let encoded = tokenizer
                 .encode(text)
                 .map_err(|error| anyhow::anyhow!("failed to encode the prompt: {error}"))?;
             request.request.prompt = crate::config::GeneratePrompt::TokenIds(encoded);
         }
+
+        // Admission is the interpreter's too.
+        //
+        // #1723 routed every package without a fused decode core through this
+        // drive, and this drive ran the declared components without ever asking
+        // the scheduler. Two things followed: the KV byte budget stopped
+        // applying to those packages at all, so an over-budget request did its
+        // work and then failed somewhere downstream for an unrelated reason;
+        // and `on_admitted` fired unconditionally, so a caller counting
+        // admissions counted requests that nothing had admitted.
+        //
+        // This sits *after* the prompt is encoded so the scheduler sizes the
+        // request it is actually going to run, and *before* any component is
+        // invoked, which is the ordering the guarantee is about.
+        let scheduler_seq = request
+            .session_id
+            .as_deref()
+            .and_then(|id| id.parse::<crate::config::SessionId>().ok())
+            .unwrap_or(0);
+        let prompt_tokens = match &request.request.prompt {
+            crate::config::GeneratePrompt::TokenIds(ids) => ids.len(),
+            crate::config::GeneratePrompt::TokenRows(rows) => {
+                rows.iter().map(Vec::len).max().unwrap_or(0)
+            }
+            crate::config::GeneratePrompt::Text(_) => 0,
+        };
+        let scheduled = self.admit_generate_request_with_scheduler(
+            scheduler_seq,
+            prompt_tokens,
+            request.request.options.max_new_tokens,
+            crate::engine::Priority::Normal,
+        )?;
+        request.request.options.max_new_tokens = scheduled.max_tokens;
+        if let Some(on_admitted) = on_admitted.as_mut() {
+            on_admitted();
+        }
+
         let options = request.request.options.clone();
+        let runtime = &*self.workflow;
         let tokenizer = runtime.package_tokenizer();
-        crate::pipeline::generation::run_declared_generation(
+        let result = crate::pipeline::generation::run_declared_generation(
             runtime, &options, tokenizer, request, None, callback,
-        )
+        );
+        // Releases the sequence's hot-tier reservation. Without it the budget
+        // and the running-batch slot leak on every interpreted request, and the
+        // second request on a tight budget is refused because the first never
+        // let go.
+        self.scheduler.complete(scheduler_seq);
+        result
     }
 
     /// Encode a workflow's declared buffered-PCM16 audio output for serving.