From 1363d477630581c47e68e6af9de2940d4049c593 Mon Sep 17 00:00:00 2001 From: xiaoyu <1259084489@qq.com> Date: Mon, 5 Oct 2026 15:09:51 +0800 Subject: [PATCH] runtime: add the worker-side engine contract MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review of the in-process serving work asked for its placement to change: the request lifecycle belongs to the worker-side runtime, not to the frontend. `docs/architecture.md` gives the frontend HTTP transport, forwarding and response delivery, and gives admission, queues and back-pressure to the runtime, and `omni-runtime` already carries the admission half as `SerialScheduler`. This adds the other half: the contract a model implements to be served. It is two required methods — say whether work may be sent, and accept it — plus two optional ones for observability: fn readiness(&self) -> Readiness; fn submit(&self, body: Vec, deadline: Instant) -> Reply; fn report(&self) -> Option { None } fn failure(&self) -> Option { None } `submit` returns a reply rather than a result, so a refusal travels the same path as an answer. A caller then has one thing to await and can still tell "there was no room" from "the runtime dropped my receiver", which two channels would make easy to confuse. The types a worker has to name — `Readiness`, `Report`, `Answer`, `EngineError` — live here with the trait, because they are what the contract is made of. Queue state is published by the worker that owns the queue: `Report` carries depth, capacity and a rejected count. Nothing here parses a decision envelope. Request bytes go in and response bytes come out, so a field the runtime has never heard of survives and a worker's error text cannot produce invalid JSON. The frontend gains the transport that serves such a worker: the two routes, the request budget measured from the headers, the body limit, the status mapping and the queue depth a reply publishes. Nothing else changes for the forwarding path. One timeout semantic, which the same review asked for. A request that cannot be answered inside its budget is `504` whatever part of the budget ran out — upload, queue or inference. `503` is reserved for a worker that cannot take work at all: still loading, failed, out of capacity, or gone. The forwarding path answers 504 for a backend timeout too, so a caller does not have to learn which mode it is talking to in order to read one. Tests are under the repository `tests/` tree per CONTRIBUTING: `tests/runtime/engine.rs` exercises the contract through its public API with two workers (an inline one and one that owns a queue), and `tests/frontend/in_process.rs` runs the transport over real sockets for readiness, byte preservation, the body limit, refusal without retry, error mapping, queue depth, and a worker that never answers. Verified on this commit: fmt and clippy -D warnings clean over the workspace, 70 tests passing across five consecutive runs, release build passes, and the existing forwarding suites are untouched and still green. --- Cargo.lock | 1 + src/frontend/Cargo.toml | 5 + src/frontend/src/engine.rs | 212 ++++++++++++++++++++++++ src/frontend/src/lib.rs | 10 ++ src/runtime/Cargo.toml | 4 + src/runtime/src/engine.rs | 140 ++++++++++++++++ src/runtime/src/lib.rs | 4 +- tests/frontend/in_process.rs | 310 +++++++++++++++++++++++++++++++++++ tests/runtime/engine.rs | 220 +++++++++++++++++++++++++ 9 files changed, 905 insertions(+), 1 deletion(-) create mode 100644 src/frontend/src/engine.rs create mode 100644 src/runtime/src/engine.rs create mode 100644 tests/frontend/in_process.rs create mode 100644 tests/runtime/engine.rs diff --git a/Cargo.lock b/Cargo.lock index 11b8ba2..06588b6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1004,6 +1004,7 @@ version = "0.1.0" dependencies = [ "axum", "futures-util", + "omni-runtime", "reqwest", "tokio", ] diff --git a/src/frontend/Cargo.toml b/src/frontend/Cargo.toml index 4ca9d80..a4c16d4 100644 --- a/src/frontend/Cargo.toml +++ b/src/frontend/Cargo.toml @@ -7,6 +7,7 @@ description = "Model-independent HTTP frontend for Jev-style decision workers" [dependencies] axum = "0.8" +omni-runtime = { path = "../runtime" } reqwest = { version = "0.12.28", default-features = false, features = ["rustls-tls", "stream"] } tokio = { version = "1", features = ["macros", "net", "rt-multi-thread", "signal"] } @@ -17,3 +18,7 @@ tokio = { version = "1", features = ["io-util", "process", "time"] } [[test]] name = "forwarding_limits" path = "../../tests/frontend/forwarding_limits.rs" + +[[test]] +name = "in_process" +path = "../../tests/frontend/in_process.rs" diff --git a/src/frontend/src/engine.rs b/src/frontend/src/engine.rs new file mode 100644 index 0000000..3a0c72b --- /dev/null +++ b/src/frontend/src/engine.rs @@ -0,0 +1,212 @@ +//! In-process serving: the `/v1/systemone` surface answered by a worker in this binary. +//! +//! The contract this serves is [`omni_runtime::engine::Engine`] — admission and the queue +//! belong to the worker-side runtime, and this module is the transport above it. It owns what +//! the transport owns and nothing else: the two routes, the request budget measured from the +//! headers, the body limit, the status mapping, and the queue depth a reply publishes. +//! +//! Nothing here parses a decision envelope. Request bytes go to the engine and response bytes +//! come back, so a field the frontend has never heard of survives and a model's error text +//! cannot produce invalid JSON. +//! +//! **One timeout semantic.** A request that cannot be answered inside its budget is `504`, +//! whatever part of the budget ran out — the upload, the queue, or the inference. `503` is +//! reserved for a worker that cannot take work at all: still loading, failed, out of capacity, +//! or gone. The forwarding path in [`crate`] answers the same way, so a caller does not have to +//! learn which mode it is talking to in order to read a timeout. + +use std::{ + future::Future, + sync::Arc, + time::{Duration, Instant}, +}; + +use axum::{ + Router, + body::{Body, to_bytes}, + extract::{Request, State}, + http::{HeaderName, HeaderValue, StatusCode, header}, + response::Response, + routing::{get, post}, +}; +use omni_runtime::engine::{Answer, Engine, EngineError, Readiness}; +use tokio::net::TcpListener; + +use crate::BoxError; + +/// Work accepted and not yet completed, as reported when a response was produced. Only sent +/// when the worker reports one, so a missing header means "not observed", not "empty". +pub const QUEUE_DEPTH: HeaderName = HeaderName::from_static("x-queue-depth"); + +/// The transport's own limits: what it will read, and how long it will wait. +#[derive(Clone, Debug)] +pub struct ServiceConfig { + /// Total request budget, covering upload, queueing and inference. + pub timeout: Duration, + pub max_body: usize, +} + +impl Default for ServiceConfig { + fn default() -> Self { + Self { + timeout: Duration::from_secs(30), + max_body: 1024 * 1024, + } + } +} + +#[derive(Clone)] +struct Service { + engine: Arc, + config: ServiceConfig, +} + +/// The in-process router: `POST /v1/systemone` and `GET /health`. +pub fn app(engine: Arc, config: ServiceConfig) -> Router { + Router::new() + .route("/v1/systemone", post(infer)) + .route("/health", get(health)) + .with_state(Service { engine, config }) +} + +/// Serves until `shutdown` resolves, then drains in-flight connections. +/// +/// Draining here covers the transport only. Stopping admission and joining the worker's own +/// thread belong to the caller, and happen after this returns. +pub async fn serve( + listener: TcpListener, + router: Router, + shutdown: impl Future + Send + 'static, +) -> Result<(), BoxError> { + axum::serve(listener, router) + .with_graceful_shutdown(shutdown) + .await?; + Ok(()) +} + +async fn health(State(service): State) -> Response { + let readiness = service.engine.readiness(); + let report = match readiness { + Readiness::Ready => service.engine.report(), + // A starting worker knows nothing yet, and a failed one has already said why. + _ => None, + }; + let (status, value) = match readiness { + Readiness::Ready => (StatusCode::OK, "ok"), + Readiness::Starting => (StatusCode::SERVICE_UNAVAILABLE, "starting"), + Readiness::Failed => (StatusCode::SERVICE_UNAVAILABLE, "failed"), + }; + let mut fields = vec![format!("\"status\":\"{value}\"")]; + if let (Readiness::Failed, Some(reason)) = (readiness, service.engine.failure()) { + fields.push(format!("\"reason\":{}", quote(&reason))); + } + if let Some(report) = report { + fields.push(format!("\"depth\":{}", report.depth)); + fields.push(format!("\"capacity\":{}", report.capacity)); + fields.push(format!("\"rejected\":{}", report.rejected)); + } + json(status, &format!("{{{}}}", fields.join(","))) +} + +async fn infer(State(service): State, request: Request) -> Response { + // The budget starts here, so a slow upload is charged to the caller rather than to the + // worker, and work that cannot finish in time is refused before it is admitted. + let deadline = Instant::now() + service.config.timeout; + if service.engine.readiness() != Readiness::Ready { + return fail(StatusCode::SERVICE_UNAVAILABLE, "model unavailable"); + } + + // Reading the body is charged to the same budget. Too large and too slow are separate + // answers so a caller can tell a big request from a late one. + let body = match tokio::time::timeout_at( + deadline.into(), + to_bytes(request.into_body(), service.config.max_body), + ) + .await + { + Ok(Ok(bytes)) => bytes.to_vec(), + Ok(Err(_)) => return fail(StatusCode::PAYLOAD_TOO_LARGE, "request body too large"), + Err(_) => { + return fail( + StatusCode::GATEWAY_TIMEOUT, + "request timed out while uploading", + ); + } + }; + if Instant::now() >= deadline { + return fail(StatusCode::GATEWAY_TIMEOUT, "request timed out"); + } + + let reply = service.engine.submit(body, deadline); + match tokio::time::timeout_at(deadline.into(), reply).await { + Ok(Ok(Ok(answer))) => with_depth(answer), + Ok(Ok(Err(error))) => error_response(error), + // The worker dropped a request it had accepted. Not bad input and not a failed + // inference, so it reads as lost capacity. + Ok(Err(_)) => fail(StatusCode::SERVICE_UNAVAILABLE, "engine stopped"), + Err(_) => fail(StatusCode::GATEWAY_TIMEOUT, "request timed out"), + } +} + +fn error_response(error: EngineError) -> Response { + let (status, message) = match error { + EngineError::InvalidRequest(detail) => (StatusCode::BAD_REQUEST, detail), + EngineError::Busy => ( + StatusCode::SERVICE_UNAVAILABLE, + "inference queue full".into(), + ), + EngineError::Unavailable => (StatusCode::SERVICE_UNAVAILABLE, "model unavailable".into()), + EngineError::InferenceFailed => { + (StatusCode::INTERNAL_SERVER_ERROR, "inference failed".into()) + } + }; + fail(status, &message) +} + +fn fail(status: StatusCode, message: &str) -> Response { + json(status, &format!("{{\"error\":{}}}", quote(message))) +} + +/// Answers with the worker's body, publishing the queue depth it reported. +fn with_depth(answer: Answer) -> Response { + let mut response = Response::new(Body::from(answer.body)); + let headers = response.headers_mut(); + headers.insert( + header::CONTENT_TYPE, + HeaderValue::from_static("application/json"), + ); + if let Ok(depth) = HeaderValue::from_str(&answer.depth.to_string()) { + headers.insert(QUEUE_DEPTH, depth); + } + response +} + +fn json(status: StatusCode, body: &str) -> Response { + let mut response = Response::new(Body::from(body.to_owned())); + *response.status_mut() = status; + response.headers_mut().insert( + header::CONTENT_TYPE, + HeaderValue::from_static("application/json"), + ); + response +} + +/// Serializes a string as a JSON string literal. Hand-rolled so this crate keeps parsing +/// nothing: a worker's error text must not be able to produce invalid JSON. +fn quote(value: &str) -> String { + let mut out = String::with_capacity(value.len() + 2); + out.push('"'); + for c in value.chars() { + match c { + '"' => out.push_str("\\\""), + '\\' => out.push_str("\\\\"), + '\n' => out.push_str("\\n"), + '\r' => out.push_str("\\r"), + '\t' => out.push_str("\\t"), + c if (c as u32) < 0x20 => out.push_str(&format!("\\u{:04x}", c as u32)), + c => out.push(c), + } + } + out.push('"'); + out +} diff --git a/src/frontend/src/lib.rs b/src/frontend/src/lib.rs index e57db90..44f9eab 100644 --- a/src/frontend/src/lib.rs +++ b/src/frontend/src/lib.rs @@ -1,4 +1,14 @@ //! Jev HTTP transport. The worker owns request parsing and inference. +//! +//! Two ways to answer `/v1/systemone`, sharing one HTTP surface: +//! +//! - [`Config`] and [`app`] forward to a worker that already speaks HTTP. The frontend owns no +//! model state. +//! - [`engine`] serves a worker linked into this binary through +//! [`omni_runtime::engine::Engine`]. Admission, the queue and readiness are the runtime's; +//! the frontend owns the transport and the status mapping. + +pub mod engine; use std::{env, error::Error, net::SocketAddr, time::Duration}; diff --git a/src/runtime/Cargo.toml b/src/runtime/Cargo.toml index 67a424d..007046f 100644 --- a/src/runtime/Cargo.toml +++ b/src/runtime/Cargo.toml @@ -15,3 +15,7 @@ tokio = { version = "1.49.0", features = ["macros", "rt-multi-thread", "time"] } [[test]] name = "scheduler" path = "../../tests/runtime/scheduler.rs" + +[[test]] +name = "engine" +path = "../../tests/runtime/engine.rs" diff --git a/src/runtime/src/engine.rs b/src/runtime/src/engine.rs new file mode 100644 index 0000000..8e40724 --- /dev/null +++ b/src/runtime/src/engine.rs @@ -0,0 +1,140 @@ +//! The worker-side contract for serving one decision at a time. +//! +//! A worker in this repository is three things: model code that turns request bytes into +//! response bytes, admission that decides what may run, and HTTP that carries the result. This +//! module defines the middle seam, so a model supplies only the first and inherits the rest. +//! +//! It lives in the runtime rather than in the frontend because admission is the runtime's: +//! `docs/architecture.md` gives the frontend transport and response delivery, and gives queues, +//! back-pressure and request bookkeeping to the worker-side runtime. A worker that implements +//! [`Engine`] can be dispatched through any admission policy here — including +//! [`crate::SerialScheduler`] — without the transport knowing which one is in use. +//! +//! Nothing here parses a decision envelope. Request bytes go in, response bytes come out, so a +//! field the runtime has never heard of survives and a model's error text cannot produce +//! invalid JSON. + +use std::time::Instant; + +use tokio::sync::oneshot; + +/// How far the worker has got. An HTTP transport publishes this as readiness; anything that +/// only needs to know whether work may be sent can compare against [`Readiness::Ready`]. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum Readiness { + /// Loading, or warming up. Work must not be sent yet. + Starting, + Ready, + /// Loading or execution failed. [`Engine::failure`] explains why, and the worker stays + /// alive so an operator can read the reason rather than guess from an exit code. + Failed, +} + +/// What the worker reports about the work it is holding. Advisory: readiness never waits on it, +/// and a field the worker does not measure is left out rather than defaulted. +#[derive(Clone, Debug, Default)] +pub struct Report { + /// Work accepted and not yet completed, running or queued. + pub depth: usize, + /// The depth at which admission starts refusing, so a reader can see how close the queue is + /// to back-pressure without knowing the configuration. + pub capacity: usize, + /// Work the worker did not complete since it started: refused because there was no room, or + /// accepted and then dropped because its caller had gone. Monotonic. + pub rejected: u64, +} + +/// The answer to one request: the bytes to return, and the queue depth at the moment the work +/// was admitted, so a transport can publish back-pressure before it becomes a refusal. +#[derive(Clone, Debug)] +pub struct Answer { + pub body: Vec, + pub depth: usize, +} + +/// Why a worker could not answer. Each variant maps to one status code and one message; that +/// mapping belongs to the transport above, which is why the distinction is made here rather +/// than left to a string. +#[derive(Debug)] +pub enum EngineError { + /// The caller's body is not a request this worker accepts. + InvalidRequest(String), + /// Out of capacity right now. The caller may retry. + Busy, + /// Not able to accept work at all: still loading, drained, or dead. + Unavailable, + /// The request was well formed and the worker failed to execute it. + InferenceFailed, +} + +impl std::fmt::Display for EngineError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::InvalidRequest(detail) => write!(f, "invalid request: {detail}"), + Self::Busy => f.write_str("no queue capacity"), + Self::Unavailable => f.write_str("engine unavailable"), + Self::InferenceFailed => f.write_str("inference failed"), + } + } +} + +impl std::error::Error for EngineError {} + +/// The outcome of one request, delivered exactly once. A refusal travels through the same +/// channel as an answer, so a caller has one thing to await and one place to read the reason: +/// +/// ```no_run +/// # use std::time::Instant; +/// # use omni_runtime::engine::Engine; +/// # async fn f(engine: &dyn Engine) { +/// match engine.submit(b"{}".to_vec(), Instant::now()).await { +/// Ok(answer) => { /* the response body is answer.body */ } +/// Err(refused) => { /* an EngineError the transport maps to a status */ } +/// } +/// # } +/// ``` +pub type Reply = oneshot::Receiver>; + +/// What a model implements to be served. +/// +/// Two methods are the whole contract: say whether work may be sent, and accept it. An +/// implementation that answers inline does both directly; one that owns a queue admits in +/// [`Engine::submit`] and reports its state through [`Engine::report`]. +pub trait Engine: Send + Sync + 'static { + /// Non-blocking. Polled on every request, so it must not lock behind queued work. + fn readiness(&self) -> Readiness; + + /// Accepts one request if there is room, and returns what will carry its outcome. + /// + /// This does not fail: a refusal is an outcome, so it is sent through the reply. Keeping + /// admission and the answer on one path is what lets a caller tell "there was no room" from + /// "the runtime dropped my receiver", which two channels would make easy to confuse. + /// + /// `deadline` is the caller's budget, covering everything from the request headers to the + /// response. Work that cannot finish inside it, or whose caller has already gone, is better + /// refused than started. + fn submit(&self, body: Vec, deadline: Instant) -> Reply; + + /// What to publish about the work being held. `None` leaves those fields out. + fn report(&self) -> Option { + None + } + + /// Set only when [`Engine::readiness`] is [`Readiness::Failed`]; surfaced verbatim so an + /// operator can read the cause without the server log. + fn failure(&self) -> Option { + None + } +} + +/// Sends `outcome` if anyone is still waiting for it. +/// +/// A worker whose caller has gone should not treat that as an error: a dropped receiver is how +/// cancellation is observed, and work the worker already captured still has to finish before +/// any device buffer it used is reused. +pub fn answer( + reply: oneshot::Sender>, + outcome: Result, +) { + let _ = reply.send(outcome); +} diff --git a/src/runtime/src/lib.rs b/src/runtime/src/lib.rs index 85c4e1d..dde9174 100644 --- a/src/runtime/src/lib.rs +++ b/src/runtime/src/lib.rs @@ -1,4 +1,6 @@ -//! Shared admission and blocking dispatch for a loaded native executor. +//! Shared admission, dispatch and the worker-side engine contract for a loaded executor. + +pub mod engine; use std::sync::Arc; diff --git a/tests/frontend/in_process.rs b/tests/frontend/in_process.rs new file mode 100644 index 0000000..b2dea95 --- /dev/null +++ b/tests/frontend/in_process.rs @@ -0,0 +1,310 @@ +//! The in-process path over real sockets: client -> omni-jev -> a worker implementing the +//! runtime's engine contract. +//! +//! The worker here is a stand-in, so none of this needs a model or a GPU. What it checks is the +//! transport's half of the contract: which status a given worker state produces, that bytes and +//! the budget are handled as documented, and that a worker which never answers is cut off. + +use std::{ + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::{Duration, Instant}, +}; + +use omni_jev::engine::{QUEUE_DEPTH, ServiceConfig, app}; +use omni_runtime::engine::{Answer, Engine, EngineError, Readiness, Reply, Report}; +use tokio::sync::oneshot; +mod common; + +use common::{client, listen}; + +const REQUEST: &str = r#"{"model":"english","state":"refund please","questions":{"q":{"type":"noul","instructions":"Ask?"}}}"#; + +/// How the stand-in answers. Each mode is one of the states the transport has to tell apart, so +/// the status codes below are the contract rather than a convention. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum Mode { + Ready, + Starting, + Failed, + /// Accepts the request and never answers it. + Silent, + /// Refuses because there is no room. + Busy, + /// Refuses because the body is not a request it accepts. + Rejects, + /// Accepts the request and fails to run it. + Fails, +} + +struct Fake { + mode: Mode, + depth: usize, + capacity: usize, + rejected: u64, + /// Requests that reached the engine at all. + seen: Arc, +} + +impl Fake { + /// A ready worker that reports the given queue depth. + fn reporting_depth(depth: usize) -> Arc { + Arc::new(Self { + mode: Mode::Ready, + depth, + ..Self::quiet() + }) + } + + fn quiet() -> Self { + Self { + mode: Mode::Ready, + depth: 0, + capacity: 4, + rejected: 0, + seen: Arc::new(AtomicUsize::new(0)), + } + } + + fn new(mode: Mode) -> Arc { + Arc::new(Self { + mode, + ..Self::quiet() + }) + } +} + +impl Engine for Fake { + fn readiness(&self) -> Readiness { + match self.mode { + Mode::Starting => Readiness::Starting, + Mode::Failed => Readiness::Failed, + _ => Readiness::Ready, + } + } + + fn submit(&self, body: Vec, _deadline: Instant) -> Reply { + self.seen.fetch_add(1, Ordering::AcqRel); + let (reply, receiver) = oneshot::channel(); + match self.mode { + Mode::Busy => { + let _ = reply.send(Err(EngineError::Busy)); + } + Mode::Rejects => { + let _ = reply.send(Err(EngineError::InvalidRequest( + "bad \"envelope\"\n".into(), + ))); + } + Mode::Fails => { + let _ = reply.send(Err(EngineError::InferenceFailed)); + } + // Keeping the sender alive means the request neither completes nor fails: only the + // caller's budget can end it. + Mode::Silent => std::mem::forget(reply), + _ => { + let depth = self.depth; + tokio::spawn(async move { + let _ = reply.send(Ok(Answer { body, depth })); + }); + } + } + receiver + } + + fn report(&self) -> Option { + Some(Report { + depth: self.depth, + capacity: self.capacity, + rejected: self.rejected, + }) + } + + fn failure(&self) -> Option { + matches!(self.mode, Mode::Failed).then(|| "no such checkpoint".to_owned()) + } +} + +/// Starts the in-process service and returns its base URL, with the worker's own counters. +async fn start(mode: Mode, config: ServiceConfig) -> (String, Arc) { + let fake = Fake::new(mode); + let seen = Arc::clone(&fake.seen); + (listen(app(fake, config)).await, seen) +} + +async fn post(url: &str, body: &str) -> reqwest::Response { + client() + .post(format!("{url}/v1/systemone")) + .body(body.to_owned()) + .send() + .await + .unwrap() +} + +async fn health(url: &str) -> (reqwest::StatusCode, String) { + let response = client().get(format!("{url}/health")).send().await.unwrap(); + let status = response.status(); + (status, response.text().await.unwrap()) +} + +#[tokio::test] +async fn readiness_decides_whether_work_is_accepted() { + // Ready answers; the other two refuse without the worker ever seeing the request, and the + // health body says which of the three it is. + for (mode, status, expected) in [ + ( + Mode::Ready, + 200, + r#"{"status":"ok","depth":0,"capacity":4,"rejected":0}"#, + ), + (Mode::Starting, 503, r#"{"status":"starting"}"#), + ( + Mode::Failed, + 503, + r#"{"status":"failed","reason":"no such checkpoint"}"#, + ), + ] { + let (url, seen) = start(mode, ServiceConfig::default()).await; + let (code, body) = health(&url).await; + assert_eq!( + (code.as_u16(), body.as_str()), + (status, expected), + "{mode:?}" + ); + + let response = post(&url, REQUEST).await; + assert_eq!(response.status().as_u16(), status, "{mode:?}"); + if mode != Mode::Ready { + assert_eq!( + seen.load(Ordering::Acquire), + 0, + "{mode:?}: an unready worker must not see the request" + ); + } + } +} + +#[tokio::test] +async fn a_decision_survives_the_transport_byte_for_byte() { + let (url, _) = start(Mode::Ready, ServiceConfig::default()).await; + let response = post(&url, REQUEST).await; + assert_eq!(response.status(), 200); + assert_eq!( + response.headers()["content-type"], + "application/json", + "the transport labels its own responses" + ); + assert_eq!(response.text().await.unwrap(), REQUEST); +} + +#[tokio::test] +async fn a_body_over_the_limit_is_refused_before_the_worker_sees_it() { + let (url, seen) = start( + Mode::Ready, + ServiceConfig { + max_body: 32, + ..Default::default() + }, + ) + .await; + assert_eq!(post(&url, &"x".repeat(64)).await.status(), 413); + assert_eq!(seen.load(Ordering::Acquire), 0); +} + +#[tokio::test] +async fn a_busy_worker_produces_503_and_is_not_retried() { + let (url, seen) = start(Mode::Busy, ServiceConfig::default()).await; + let response = post(&url, REQUEST).await; + assert_eq!(response.status(), 503); + assert_eq!( + response.text().await.unwrap(), + r#"{"error":"inference queue full"}"# + ); + assert_eq!( + seen.load(Ordering::Acquire), + 1, + "a refusal is an answer, not a reason to ask again" + ); +} + +#[tokio::test] +async fn worker_errors_map_to_one_status_each() { + for (mode, status, expected) in [ + ( + Mode::Rejects, + 400, + "{\"error\":\"bad \\\"envelope\\\"\\n\"}", + ), + (Mode::Fails, 500, r#"{"error":"inference failed"}"#), + ] { + let (url, _) = start(mode, ServiceConfig::default()).await; + let response = post(&url, REQUEST).await; + assert_eq!(response.status().as_u16(), status, "{mode:?}"); + // The worker's text is passed through as JSON, so it has to survive being quoted; the + // fake's message contains a quote and a newline on purpose. + assert_eq!(response.text().await.unwrap(), expected, "{mode:?}"); + } +} + +#[tokio::test] +async fn the_reply_publishes_the_depth_the_worker_reported() { + // The worker, not the transport, measures its own queue, so the header carries whatever the + // reply said. A depth of zero is included rather than treated as absent. + let url = listen(app(Fake::reporting_depth(3), ServiceConfig::default())).await; + + let response = post(&url, REQUEST).await; + assert_eq!(response.status(), 200); + assert_eq!( + response.headers()[QUEUE_DEPTH], + "3", + "back-pressure has to be visible before it becomes a refusal" + ); +} + +#[tokio::test] +async fn a_worker_that_never_answers_is_cut_off_by_the_budget() { + let (url, _) = start( + Mode::Silent, + ServiceConfig { + timeout: Duration::from_millis(150), + ..Default::default() + }, + ) + .await; + + let started = Instant::now(); + let response = post(&url, REQUEST).await; + assert_eq!( + response.status(), + 504, + "a request that cannot be answered inside its budget is a gateway timeout" + ); + assert!( + started.elapsed() < Duration::from_secs(2), + "the budget, not the client, has to end the wait: took {:?}", + started.elapsed() + ); + assert_eq!( + response.text().await.unwrap(), + r#"{"error":"request timed out"}"# + ); +} + +/// A 503 is for a worker that cannot take work; a 504 is for work that ran out of budget. The +/// two must not blur, or a caller cannot tell "retry later" from "this took too long". +#[tokio::test] +async fn timeout_and_unavailability_are_different_answers() { + let (slow, _) = start( + Mode::Silent, + ServiceConfig { + timeout: Duration::from_millis(100), + ..Default::default() + }, + ) + .await; + assert_eq!(post(&slow, REQUEST).await.status(), 504); + + let (loading, _) = start(Mode::Starting, ServiceConfig::default()).await; + assert_eq!(post(&loading, REQUEST).await.status(), 503); +} diff --git a/tests/runtime/engine.rs b/tests/runtime/engine.rs new file mode 100644 index 0000000..48c5dd3 --- /dev/null +++ b/tests/runtime/engine.rs @@ -0,0 +1,220 @@ +//! The worker-side engine contract: what an implementation must do, and what a caller may +//! assume. Everything here goes through the public API, because that is all a worker or a +//! transport sees. + +use std::{ + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::{Duration, Instant}, +}; + +use omni_runtime::engine::{Answer, Engine, EngineError, Readiness, Reply, Report}; +use tokio::sync::oneshot; + +/// The smallest worker: ready immediately, answers with its own body, and reports nothing. +struct Echo { + seen: Arc, +} + +impl Echo { + fn new() -> Arc { + Arc::new(Self { + seen: Arc::new(AtomicUsize::new(0)), + }) + } +} + +impl Engine for Echo { + fn readiness(&self) -> Readiness { + Readiness::Ready + } + + fn submit(&self, body: Vec, _deadline: Instant) -> Reply { + self.seen.fetch_add(1, Ordering::AcqRel); + let (reply, receiver) = oneshot::channel(); + if body.is_empty() { + let _ = reply.send(Err(EngineError::InvalidRequest( + "empty request body".into(), + ))); + } else { + let _ = reply.send(Ok(Answer { body, depth: 0 })); + } + receiver + } +} + +/// A worker that owns a queue: it admits nothing beyond its capacity, and reports the depth it +/// is holding. This is the shape the trait exists for. +struct Queued { + capacity: usize, + depth: AtomicUsize, + rejected: AtomicUsize, + held: std::sync::Mutex>>>, +} + +impl Queued { + fn new(capacity: usize) -> Arc { + Arc::new(Self { + capacity, + depth: AtomicUsize::new(0), + rejected: AtomicUsize::new(0), + held: std::sync::Mutex::new(Vec::new()), + }) + } + + /// Answers everything admitted so far, releasing the depth it was holding. + fn release_all(&self) { + for reply in self.held.lock().unwrap().drain(..) { + self.depth.fetch_sub(1, Ordering::AcqRel); + let _ = reply.send(Ok(Answer { + body: b"answered".to_vec(), + depth: 0, + })); + } + } +} + +impl Engine for Queued { + fn readiness(&self) -> Readiness { + Readiness::Ready + } + + fn submit(&self, _body: Vec, _deadline: Instant) -> Reply { + let (reply, receiver) = oneshot::channel(); + // Admission and the answer travel on one path, so a refusal is sent rather than + // returned: the caller has a single thing to await either way. + if self.depth.load(Ordering::Acquire) >= self.capacity { + self.rejected.fetch_add(1, Ordering::Relaxed); + let _ = reply.send(Err(EngineError::Busy)); + return receiver; + } + self.depth.fetch_add(1, Ordering::AcqRel); + self.held.lock().unwrap().push(reply); + receiver + } + + fn report(&self) -> Option { + Some(Report { + depth: self.depth.load(Ordering::Acquire), + capacity: self.capacity, + rejected: self.rejected.load(Ordering::Relaxed) as u64, + }) + } +} + +#[tokio::test] +async fn an_answer_arrives_on_the_reply() { + let engine: Arc = Echo::new(); + let answer = engine + .submit(b"hello".to_vec(), Instant::now()) + .await + .expect("the sender is held until it answers") + .expect("a non-empty body is answered"); + assert_eq!(answer.body, b"hello"); +} + +#[tokio::test] +async fn a_refusal_arrives_on_the_same_reply() { + let engine: Arc = Echo::new(); + let refused = engine + .submit(Vec::new(), Instant::now()) + .await + .expect("the refusal is the message, not a dropped channel") + .expect_err("an empty body is refused"); + assert!( + matches!(refused, EngineError::InvalidRequest(_)), + "{refused:?}" + ); +} + +#[tokio::test] +async fn a_worker_without_observability_reports_nothing() { + let echo = Echo::new(); + let seen = Arc::clone(&echo.seen); + let engine: Arc = echo; + assert_eq!(engine.readiness(), Readiness::Ready); + assert!(engine.report().is_none()); + assert!(engine.failure().is_none()); + assert_eq!( + seen.load(Ordering::Acquire), + 0, + "asking a worker how it is must not run any work" + ); +} + +/// The trait has to carry a queue, not just a pass-through: admission, the depth it reports, +/// and the fact that a refusal does not release a slot. +#[tokio::test] +async fn a_queued_worker_admits_to_capacity_and_says_so() { + let engine = Queued::new(2); + let first = engine.submit(b"one".to_vec(), Instant::now()); + let second = engine.submit(b"two".to_vec(), Instant::now()); + assert_eq!( + engine + .report() + .map(|report| (report.depth, report.capacity)), + Some((2, 2)), + "both accepted requests are outstanding" + ); + + let third = engine.submit(b"three".to_vec(), Instant::now()); + assert!( + matches!(third.await, Ok(Err(EngineError::Busy))), + "the third has nowhere to go and is told so on its own reply" + ); + let report = engine.report().expect("this worker reports its queue"); + assert_eq!(report.depth, 2, "a refused request must not be counted"); + assert_eq!(report.rejected, 1); + + // A refusal is not a lost slot: answering the accepted work frees the queue again, and the + // next request is admitted rather than refused. + engine.release_all(); + assert!(matches!(first.await, Ok(Ok(_)))); + assert!(matches!(second.await, Ok(Ok(_)))); + assert_eq!(engine.report().map(|report| report.depth), Some(0)); + + let fourth = engine.submit(b"fourth".to_vec(), Instant::now()); + assert_eq!( + engine.report().map(|report| report.depth), + Some(1), + "the freed slot is usable again" + ); + engine.release_all(); + assert!(matches!(fourth.await, Ok(Ok(_)))); +} + +/// A dropped receiver is how cancellation is observed. Dropping it must not disturb the +/// worker's own accounting, because work it already admitted still has to finish. +#[tokio::test] +async fn dropping_a_reply_does_not_disturb_the_queue() { + let engine = Queued::new(2); + let abandoned = engine.submit(b"one".to_vec(), Instant::now()); + drop(abandoned); + assert_eq!(engine.report().map(|report| report.depth), Some(1)); + engine.release_all(); + assert_eq!(engine.report().map(|report| report.depth), Some(0)); +} + +#[test] +fn an_error_says_what_it_is() { + assert_eq!( + EngineError::InvalidRequest("no model field".into()).to_string(), + "invalid request: no model field" + ); + assert_eq!(EngineError::Busy.to_string(), "no queue capacity"); + assert_eq!(EngineError::Unavailable.to_string(), "engine unavailable"); + assert_eq!(EngineError::InferenceFailed.to_string(), "inference failed"); +} + +/// The transport reads a deadline off the contract; it has to be the caller's own instant. +#[tokio::test] +async fn a_deadline_is_an_instant_the_caller_owns() { + let engine = Queued::new(1); + let deadline = Instant::now() + Duration::from_millis(50); + let reply = engine.submit(b"one".to_vec(), deadline); + assert!(deadline > Instant::now(), "the budget is in the future"); + engine.release_all(); + assert!(matches!(reply.await, Ok(Ok(_)))); +}