Return overload response for KV admission refusal

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

@@ -39,6 +39,17 @@ pub enum ContinuousBatchEvent {
     },
 }
 
+#[derive(Debug)]
+pub enum ContinuousBatchAdmission {
+    Assigned {
+        handle: ContinuousBatchHandle,
+    },
+    Rejected {
+        handle: ContinuousBatchHandle,
+        error: anyhow::Error,
+    },
+}
+
 struct BatchRow {
     result_index: usize,
     physical_row: usize,
@@ -105,6 +116,7 @@ pub struct ContinuousBatchManager<'a> {
     static_max_len: usize,
     queue: VecDeque<PendingContinuousRequest>,
     rows: Vec<Option<ContinuousBatchRow>>,
+    admissions: VecDeque<ContinuousBatchAdmission>,
     events: VecDeque<ContinuousBatchEvent>,
     next_handle: usize,
 }
@@ -132,6 +144,7 @@ impl<'a> ContinuousBatchManager<'a> {
             static_max_len,
             queue: VecDeque::new(),
             rows: (0..max_batch).map(|_| None).collect(),
+            admissions: VecDeque::new(),
             events: VecDeque::new(),
             next_handle: 0,
         })
@@ -161,6 +174,8 @@ impl<'a> ContinuousBatchManager<'a> {
         let max_context = self.max_context_for_request(&options);
         let chain = build_processor_chain(&options, Some(self.tokenizer))?;
         if reached_context_limit(prompt_tokens.len(), max_context) {
+            self.admissions
+                .push_back(ContinuousBatchAdmission::Assigned { handle });
             ensure_constrained_finish(&options, "", FinishReason::Length)?;
             self.events.push_back(ContinuousBatchEvent::Finished {
                 handle,
@@ -180,7 +195,7 @@ impl<'a> ContinuousBatchManager<'a> {
 
     /// Advance all rows with pending logits by one generated token.
     pub fn step(&mut self) -> anyhow::Result<()> {
-        self.admit_available_rows()?;
+        self.admit_available_rows();
         let ready_rows = self
             .rows
             .iter()
@@ -205,15 +220,35 @@ impl<'a> ContinuousBatchManager<'a> {
             }
         }
 
-        self.admit_available_rows()?;
         self.decode_next_pending_rows()
     }
 
+    /// Assign queued requests to currently available decode rows.
+    pub fn admit_pending(&mut self) {
+        self.admit_available_rows();
+    }
+
     /// Drain token/result events emitted by previous `submit` or `step` calls.
     pub fn poll(&mut self) -> Vec<ContinuousBatchEvent> {
         self.events.drain(..).collect()
     }
 
+    pub fn poll_admissions(&mut self) -> Vec<ContinuousBatchAdmission> {
+        self.admissions.drain(..).collect()
+    }
+
+    pub fn cancel_pending(&mut self, handle: ContinuousBatchHandle) -> bool {
+        let Some(index) = self
+            .queue
+            .iter()
+            .position(|pending| pending.handle == handle)
+        else {
+            return false;
+        };
+        self.queue.remove(index);
+        true
+    }
+
     pub fn max_batch(&self) -> usize {
         self.rows.len()
     }
@@ -239,15 +274,22 @@ impl<'a> ContinuousBatchManager<'a> {
         Some(configured.map_or(self.static_max_len, |limit| limit.min(self.static_max_len)))
     }
 
-    fn admit_available_rows(&mut self) -> anyhow::Result<()> {
+    fn admit_available_rows(&mut self) {
         while !self.queue.is_empty() {
             let Some(row_index) = self.rows.iter().position(|row| row.is_none()) else {
                 break;
             };
             let pending = self.queue.pop_front().expect("queue checked non-empty");
-            self.admit_pending_into_row(pending, row_index)?;
+            let handle = pending.handle;
+            match self.admit_pending_into_row(pending, row_index) {
+                Ok(()) => self
+                    .admissions
+                    .push_back(ContinuousBatchAdmission::Assigned { handle }),
+                Err(error) => self
+                    .admissions
+                    .push_back(ContinuousBatchAdmission::Rejected { handle, error }),
+            }
         }
-        Ok(())
     }
 
     fn admit_pending_into_row(
@@ -1003,6 +1045,107 @@ fn collect_batch_results(
 mod tests {
     use super::*;
     use crate::sampling::sample_categorical;
+    use onnx_genai_ort::decode::BatchedDecodeSession;
+    use onnx_genai_ort::{OrtError, Value};
+    use std::path::Path;
+
+    struct RejectAssignDecode;
+
+    impl<'a> BatchedDecodeSession<'a> for RejectAssignDecode {
+        fn batch_size(&self) -> usize {
+            1
+        }
+
+        fn max_len(&self) -> usize {
+            32
+        }
+
+        fn row_len(&self, _row: usize) -> onnx_genai_ort::Result<usize> {
+            Ok(0)
+        }
+
+        fn active_rows(&self) -> Vec<usize> {
+            Vec::new()
+        }
+
+        fn deactivate_row(&mut self, _row: usize) -> onnx_genai_ort::Result<()> {
+            Ok(())
+        }
+
+        fn assign_row(&mut self, _row: usize) -> onnx_genai_ort::Result<()> {
+            Err(OrtError::InvalidArgument(
+                "deliberate row assignment failure".to_string(),
+            ))
+        }
+
+        fn step_select(
+            &mut self,
+            _next_token_ids: &[i64],
+            _position_ids: &[i64],
+            _advance_rows: &[bool],
+        ) -> onnx_genai_ort::Result<Value> {
+            unreachable!("a rejected row must never decode")
+        }
+
+        fn step_active(
+            &mut self,
+            _next_token_ids: &[i64],
+            _position_ids: &[i64],
+        ) -> onnx_genai_ort::Result<Value> {
+            unreachable!("a rejected row must never decode")
+        }
+    }
+
+    struct AcceptAssignDecode;
+
+    impl<'a> BatchedDecodeSession<'a> for AcceptAssignDecode {
+        fn batch_size(&self) -> usize {
+            1
+        }
+
+        fn max_len(&self) -> usize {
+            32
+        }
+
+        fn row_len(&self, _row: usize) -> onnx_genai_ort::Result<usize> {
+            Ok(0)
+        }
+
+        fn active_rows(&self) -> Vec<usize> {
+            Vec::new()
+        }
+
+        fn deactivate_row(&mut self, _row: usize) -> onnx_genai_ort::Result<()> {
+            Ok(())
+        }
+
+        fn assign_row(&mut self, _row: usize) -> onnx_genai_ort::Result<()> {
+            Ok(())
+        }
+
+        fn step_select(
+            &mut self,
+            _next_token_ids: &[i64],
+            _position_ids: &[i64],
+            _advance_rows: &[bool],
+        ) -> onnx_genai_ort::Result<Value> {
+            unreachable!("admission must not require backend generation")
+        }
+
+        fn step_active(
+            &mut self,
+            _next_token_ids: &[i64],
+            _position_ids: &[i64],
+        ) -> onnx_genai_ort::Result<Value> {
+            unreachable!("admission must not require backend generation")
+        }
+    }
+
+    fn test_tokenizer() -> Tokenizer {
+        let path = Path::new(env!("CARGO_MANIFEST_DIR"))
+            .join("../../tests/fixtures/tiny-llm/tokenizer.json");
+        Tokenizer::from_file(path).expect("load tiny tokenizer")
+    }
 
     #[test]
     fn per_row_sampling_is_seedable_and_independent() {
@@ -1021,4 +1164,65 @@ mod tests {
         assert_eq!(sequence(1), sequence(1));
         assert_ne!(sequence(0), sequence(1));
     }
+
+    #[test]
+    fn row_assignment_failure_is_reported_for_the_submitted_handle() {
+        let tokenizer = test_tokenizer();
+        let mut manager =
+            ContinuousBatchManager::new(Box::new(RejectAssignDecode), &tokenizer, None, 1).unwrap();
+        let handle = manager
+            .submit(GenerateRequest::new(GeneratePrompt::TokenIds(vec![1])))
+            .unwrap();
+
+        assert!(manager.poll_admissions().is_empty());
+        manager.step().unwrap();
+        let admissions = manager.poll_admissions();
+        assert_eq!(admissions.len(), 1);
+        match &admissions[0] {
+            ContinuousBatchAdmission::Rejected {
+                handle: rejected,
+                error,
+            } => {
+                assert_eq!(*rejected, handle);
+                assert!(
+                    error
+                        .to_string()
+                        .contains("deliberate row assignment failure")
+                );
+            }
+            ContinuousBatchAdmission::Assigned { .. } => panic!("failed row was acknowledged"),
+        }
+        assert!(!manager.has_pending_work());
+    }
+
+    #[test]
+    fn successful_row_admission_is_observable_before_backend_generation() {
+        let tokenizer = test_tokenizer();
+        let mut manager =
+            ContinuousBatchManager::new(Box::new(AcceptAssignDecode), &tokenizer, None, 1).unwrap();
+        let handle = manager
+            .submit(GenerateRequest::new(GeneratePrompt::TokenIds(vec![1])))
+            .unwrap();
+
+        manager.admit_pending();
+        let admissions = manager.poll_admissions();
+        assert!(matches!(
+            admissions.as_slice(),
+            [ContinuousBatchAdmission::Assigned { handle: assigned }] if *assigned == handle
+        ));
+    }
+
+    #[test]
+    fn pending_request_can_be_cancelled_without_an_admission_event() {
+        let tokenizer = test_tokenizer();
+        let mut manager =
+            ContinuousBatchManager::new(Box::new(RejectAssignDecode), &tokenizer, None, 1).unwrap();
+        let handle = manager
+            .submit(GenerateRequest::new(GeneratePrompt::TokenIds(vec![1])))
+            .unwrap();
+
+        assert!(manager.cancel_pending(handle));
+        assert!(!manager.has_pending_work());
+        assert!(manager.poll_admissions().is_empty());
+    }
 }

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

@@ -491,6 +491,13 @@ pub(crate) fn governor_kv_config(
     Ok(ModelKvConfig::known(page_size_bytes, tokens_per_page))
 }
 
+#[cfg_attr(
+    not(feature = "native-backend"),
+    expect(
+        dead_code,
+        reason = "native KV admission is used only by the native backend"
+    )
+)]
 pub(crate) fn governor_native_kv_config(
     kv_model: Option<&KvModelInfo>,
     config: &EngineConfig,