diff --git a/.agents/skills/add-new-model/SKILL.md b/.agents/skills/add-new-model/SKILL.md index cd065faf..5aa277a1 100644 --- a/.agents/skills/add-new-model/SKILL.md +++ b/.agents/skills/add-new-model/SKILL.md @@ -44,16 +44,18 @@ redesign serving infrastructure or introduce a general model framework. | `src/models//` | Separate model-specific processors and batch adapters, plus executor checkpoint loading, weights, forward passes, learned heads, device state, and kernel selection. A native Rust worker may live in `native/`, following existing models. | | `src/models//` | Model-family execution shared by implementations that actually need it. Reuse compatible code; extract a common implementation when there is a second user. | | `src/backends/cuda/` or `src/backends/metal/` | Hardware operations and kernels, native host glue, build scripts, and backend integration. GPU processing transforms or tensor packing retain their processor/batcher ownership. | +| `src/runtime/` | Shared native admission and dispatch policy. Reuse the serial scheduler for single-executor work; keep model layouts and device state in model adapters/executors. | | `src/frontend/` | Transport and existing worker adapters when integration requires them. Keep model-specific input/output semantics in processors rather than the transport layer. | | Root `Cargo.toml` and affected crate manifests | Workspace membership, dependencies, binaries, and explicit test registration. Update `Cargo.lock` when dependencies change. | | `recipe//` | Checkpoint preparation/export scripts, setup and launch instructions, and example requests. Reusable runtime implementation belongs in `src/`. | | `tests//` and `tests//` | All test bodies, test helpers, and fixtures, including contract, tokenizer, checkpoint, HTTP, and kernel tests. | -The shared processing/scheduling runtime is planned. Integrate with available -worker interfaces while separating processors and batch adapters from forward -code; adding a model alone does not require implementing that runtime. When -shared scheduling is implemented, it owns queues and admission policy, while -models declare supported layouts and compatibility constraints. The native +The [native runtime](../../../src/runtime/README.md) provides serial FIFO admission +and blocking dispatch per loaded executor. Integrate with available worker +interfaces while separating processors and batch adapters from forward code. +Models declare their execution unit, layouts and compatibility constraints; +shared runtime code owns admission policy. Processing orchestration, queue/token +budgets and dynamic batching remain planned. The native target uses Rust host orchestration and bindings with CUDA/Metal device kernels. **Do not add test code or fixtures under `src/`.** This includes inline Rust test diff --git a/.agents/skills/precheck-pr/SKILL.md b/.agents/skills/precheck-pr/SKILL.md index b9cec776..6cdf66cf 100644 --- a/.agents/skills/precheck-pr/SKILL.md +++ b/.agents/skills/precheck-pr/SKILL.md @@ -33,9 +33,9 @@ user instructions for those actions. - Check ownership against the architecture contracts: transport, independent processing and scheduling, model-specific adapters and executors, and hardware operations. In the native target, Rust orchestrates host work and dispatch; - CUDA/Metal execute device kernels. Shared runtime layers remain planned until - implemented; existing worker pipelines are not proof of dynamic batching. - Consult the relevant `src/frontend/`, `src/models/`, `src/backends/`, or + CUDA/Metal execute device kernels. Serial runtime admission and dispatch are + implemented; further orchestration and dynamic batching remain planned. + Consult the relevant `src/frontend/`, `src/runtime/`, `src/models/`, `src/backends/`, or `recipe/` documentation for the changed component's contract and current status. - For frontend changes, inspect affected request validation, error handling, cancellation, and response behavior, including the model-worker boundary. diff --git a/.agents/skills/self-review/SKILL.md b/.agents/skills/self-review/SKILL.md index e40bbf50..b083550a 100644 --- a/.agents/skills/self-review/SKILL.md +++ b/.agents/skills/self-review/SKILL.md @@ -18,7 +18,10 @@ semantics and distinguish target layers from current worker capabilities. For processing or batching changes, trace prepared inputs through executor layouts and output reconstruction. Apply the architecture contract's batching, -numerical, and lifetime invariants to the affected model. +numerical, and lifetime invariants to the affected model. For shared runtime +admission changes, check FIFO execution units, cancellation before dispatch, +permit/resource retention after dispatch, and release on errors or panics against +the [runtime contract](../../../src/runtime/README.md). ## Conditional A/B checks diff --git a/.agents/skills/system1-omni-review/SKILL.md b/.agents/skills/system1-omni-review/SKILL.md index fca9df73..3f8d8c17 100644 --- a/.agents/skills/system1-omni-review/SKILL.md +++ b/.agents/skills/system1-omni-review/SKILL.md @@ -20,10 +20,10 @@ baseline; target code and configuration determine current paths and commands. ## Review contracts -- Check separate transport, processing, scheduling, model-executor, and backend ownership against the architecture contracts. Model-specific processors and batch adapters preserve semantics; shared runtime code owns orchestration and scheduling when implemented. Rust owns native host orchestration and bindings/dispatch; CUDA/Metal supply device operations, including optional processing transforms and packing. Shared runtime layers are planned, so assess current workers against their implemented contracts. CUDA and Metal need not have identical internals. A directory or pass-through HTTP support does not prove native support or batching. +- Check separate transport, processing, scheduling, model-executor, and backend ownership against the architecture contracts. Model-specific processors and batch adapters preserve semantics; shared runtime code owns orchestration and scheduling when implemented. Rust owns native host orchestration and bindings/dispatch; CUDA/Metal supply device operations, including optional processing transforms and packing. Serial runtime admission and dispatch are implemented; further orchestration and batching remain planned. CUDA and Metal need not have identical internals. A directory or pass-through HTTP support does not prove native support or batching. - For `src/frontend/`, inspect exact envelope/body and end-to-end header preservation, hop-by-hop filtering, backend path prefix/query handling, timeout/status mapping, health forwarding, cancellation and shutdown. The configured backend bypasses proxies, follows no redirects and retries no requests; changes need transport regression evidence, not a model benchmark. - For models, trace serialization → prompt/template → token IDs → one-pass prefill → scoring → response. Preserve the specific model's contract; do not impose Cua-S1's 26-letter choice restriction on a new model that implements score/noul. Verify pinned checkpoint/adapter/tokenizer/configuration identity, shapes/dtypes and readiness after real initialization. -- For processing or batching changes, apply the architecture contract's compatibility, reconstruction, numerical, and lifetime invariants to the affected model. Require executor/kernel evidence for batch capability. +- For processing or batching changes, apply the architecture contract's compatibility, reconstruction, numerical, and lifetime invariants to the affected model. For admission changes, inspect the [runtime contract](../../../src/runtime/README.md), execution-unit granularity, cancellation before/after dispatch and permit/resource lifetimes. Require executor/kernel evidence for batch capability. - For native/FFI/CUDA changes, inspect Rust and C ABI in lockstep, buffer and graph lifetimes, stream/device ownership, sizes/overflow, BF16 rounding locations, and error propagation/recovery. Compile success and ignored GPU tests do not validate kernels. Test changed inputs, reuse, growth, eviction and failure recovery when graph/cache behavior changes. - For performance or numerical claims, freeze baseline/candidate, hardware/toolchain, checkpoint, input manifest, precision, variables, tolerances and budget before comparing. Separate quality/parity from speed, native inference from frontend overhead, and cold/capture costs from warmed latency. Preserve errors and denominators; never relax tolerances or discard failures to produce a winner. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 2e5a7d49..4eae78b7 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -14,16 +14,17 @@ implementations need it. Keep processors and batch adapters separate from the model's forward implementation, even when they live in the same model directory. The README's target architecture shares processing orchestration and scheduling across -models; these layers are planned, while current worker pipelines remain -model-specific. Model executors own weights, forward passes, learned heads, -and device state. +models. Native workers already reuse +[serial admission and dispatch](src/runtime/README.md); processing orchestration +and GPU batching remain planned. Model executors own weights, forward passes, +learned heads, and device state. For native implementations, Rust owns host processing, scheduling, model orchestration, and backend bindings/dispatch. CUDA/Metal provide device operations used by models or processing modules. Document processor inputs, executor layouts and batch constraints, output reconstruction, and state/buffer lifetimes alongside the existing API and numerical contract. Keep current worker -integration usable while shared runtime components remain unimplemented. +integration usable while further shared runtime components remain unimplemented. All test bodies, test helpers and fixtures belong in the repository-level `tests/` tree. Do not add inline test bodies or crate-local `tests/` directories diff --git a/Cargo.lock b/Cargo.lock index 715fbd31..354a6120 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -951,6 +951,7 @@ dependencies = [ "half", "memmap2", "omni-qwen3-5-native", + "omni-runtime", "safetensors 0.8.0", "serde_json", "tokenizers", @@ -988,6 +989,7 @@ dependencies = [ "anyhow", "axum", "omni-qwen3-5-native", + "omni-runtime", "serde_json", "tokenizers", "tokio", @@ -1006,6 +1008,14 @@ dependencies = [ "serde_json", ] +[[package]] +name = "omni-runtime" +version = "0.1.0" +dependencies = [ + "anyhow", + "tokio", +] + [[package]] name = "once_cell" version = "1.21.4" diff --git a/Cargo.toml b/Cargo.toml index e7f33264..de28591e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,3 +1,3 @@ [workspace] -members = ["src/frontend", "src/models/cua_s1/native", "src/models/qwen3_5/native", "src/models/open_jev/native", "src/models/laya"] +members = ["src/frontend", "src/runtime", "src/models/cua_s1/native", "src/models/qwen3_5/native", "src/models/open_jev/native", "src/models/laya"] resolver = "3" diff --git a/README.md b/README.md index 7b34c2d4..92ee71d0 100644 --- a/README.md +++ b/README.md @@ -47,8 +47,8 @@ shared CUDA kernels in this repository. model workers. - **Model-owned execution.** Model executors own weights, forward passes, learned heads, device state, and kernel selection. Native workers have separate - processing and executor modules; shared processing orchestration and scheduling - are planned. + processing and executor modules, with shared FIFO admission and blocking + dispatch in the [native runtime](src/runtime/README.md). - **Native CUDA workers.** The Cua-S1 4B 0.2 `text` adapter and Open-Jev-27B-v1.1 run as native workers with shared CUDA kernels. Cua-S1 also has a Python worker that serves as the correctness reference. @@ -65,10 +65,11 @@ Share processing and scheduling; let each model own its execution. ![System1-Omni target architecture: Rust frontend, independent processing and batching layers, model executors, and CUDA and Metal backends](docs/assets/architecture.svg) -The diagram shows the **target architecture**, not an implemented shared runtime. -Today the frontend forwards HTTP requests to separately running workers, whose -model-specific handlers coordinate independent processors and executors. Shared -processing orchestration, scheduling, and dynamic batching are planned. +The diagram shows the **target architecture**. Today the frontend forwards HTTP +requests to separately running workers, whose handlers coordinate independent +processors and executors. Both native workers use shared FIFO admission and +blocking dispatch per loaded executor. Processing orchestration, batch budgets, +compatibility grouping and dynamic batching remain planned. | Layer | Responsibility | Native target implementation | | --- | --- | --- | @@ -105,14 +106,15 @@ repository root. | Directory | Responsibility | | --- | --- | | [`src/frontend/`](src/frontend/) | Rust serving code, Python worker adapters, and the small engine interface. | +| [`src/runtime/`](src/runtime/README.md) | Shared native execution admission and blocking dispatch. | | [`src/models/`](src/models/) | Model contracts, existing worker pipelines, and model executors, including the shared Qwen3.5/3.8 prefill implementation. | | [`src/backends/cuda/`](src/backends/cuda/) | NVIDIA GPU operations and kernel integration. | | [`src/backends/metal/`](src/backends/metal/) | Apple GPU operations and kernel integration. | | [`recipe/`](recipe/) | Model setup instructions, launch commands, configuration examples, and example requests. | | [`docs/`](docs/) | Project documentation and architecture assets. | -The frontend, both native workers, their shared Qwen3.5/3.8 prefill -implementation and the Laya checkpoint reader are Cargo workspace members. +The frontend, native runtime, both native workers, their shared Qwen3.5/3.8 +prefill implementation and the Laya checkpoint reader are Cargo workspace members. The other model and backend directories currently document planned work; they do not prescribe process boundaries. @@ -158,8 +160,9 @@ requests and a matched comparison with raw HF Transformers and OpenJev-Fast. ## Roadmap The current focus is the native Cua-S1 and Open-Jev CUDA workers and the serving -benchmark harness. Planned work includes the shared processing and scheduling -layer shown above, bounded dynamic batching with batch-capable executors, +benchmark harness. Planned work extends the shared runtime with processing +orchestration, admission budgets, compatibility grouping and bounded dynamic +batching with batch-capable executors, the in-repository LAYA model engine, additional model engines and GPU backends including Metal, and per-model performance measurements as implementations are added and validated. diff --git a/docs/architecture.md b/docs/architecture.md index 720c3ecb..b1988e80 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -11,8 +11,10 @@ to separately running workers. Cua-S1 and Open-Jev have native Rust/CUDA workers that share the [Qwen3.5/3.8 executor](../src/models/qwen3_5/native/). Their model-specific workers coordinate independent processor and executor modules through `prepare` → `execute` → `finish`. The shared Qwen executor accepts one -prompt per forward call; shared processing orchestration, scheduling, and -dynamic batching are planned. +prompt per forward call. Both workers 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. The native workers currently compute their decision heads on the CPU after downloading the final hidden state. GPU head execution belongs to the target @@ -22,9 +24,9 @@ also planned; Python workers retain their documented reference/serving roles. ## Native worker boundaries Both native workers separate `processing.rs` from `executor.rs`; `engine.rs` -assembles them from the checkpoint, and the HTTP handler coordinates the three -stages. Preparation validates the entire request before any forward call and -returns executor inputs plus a response context. The context retains question +assembles them with a `SerialScheduler` per loaded executor, and the HTTP handler +coordinates the three stages. Preparation validates the entire request before +any forward call and returns executor inputs plus a response context. The context retains question and candidate identity, usage, and response metadata outside the executor. | Worker | Prepared executor inputs | Executor outputs | Response finishing | @@ -32,9 +34,12 @@ and candidate identity, usage, and response metadata outside the executor. | Cua-S1 | One unpadded token-ID vector and option count per question, in request order. | One FP32 answer-letter logit vector per question. | Per-question softmax, choice/confidence, ordered answers, and token usage. | | Open-Jev | Token-ID vectors grouped by question, then independent candidate, in request order. | One FP32 learned scalar per candidate in the same grouping. | Add the `noul` false logit of zero, calibrate across each complete question, and restore typed answers, usage, and metadata. | -These input collections are serial work, not GPU batches. Cua-S1 retains its -per-question model lock; Open-Jev retains its request-wide lock. Executors own -the loaded Qwen model and CPU head weights, preserving FP64 accumulation and +These input collections are serial work, not GPU batches. Shared runtime +admission precedes blocking dispatch: Cua-S1 admits one question forward at a +time; Open-Jev admits one complete request. Cua-S1's CPU letter projection stays +outside admission; Open-Jev's scalar heads remain inside its request unit. The +model mutexes guard mutable state, retaining per-question/request granularity. +Executors own the loaded Qwen model and CPU head weights, preserving FP64 accumulation and the existing FP32 rounding and bias order. Finishing checks output cardinality before reconstruction. HTTP validation, error status/body conventions, and real warmup before readiness remain model-specific and unchanged. @@ -82,6 +87,13 @@ probabilities must not be normalized across unrelated questions or requests. ## Scheduling and batching contracts +The implemented serial scheduler limits admitted execution to one unit per loaded +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. + - Batch only work accepted by the same loaded executor, with compatible checkpoint/adapter identity, device, dtype, and input layout. Models declare additional constraints and limits; the scheduler owns queue and admission @@ -107,15 +119,20 @@ ownership even when their kernels live under `src/backends/`. Keep Rust declarations and the native ABI in sync, and rebuild all consumers when the ABI changes. Buffer handles and captured graphs must outlive queued device operations; synchronize before reading host results or reclaiming -storage. Cancellation or timeout does not authorize freeing buffers still in -use by submitted work. Keep readiness tied to successful initialization and -the worker's real warmup; extraction into shared layers preserves that behavior. +storage. Cancellation while waiting for admission removes the queued caller +without dispatch. Once dispatched, the runtime retains admission and captured executor +resources until the blocking closure completes, even if its caller is cancelled. +The executor must synchronize device work before returning; cancellation or +timeout does not authorize freeing buffers still in use by submitted work. Keep +readiness tied to successful initialization and the worker's real warmup; +extraction into shared layers preserves that behavior. ## Integration and validation Keep model-specific processors and batch adapters separate from forward implementations, even when colocated under `src/models//`. Reuse the -current worker interface while shared runtime components are unimplemented. +current worker interface and serial runtime while further runtime components +remain unimplemented. Extract common orchestration when implementing that layer, with explicit scope and current consumers; adding a model alone does not require a new framework. diff --git a/mkdocs.yml b/mkdocs.yml index 4365df2c..f4632ee2 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -27,6 +27,7 @@ exclude_docs: | !/docs/assets/** !/recipe/**/*.md !/src/frontend/README.md + !/src/runtime/README.md !/assets/** **/.*/** @@ -49,6 +50,7 @@ nav: - Architecture and contracts: docs/architecture.md - Supported models and hardware: docs/supported-models.md - Frontend: src/frontend/README.md + - Native runtime: src/runtime/README.md - Recipes: - recipe/README.md - Laya text worker: recipe/laya/README.md diff --git a/src/models/cua_s1/README.md b/src/models/cua_s1/README.md index c6cb89be..36530206 100644 --- a/src/models/cua_s1/README.md +++ b/src/models/cua_s1/README.md @@ -11,10 +11,12 @@ with independent [processing](native/src/processing.rs) and [executor](native/src/executor.rs) modules. Preparation returns token IDs and an option count per question, plus a response context for identity/order and usage. The executor returns FP32 letter logits; response finishing owns per-question -softmax and the choice/confidence interpretation. Shared orchestration and -scheduling belong to the planned worker runtime. Execution remains one prompt -per question with the existing per-question lock and CPU answer-letter projection -after CUDA prefill. The contracts below describe its existing behavior. +softmax and the choice/confidence interpretation. The engine owns a +[shared serial scheduler](../../runtime/README.md) that admits each question's +forward before blocking dispatch. Execution remains one prompt per question +with a model-state mutex and CPU answer-letter projection after CUDA prefill, +outside runtime admission. Processing orchestration and GPU batching remain +planned. The contracts below describe its existing behavior. ## Pinned revisions diff --git a/src/models/cua_s1/native/Cargo.toml b/src/models/cua_s1/native/Cargo.toml index 9da20a38..a8eb3ae1 100644 --- a/src/models/cua_s1/native/Cargo.toml +++ b/src/models/cua_s1/native/Cargo.toml @@ -15,6 +15,7 @@ axum = "0.8.8" half = "2.7.1" # the CUDA kernels live in libqwen3_5_cuda.so, loaded at run time omni-qwen3-5-native = { path = "../../qwen3_5/native" } +omni-runtime = { path = "../../../runtime" } memmap2 = "0.9.9" safetensors = "0.8.0" serde_json = { version = "1.0.149", features = ["float_roundtrip", "preserve_order"] } diff --git a/src/models/cua_s1/native/src/engine.rs b/src/models/cua_s1/native/src/engine.rs index 9a7b25f2..21e2fd18 100644 --- a/src/models/cua_s1/native/src/engine.rs +++ b/src/models/cua_s1/native/src/engine.rs @@ -3,6 +3,7 @@ use std::path::Path; use anyhow::{Context, Result, ensure}; +use omni_runtime::SerialScheduler; use crate::contract::LETTERS; use crate::executor::Executor; @@ -10,6 +11,7 @@ use crate::processing::Processor; pub struct Engine { pub processor: Processor, + pub scheduler: SerialScheduler, pub executor: Executor, } @@ -34,6 +36,7 @@ impl Engine { let executor = Executor::load(dir, library, &ids).await?; Ok(Self { processor, + scheduler: SerialScheduler::default(), executor, }) } diff --git a/src/models/cua_s1/native/src/executor.rs b/src/models/cua_s1/native/src/executor.rs index a3b19c40..3db4c425 100644 --- a/src/models/cua_s1/native/src/executor.rs +++ b/src/models/cua_s1/native/src/executor.rs @@ -4,6 +4,7 @@ use std::path::Path; use std::sync::{Arc, Mutex}; use anyhow::{Context, Result, ensure}; +use omni_runtime::SerialScheduler; use serde_json::Value as Json; use crate::model::Model; @@ -33,19 +34,24 @@ impl Executor { }) } - /// Serial prefill, retaining the existing per-question model lock. + /// Admit each question through the worker's shared runtime scheduler. /// Returns FP32 letter logits in input order; softmax belongs to processing. - pub async fn execute(&self, inputs: Vec) -> Result>> { + pub async fn execute( + &self, + scheduler: &SerialScheduler, + inputs: Vec, + ) -> Result>> { let mut rows = Vec::with_capacity(inputs.len()); for input in inputs { let model = self.model.clone(); - let last = tokio::task::spawn_blocking(move || { - model - .lock() - .map_err(|_| anyhow::anyhow!("poisoned"))? - .forward(&input.ids) - }) - .await??; + let last = scheduler + .run(move || { + model + .lock() + .map_err(|_| anyhow::anyhow!("poisoned"))? + .forward(&input.ids) + }) + .await?; rows.push(letter_logits(&self.letters, &last, input.n_options)?); } Ok(rows) diff --git a/src/models/cua_s1/native/src/main.rs b/src/models/cua_s1/native/src/main.rs index 3e464502..f2912b9f 100644 --- a/src/models/cua_s1/native/src/main.rs +++ b/src/models/cua_s1/native/src/main.rs @@ -35,7 +35,10 @@ async fn decide(engine: &Engine, raw: &[u8]) -> Response { Err(e) => return reply(e.status, json!({"detail": e.message})), }; let result = async { - let rows = engine.executor.execute(prepared.inputs).await?; + let rows = engine + .executor + .execute(&engine.scheduler, prepared.inputs) + .await?; prepared.context.finish(rows) } .await; diff --git a/src/models/open_jev/README.md b/src/models/open_jev/README.md index 7099582a..a6b22822 100644 --- a/src/models/open_jev/README.md +++ b/src/models/open_jev/README.md @@ -25,5 +25,6 @@ calibration, usage, and metadata. The executor owns the trained head and returns one FP32 scalar per candidate. Response finishing adds the `noul` false baseline and applies calibrated normalization across each complete question. The worker retains request-wide model locking and independent single-prompt execution, with -the scalar head on the CPU after CUDA prefill. Shared runtime scheduling and -dynamic batching remain planned. +the scalar head on the CPU after CUDA prefill. The engine owns a +[shared serial scheduler](../../runtime/README.md) that admits the complete request +before blocking dispatch. GPU batching and batch budgets remain planned. diff --git a/src/models/open_jev/native/Cargo.toml b/src/models/open_jev/native/Cargo.toml index 5384417d..346440c6 100644 --- a/src/models/open_jev/native/Cargo.toml +++ b/src/models/open_jev/native/Cargo.toml @@ -9,6 +9,7 @@ description = "Native Rust/CUDA worker for Open-Jev-27B-v1.1" anyhow = "1.0.100" axum = "0.8.8" omni-qwen3-5-native = { path = "../../qwen3_5/native" } +omni-runtime = { path = "../../../runtime" } serde_json = { version = "1.0.149", features = ["float_roundtrip", "preserve_order"] } tokenizers = { version = "=0.22.2", default-features = false, features = ["onig"] } tokio = { version = "1.49.0", features = ["macros", "net", "rt-multi-thread", "sync"] } diff --git a/src/models/open_jev/native/src/engine.rs b/src/models/open_jev/native/src/engine.rs index fb00e63e..5e5d3d66 100644 --- a/src/models/open_jev/native/src/engine.rs +++ b/src/models/open_jev/native/src/engine.rs @@ -3,6 +3,7 @@ use std::path::Path; use anyhow::{Context, Result, ensure}; +use omni_runtime::SerialScheduler; use serde_json::Value; use crate::contract::MODEL_ID; @@ -13,6 +14,7 @@ pub use crate::contract::{BASE_REVISION, CHECKPOINT_REVISION}; pub struct Engine { pub processor: Processor, + pub scheduler: SerialScheduler, pub executor: Executor, } @@ -52,6 +54,7 @@ impl Engine { let executor = Executor::load(dir, library, head).await?; Ok(Self { processor, + scheduler: SerialScheduler::default(), executor, }) } diff --git a/src/models/open_jev/native/src/executor.rs b/src/models/open_jev/native/src/executor.rs index 195473ac..9504c02e 100644 --- a/src/models/open_jev/native/src/executor.rs +++ b/src/models/open_jev/native/src/executor.rs @@ -5,6 +5,7 @@ use std::sync::{Arc, Mutex}; use anyhow::{Context, Result, ensure}; use omni_qwen3_5_native::model::{Config, Model}; +use omni_runtime::SerialScheduler; use serde_json::Value; #[derive(Clone)] @@ -79,24 +80,29 @@ impl Executor { } /// Inputs and outputs are grouped by question, then candidate, in prepared order. - /// Each candidate is one unpadded forward call; the model lock spans this request. - pub async fn execute(&self, ids: Vec>>) -> Result>> { + /// Admit one whole request; the model lock spans its independent candidate calls. + pub async fn execute( + &self, + scheduler: &SerialScheduler, + ids: Vec>>, + ) -> Result>> { let model = self.model.clone(); let head = self.head.clone(); - tokio::task::spawn_blocking(move || { - let mut model = model - .lock() - .map_err(|_| anyhow::anyhow!("poisoned model"))?; - ids.iter() - .map(|candidates| { - candidates - .iter() - .map(|ids| head.score(model.forward(ids)?)) - .collect::>>() - }) - .collect::>>() - }) - .await? + scheduler + .run(move || { + let mut model = model + .lock() + .map_err(|_| anyhow::anyhow!("poisoned model"))?; + ids.iter() + .map(|candidates| { + candidates + .iter() + .map(|ids| head.score(model.forward(ids)?)) + .collect::>>() + }) + .collect::>>() + }) + .await } } diff --git a/src/models/open_jev/native/src/main.rs b/src/models/open_jev/native/src/main.rs index 4dbfe2df..7352bff5 100644 --- a/src/models/open_jev/native/src/main.rs +++ b/src/models/open_jev/native/src/main.rs @@ -27,7 +27,10 @@ async fn decide(engine: &Engine, raw: &[u8]) -> Response { Err(e) => return error(StatusCode::UNPROCESSABLE_ENTITY, e), }; let result = async { - let rows = engine.executor.execute(prepared.inputs).await?; + let rows = engine + .executor + .execute(&engine.scheduler, prepared.inputs) + .await?; prepared.context.finish(rows) } .await; diff --git a/src/runtime/Cargo.toml b/src/runtime/Cargo.toml new file mode 100644 index 00000000..67a424da --- /dev/null +++ b/src/runtime/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "omni-runtime" +version = "0.1.0" +edition = "2024" +publish = false +description = "Shared native worker execution admission and dispatch" + +[dependencies] +anyhow = "1.0.100" +tokio = { version = "1.49.0", features = ["rt", "sync"] } + +[dev-dependencies] +tokio = { version = "1.49.0", features = ["macros", "rt-multi-thread", "time"] } + +[[test]] +name = "scheduler" +path = "../../tests/runtime/scheduler.rs" diff --git a/src/runtime/README.md b/src/runtime/README.md new file mode 100644 index 00000000..c06c3036 --- /dev/null +++ b/src/runtime/README.md @@ -0,0 +1,48 @@ +# Native execution runtime + +`omni-runtime` provides `SerialScheduler` for the Cua-S1 and Open-Jev native +workers. Their engines assemble a processor, one scheduler for the loaded +executor, and the executor. HTTP handlers still coordinate +`prepare` → `execute` → `finish`. + +The scheduler owns FIFO admission and blocking dispatch. A caller waits +asynchronously for one execution permit before entering Tokio's blocking pool. +Cloning a scheduler shares that admission queue; separate loaded executors have +separate queues. The model mutex remains a guard for Rust's mutable executor +state. Waiting requests no longer occupy blocking threads waiting for that mutex. + +| Worker | Admitted unit | Work outside admission | +| --- | --- | --- | +| 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. | + +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 +model identifiers, question/candidate mappings, tensor layouts or backend APIs. +The existing numerical, ordering, whole-request failure and readiness contracts +remain in the processors, executors and worker handlers. + +Cancellation while waiting removes the caller from admission without dispatch. +After dispatch, the blocking closure retains the permit and its captured +resources until completion, even if the caller disconnects or stops waiting. +Executors must finish queued device work before returning; the current Qwen +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 still accepts one prompt per forward; this scheduler does not pack +inputs 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: + +```sh +cargo test --locked -p omni-runtime +``` + +See the [architecture contracts](../../docs/architecture.md) and +[benchmark protocol](../../benchmarks/README.md) when changing admission policy. diff --git a/src/runtime/src/lib.rs b/src/runtime/src/lib.rs new file mode 100644 index 00000000..85c4e1db --- /dev/null +++ b/src/runtime/src/lib.rs @@ -0,0 +1,46 @@ +//! Shared admission and blocking dispatch for a loaded native executor. + +use std::sync::Arc; + +use anyhow::{Context, Result}; +use tokio::sync::Semaphore; + +/// FIFO admission of one execution unit at a time. Clones share the same queue. +/// Each loaded executor gets its own scheduler; models define the unit of work. +#[derive(Clone)] +pub struct SerialScheduler { + admission: Arc, +} + +impl Default for SerialScheduler { + fn default() -> Self { + Self { + admission: Arc::new(Semaphore::new(1)), + } + } +} + +impl SerialScheduler { + /// Wait asynchronously before dispatching blocking model work. + /// Cancellation while waiting removes the work from admission. Once dispatched, + /// the work retains its permit and captured resources until the closure returns, + /// even when its caller stops waiting. Device work must complete before return. + pub async fn run(&self, work: F) -> Result + where + F: FnOnce() -> Result + Send + 'static, + T: Send + 'static, + { + let permit = self + .admission + .clone() + .acquire_owned() + .await + .expect("private admission semaphore is never closed"); + tokio::task::spawn_blocking(move || { + let _permit = permit; + work() + }) + .await + .context("execution task failed")? + } +} diff --git a/tests/runtime/scheduler.rs b/tests/runtime/scheduler.rs new file mode 100644 index 00000000..b5999340 --- /dev/null +++ b/tests/runtime/scheduler.rs @@ -0,0 +1,186 @@ +use std::future::{Future, poll_fn}; +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex, mpsc}; +use std::task::Poll; +use std::time::Duration; + +use anyhow::Result; +use omni_runtime::SerialScheduler; +use tokio::sync::oneshot; +use tokio::task::JoinHandle; + +async fn deadline(future: impl Future) -> T { + tokio::time::timeout(Duration::from_secs(5), future) + .await + .expect("scheduler made no progress") +} + +async fn poll_waiting(mut future: Pin<&mut impl Future>) { + poll_fn(|cx| { + assert!(future.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; +} + +async fn blocked(scheduler: SerialScheduler) -> (JoinHandle>, mpsc::Sender<()>) { + let (started, ready) = oneshot::channel(); + let (release, wait) = mpsc::channel(); + let task = tokio::spawn(async move { + scheduler + .run(move || { + started.send(()).unwrap(); + wait.recv().unwrap(); + Ok(()) + }) + .await + }); + deadline(ready).await.unwrap(); + (task, release) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn queued_units_run_serially_in_fifo_order() { + let scheduler = SerialScheduler::default(); + let (first, release) = blocked(scheduler.clone()).await; + let order = Arc::new(Mutex::new(Vec::new())); + let mut second = Box::pin(scheduler.run({ + let order = order.clone(); + move || { + order.lock().unwrap().push(2); + Ok(20) + } + })); + let mut third = Box::pin(scheduler.run({ + let order = order.clone(); + move || { + order.lock().unwrap().push(3); + Ok(30) + } + })); + poll_waiting(second.as_mut()).await; + poll_waiting(third.as_mut()).await; + // A spare blocking thread exists, but neither queued unit may execute yet. + assert!( + tokio::time::timeout(Duration::from_millis(100), second.as_mut()) + .await + .is_err() + ); + assert!(order.lock().unwrap().is_empty()); + release.send(()).unwrap(); + deadline(first).await.unwrap().unwrap(); + let (a, b) = deadline(async { tokio::join!(second, third) }).await; + assert_eq!((a.unwrap(), b.unwrap()), (20, 30)); + assert_eq!(*order.lock().unwrap(), [2, 3]); +} + +#[test] +fn cancelling_queued_work_never_dispatches_it() { + // One blocking thread makes premature dispatch observable after the held job. + tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .max_blocking_threads(1) + .enable_time() + .build() + .unwrap() + .block_on(async { + let scheduler = SerialScheduler::default(); + let (first, release) = blocked(scheduler.clone()).await; + let calls = Arc::new(AtomicUsize::new(0)); + let mut cancelled = Box::pin(scheduler.run({ + let calls = calls.clone(); + move || { + calls.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + })); + poll_waiting(cancelled.as_mut()).await; + drop(cancelled); + release.send(()).unwrap(); + deadline(first).await.unwrap().unwrap(); + assert_eq!(deadline(scheduler.run(|| Ok(42))).await.unwrap(), 42); + assert_eq!(calls.load(Ordering::SeqCst), 0); + }); +} + +struct DeviceState(Arc); + +impl Drop for DeviceState { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn cancelling_dispatched_work_retains_admission_and_resources_until_completion() { + let scheduler = SerialScheduler::default(); + let dropped = Arc::new(AtomicBool::new(false)); + let state = Arc::new(DeviceState(dropped.clone())); + let (started, ready) = oneshot::channel(); + let (release, wait) = mpsc::channel(); + let first = tokio::spawn({ + let scheduler = scheduler.clone(); + let state = state.clone(); + async move { + scheduler + .run(move || { + started.send(()).unwrap(); + wait.recv().unwrap(); + drop(state); + Ok(()) + }) + .await + } + }); + deadline(ready).await.unwrap(); + first.abort(); + assert!(deadline(first).await.unwrap_err().is_cancelled()); + drop(state); + assert!(!dropped.load(Ordering::SeqCst)); + let mut next = Box::pin(scheduler.run({ + let dropped = dropped.clone(); + move || { + assert!(dropped.load(Ordering::SeqCst)); + Ok(()) + } + })); + assert!( + tokio::time::timeout(Duration::from_millis(100), next.as_mut()) + .await + .is_err() + ); + assert!(!dropped.load(Ordering::SeqCst)); + release.send(()).unwrap(); + deadline(next).await.unwrap(); + assert!(dropped.load(Ordering::SeqCst)); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn errors_and_panics_release_admission() { + let scheduler = SerialScheduler::default(); + let error = scheduler + .run(|| -> Result<()> { anyhow::bail!("forward failed") }) + .await; + assert_eq!(error.unwrap_err().to_string(), "forward failed"); + let panic = scheduler + .run(|| -> Result<()> { panic!("executor panicked") }) + .await; + assert!( + panic + .unwrap_err() + .to_string() + .contains("execution task failed") + ); + assert_eq!(deadline(scheduler.run(|| Ok(42))).await.unwrap(), 42); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn loaded_executors_have_independent_admission() { + let first_scheduler = SerialScheduler::default(); + let second_scheduler = SerialScheduler::default(); + let (first, release) = blocked(first_scheduler).await; + assert_eq!(deadline(second_scheduler.run(|| Ok(42))).await.unwrap(), 42); + release.send(()).unwrap(); + deadline(first).await.unwrap().unwrap(); +}