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(_)))); +}