diff --git a/docs/architecture.md b/docs/architecture.md index 34e4bd27..33ec4043 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -20,11 +20,14 @@ uses a separate Hopper CUDA backend for one complete padded request. These worke coordinate independent processors and executors through `prepare` → `execute` → `finish` and use the [native runtime](../src/runtime/README.md) for FIFO admission and blocking dispatch -per loaded executor. Shared processing orchestration, batch budgets, -compatibility grouping and dynamic batching are planned. +per loaded executor. The [Decider native worker](../recipe/decider/README.md) +also uses this runtime and can opt into outstanding-request and processed-token +budgets plus HTTP request deadlines. Shared processing orchestration, +cross-request batch assembly, compatibility grouping and dynamic batching are planned. Cua-S1 and Open-Jev compute their decision heads on the CPU after downloading -the final hidden state. JEMM uses the shared Qwen executor with separate unmerged +the final hidden state. Decider projects its selected LM-head rows on CUDA. +JEMM uses the shared Qwen executor with separate unmerged FP32 adapters and framework-reference arithmetic. It computes selected BF16 label-head logits on CUDA, then calibrates typed answers on the CPU. JEMM admits one complete text or 1–4-image request and synchronizes both model streams before @@ -154,7 +157,14 @@ executor, with FIFO waiting. Engines own scheduler instances; model-specific execution adapters declare the unit and submit owned closures. Waiting is async, so queued requests do not occupy blocking threads waiting for a model lock. Schedulers are local to each worker/executor; this does not coordinate separate -processes. Queue-length limits and token budgets are not implemented yet. +processes. Workers can separately opt into the runtime's `BoundedAdmission` for +outstanding-request and caller-defined cost budgets. Decider's HTTP adapter +reserves requests before upload/preparation and complete prepared-row tokens +before execution, then carries those reservations in `ExecutionContext` through +blocking work. The worker selects configuration and the processor defines token +cost; shared runtime primitives own capacity accounting, cancellation/deadline +checks and FIFO execution admission. These request budgets do not add +cross-request batching or replace the scheduler's single execution permit. - Batch only work accepted by the same loaded executor, with compatible checkpoint/adapter identity, device, dtype, and input layout. Models declare diff --git a/recipe/README.md b/recipe/README.md index e771de81..27778139 100644 --- a/recipe/README.md +++ b/recipe/README.md @@ -16,14 +16,16 @@ For a first real decision, follow the [complete CPU walkthrough](../docs/getting - [Cua-S1 4B 0.2 native text worker](cua_s1/native.md): build the CUDA library and the Rust worker, export BF16 weights with the original FP32 adapter and start the worker. +- [Decider-2B v11 native text worker](decider/README.md): load verified pinned + weights, build the Rust/CUDA worker and configure optional batching, Graph, + prefix reuse, request/token capacity and HTTP deadlines; + [validation](decider/validation.md). - [Cua-S1 native screenshot worker](cua_s1/native_screenshot.md): export the multimodal language checkpoint and serve one-image decisions with Rust/CUDA. - [Open-Jev native text worker](open_jev/native.md): export the merged text backbone and trained decision head of Open-Jev-27B-v1.1 or Open-Jev-9B, then serve with Rust and CUDA. [Open-Jev-9B validation](open_jev/validation-9b.md) compares the 9B worker with the reference. -- [Decider-2B v11 native worker](decider/README.md): pinned BF16 text decisions, - optional request-local batching, Graph and prefix reuse; [validation](decider/validation.md). - [JEMM native worker](jemm/README.md): pinned text/image decisions and export; known multi-image numerical limitation is documented in the model contract. - [JEV-27B-VL experimental worker](jev_vl/README.md): native text and preencoded diff --git a/recipe/decider/README.md b/recipe/decider/README.md index b572586f..6f274d3a 100644 --- a/recipe/decider/README.md +++ b/recipe/decider/README.md @@ -36,6 +36,9 @@ Worker terminal: ```sh DECIDER_GRAPH=0 \ +DECIDER_MAX_PENDING_REQUESTS=8 \ +DECIDER_MAX_PENDING_TOKENS=1048576 \ +DECIDER_REQUEST_TIMEOUT_MS=5000 \ DECIDER_MODEL=/path/to/models/decider-2b-v11 \ DECIDER_CUDA_LIB="$PWD/target/release/libqwen3_5_cuda.so" \ DECIDER_HOST=127.0.0.1 DECIDER_PORT=8000 \ @@ -85,9 +88,9 @@ with empty answers and zero input/output usage. No partial answers are returned. Supported inputs are plain text/JSON state, independent Choice (2–255 options), Noul and isolated Score (2–10 levels). State-prefix truncation is 32,768 tokens; complete rows can include question suffixes up to 36,864 tokens. At most 1,024 -expanded rows and 1,048,576 processed tokens are accepted per request. These bounds -do not bound all waiting requests together; runtime pending-queue limits and -cross-request batching remain separate work. +expanded rows and 1,048,576 processed tokens are accepted per request. These processor bounds +are separate from the optional pending-request, processed-token and deadline +controls described below. Cross-request batching remains unsupported. The pinned worker requires released calibration. Images, video, chat/schema-first, packed questions, cross-request prefix caching and quantized/CPU/Metal execution @@ -166,7 +169,7 @@ worker lifetime and keeps the completed eager result; health then reports `eager Launch/inference failures still retire the worker. Capture and startup costs must be measured separately from warm replay. -## Optional request-local shared prefixes (pending integration) +## Optional request-local shared prefixes This integration depends on the pending Qwen continuation/fixed-GEMM/executor stack (#98/#99). Build this integrated branch with ABI9; older libraries are @@ -218,3 +221,85 @@ DECIDER_MODEL=/path/to/models/decider-2b-v11 \ DECIDER_CUDA_LIB="$PWD/target/release/libqwen3_5_cuda.so" \ target/release/omni-decider ``` + +## Capacity and deadline semantics + +These optional HTTP limits apply to eager, packed, Graph, fixed, shared and auto +execution. They do not apply to startup warmup or direct `Engine::predict` calls. +The example values in the startup command are independent configuration examples, +not measured deployment recommendations. Omit a value to leave that limit unset; +explicit values must be positive integers and invalid values fail before loading. +Graph remains incompatible with fixed/shared/auto; these limits do not change +that mode constraint. + +| Variable | What it bounds | +| --- | --- | +| `DECIDER_MAX_PENDING_REQUESTS` | Outstanding HTTP requests, including body upload, preparation, execution queueing, execution and response reconstruction. A slot is reserved before reading the body. | +| `DECIDER_MAX_PENDING_TOKENS` | The sum of complete prepared-row token lengths across admitted requests. Reservation follows preprocessing; repeated state and isolated Score rows count in full. This differs from API `usage.input_tokens` and does not change it. | +| `DECIDER_REQUEST_TIMEOUT_MS` | One absolute deadline from request admission through response reconstruction, including slow body upload and waiting. It does not cover startup, `/health` or network delivery of the response. | + +Request/token capacity uses shared +[`BoundedAdmission`](../../src/runtime/README.md#optional-request-capacity-and-execution-control). +The separate FIFO scheduler still admits one complete Decider execution request +at a time. Request capacity also bounds preprocessing; a token-only deployment +must prepare a request before it can determine its token cost. + +Temporary request/token contention returns `503`, `Content-Type: application/json`, +`Retry-After: 1`, and `{"detail":"inference capacity exhausted"}`. Use bounded +retries with backoff and an overall client deadline for this response. + +A single request whose processed-token cost exceeds the entire token ceiling +returns JSON `422` with `{"detail":"request exceeds processed-token capacity"}` +and no `Retry-After` header. It is rejected before execution and cannot succeed +unchanged under that configuration: reduce its prepared work or raise the +configured budget and restart the worker. A request costing exactly the ceiling +is accepted when that token capacity is free. + +Deadline expiry returns JSON `504` with `{"detail":"request timed out"}`. +The deadline is strict: work that completes after it may still yield `504` +instead of a successful response. Timeout/disconnect signals cancellation and +prevents work that has not started from running at the next context check; +it does not preempt already-running CPU/CUDA work. Detached preparation and +execution retain their capacity until they finish, and submitted device work +must synchronize before its resources are released. A `504` therefore does +not imply an immediately free request slot. + +Existing content-type, payload and validation checks remain in force: JSON is +required (`415` otherwise), the request body limit is 8 MiB (`413`), and processor +validation errors return `422`. A not-ready worker returns `503` with +`{"detail":"model unavailable"}` rather than the capacity-refusal response. + +Within-request packing is configured separately with `DECIDER_BATCH_MAX_ROWS` +(default `1`, allowed `1..=4`) and `DECIDER_BATCH_MAX_TOKENS` (default `4096`, +allowed `1..=4096`). These select packed groups within an admitted request; +they do not replace the outstanding-token budget. Rows above the packing token +budget run alone, subject to the processor/executor's separate input limits. + +## Health and validation + +`GET /health` is outside request admission and returns `200` for a ready +executor or `503` for an unavailable one. Existing model, checkpoint, execution, +batch and calibration metadata remain present. When any HTTP runtime limit is +set, `runtime_limits` reports `requests`, `processed_tokens` and `timeout_ms`; +unset members are `null`. + +Only a configured request-count bound adds top-level `depth`, `capacity` and +`rejected`. `depth` includes all outstanding request reservations, including +blocking work retained after a timeout; `capacity` is the request-count ceiling. +`rejected` counts request-capacity and token-capacity refusals, including requests +larger than the entire token ceiling. It excludes timeouts and processor +validation failures. Token-only or timeout-only deployments do not publish +these three fields. Token occupancy is not exposed by this health interface. + +The configuration and HTTP control checks use no checkpoint or GPU: + +```sh +cargo test --locked -p omni-runtime +cargo test --locked -p omni-decider-native --test runtime +cargo test --locked -p omni-decider-native --lib serve::runtime_tests +``` + +These checks cover admission, cancellation, capacity retention and synthetic +HTTP deadlines. They do not establish model parity, performance, GPU-memory +limits or other-hardware support. Device/checkpoint qualification remains a +separate check with the pinned artifacts and matching rebuilt CUDA library. diff --git a/src/models/decider/native/Cargo.toml b/src/models/decider/native/Cargo.toml index ca279132..d5d30849 100644 --- a/src/models/decider/native/Cargo.toml +++ b/src/models/decider/native/Cargo.toml @@ -16,7 +16,7 @@ safetensors = "0.8.0" omni-qwen3-5-native = { path = "../../qwen3_5/native" } omni-runtime = { path = "../../../runtime" } axum = "0.8.8" -tokio = { version = "1.49.0", features = ["macros", "net", "rt-multi-thread", "sync", "signal"] } +tokio = { version = "1.49.0", features = ["macros", "net", "rt-multi-thread", "sync", "signal", "time"] } [[bin]] name = "omni-decider" @@ -49,3 +49,14 @@ path = "../../../../tests/decider/prefix.rs" [[example]] name = "decider-run" path = "../../../../tests/decider/runner.rs" + +[[test]] +name = "runtime" +path = "../../../../tests/decider/runtime.rs" + +[[test]] +name = "startup" +path = "../../../../tests/decider/startup.rs" + +[dev-dependencies] +tokio = { version = "1.49.0", features = ["io-util"] } diff --git a/src/models/decider/native/src/engine.rs b/src/models/decider/native/src/engine.rs index 3d89dd00..732ec7c1 100644 --- a/src/models/decider/native/src/engine.rs +++ b/src/models/decider/native/src/engine.rs @@ -1,7 +1,10 @@ //! Assemble verified artifacts, model-owned processing, and serial eager execution. use crate::{ - Config, Kind, Limits, Processor, batching::BatchLimits, checkpoint::Checkpoint, + Config, Kind, Limits, Processor, + batching::BatchLimits, + checkpoint::Checkpoint, executor::Executor, + runtime::{Runtime, RuntimeConfig}, }; use anyhow::Result; use omni_runtime::SerialScheduler; @@ -12,17 +15,19 @@ pub struct Engine { pub processor: Arc, pub executor: Executor, pub scheduler: SerialScheduler, + pub runtime: Runtime, metadata: Value, } impl Engine { pub async fn load(dir: &Path, library: &Path) -> Result { let graph = crate::options::graph_env()?; - Self::load_with_modes( + Self::load_with_modes_and_runtime( dir, library, BatchLimits::from_env()?, graph, crate::options::prefix_env(graph)?, + RuntimeConfig::from_env()?, ) .await } @@ -31,7 +36,7 @@ impl Engine { library: &Path, batching: BatchLimits, ) -> Result { - Self::load_with_options(dir, library, batching, false).await + Self::load_with_runtime(dir, library, batching, RuntimeConfig::from_env()?).await } pub async fn load_with_options( dir: &Path, @@ -54,11 +59,48 @@ impl Engine { batching: BatchLimits, graph: bool, prefix_mode: crate::prefix::PrefixMode, + ) -> Result { + Self::load_with_modes_and_runtime( + dir, + library, + batching, + graph, + prefix_mode, + RuntimeConfig::from_env()?, + ) + .await + } + /// Explicit runtime configuration with the existing independent eager defaults. + pub async fn load_with_runtime( + dir: &Path, + library: &Path, + batching: BatchLimits, + config: RuntimeConfig, + ) -> Result { + Self::load_with_modes_and_runtime( + dir, + library, + batching, + false, + crate::prefix::PrefixMode::Off, + config, + ) + .await + } + /// Select execution and HTTP policy together without reading environment switches. + pub async fn load_with_modes_and_runtime( + dir: &Path, + library: &Path, + batching: BatchLimits, + graph: bool, + prefix_mode: crate::prefix::PrefixMode, + config: RuntimeConfig, ) -> Result { anyhow::ensure!( !graph || prefix_mode == crate::prefix::PrefixMode::Off, "shared/fixed execution requires Graph off" ); + let runtime = Runtime::new(config)?; let directory = dir.to_owned(); let (processor,checkpoint,mut metadata) = tokio::task::spawn_blocking(move || -> Result<_> { let processor = Processor::load(&directory,Limits::default())?; @@ -86,6 +128,7 @@ impl Engine { processor: Arc::new(processor), executor, scheduler: SerialScheduler::default(), + runtime, metadata, }; // Loading returns only after a real complete model decision, before any socket binds. @@ -120,6 +163,16 @@ impl Engine { } else { "unavailable" }); + let config = self.runtime.config(); + if config.max_requests.is_some() || config.max_tokens.is_some() || config.timeout.is_some() + { + metadata["runtime_limits"] = json!({"requests":config.max_requests,"processed_tokens":config.max_tokens,"timeout_ms":config.timeout.map(|timeout|timeout.as_millis())}); + } + if let Some(report) = self.runtime.report() { + metadata["depth"] = json!(report.depth); + metadata["capacity"] = json!(report.capacity); + metadata["rejected"] = json!(report.rejected); + } metadata } } diff --git a/src/models/decider/native/src/executor.rs b/src/models/decider/native/src/executor.rs index 5d73735b..dfc344be 100644 --- a/src/models/decider/native/src/executor.rs +++ b/src/models/decider/native/src/executor.rs @@ -12,6 +12,7 @@ use omni_qwen3_5_native::{ model::{GraphStats, Model}, }; use omni_runtime::SerialScheduler; +use omni_runtime::control::ExecutionContext; use std::{ ffi::c_void, path::Path, @@ -293,6 +294,16 @@ impl Executor { scheduler: &SerialScheduler, rows: Vec, ) -> Result>> { + self.execute_with_context(scheduler, rows, ExecutionContext::default()) + .await + } + pub async fn execute_with_context( + &self, + scheduler: &SerialScheduler, + rows: Vec, + context: ExecutionContext, + ) -> Result>> { + context.check()?; validate_rows(&rows, &self.labels)?; ensure!(self.is_ready(), "executor unavailable"); if rows.is_empty() { @@ -305,7 +316,7 @@ impl Executor { ); let prefix_stats = self.prefix_stats.clone(); scheduler - .run(move || { + .run_with_context(context, move || { let mut guard = match loaded.lock() { Ok(guard) => guard, Err(_) => { diff --git a/src/models/decider/native/src/lib.rs b/src/models/decider/native/src/lib.rs index 3fd6067d..128cfe0f 100644 --- a/src/models/decider/native/src/lib.rs +++ b/src/models/decider/native/src/lib.rs @@ -35,6 +35,7 @@ pub mod batching; mod checkpoint; pub mod engine; pub mod executor; +pub mod runtime; pub mod serve; pub mod options; diff --git a/src/models/decider/native/src/runtime.rs b/src/models/decider/native/src/runtime.rs new file mode 100644 index 00000000..5a092375 --- /dev/null +++ b/src/models/decider/native/src/runtime.rs @@ -0,0 +1,150 @@ +//! Opt-in request and processed-token budgets for one loaded Decider worker. +use anyhow::{Context, Result, ensure}; +use omni_runtime::{admission::BoundedAdmission, control::ExecutionContext, engine::Report}; +use std::time::{Duration, Instant}; + +#[derive(Clone, Copy, Debug, Default)] +/// All limits are opt-in. Unset environment variables retain unbounded serving. +/// `DECIDER_MAX_PENDING_REQUESTS` includes upload, preparation, queueing and execution. +/// `DECIDER_MAX_PENDING_TOKENS` counts complete prepared rows, not billable usage. +/// `DECIDER_REQUEST_TIMEOUT_MS` covers upload through response reconstruction. +pub struct RuntimeConfig { + pub max_requests: Option, + pub max_tokens: Option, + pub timeout: Option, +} +impl RuntimeConfig { + pub fn from_env() -> Result { + let read = |name: &str| match std::env::var(name) { + Ok(value) => Ok(Some(value)), + Err(std::env::VarError::NotPresent) => Ok(None), + Err(error) => Err(anyhow::anyhow!("{name}: {error}")), + }; + Self::from_values( + read("DECIDER_MAX_PENDING_REQUESTS")?.as_deref(), + read("DECIDER_MAX_PENDING_TOKENS")?.as_deref(), + read("DECIDER_REQUEST_TIMEOUT_MS")?.as_deref(), + ) + } + pub fn from_values( + requests: Option<&str>, + tokens: Option<&str>, + timeout: Option<&str>, + ) -> Result { + let parse = |name: &str, value: Option<&str>| -> Result> { + value + .map(|value| { + let value: usize = value.parse().with_context(|| format!("invalid {name}"))?; + ensure!(value > 0, "{name} must be positive"); + Ok(value) + }) + .transpose() + }; + let config = Self { + max_requests: parse("DECIDER_MAX_PENDING_REQUESTS", requests)?, + max_tokens: parse("DECIDER_MAX_PENDING_TOKENS", tokens)?, + timeout: parse("DECIDER_REQUEST_TIMEOUT_MS", timeout)? + .map(|n| Duration::from_millis(n as u64)), + }; + config.validate()?; + Ok(config) + } + fn validate(&self) -> Result<()> { + ensure!( + self.max_requests != Some(0) && self.max_tokens != Some(0), + "runtime budgets must be positive" + ); + if let Some(timeout) = self.timeout { + ensure!( + !timeout.is_zero() && Instant::now().checked_add(timeout).is_some(), + "request timeout must be positive and fit the monotonic clock" + ); + } + Ok(()) + } +} + +#[derive(Debug)] +pub struct Overloaded; +impl std::fmt::Display for Overloaded { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("inference capacity exhausted") + } +} +impl std::error::Error for Overloaded {} + +#[derive(Debug)] +pub struct OversizedRequest; +impl std::fmt::Display for OversizedRequest { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("request exceeds processed-token capacity") + } +} +impl std::error::Error for OversizedRequest {} + +pub struct Runtime { + config: RuntimeConfig, + requests: Option, + tokens: Option, +} +impl Runtime { + pub fn new(config: RuntimeConfig) -> Result { + config.validate()?; + Ok(Self { + config, + requests: config.max_requests.map(|n| BoundedAdmission::new(n, None)), + tokens: config + .max_tokens + .map(|n| BoundedAdmission::new(usize::MAX, Some(n))), + }) + } + pub fn config(&self) -> RuntimeConfig { + self.config + } + /// Reserve before upload and preparation. The context follows all blocking work. + pub fn admit(&self, started: Instant) -> Result { + let context = ExecutionContext::new(self.config.timeout.map(|timeout| started + timeout)); + match &self.requests { + Some(admission) => Ok(context.retain(admission.try_admit(1).ok_or(Overloaded)?)), + None => Ok(context), + } + } + /// Cost is the sum of complete rows, including repeated state and Score expansion. + pub fn reserve_tokens( + &self, + context: ExecutionContext, + tokens: usize, + ) -> Result { + context.check()?; + match &self.tokens { + Some(admission) => { + let permit = admission.try_admit(tokens).ok_or_else(|| { + // Admission records every refusal; the HTTP adapter distinguishes + // requests that can fit after contention clears from those that cannot. + if self + .config + .max_tokens + .is_some_and(|ceiling| tokens > ceiling) + { + anyhow::Error::new(OversizedRequest) + } else { + anyhow::Error::new(Overloaded) + } + })?; + Ok(context.retain(permit)) + } + None => Ok(context), + } + } + /// Request-depth metrics exist only when a request-count bound is configured. + /// Token-only limits have no finite request capacity to report. + pub fn report(&self) -> Option { + self.requests.as_ref().map(|admission| { + let mut report = admission.report(); + if let Some(tokens) = &self.tokens { + report.rejected = report.rejected.saturating_add(tokens.report().rejected); + } + report + }) + } +} diff --git a/src/models/decider/native/src/serve.rs b/src/models/decider/native/src/serve.rs index 3b7dcbe3..b00db09c 100644 --- a/src/models/decider/native/src/serve.rs +++ b/src/models/decider/native/src/serve.rs @@ -1,15 +1,21 @@ //! Model worker HTTP transport; Decider interpretation belongs to Processor. -use crate::{Limits, engine::Engine}; +use crate::{ + Limits, + engine::Engine, + runtime::{Overloaded, OversizedRequest, Runtime}, +}; use axum::{ Json, Router, body::Bytes, - extract::{DefaultBodyLimit, State, rejection::BytesRejection}, - http::{HeaderMap, StatusCode}, + extract::{DefaultBodyLimit, FromRequest, Request, State}, + http::StatusCode, response::{IntoResponse, Response}, routing::{get, post}, }; +use omni_runtime::control::{ExecutionContext, ExecutionStopped}; use serde_json::json; use std::sync::Arc; +use std::{future::Future, time::Instant}; fn error(status: StatusCode, message: impl ToString) -> Response { (status, Json(json!({"detail":message.to_string()}))).into_response() @@ -22,12 +28,9 @@ async fn health(State(engine): State>) -> Response { }; (status, Json(engine.health())).into_response() } -async fn decide( - State(engine): State>, - headers: HeaderMap, - body: Result, -) -> Response { - let content_type = headers +async fn decide(State(engine): State>, request: Request) -> Response { + let content_type = request + .headers() .get("content-type") .and_then(|h| h.to_str().ok()) .unwrap_or("") @@ -41,18 +44,84 @@ async fn decide( "Content-Type must be application/json", ); } - let raw = match body { - Ok(raw) => raw, - Err(e) => return error(e.status(), e.body_text()), - }; if !engine.executor.is_ready() { return error(StatusCode::SERVICE_UNAVAILABLE, "model unavailable"); } - // Tokenization/row construction runs outside admission and the model mutex. + within_budget(&engine.runtime, request, |raw, context| { + infer(&engine, raw, context) + }) + .await +} + +async fn within_budget(runtime: &Runtime, request: Request, infer: F) -> Response +where + F: FnOnce(Bytes, ExecutionContext) -> Fut, + Fut: Future, +{ + let context = match runtime.admit(Instant::now()) { + Ok(context) => context, + Err(error) => return runtime_error(error, StatusCode::INTERNAL_SERVER_ERROR), + }; + let _cancel = context.cancel_on_drop(); + let deadline = context.deadline(); + let work = async move { + if let Err(error) = context.check() { + return runtime_error(error.into(), StatusCode::INTERNAL_SERVER_ERROR); + } + let raw = match Bytes::from_request(request, &()).await { + Ok(raw) => raw, + Err(error) => return self::error(error.status(), error.body_text()), + }; + if let Err(error) = context.check() { + return runtime_error(error.into(), StatusCode::INTERNAL_SERVER_ERROR); + } + infer(raw, context).await + }; + match deadline { + Some(deadline) => match tokio::time::timeout_at(deadline.into(), work).await { + Ok(response) => response, + Err(_) => error(StatusCode::GATEWAY_TIMEOUT, "request timed out"), + }, + None => work.await, + } +} + +fn runtime_error(error: anyhow::Error, fallback: StatusCode) -> Response { + if error.is::() { + let mut response = self::error( + StatusCode::SERVICE_UNAVAILABLE, + "inference capacity exhausted", + ); + response + .headers_mut() + .insert("retry-after", axum::http::HeaderValue::from_static("1")); + response + } else if error.is::() { + self::error(StatusCode::UNPROCESSABLE_ENTITY, error) + } else if error.is::() { + self::error(StatusCode::GATEWAY_TIMEOUT, "request timed out") + } else if fallback == StatusCode::UNPROCESSABLE_ENTITY { + self::error(fallback, error) + } else { + self::error(fallback, "model inference failed") + } +} + +async fn infer(engine: &Engine, raw: Bytes, context: ExecutionContext) -> Response { + // Request capacity covers preparation; processed-token admission follows it. + // Tokenization and row construction stay outside the model mutex. let processor = engine.processor.clone(); - let prepared = match tokio::task::spawn_blocking(move || processor.prepare(&raw)).await { + let preparing = context.clone(); + let prepared = match tokio::task::spawn_blocking(move || { + preparing.check()?; + let prepared = processor.prepare(&raw)?; + preparing.check()?; + Ok::<_, anyhow::Error>(prepared) + }) + .await + { Ok(Ok(prepared)) => prepared, - Ok(Err(e)) => return error(StatusCode::UNPROCESSABLE_ENTITY, e), + Ok(Err(e)) => return runtime_error(e, StatusCode::UNPROCESSABLE_ENTITY), Err(e) => { eprintln!("Decider preparation failed: {e}"); return error( @@ -61,22 +130,32 @@ async fn decide( ); } }; + let context = match engine + .runtime + .reserve_tokens(context, prepared.processed_tokens()) + { + Ok(context) => context, + Err(error) => return runtime_error(error, StatusCode::INTERNAL_SERVER_ERROR), + }; let result = engine .executor - .execute(&engine.scheduler, prepared.rows) + .execute_with_context(&engine.scheduler, prepared.rows, context.clone()) .await .and_then(|rows| prepared.context.finish(rows)); match result { - Ok(body) => Json(body).into_response(), + Ok(body) => match context.check() { + Ok(()) => Json(body).into_response(), + Err(error) => runtime_error(error.into(), StatusCode::INTERNAL_SERVER_ERROR), + }, Err(e) => { eprintln!("Decider inference failed: {e:#}"); - error( + runtime_error( + e, if engine.executor.is_ready() { StatusCode::INTERNAL_SERVER_ERROR } else { StatusCode::SERVICE_UNAVAILABLE }, - "model inference failed", ) } } @@ -88,3 +167,7 @@ pub fn router(engine: Arc) -> Router { .layer(DefaultBodyLimit::max(Limits::default().max_request_bytes)) .with_state(engine) } + +#[cfg(test)] +#[path = "../../../../../tests/decider/serve_runtime.rs"] +mod runtime_tests; diff --git a/src/runtime/Cargo.toml b/src/runtime/Cargo.toml index 007046f2..6fa06401 100644 --- a/src/runtime/Cargo.toml +++ b/src/runtime/Cargo.toml @@ -19,3 +19,11 @@ path = "../../tests/runtime/scheduler.rs" [[test]] name = "engine" path = "../../tests/runtime/engine.rs" + +[[test]] +name = "admission" +path = "../../tests/runtime/admission.rs" + +[[test]] +name = "control" +path = "../../tests/runtime/control.rs" diff --git a/src/runtime/README.md b/src/runtime/README.md index d38a06d8..b5e9124c 100644 --- a/src/runtime/README.md +++ b/src/runtime/README.md @@ -1,7 +1,7 @@ # Native execution runtime -`omni-runtime` provides `SerialScheduler` for the Cua-S1, Open-Jev, Decider and Laya native -workers. Their engines assemble a processor, one scheduler for the loaded +`omni-runtime` provides `SerialScheduler` for the Cua-S1, Open-Jev, Laya and +Decider native workers. Their engines assemble a processor, one scheduler for the loaded executor, and the executor. HTTP handlers still coordinate `prepare` → `execute` → `finish`. @@ -15,8 +15,8 @@ state. Waiting requests no longer occupy blocking threads waiting for that mutex | --- | --- | --- | | Cua-S1 | One question's unpadded Qwen forward. Release before admitting its next question. | Request preparation, CPU letter projection and response finishing. | | Open-Jev | One complete request's independent candidate forwards and CPU scalar heads. | Request preparation and calibrated response finishing. | -| Decider | One complete request's backbone rows, selected CUDA label head and synchronized readback. | Request preparation and whole-question calibrated response finishing. | | Laya | One complete padded request, including GPU scorer/action head and copyback. | Request preparation/padding and calibrated response finishing. | +| Decider | One complete request's independent rows, including within-request packing and LM-head projection. | Request preparation and calibrated response finishing. An optional request reservation covers both; token reservation begins after preparation. | Models retain weights, heads, device state, scratch buffers and graph caches. The scheduler takes an owned closure and returns its result; it does not know @@ -32,20 +32,63 @@ forward synchronizes the final hidden-state download. Errors and panics release admission when the blocking task finishes. Model mutex poisoning retains its existing failure behavior. -This implements serial admission per executor within a worker process. Admitted -concurrency is one. Queue-length limits, token budgets, compatibility grouping, -shared processing orchestration and GPU batching remain planned. The Qwen -executor runs single prompts, bounded packed prefill (Open-Jev-27B-v1.1's -candidates within a request and Decider's complete rows) or prompts that share -prefixes in one call (Decider opt-in); this scheduler does not pack inputs or coordinate separately -running workers. +## Optional request capacity and execution control + +Two admission layers have different lifetimes. `SerialScheduler` owns the FIFO +execution queue and permits one execution unit per loaded executor. It waits +instead of refusing work and does not itself bound the number of waiting +requests. A worker can additionally use the shared `admission::BoundedAdmission` +to refuse requests when an outstanding-request or caller-defined cost budget is +full. `try_admit(cost)` returns an owned `Permit` or refuses without reserving; +costs above the entire budget are refused, never clamped. Dropping the permit +releases its reservation. Neither layer coordinates separately running workers. + +`control::ExecutionContext` carries an optional absolute deadline, shared +cancellation state and retained admission permits. Context clones share each +retained reservation; capacity is released after its last owner drops. The +awaiting caller holds a `CancellationGuard` from `cancel_on_drop()` so abandoning +the request signals cancellation while detached blocking work retains its +context and capacity. `check()` reports `ExecutionStopped::Deadline` or +`ExecutionStopped::Cancelled`; it does not interrupt running work. + +`SerialScheduler::run_with_context` checks the context before queueing and again +on the blocking thread immediately before calling the work closure, including +when that closure waited in Tokio's blocking pool. The context and scheduler +permit remain owned through completion. This method does not install an async +timeout or preempt device work; the HTTP adapter enforces the response deadline, +and the executor must synchronize submitted device work before returning. +The existing `SerialScheduler::run` remains available without request control. + +The Decider HTTP worker opts into both layers. Its worker adapter reads the +environment, reserves a request slot before upload/preparation, and reserves +processed tokens after preparation. The model processor defines that token cost; +the shared admission type enforces the budgets, and the scheduler still owns +FIFO execution admission. Context clones retain reservations through blocking +preparation and execution even after a timeout or disconnect. Response +reconstruction remains in the model adapter. Internal `Engine::predict`, +including startup warmup, uses the direct inference path without HTTP limits. +See the [Decider recipe](../../recipe/decider/README.md) for configuration, +health fields and overload/deadline responses. Other workers keep their current +admission behavior unless they explicitly opt in. + +Shared processing orchestration, compatibility grouping, cross-request dynamic +batching and a shared worker drain lifecycle remain planned. Models may pack +supported inputs within one request; this scheduler does not assemble GPU batches. +The Qwen executor supports single prompts, bounded packed prefill (Open-Jev +candidates or Decider complete rows), and request-local shared prefixes +(Decider opt-in). Graph, fixed and shared execution retain the same complete +request admission unit and context lifetime; the scheduler does not construct +GPU batches or coordinate separately running workers. CPU tests cover FIFO serialization, queued cancellation, retention of admission -and resources after dispatched cancellation, errors/panics, and independent -loaded executors: +and resources after dispatched cancellation, errors/panics, independent loaded +executors, bounded request/cost admission, and context checks under blocking-pool +backlog. Decider adds configuration and HTTP upload/preparation deadline checks: ```sh cargo test --locked -p omni-runtime +cargo test --locked -p omni-decider-native --test runtime +cargo test --locked -p omni-decider-native --lib serve::runtime_tests ``` See the [architecture contracts](../../docs/architecture.md) and diff --git a/src/runtime/src/admission.rs b/src/runtime/src/admission.rs new file mode 100644 index 00000000..58995aac --- /dev/null +++ b/src/runtime/src/admission.rs @@ -0,0 +1,220 @@ +//! Optional capacity limits for a worker that wants them. +//! +//! `SerialScheduler` admits one unit of work at a time and never refuses: a hundred requests +//! queue behind each other and the only bound is the caller's patience. That is the right +//! default for the workers here today, and this type is what a worker opts into when it wants a +//! bound instead. +//! +//! **Units are the caller's, and they are explicit.** Admission counts two things: how many +//! requests are outstanding, and how much each one costs. A request counts as one; the cost is +//! whatever the caller passes, so a prefill worker can admit by token count, a worker with +//! large payloads by bytes, and a worker that cares about neither by `1`. Nothing here guesses +//! which of those a model means. +//! +//! **The permit is owned, not borrowed.** It holds a share of the admission, so it can be moved +//! into the blocking work it was granted for — which is what the runtime's executor already +//! does — and the capacity is released when that work finishes with it rather than when the +//! request that asked for it goes away. +//! +//! This is admission only. It decides whether work may start and hands back a permit that +//! releases when the work is done with; it does not own a thread, run anything, or move a +//! request between states. Readiness, worker-thread ownership and drain are separate. + +use std::sync::{Arc, Mutex}; + +use crate::engine::Report; + +/// A bound a worker may opt into. Clones share it, so a queue can hold one and so can the +/// handlers that feed it. +/// +/// ``` +/// use omni_runtime::admission::BoundedAdmission; +/// +/// let admission = BoundedAdmission::new(2, None); +/// let first = admission.try_admit(1).expect("one of two"); +/// let second = admission.try_admit(1).expect("two of two"); +/// assert!(admission.try_admit(1).is_none(), "the third has nowhere to go"); +/// +/// // A permit releases its capacity when the work is done with. +/// drop(first); +/// assert!(admission.try_admit(1).is_some()); +/// ``` +/// +/// It wraps a shared inner rather than *being* an [`Arc`] alias: an alias would inherit `Arc`'s +/// associated functions, so `BoundedAdmission::new` would resolve to `Arc::new`. +#[derive(Clone, Debug)] +pub struct BoundedAdmission(Arc); + +/// The bound and the counters behind one [`BoundedAdmission`]. +#[derive(Debug)] +struct Admission { + /// Requests that may be outstanding at once. Zero would admit nothing, so it is refused at + /// construction rather than silently rejecting every request. + requests: usize, + /// Total cost that may be outstanding, when the caller wants to bound something other than + /// a request count. `None` bounds requests only. + units: Option, + held: Mutex, +} + +/// What admission is currently holding, and what it has refused. +#[derive(Debug, Default)] +struct Held { + outstanding: usize, + units: usize, + rejected: u64, +} + +/// One admitted piece of work's reservation, released when its [`Permit`] drops. +/// +/// A permit keeps the caller's admitted cost in `units`, which [`Permit::units`] +/// reports unchanged. Its `charged` field records what advanced the shared unit +/// total: the cost when a unit budget exists, and zero for request-only admission. +/// Release subtracts that charge exactly once; oversized requests are refused, +/// never clamped. +#[derive(Clone, Copy, Debug)] +pub struct Reservation { + /// The cost the work was admitted at, which is what [`Permit::units`] reports. + units: usize, + /// What the shared unit total was advanced by, so release subtracts exactly that. Zero when + /// no unit budget is configured: there is then no total, and charging it would overflow for + /// a cost near `usize::MAX` while bounding nothing. + charged: usize, +} + +impl BoundedAdmission { + /// A bound of `requests` at a time, optionally also of `units` at a time. + /// + /// # Panics + /// + /// If `requests` is zero, or if `units` is `Some(0)`. Both would admit nothing at all, + /// which is a configuration mistake rather than a limit worth enforcing, and it would show + /// up as every request refused rather than as the mistake it is. + pub fn new(requests: usize, units: Option) -> Self { + assert!( + requests > 0, + "admission with no request capacity admits nothing" + ); + assert!( + units != Some(0), + "admission with no unit capacity admits nothing" + ); + Self(Arc::new(Admission { + requests, + units, + held: Mutex::new(Held::default()), + })) + } + + /// Admits one piece of work costing `units`, or refuses it because there is no room. + /// + /// The permit is the reservation: capacity is held until it is dropped, which is what + /// makes "outstanding" mean "accepted and not yet finished with" rather than "accepted". + /// A refusal is counted and nothing is reserved. + /// + /// A request costing more than the unit ceiling is refused rather than admitted and + /// overrunning it, so a reservation is never larger than the ceiling and the held total + /// cannot reach `usize::MAX` by accumulation. + /// + /// The permit owns a share of this admission, so it can outlive the handle that issued it + /// and be moved into the blocking work it was granted for. + pub fn try_admit(&self, units: usize) -> Option { + let mut held = self.0.lock(); + if held.outstanding >= self.0.requests { + held.rejected += 1; + return None; + } + let charged = match self.0.units { + Some(ceiling) => { + // A request larger than the whole ceiling is refused rather than clamped, so a + // reservation is never bigger than the budget it was admitted under. What is + // left of the ceiling is then known to fit, because `held.units <= ceiling` + // holds as an invariant once every reservation is bounded this way. + if units > ceiling || held.units > ceiling - units { + held.rejected += 1; + return None; + } + units + } + // No unit budget, so there is no total to keep and nothing is charged. Adding the + // caller's cost here would overflow for a cost near `usize::MAX` while protecting + // nothing: the request bound is the only limit this configuration asked for. + None => 0, + }; + held.outstanding += 1; + held.units += charged; + Some(Permit { + admission: Arc::clone(&self.0), + reserved: Reservation { units, charged }, + }) + } + + /// What a worker would publish as its queue state. `capacity` is the request bound, which + /// is the one every caller of this type has. + pub fn report(&self) -> Report { + let held = self.0.lock(); + Report { + depth: held.outstanding, + capacity: self.0.requests, + rejected: held.rejected, + } + } +} + +impl Admission { + /// Releases one reservation. Called when its [`Permit`] drops. + /// + /// Guarded rather than saturating: a reservation larger than what is held would mean two + /// permits for one admission, and quietly underflowing would hide that. The public API + /// cannot produce one, because a reservation is refused unless it fits under the ceiling. + fn release(&self, reservation: Reservation) { + let mut held = self.lock(); + held.outstanding = held + .outstanding + .checked_sub(1) + .expect("a permit releases exactly one admission"); + if self.units.is_some() { + held.units = held + .units + .checked_sub(reservation.charged) + .expect("a permit releases the units it was admitted with"); + } + } + + fn lock(&self) -> std::sync::MutexGuard<'_, Held> { + // A panic while holding this leaves the counters readable and consistent, so a + // poisoned lock is not a reason to refuse every later request. + self.held.lock().unwrap_or_else(|error| error.into_inner()) + } +} + +/// One admitted piece of work. Holding it holds its capacity; dropping it releases both. +/// +/// It owns a share of its admission rather than borrowing it, so it can be moved into the +/// blocking work it was granted for. The runtime's own executor does exactly that; a borrow +/// would not compile there, because that work has to be `'static`. +/// +/// This is deliberately not `Clone` and not `Copy`: two permits for one admission would release +/// capacity twice, which is the whole failure this type exists to prevent. +#[derive(Debug)] +pub struct Permit { + admission: Arc, + reserved: Reservation, +} + +impl Permit { + /// The cost this work was admitted at, for a caller that wants to report or account it. + /// + /// This is the cost the caller passed, unchanged: nothing is clamped, because a request too + /// large for the unit ceiling is refused rather than shaved to fit. When no unit budget is + /// configured the cost is still reported, even though no unit total was kept for it. + pub fn units(&self) -> usize { + self.reserved.units + } +} + +impl Drop for Permit { + fn drop(&mut self) { + self.admission.release(self.reserved); + } +} diff --git a/src/runtime/src/control.rs b/src/runtime/src/control.rs new file mode 100644 index 00000000..28b24304 --- /dev/null +++ b/src/runtime/src/control.rs @@ -0,0 +1,73 @@ +//! Request deadlines, cancellation and owned admission across blocking work. +use std::{ + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, + time::Instant, +}; + +use crate::admission::Permit; + +#[derive(Clone, Default)] +pub struct ExecutionContext { + deadline: Option, + cancelled: Arc, + permits: Vec>, +} + +#[derive(Debug, PartialEq, Eq)] +pub enum ExecutionStopped { + Deadline, + Cancelled, +} + +impl std::fmt::Display for ExecutionStopped { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + Self::Deadline => "request deadline expired", + Self::Cancelled => "request cancelled", + }) + } +} +impl std::error::Error for ExecutionStopped {} + +impl ExecutionContext { + pub fn new(deadline: Option) -> Self { + Self { + deadline, + ..Self::default() + } + } + /// Keep capacity alive in every clone, including detached blocking tasks. + pub fn retain(mut self, permit: Permit) -> Self { + self.permits.push(Arc::new(permit)); + self + } + pub fn deadline(&self) -> Option { + self.deadline + } + pub fn check(&self) -> Result<(), ExecutionStopped> { + if self + .deadline + .is_some_and(|deadline| Instant::now() >= deadline) + { + return Err(ExecutionStopped::Deadline); + } + if self.cancelled.load(Ordering::Acquire) { + return Err(ExecutionStopped::Cancelled); + } + Ok(()) + } + /// The awaiting caller owns this guard; blocking work owns context clones. + pub fn cancel_on_drop(&self) -> CancellationGuard { + CancellationGuard(self.cancelled.clone()) + } +} + +pub struct CancellationGuard(Arc); +impl Drop for CancellationGuard { + fn drop(&mut self) { + self.0.store(true, Ordering::Release); + } +} diff --git a/src/runtime/src/lib.rs b/src/runtime/src/lib.rs index dde9174f..1f0e7d2b 100644 --- a/src/runtime/src/lib.rs +++ b/src/runtime/src/lib.rs @@ -1,5 +1,7 @@ //! Shared admission, dispatch and the worker-side engine contract for a loaded executor. +pub mod admission; +pub mod control; pub mod engine; use std::sync::Arc; @@ -45,4 +47,24 @@ impl SerialScheduler { .await .context("execution task failed")? } + /// Retain request capacity through device synchronization. Check again on the + /// blocking thread: dispatch itself can wait behind unrelated blocking work. + pub async fn run_with_context( + &self, + context: control::ExecutionContext, + work: F, + ) -> Result + where + F: FnOnce() -> Result + Send + 'static, + T: Send + 'static, + { + context.check()?; + self.run(move || { + context.check()?; + let result = work(); + drop(context); + result + }) + .await + } } diff --git a/tests/decider/runtime.rs b/tests/decider/runtime.rs new file mode 100644 index 00000000..2ae2273a --- /dev/null +++ b/tests/decider/runtime.rs @@ -0,0 +1,69 @@ +use omni_decider_native::runtime::{Runtime, RuntimeConfig}; +use std::time::{Duration, Instant}; + +#[test] +fn default_policy_is_unbounded_and_has_no_deadline() { + let runtime = Runtime::new(RuntimeConfig::from_values(None, None, None).unwrap()).unwrap(); + let request = runtime.admit(Instant::now()).unwrap(); + assert!(request.deadline().is_none()); + assert!(runtime.report().is_none()); +} +#[test] +fn invalid_limits_are_configuration_errors() { + for values in [ + (Some("0"), None, None), + (None, Some("0"), None), + (None, None, Some("0")), + (Some("-1"), None, None), + (None, None, Some("nan")), + (None, None, Some("18446744073709551616")), + ] { + assert!( + RuntimeConfig::from_values(values.0, values.1, values.2) + .and_then(Runtime::new) + .is_err() + ); + } +} +#[test] +fn request_and_processed_token_limits_are_independent_and_recover() { + let runtime = + Runtime::new(RuntimeConfig::from_values(Some("2"), Some("100"), Some("1000")).unwrap()) + .unwrap(); + let now = Instant::now(); + let first = runtime.admit(now).unwrap(); + assert_eq!(first.deadline(), Some(now + Duration::from_secs(1))); + let first = runtime.reserve_tokens(first, 60).unwrap(); + let second = runtime.admit(now).unwrap(); + assert!(runtime.admit(now).is_err()); + assert!(runtime.reserve_tokens(second, 41).is_err()); + assert_eq!(runtime.report().unwrap().depth, 1); + assert_eq!(runtime.report().unwrap().rejected, 2); + let second = runtime + .reserve_tokens(runtime.admit(now).unwrap(), 40) + .unwrap(); + drop(first); + drop(second); + assert_eq!(runtime.report().unwrap().depth, 0); + let all = runtime + .reserve_tokens(runtime.admit(now).unwrap(), 100) + .unwrap(); + drop(all); + assert!( + runtime + .reserve_tokens(runtime.admit(now).unwrap(), 101) + .is_err() + ); + assert_eq!(runtime.report().unwrap().depth, 0); +} +#[test] +fn deadline_before_preparation_is_explicit_and_does_not_hold_capacity() { + let runtime = + Runtime::new(RuntimeConfig::from_values(Some("1"), None, Some("1")).unwrap()).unwrap(); + let context = runtime + .admit(Instant::now() - Duration::from_secs(1)) + .unwrap(); + assert!(context.check().is_err()); + drop(context); + assert_eq!(runtime.report().unwrap().depth, 0); +} diff --git a/tests/decider/serve_runtime.rs b/tests/decider/serve_runtime.rs new file mode 100644 index 00000000..34b5ff8c --- /dev/null +++ b/tests/decider/serve_runtime.rs @@ -0,0 +1,189 @@ +use super::*; +use crate::runtime::{Runtime, RuntimeConfig}; +use axum::body::Body; +use std::time::{Duration, Instant}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; + +fn policy(requests: Option<&str>, timeout: Option<&str>) -> Runtime { + Runtime::new(RuntimeConfig::from_values(requests, None, timeout).unwrap()).unwrap() +} +#[tokio::test] +async fn overload_is_json_503_with_retry_after_and_never_reads_the_body() { + let runtime = policy(Some("1"), None); + let held = runtime.admit(Instant::now()).unwrap(); + let response = within_budget(&runtime, Request::new(Body::empty()), |_, _| async { + panic!("refused request reached preparation"); + #[allow(unreachable_code)] + error(StatusCode::OK, "") + }) + .await; + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + assert_eq!(response.headers()["retry-after"], "1"); + drop(held); + assert_eq!(runtime.report().unwrap().depth, 0); +} +#[tokio::test] +async fn upload_is_inside_the_same_deadline_and_releases_capacity() { + let runtime = Arc::new(policy(Some("1"), Some("80"))); + let app = Router::new() + .route( + "/", + post( + |State(runtime): State>, request: Request| async move { + within_budget(&runtime, request, |_, _| async { + error(StatusCode::OK, "done") + }) + .await + }, + ), + ) + .with_state(runtime.clone()); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + let mut socket = tokio::net::TcpStream::connect(address).await.unwrap(); + socket.write_all(b"POST / HTTP/1.1\r\nHost: localhost\r\nContent-Length: 100\r\nConnection: close\r\n\r\nx").await.unwrap(); + let mut response = String::new(); + tokio::time::timeout(Duration::from_secs(3), socket.read_to_string(&mut response)) + .await + .unwrap() + .unwrap(); + assert!(response.starts_with("HTTP/1.1 504"), "{response}"); + assert_eq!(runtime.report().unwrap().depth, 0); + server.abort(); +} +#[tokio::test] +async fn deadline_covers_preparation_and_returns_504() { + let runtime = policy(Some("1"), Some("20")); + let (release, wait) = std::sync::mpsc::channel(); + let response = within_budget( + &runtime, + Request::new(Body::empty()), + |_, context| async move { + let context = context.clone(); + let outcome = tokio::task::spawn_blocking(move || { + wait.recv().unwrap(); + context.check() + }) + .await + .unwrap(); + match outcome { + Ok(()) => error(StatusCode::OK, ""), + Err(stopped) => runtime_error(stopped.into(), StatusCode::INTERNAL_SERVER_ERROR), + } + }, + ) + .await; + assert_eq!(response.status(), StatusCode::GATEWAY_TIMEOUT); + assert_eq!(runtime.report().unwrap().depth, 1); + release.send(()).unwrap(); + tokio::time::timeout(Duration::from_secs(2), async { + while runtime.report().unwrap().depth > 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); +} +#[tokio::test] +async fn successful_response_and_existing_body_rejections_are_preserved() { + let runtime = policy(Some("1"), Some("1000")); + let response = within_budget( + &runtime, + Request::new(Body::from("ok")), + |raw, context| async move { + assert_eq!(raw.as_ref(), b"ok"); + context.check().unwrap(); + error( + StatusCode::UNPROCESSABLE_ENTITY, + "original validation error", + ) + }, + ) + .await; + assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY); + assert_eq!(runtime.report().unwrap().depth, 0); + let response = within_budget( + &runtime, + Request::new(Body::from(vec![0; 3 * 1024 * 1024])), + |_, _| async { + panic!("oversize reached preparation"); + #[allow(unreachable_code)] + error(StatusCode::OK, "") + }, + ) + .await; + assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE); + assert_eq!(runtime.report().unwrap().depth, 0); +} + +#[tokio::test] +async fn token_ceiling_refusal_is_json_422_without_retry_after() { + let runtime = + Runtime::new(RuntimeConfig::from_values(Some("2"), Some("100"), None).unwrap()).unwrap(); + // A permanently oversized request is invalid both when idle and under contention. + for held_tokens in [0, 60] { + let held = runtime + .reserve_tokens(runtime.admit(Instant::now()).unwrap(), held_tokens) + .unwrap(); + let response = within_budget(&runtime, Request::new(Body::empty()), |_, context| async { + match runtime.reserve_tokens(context, 101) { + Ok(_) => panic!("oversized request reached execution"), + Err(error) => runtime_error(error, StatusCode::INTERNAL_SERVER_ERROR), + } + }) + .await; + assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY); + assert!(!response.headers().contains_key("retry-after")); + assert_eq!(response.headers()["content-type"], "application/json"); + let body = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + assert_eq!( + serde_json::from_slice::(&body).unwrap(), + json!({"detail": "request exceeds processed-token capacity"}) + ); + assert_eq!(runtime.report().unwrap().depth, 1); + drop(held); + assert_eq!(runtime.report().unwrap().depth, 0); + } + assert_eq!(runtime.report().unwrap().rejected, 2); + let whole = runtime + .reserve_tokens(runtime.admit(Instant::now()).unwrap(), 100) + .unwrap(); + drop(whole); +} + +#[tokio::test] +async fn temporary_token_contention_is_json_503_with_retry_after_and_recovers() { + let runtime = + Runtime::new(RuntimeConfig::from_values(Some("2"), Some("100"), None).unwrap()).unwrap(); + let held = runtime + .reserve_tokens(runtime.admit(Instant::now()).unwrap(), 60) + .unwrap(); + let response = within_budget(&runtime, Request::new(Body::empty()), |_, context| async { + match runtime.reserve_tokens(context, 41) { + Ok(_) => panic!("overloaded request reached execution"), + Err(error) => runtime_error(error, StatusCode::INTERNAL_SERVER_ERROR), + } + }) + .await; + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + assert_eq!(response.headers()["retry-after"], "1"); + assert_eq!(response.headers()["content-type"], "application/json"); + let body = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + assert_eq!( + serde_json::from_slice::(&body).unwrap(), + json!({"detail": "inference capacity exhausted"}) + ); + assert_eq!(runtime.report().unwrap().depth, 1); + assert_eq!(runtime.report().unwrap().rejected, 1); + drop(held); + let retried = runtime + .reserve_tokens(runtime.admit(Instant::now()).unwrap(), 41) + .unwrap(); + drop(retried); + assert_eq!(runtime.report().unwrap().depth, 0); +} diff --git a/tests/decider/startup.rs b/tests/decider/startup.rs new file mode 100644 index 00000000..515a3d16 --- /dev/null +++ b/tests/decider/startup.rs @@ -0,0 +1,150 @@ +//! Combined worker switches must be validated before checkpoint or CUDA loading. +use std::process::Command; + +fn rejects_before_loading(settings: &[(&str, &str)], expected: &str) { + let missing = std::env::temp_dir().join(format!( + "decider-startup-missing-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + assert!(!missing.exists()); + let mut command = Command::new(env!("CARGO_BIN_EXE_omni-decider")); + for key in [ + "DECIDER_GRAPH", + "DECIDER_PREFIX", + "DECIDER_FIXED", + "DECIDER_BATCH_MAX_ROWS", + "DECIDER_BATCH_MAX_TOKENS", + "DECIDER_MAX_PENDING_REQUESTS", + "DECIDER_MAX_PENDING_TOKENS", + "DECIDER_REQUEST_TIMEOUT_MS", + "DECIDER_PORT", + ] { + command.env_remove(key); + } + let output = command + .env("DECIDER_MODEL", missing) + .env("DECIDER_CUDA_LIB", "unavailable-cuda-library") + .env("CUA_S1_GRAPH", "1") + .envs(settings.iter().copied()) + .output() + .unwrap(); + let stderr = String::from_utf8(output.stderr).unwrap(); + assert!(!output.status.success()); + assert!(stderr.contains(expected), "{settings:?}: {stderr}"); + assert!(!stderr.contains("No such file"), "{stderr}"); +} + +#[test] +fn graph_and_prefix_validation_survives_runtime_integration() { + for (settings, expected) in [ + (vec![("DECIDER_GRAPH", "2")], "DECIDER_GRAPH must be 0 or 1"), + ( + vec![("DECIDER_PREFIX", "invalid")], + "DECIDER_PREFIX must be", + ), + (vec![("DECIDER_FIXED", "invalid")], "DECIDER_FIXED must be"), + ( + vec![("DECIDER_PREFIX", "1"), ("DECIDER_FIXED", "1")], + "select only one", + ), + ( + vec![("DECIDER_GRAPH", "1"), ("DECIDER_PREFIX", "auto")], + "shared/fixed execution requires DECIDER_GRAPH=0", + ), + ] { + rejects_before_loading(&settings, expected); + } +} + +#[test] +fn runtime_validation_survives_every_execution_mode() { + for mode in [ + vec![], + vec![("DECIDER_GRAPH", "1")], + vec![("DECIDER_PREFIX", "1")], + vec![("DECIDER_PREFIX", "auto")], + vec![("DECIDER_FIXED", "1")], + ] { + for (setting, expected) in [ + ( + "DECIDER_MAX_PENDING_REQUESTS", + "DECIDER_MAX_PENDING_REQUESTS must be positive", + ), + ( + "DECIDER_MAX_PENDING_TOKENS", + "DECIDER_MAX_PENDING_TOKENS must be positive", + ), + ( + "DECIDER_REQUEST_TIMEOUT_MS", + "DECIDER_REQUEST_TIMEOUT_MS must be positive", + ), + ] { + let mut settings = mode.clone(); + settings.push((setting, "0")); + rejects_before_loading(&settings, expected); + } + } +} + +#[tokio::test] +async fn explicit_mode_constructor_rejects_graph_with_prefix_before_loading() { + use omni_decider_native::{ + batching::BatchLimits, engine::Engine, prefix::PrefixMode, runtime::RuntimeConfig, + }; + for mode in [PrefixMode::Fixed, PrefixMode::Shared, PrefixMode::Auto] { + let result = Engine::load_with_modes_and_runtime( + std::path::Path::new("unavailable-model"), + std::path::Path::new("unavailable-cuda-library"), + BatchLimits::default(), + true, + mode, + RuntimeConfig::default(), + ) + .await; + assert!( + result + .err() + .unwrap() + .to_string() + .contains("shared/fixed execution requires Graph off") + ); + } +} + +#[tokio::test] +async fn explicit_mode_constructor_validates_runtime_before_loading() { + use omni_decider_native::{ + batching::BatchLimits, engine::Engine, prefix::PrefixMode, runtime::RuntimeConfig, + }; + for (graph, mode) in [ + (false, PrefixMode::Off), + (true, PrefixMode::Off), + (false, PrefixMode::Fixed), + (false, PrefixMode::Shared), + (false, PrefixMode::Auto), + ] { + let result = Engine::load_with_modes_and_runtime( + std::path::Path::new("unavailable-model"), + std::path::Path::new("unavailable-cuda-library"), + BatchLimits::default(), + graph, + mode, + RuntimeConfig { + max_requests: Some(0), + ..RuntimeConfig::default() + }, + ) + .await; + assert!( + result + .err() + .unwrap() + .to_string() + .contains("runtime budgets must be positive") + ); + } +} diff --git a/tests/runtime/admission.rs b/tests/runtime/admission.rs new file mode 100644 index 00000000..f7e4ae8f --- /dev/null +++ b/tests/runtime/admission.rs @@ -0,0 +1,370 @@ +//! Optional capacity limits: what they admit, what they refuse, and what a permit means. +//! +//! Everything goes through the public API, because that is all a worker sees. + +use std::{ + mem::ManuallyDrop, + sync::{ + Arc, Barrier, + atomic::{AtomicUsize, Ordering}, + }, + thread, +}; + +use omni_runtime::admission::BoundedAdmission; + +#[test] +fn request_only_admission_releases_large_cost_first() { + request_only_admission_preserves_costs_and_recovers(true); +} + +#[test] +fn request_only_admission_releases_small_cost_first() { + request_only_admission_preserves_costs_and_recovers(false); +} + +fn request_only_admission_preserves_costs_and_recovers(large_first: bool) { + let admission = BoundedAdmission::new(2, None); + let large = admission.try_admit(usize::MAX).expect("first request"); + let small = admission.try_admit(1).expect("second request"); + assert_eq!(large.units(), usize::MAX); + assert_eq!(small.units(), 1); + assert_eq!(admission.report().depth, 2); + assert!(admission.try_admit(0).is_none(), "request capacity is full"); + assert_eq!(admission.report().rejected, 1); + + let (first, last) = if large_first { + (large, small) + } else { + (small, large) + }; + // Keep a broken release from triggering a second destructor panic while unwinding. + let last = ManuallyDrop::new(last); + drop(first); + assert_eq!(admission.report().depth, 1); + drop(ManuallyDrop::into_inner(last)); + assert_eq!(admission.report().depth, 0); + let again = admission.try_admit(usize::MAX).expect("capacity recovered"); + assert_eq!(again.units(), usize::MAX); + drop(again); + assert_eq!(admission.report().depth, 0); + assert_eq!(admission.report().rejected, 1); +} + +#[test] +fn a_request_bound_refuses_once_it_is_full() { + let admission = BoundedAdmission::new(2, None); + let first = admission.try_admit(1).expect("the first of two"); + let second = admission.try_admit(1).expect("the second of two"); + assert!( + admission.try_admit(1).is_none(), + "the third has nowhere to go" + ); + + let report = admission.report(); + assert_eq!(report.depth, 2, "a refused request must not be counted"); + assert_eq!(report.capacity, 2); + assert_eq!(report.rejected, 1); + + // A refusal is not a lost slot: releasing one makes room for the next. + drop(first); + assert_eq!(admission.report().depth, 1); + assert!(admission.try_admit(1).is_some()); + drop(second); +} + +/// The point of a permit: capacity means "accepted and not yet finished with". A worker that +/// hands work to a thread and forgets to release would be bounded forever, so release is tied +/// to the permit's lifetime rather than to a call someone has to remember. +#[test] +fn capacity_is_held_until_the_permit_is_dropped() { + let admission = BoundedAdmission::new(1, None); + let permit = admission.try_admit(1).expect("one of one"); + assert!(admission.try_admit(1).is_none(), "still held"); + drop(permit); + assert_eq!(admission.report().depth, 0); + assert!(admission.try_admit(1).is_some(), "released"); +} + +/// Units are the caller's, which is what lets a prefill worker bound tokens and a worker with +/// large payloads bound bytes without this type knowing which it is looking at. +#[test] +fn a_unit_bound_admits_by_cost_rather_than_by_count() { + // Room for three requests, or 100 units, whichever runs out first. + let admission = BoundedAdmission::new(3, Some(100)); + + let big = admission.try_admit(60).expect("60 of 100 units"); + assert_eq!(admission.report().depth, 1, "one request, 60 units"); + + let small = admission.try_admit(40).expect("the remaining 40 units"); + assert!( + admission.try_admit(1).is_none(), + "the unit ceiling is reached even though only two requests are outstanding" + ); + + // Releasing the large one frees its units, not just its request slot. + drop(big); + assert!(admission.try_admit(50).is_some(), "60 units came back"); + drop(small); +} + +/// A request that costs more than the whole ceiling is refused rather than admitted and +/// overrunning it — the case a naive `depth < capacity` check gets wrong. +#[test] +fn one_request_larger_than_the_ceiling_is_refused() { + let admission = BoundedAdmission::new(10, Some(100)); + assert!(admission.try_admit(101).is_none()); + assert_eq!(admission.report().depth, 0, "nothing was reserved"); + assert!( + admission.try_admit(100).is_some(), + "exactly the ceiling fits" + ); +} + +#[test] +fn the_request_count_is_reported_as_capacity() { + let admission = BoundedAdmission::new(7, Some(1_000)); + assert_eq!(admission.report().capacity, 7); + assert_eq!(admission.report().depth, 0); + assert_eq!(admission.report().rejected, 0); +} + +/// A zero bound admits nothing, which is a configuration mistake rather than a limit. It is +/// refused where it is written instead of showing up as every request being turned away. +#[test] +#[should_panic(expected = "admits nothing")] +fn a_zero_request_bound_is_refused_at_construction() { + let _ = BoundedAdmission::new(0, None); +} + +#[test] +#[should_panic(expected = "admits nothing")] +fn a_zero_unit_bound_is_refused_at_construction() { + let _ = BoundedAdmission::new(1, Some(0)); +} + +/// Clones share the bound, so the queue and the handlers feeding it cannot disagree about how +/// much room is left. +#[test] +fn clones_share_one_bound() { + let admission = BoundedAdmission::new(1, None); + let shared = admission.clone(); + let permit = shared.try_admit(1).expect("one of one"); + assert!(admission.try_admit(1).is_none(), "the clone held the slot"); + drop(permit); + assert!(admission.try_admit(1).is_some()); + assert_eq!(admission.report().rejected, 1, "both see the same count"); +} + +/// Regression: every admission is released by exactly what it reserved, in any order. The +/// permit stores the units it actually reserved rather than the units it was asked for, so the +/// counter cannot be driven negative — an earlier version stored the requested cost and, where +/// the two differed, released past zero and panicked on a later release. +#[test] +fn a_reservation_is_released_by_what_was_reserved_not_by_what_was_asked() { + // Room for many requests and exactly 100 units, so the unit ceiling is what binds. + let admission = BoundedAdmission::new(10, Some(100)); + + let a = admission.try_admit(10).expect("10 of 100"); + let b = admission.try_admit(60).expect("60 more"); + let c = admission.try_admit(30).expect("the remaining 30"); + assert!( + admission.try_admit(1).is_none(), + "the ceiling is exactly reached" + ); + + // Release order is irrelevant when what goes in is what comes out. + drop(b); + let again = admission.try_admit(60).expect("b's units came back"); + drop(a); + drop(c); + assert_eq!( + admission.report().depth, + 1, + "only the re-admitted 60 is held" + ); + drop(again); + assert_eq!(admission.report().depth, 0, "and it released its own units"); + assert!( + admission.try_admit(100).is_some(), + "the ceiling is whole again" + ); +} + +/// A request costing more than the ceiling is refused outright, which is what keeps the held +/// total below `usize::MAX`: a reservation larger than the ceiling it was admitted under would +/// be released against the others and unbalance the counter. +#[test] +fn a_request_larger_than_the_unit_ceiling_is_refused_however_large_it_is() { + let admission = BoundedAdmission::new(10, Some(64)); + + for cost in [65, 1_000, usize::MAX] { + assert!( + admission.try_admit(cost).is_none(), + "a request costing {cost} must be refused, not clamped" + ); + } + assert_eq!(admission.report().depth, 0, "nothing was reserved"); + assert_eq!(admission.report().rejected, 3, "all three were counted"); + + // The largest request that fits is still admitted, and released with its own units. + let full = admission.try_admit(64).expect("exactly the ceiling"); + assert_eq!(full.units(), 64); + assert!( + admission.try_admit(1).is_none(), + "and it took the whole ceiling" + ); + drop(full); + assert_eq!(admission.report().depth, 0); + assert!(admission.try_admit(64).is_some(), "the ceiling is reusable"); +} + +/// The ceiling is the whole budget, so one request may take all of it and no more. +#[test] +fn the_ceiling_is_the_largest_single_reservation() { + let admission = BoundedAdmission::new(10, Some(usize::MAX)); + + let all = admission.try_admit(usize::MAX).expect("the whole ceiling"); + assert!( + admission.try_admit(1).is_none(), + "even one unit is refused: the ceiling is taken" + ); + drop(all); + assert_eq!(admission.report().depth, 0); + assert!( + admission.try_admit(usize::MAX).is_some(), + "and the whole ceiling comes back" + ); +} + +/// Concurrent admitters can never together hold more than the request bound, and every +/// admission is released exactly once. A counter that drifted one way would let the live count +/// pass the bound; drifting the other way would leave permits outstanding at the end. +#[test] +fn concurrent_admissions_never_exceed_the_bound_and_all_come_back() { + const THREADS: usize = 32; + const ROUNDS: usize = 100; + const CAPACITY: usize = 4; + + for _ in 0..ROUNDS { + let admission = BoundedAdmission::new(CAPACITY, None); + let barrier = Arc::new(Barrier::new(THREADS)); + let live = Arc::new(AtomicUsize::new(0)); + let peak = Arc::new(AtomicUsize::new(0)); + + let handles: Vec<_> = (0..THREADS) + .map(|_| { + let admission = admission.clone(); + let barrier = Arc::clone(&barrier); + let live = Arc::clone(&live); + let peak = Arc::clone(&peak); + thread::spawn(move || { + barrier.wait(); + if let Some(_permit) = admission.try_admit(1) { + let now = live.fetch_add(1, Ordering::SeqCst) + 1; + peak.fetch_max(now, Ordering::SeqCst); + live.fetch_sub(1, Ordering::SeqCst); + } + }) + }) + .collect(); + for handle in handles { + handle.join().expect("no admitting thread may panic"); + } + + assert_eq!( + live.load(Ordering::SeqCst), + 0, + "every admission was released" + ); + assert!( + peak.load(Ordering::SeqCst) <= CAPACITY, + "at most the capacity is ever held at once" + ); + assert_eq!(admission.report().depth, 0, "every admission was released"); + assert_eq!( + admission.report().capacity, + CAPACITY, + "the bound outlives the contention" + ); + } +} + +/// A panic while the counters are locked must not refuse every later request: the counters are +/// still consistent, so the lock is recovered rather than poisoned for good. +#[test] +fn a_poisoned_lock_keeps_admitting() { + let admission = BoundedAdmission::new(1, None); + + let held = admission.clone(); + let panicked = thread::spawn(move || { + let _permit = held.try_admit(1).expect("one of one"); + panic!("admit and then die holding the lock"); + }) + .join(); + assert!(panicked.is_err(), "the helper thread was meant to panic"); + + let report = admission.report(); + assert_eq!( + report.depth, 0, + "the panicking thread's permit was dropped as it unwound, so its slot came back" + ); + assert!( + admission.try_admit(1).is_some(), + "the recovered lock admits" + ); +} + +/// A permit is owned, so it can be moved into the work it was granted for instead of being tied +/// to the task that asked for it. This is the shape the runtime's executor uses, and the reason +/// the permit is not a borrow of the admission. +#[test] +fn a_permit_can_be_moved_into_the_work_it_was_granted_for() { + let admission = BoundedAdmission::new(1, None); + let permit = admission.try_admit(1).expect("one of one"); + assert!(admission.try_admit(1).is_none(), "held by the permit"); + + let held = admission.clone(); + thread::spawn(move || { + // The work owns the permit and releases the capacity when it finishes with it. + let _permit = permit; + assert!( + held.try_admit(1).is_none(), + "still held while the work runs" + ); + }) + .join() + .expect("the work thread may not panic"); + + assert_eq!(admission.report().depth, 0, "the work released its slot"); + assert!(admission.try_admit(1).is_some(), "and it can be reused"); +} + +/// Regression: a request-only bound has no unit budget, so nothing is accumulated. Adding the +/// caller's cost anyway overflowed for a cost near `usize::MAX`: debug builds panicked on +/// admission, release builds wrapped the total and then panicked in `release`. +#[test] +fn a_request_only_bound_does_not_accumulate_units() { + for reverse in [false, true] { + let admission = BoundedAdmission::new(2, None); + let huge = admission.try_admit(usize::MAX).expect("one of two"); + let one = admission.try_admit(1).expect("two of two"); + assert_eq!(admission.report().depth, 2); + assert_eq!(huge.units(), usize::MAX, "the caller's cost is preserved"); + assert_eq!(one.units(), 1, "and so is the small one's"); + + // Both drop orders have to release cleanly and leave the same depth. + if reverse { + drop(one); + drop(huge); + } else { + drop(huge); + drop(one); + } + assert_eq!(admission.report().depth, 0, "every admission came back"); + assert!( + admission.try_admit(usize::MAX).is_some(), + "and the bound admits again afterwards" + ); + } +} diff --git a/tests/runtime/control.rs b/tests/runtime/control.rs new file mode 100644 index 00000000..38521678 --- /dev/null +++ b/tests/runtime/control.rs @@ -0,0 +1,235 @@ +use anyhow::Result; +use omni_runtime::{SerialScheduler, admission::BoundedAdmission, control::ExecutionContext}; +use std::{ + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + mpsc, + }, + time::{Duration, Instant}, +}; +use tokio::sync::oneshot; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn expired_queued_work_does_not_execute() { + let scheduler = SerialScheduler::default(); + let (started, ready) = oneshot::channel(); + let (release, wait) = mpsc::channel(); + let first = tokio::spawn({ + let scheduler = scheduler.clone(); + async move { + scheduler + .run(move || { + started.send(()).unwrap(); + wait.recv().unwrap(); + Ok(()) + }) + .await + } + }); + ready.await.unwrap(); + let calls = Arc::new(AtomicUsize::new(0)); + let context = ExecutionContext::new(Some(Instant::now() + Duration::from_millis(40))); + let queued = tokio::spawn({ + let scheduler = scheduler.clone(); + let calls = calls.clone(); + async move { + scheduler + .run_with_context(context, move || { + calls.fetch_add(1, Ordering::SeqCst); + Ok(()) + }) + .await + } + }); + tokio::time::sleep(Duration::from_millis(80)).await; + release.send(()).unwrap(); + first.await.unwrap().unwrap(); + assert!( + queued + .await + .unwrap() + .unwrap_err() + .to_string() + .contains("deadline") + ); + assert_eq!(calls.load(Ordering::SeqCst), 0); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn cancellation_after_dispatch_keeps_both_reservations_until_work_finishes() { + let admission = BoundedAdmission::new(1, Some(10)); + let tokens = BoundedAdmission::new(1, Some(20)); + let context = ExecutionContext::new(None) + .retain(admission.try_admit(3).unwrap()) + .retain(tokens.try_admit(20).unwrap()); + let (started, ready) = oneshot::channel(); + let (release, wait) = mpsc::channel(); + let scheduler = SerialScheduler::default(); + let task = tokio::spawn({ + let scheduler = scheduler.clone(); + async move { + let _cancel = context.cancel_on_drop(); + scheduler + .run_with_context(context, move || { + started.send(()).unwrap(); + wait.recv().unwrap(); + Ok(()) + }) + .await + } + }); + ready.await.unwrap(); + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + assert!(admission.try_admit(1).is_none()); + assert!(tokens.try_admit(1).is_none()); + release.send(()).unwrap(); + scheduler_barrier(&scheduler).await; + assert!(admission.try_admit(10).is_some()); + assert!(tokens.try_admit(20).is_some()); +} + +async fn scheduler_barrier(scheduler: &SerialScheduler) { + tokio::time::timeout(Duration::from_secs(5), scheduler.run(|| Ok(()))) + .await + .unwrap() + .unwrap(); +} + +#[test] +fn blocking_pool_backlog_checks_cancellation_before_starting_work() { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .max_blocking_threads(1) + .enable_time() + .build() + .unwrap() + .block_on(async { + let (started, ready) = oneshot::channel(); + let (release, wait) = mpsc::channel(); + let blocker = tokio::task::spawn_blocking(move || { + started.send(()).unwrap(); + wait.recv().unwrap(); + }); + ready.await.unwrap(); + let calls = Arc::new(AtomicUsize::new(0)); + let context = ExecutionContext::new(None); + let cancel = context.cancel_on_drop(); + let scheduler = SerialScheduler::default(); + let mut work = Box::pin(scheduler.run_with_context(context, { + let calls = calls.clone(); + move || { + calls.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + })); + std::future::poll_fn(|cx| { + assert!(std::future::Future::poll(work.as_mut(), cx).is_pending()); + std::task::Poll::Ready(()) + }) + .await; + drop(cancel); + drop(work); + release.send(()).unwrap(); + blocker.await.unwrap(); + scheduler_barrier(&scheduler).await; + assert_eq!(calls.load(Ordering::SeqCst), 0); + }); +} + +#[tokio::test] +async fn errors_and_panics_release_owned_capacity() { + let admission = BoundedAdmission::new(1, None); + let scheduler = SerialScheduler::default(); + for panic in [false, true] { + let context = ExecutionContext::new(None).retain(admission.try_admit(1).unwrap()); + let result: Result<()> = scheduler + .run_with_context(context, move || { + if panic { + panic!("expected panic") + } + anyhow::bail!("expected error") + }) + .await; + assert!(result.is_err()); + assert_eq!(admission.report().depth, 0); + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn cancellation_while_waiting_releases_request_capacity_without_dispatch() { + let scheduler = SerialScheduler::default(); + let (started, ready) = oneshot::channel(); + let (release, wait) = mpsc::channel(); + let first = tokio::spawn({ + let scheduler = scheduler.clone(); + async move { + scheduler + .run(move || { + started.send(()).unwrap(); + wait.recv().unwrap(); + Ok(()) + }) + .await + } + }); + ready.await.unwrap(); + let admission = BoundedAdmission::new(1, None); + let context = ExecutionContext::new(None).retain(admission.try_admit(1).unwrap()); + let cancel = context.cancel_on_drop(); + let mut queued = Box::pin(scheduler.run_with_context(context, || -> Result<()> { + panic!("cancelled work executed") + })); + std::future::poll_fn(|cx| { + assert!(std::future::Future::poll(queued.as_mut(), cx).is_pending()); + std::task::Poll::Ready(()) + }) + .await; + assert_eq!(admission.report().depth, 1); + drop(cancel); + drop(queued); + assert_eq!(admission.report().depth, 0); + release.send(()).unwrap(); + first.await.unwrap().unwrap(); + scheduler_barrier(&scheduler).await; +} + +#[test] +fn blocking_pool_backlog_checks_deadline_before_starting_work() { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .max_blocking_threads(1) + .enable_time() + .build() + .unwrap() + .block_on(async { + let (started, ready) = oneshot::channel(); + let (release, wait) = mpsc::channel(); + let blocker = tokio::task::spawn_blocking(move || { + started.send(()).unwrap(); + wait.recv().unwrap(); + }); + ready.await.unwrap(); + let scheduler = SerialScheduler::default(); + let calls = Arc::new(AtomicUsize::new(0)); + let context = ExecutionContext::new(Some(Instant::now() + Duration::from_millis(30))); + let mut work = Box::pin(scheduler.run_with_context(context, { + let calls = calls.clone(); + move || { + calls.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + })); + std::future::poll_fn(|cx| { + assert!(std::future::Future::poll(work.as_mut(), cx).is_pending()); + std::task::Poll::Ready(()) + }) + .await; + tokio::time::sleep(Duration::from_millis(60)).await; + release.send(()).unwrap(); + blocker.await.unwrap(); + assert!(work.await.unwrap_err().to_string().contains("deadline")); + assert_eq!(calls.load(Ordering::SeqCst), 0); + }); +}