diff --git a/README.md b/README.md index cbf99e6..bf01262 100644 --- a/README.md +++ b/README.md @@ -26,8 +26,10 @@ System1-Omni models, designed around a Rust frontend, model-owned execution, and high-performance CUDA and Metal backends. -The Rust frontend forwards requests to a separately running model worker. The -Cua-S1 4B 0.2 `text` adapter and Open-Jev-27B-v1.1 have native workers using +The Rust frontend forwards requests to a separately running model worker, and +defines the small engine interface a model implements to serve in-process +instead, behind a bounded queue with readiness and shutdown the transport owns. +The Cua-S1 4B 0.2 `text` adapter and Open-Jev-27B-v1.1 have native workers using shared CUDA kernels in this repository. ## News diff --git a/src/frontend/README.md b/src/frontend/README.md index 9569594..ee5b268 100644 --- a/src/frontend/README.md +++ b/src/frontend/README.md @@ -1,8 +1,14 @@ # Rust frontend -An Axum/Tokio server that forwards requests to a separately running model worker -using Reqwest. The worker handles validation, media loading, preprocessing and -inference. +An Axum/Tokio server for the `/v1/systemone` decision API. It runs in one of two +modes, and the HTTP surface is the same in both: + +- **Forwarding** (default). Requests are passed to a separately running model worker + with Reqwest. The worker handles validation, media loading, preprocessing and + inference, and this process owns no model state. +- **Native.** A model linked into this binary answers the requests in-process. The + transport owns the request lifecycle — budget, queue admission, readiness, + shutdown — and the model owns the bytes. ## Run @@ -20,8 +26,38 @@ must be an IP address and port. The backend URL accepts a path prefix, such as `http://localhost:8000/worker`, but no credentials, query or fragment. Backend connections bypass system HTTP proxies. +A model is linked in by selecting an engine. `passthrough` is a stand-in that echoes +every valid request: + +```sh +OMNI_SYSTEMONE_ENGINE=passthrough \ +OMNI_JEV_BIND=127.0.0.1:8080 \ + ./target/release/omni-jev +``` + +It exists to check a deployment's wiring — bind address, readiness probe, queue size, +signal handling — before the model is available, and to exercise the lifecycle without +hardware. It is not a model and is not a production path. + +## Configuration + +| Variable | Default | Meaning | +| --- | --- | --- | +| `OMNI_JEV_BIND` | `127.0.0.1:8080` | Listen address, in both modes. | +| `OMNI_JEV_BACKEND_URL` | `http://127.0.0.1:8000` | Forwarding mode only. | +| `OMNI_SYSTEMONE_ENGINE` | unset | Unset forwards; `passthrough` serves the stand-in. An unknown value is a startup error. | +| `OMNI_SYSTEMONE_CHECKPOINT` | `/checkpoints` | Where the engine's weights are. | +| `OMNI_SYSTEMONE_QUEUE` | `32` | Requests accepted and not yet completed. | +| `OMNI_SYSTEMONE_TIMEOUT_MS` | `30000` | Request budget, covering upload, queueing and inference. | +| `OMNI_SYSTEMONE_MAX_BODY` | `1048576` | Largest request body accepted. | +| `OMNI_SYSTEMONE_DRAIN_MS` | `10000` | How long shutdown waits for accepted work. | + +A malformed value is a startup error rather than a silent fallback to the default. + ## HTTP interface +### Forwarding mode + - `POST /v1/systemone` forwards the `model`, `state` and `questions` envelope unchanged. Workers return `choice`, `score` or `noul` decisions; see the [Jev API reference](https://docs.typesafe.ai/api). @@ -37,6 +73,76 @@ Text, image, audio, video and mixed payloads pass through as bytes. Actual infer support depends on the worker. The [Laya recipe](../../recipe/laya/README.md) verifies text decisions against a real backend. +### Native mode + +- `POST /v1/systemone` passes the body to the engine and returns the body it produced. + Nothing is parsed or re-serialized, so a field the frontend has never heard of + survives, and a model's error text cannot produce invalid JSON. +- `GET /health` is the readiness probe. It answers `200` with `{"status":"ok"}` only + after the engine has loaded **and** warmed up; before that it is `503` with + `{"status":"starting"}`. An engine that failed after startup answers `503` with + `{"status":"failed","reason":…}` and the process keeps running, so the failure is + readable by an operator instead of vanishing with the process. When the engine + reports one, `depth`, `capacity` and `rejected` are included; an engine that reports + nothing gets none of them rather than a zero it did not measure. `rejected` counts work + the engine did not complete: refused because there was no room, or accepted and then + dropped. +- Successful responses carry `x-queue-depth`: how many requests were outstanding when + this one was dequeued, so it includes this one. Backpressure is visible before it + turns into a refusal. +- Status codes: `400` invalid request, `413` body over `OMNI_SYSTEMONE_MAX_BODY`, + `500` the engine failed to run an accepted request, `503` not ready / queue full / + request budget expired / the engine stopped, `504` the engine did not answer inside + the budget. + +Lifecycle, which the acceptance criteria for a persistent server depend on: + +- One worker thread owns the engine, so a runtime that binds to its creating thread is + created by the thread that keeps it. Requests run one at a time, in arrival order, + with no cross-request batching. +- The queue is bounded and refuses rather than grows: once `OMNI_SYSTEMONE_QUEUE` + requests are outstanding — running **or** queued, not only waiting in the channel — + further requests get `503` immediately. The slot is taken, and the limit checked, under + the admission lock and before the request is handed over, so a request that is running + still counts against it and the depth cannot release a slot that was never reserved. +- The budget starts when the request headers arrive. A body that arrives too slowly, a + request that waits too long, or one whose client has already gone is dropped without + reaching the engine, so no GPU work is spent on an answer nobody will read. Work + already submitted is never cancelled. `Request::deadline_expired` reports the same + condition to a model that wants to check before starting something expensive. +- `SIGTERM` or `SIGINT` closes admission immediately and gives the rest of the shutdown + the same budget: responses in flight get up to `OMNI_SYSTEMONE_DRAIN_MS` to finish, + accepted work is drained within it, and the process exits when the budget is spent. A + request that is still uploading at that point is not going to finish, and holding the + process open for its whole request budget would mean the drain budget bounded nothing. + A drain that could not finish is reported rather than hidden. +- An engine that fails one request is retired: the caller gets `500`, `/health` turns + `503` with the reason, and no further request is offered to an engine whose state can + no longer be vouched for. A malformed request does not retire anything. + +A model supplies loading, warmup and one function from a request body to a response +body; the queue, readiness and shutdown behaviour above come from the transport: + +```rust +let running = omni_jev::worker::spawn( + checkpoint, + omni_jev::worker::Options::default(), + Model::load, + |model| model.warmup(), + |request| match model_decide(request.engine, request.body) { + Ok(body) => Ok(omni_jev::engine::Answer { + body, + ..request.answer() + }), + Err(detail) => Err(omni_jev::worker::Failure::InvalidRequest(detail)), + }, +)?; +``` + +Shutting down is `running.stop_accepting()` followed by `running.run_until_drained(deadline)`, +where the deadline comes from `running.drain_budget()` measured when the signal arrived — +one deadline for the whole sequence rather than one per step. + ## Checks From the repository root: @@ -51,6 +157,8 @@ cargo clippy --workspace --locked --all-targets -- -D warnings cargo test --workspace --locked ``` -Tests use local mock workers; no model weights or GPU are needed. They cover -multimodal byte preservation, authorization, connection reuse, large uploads, -backend errors, timeouts, health and binary startup/shutdown. +Tests use local mock workers and a fake engine; no model weights and no GPU are needed. +They cover multimodal byte preservation, authorization, connection reuse, large uploads, +backend errors, timeouts, health, and — for the native path — readiness, queue +admission, expired and cancelled requests, error mapping, panic isolation, drain, and +binary startup and shutdown. diff --git a/src/frontend/src/engine.rs b/src/frontend/src/engine.rs new file mode 100644 index 0000000..bc2a653 --- /dev/null +++ b/src/frontend/src/engine.rs @@ -0,0 +1,292 @@ +//! In-process serving: the same `/v1/systemone` surface, answered by a linked engine. +//! +//! The forwarding path in [`crate`] stays available for workers that already speak HTTP. +//! This module is for a model in the same binary: the transport owns admission, the +//! request budget and readiness, and the model owns the bytes. Nothing here parses the +//! decision envelope, so a field this crate has never heard of survives. + +use std::{ + fmt, + 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 tokio::{net::TcpListener, sync::oneshot}; + +use crate::BoxError; + +/// Requests accepted and not yet completed, as reported when a response was produced. Only +/// sent when the engine reports one, so a missing header means "not observed", not "empty". +pub const QUEUE_DEPTH: HeaderName = HeaderName::from_static("x-queue-depth"); + +/// How far the engine has got. `GET /health` is the readiness probe a load balancer or a +/// benchmark script polls before sending traffic, and `Starting` is what it sees while +/// weights load and the engine warms up. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum Readiness { + Starting, + Ready, + /// Startup failed. The process keeps running and `/health` carries the reason, so the + /// failure is readable by an operator instead of vanishing with the process. + Failed, +} + +/// What an engine tells `/health` about itself. Advisory: readiness never waits on it. +#[derive(Clone, Copy, Debug, Default)] +pub struct Report { + /// Requests accepted and not yet completed, including the one running. + pub depth: usize, + /// The depth at which admission starts refusing. Reported so an operator can see how + /// close the queue is to refusing without knowing the configuration. + pub capacity: usize, + /// Work the engine did not complete, since it started: refused because there was no + /// room, and accepted but then dropped. Monotonic. + pub rejected: u64, +} + +/// The reply the transport sends. `depth` is how many requests were outstanding when this +/// one was dequeued, so it includes this one; it becomes the `x-queue-depth` header, which +/// makes backpressure visible before it turns into a refusal. +pub struct Answer { + pub body: Vec, + pub depth: usize, +} + +/// A closed receiver means the client is gone: a queued request must then be skipped +/// rather than executed, and an engine that has already started work must still finish it +/// before reusing GPU buffers. +pub type Reply = oneshot::Receiver>; + +/// Why an engine could not answer. The transport maps each variant to one status code and +/// one message, so the mapping is a property of this crate rather than of a model. +#[derive(Debug)] +pub enum EngineError { + /// The caller's body is not a request this engine 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 engine failed to execute it. + InferenceFailed, +} + +impl fmt::Display for EngineError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> 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 model side of the transport. Implementations are usually a [`crate::worker::Handle`], +/// which forwards to the thread that owns the engine. +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. `Ok` means the engine owns it and the transport waits for the + /// reply until `deadline`; [`EngineError::Busy`] means nothing was accepted. + fn submit(&self, body: Vec, deadline: Instant) -> Result; + + /// What `/health` publishes while ready. `None` leaves those fields out. + fn report(&self) -> Option { + None + } + + /// Set only when [`Engine::readiness`] is [`Readiness::Failed`], and surfaced verbatim + /// so an operator can read the cause without the server log. + fn failure(&self) -> Option { + None + } +} + +/// Admission policy: a bounded queue, and a budget that starts when the headers arrive. +#[derive(Clone, Debug)] +pub struct ServiceConfig { + /// Total request budget: upload, queueing and inference together. + 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 native 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 engine's +/// own worker happen after this returns, and belong to the caller. +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 engine 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 + // engine, and work that cannot finish in time is refused before it queues. + 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::SERVICE_UNAVAILABLE, + "request deadline expired while uploading", + ); + } + }; + if Instant::now() >= deadline { + return fail(StatusCode::SERVICE_UNAVAILABLE, "request deadline expired"); + } + + let reply = match service.engine.submit(body, deadline) { + Ok(reply) => reply, + Err(error) => return error_response(error), + }; + match tokio::time::timeout_at(deadline.into(), reply).await { + Ok(Ok(Ok(answer))) => with_depth(answer), + Ok(Ok(Err(error))) => error_response(error), + // The engine dropped a request it had already 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, "inference 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 engine'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 model'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 0de9ba5..8be3bd0 100644 --- a/src/frontend/src/lib.rs +++ b/src/frontend/src/lib.rs @@ -1,4 +1,16 @@ //! Jev HTTP transport. The worker owns request parsing and inference. +//! +//! Two ways to answer `/v1/systemone`: +//! +//! - [`Config`] and [`app`] forward to a worker that already speaks HTTP. The frontend +//! owns no model state, which is how the Jev path works today. +//! - [`engine`] and [`worker`] serve a model linked into this binary. The transport owns +//! the request lifecycle (budget, queue admission, readiness, shutdown) and the model +//! owns the bytes. + +pub mod engine; +pub mod passthrough; +pub mod worker; use std::{env, error::Error, net::SocketAddr, time::Duration}; diff --git a/src/frontend/src/main.rs b/src/frontend/src/main.rs index 1d47857..39a2cce 100644 --- a/src/frontend/src/main.rs +++ b/src/frontend/src/main.rs @@ -1,20 +1,226 @@ -use omni_jev::{BoxError, Config}; +//! The `omni-jev` binary. Two modes, selected by `OMNI_SYSTEMONE_ENGINE`. +//! +//! Without it, the process forwards to a worker over HTTP and owns no model state. With +//! `OMNI_SYSTEMONE_ENGINE=passthrough` it serves a linked engine in-process, which is the +//! shape a real model takes: the engine is built on a worker thread, `/health` follows its +//! warmup, and SIGTERM stops admission and drains accepted work before the process exits. + +use std::{ + env, + path::PathBuf, + process::ExitCode, + str::FromStr, + sync::Arc, + time::{Duration, Instant}, +}; + +use omni_jev::{ + BoxError, Config, + engine::{self, Answer, Engine, ServiceConfig}, + passthrough, worker, +}; use tokio::net::TcpListener; #[tokio::main] -async fn main() -> Result<(), BoxError> { - let config = Config::from_env()?; +async fn main() -> ExitCode { + match run().await { + Ok(()) => ExitCode::SUCCESS, + Err(error) => { + eprintln!("omni-jev: {error}"); + ExitCode::FAILURE + } + } +} + +async fn run() -> Result<(), BoxError> { + let bind = setting("OMNI_JEV_BIND", Config::DEFAULT_BIND.to_owned())?; + // An unrecognized engine name is a configuration error, not a fallback: silently + // forwarding would send inference somewhere the operator did not ask for. + match env::var("OMNI_SYSTEMONE_ENGINE") { + Err(env::VarError::NotPresent) => return forward(&bind).await, + Ok(name) if name.is_empty() => return forward(&bind).await, + Ok(name) if name == "passthrough" => {} + Ok(name) => return Err(format!("unknown OMNI_SYSTEMONE_ENGINE {name:?}").into()), + Err(error) => return Err(format!("OMNI_SYSTEMONE_ENGINE: {error}").into()), + } + + let queue = setting("OMNI_SYSTEMONE_QUEUE", 32)?; + let service = ServiceConfig { + timeout: Duration::from_millis(setting("OMNI_SYSTEMONE_TIMEOUT_MS", 30_000)?), + max_body: setting("OMNI_SYSTEMONE_MAX_BODY", 1024 * 1024)?, + }; + let running = worker::spawn( + PathBuf::from(setting( + "OMNI_SYSTEMONE_CHECKPOINT", + "/checkpoints".to_owned(), + )?), + worker::Options { + queue_capacity: queue, + drain: Duration::from_millis(setting("OMNI_SYSTEMONE_DRAIN_MS", 10_000)?), + }, + // The engine is built here, on the worker thread, which is where a CUDA context + // would be created and where it has to stay. + passthrough::Passthrough::load, + |_engine| Ok(()), + |request| match request.body { + [] => Err(worker::Failure::InvalidRequest("empty request body".into())), + body => Ok(Answer { + body: body.to_vec(), + ..request.answer() + }), + }, + )?; + let app = engine::app(Arc::clone(&running.handle) as Arc, service); + let listener = TcpListener::bind(&bind).await?; + // The `spawn` above returned only after loading and warmup, so discovery by this line + // means `/health` answers 200 rather than a race with startup. + eprintln!( + "omni-jev listening on {} (engine passthrough, queue {queue})", + listener.local_addr()? + ); + + // No timer runs while serving: the budget belongs to shutdown, and arming it at startup + // would retire a healthy process once it elapsed. + let drain = running.drain_budget(); + if !serve_until_stopped(listener, app, running, drain).await? { + return Err("engine did not drain before shutdown".into()); + } + Ok(()) +} + +/// How long the transport gets to write out what it is still holding once the engine has +/// stopped. An idle keep-alive connection is still an open connection and graceful shutdown +/// waits for it, so this is what stops that wait from running to the whole drain budget. +const CLOSE_GRACE: Duration = Duration::from_millis(250); + +/// Serves until the process is asked to stop, then shuts down. Returns false if the engine +/// could not finish within the drain budget. +/// +/// One `serve` call runs for the life of the process, and everything that has to happen in +/// order once a signal arrives happens in its shutdown future: close admission, take the one +/// deadline everything shares, drain the worker, then let the transport write out what it is +/// still holding. +/// +/// **One absolute deadline, taken once at the signal, and no period added on top of it.** The +/// worker drains until that instant; the transport then gets whatever remains of it, up to a +/// bounded grace. A grace measured from the end of the drain instead would make the process +/// live `drain + grace` however little work was left. And the transport cannot be asked to +/// stop — it waits for open connections for as long as they are open, including idle +/// keep-alive ones that would never close on their own — so the sequence is bounded by no +/// longer waiting on it. +async fn serve_until_stopped( + listener: TcpListener, + app: axum::Router, + running: worker::Running, + drain: Duration, +) -> Result { + let (outcome, drained) = tokio::sync::oneshot::channel::(); + // Tells the outer wait that the engine is done, and when the transport's own grace period + // ends. An oneshot rather than a `Notify`, which stores a permit and would let the wait + // finish before the window had actually been served. + let (closed, finish) = tokio::sync::oneshot::channel::(); + let admission = Arc::clone(&running.handle); + let shutdown = async move { + stop_signal().await; + // Admission closes first: a request still uploading when shutdown begins could + // otherwise submit work that keeps the engine busy for its whole request budget, + // long after the deadline below has passed. + admission.stop_accepting(); + let deadline = Instant::now() + drain; + // Off the runtime thread: this waits on the engine thread, and blocking a worker here + // would stall the transport draining alongside it. + let done = tokio::task::spawn_blocking(move || running.run_until_drained(deadline)) + .await + .unwrap_or(false); + let _ = outcome.send(done); + // The transport's window is a bounded grace, clamped to the deadline itself. Both + // halves matter: without the bound an idle keep-alive connection holds the process to + // the whole drain budget, and without the clamp the process lives `drain + grace` + // however little work was left. + let _ = closed.send(deadline.min(Instant::now() + CLOSE_GRACE)); + // Nothing left to do but let the transport finish. The caller stops waiting at the + // instant above whether or not this resolves, so this needs no bound of its own. + std::future::pending::<()>().await; + }; + + let mut serving = std::pin::pin!(engine::serve(listener, app, shutdown)); + tokio::select! { + // Both branches run only once a signal is in: `serve` cannot resolve before its + // shutdown future does, and that future waits for the signal. + waited = &mut serving => waited?, + ends = finish => { + let now = Instant::now(); + let ends = ends.unwrap_or_else(|_| (now + CLOSE_GRACE).min(now + drain)); + // Polled rather than awaited once: the transport writes out what it is holding and + // is done, so this returns as soon as it does. Waiting a fixed period here instead + // would make every shutdown last that period, however little was left. + let mut serving = &mut serving; + let settled = std::future::poll_fn(|cx| { + if Instant::now() >= ends { + return std::task::Poll::Ready(false); + } + match std::future::Future::poll(std::pin::Pin::new(&mut serving), cx) { + std::task::Poll::Ready(Ok(())) => std::task::Poll::Ready(true), + std::task::Poll::Ready(Err(error)) => { + panic!("engine::serve failed during shutdown: {error}") + } + std::task::Poll::Pending => { + // Wake when the window closes, so a transport that never finishes + // cannot outlive it. + cx.waker().wake_by_ref(); + std::task::Poll::Pending + } + } + }) + .await; + if !settled { + eprintln!("omni-jev: shutdown did not finish within {drain:?}, exiting"); + } + } + } + Ok(drained.await.unwrap_or(false)) +} + +/// Forwards to an HTTP worker, owning no model state. +async fn forward(bind: &str) -> Result<(), BoxError> { + let backend = setting( + "OMNI_JEV_BACKEND_URL", + Config::DEFAULT_BACKEND_URL.to_owned(), + )?; + let config = Config::new(bind, &backend)?; let app = omni_jev::app(&config)?; - let listener = TcpListener::bind(config.bind).await?; + let listener = TcpListener::bind(&config.bind).await?; + // This line is the documented way to find a port bound with `:0`, and the way a + // supervisor knows the process is up. It stays first and stays unchanged. eprintln!("omni-jev listening on {}", listener.local_addr()?); eprintln!("forwarding to {}", config.backend_url); - axum::serve(listener, app) - .with_graceful_shutdown(shutdown_signal()) - .await?; + engine::serve(listener, app, shutdown_signal()).await?; Ok(()) } +/// Reads a setting, failing on a malformed value rather than quietly using the default: a +/// typo in `OMNI_SYSTEMONE_TIMEOUT_MS` must not become a timeout nobody asked for. +fn setting(name: &str, default: T) -> Result +where + T::Err: std::fmt::Display, +{ + match env::var(name) { + Err(env::VarError::NotPresent) => Ok(default), + Err(error) => Err(format!("{name}: {error}").into()), + Ok(value) if value.is_empty() => Ok(default), + Ok(value) => value + .parse() + .map_err(|error| format!("{name}: {error}").into()), + } +} + +/// Resolves when the process is asked to stop. The forwarding path owns no worker, so it +/// has nothing to admit or drain. async fn shutdown_signal() { + stop_signal().await +} + +async fn stop_signal() { #[cfg(unix)] let terminate = async { use tokio::signal::unix::{SignalKind, signal}; diff --git a/src/frontend/src/passthrough.rs b/src/frontend/src/passthrough.rs new file mode 100644 index 0000000..24addf6 --- /dev/null +++ b/src/frontend/src/passthrough.rs @@ -0,0 +1,20 @@ +//! A stand-in engine that echoes its request. Not a model. +//! +//! This exists so the native path can be run and tested end to end without a GPU, a +//! checkpoint or CUDA: `OMNI_SYSTEMONE_ENGINE=passthrough` serves the same transport, +//! queue, readiness and shutdown behaviour a real engine gets, and answers every valid +//! request with its own body. That is how a deployment's wiring — bind address, readiness +//! probe, queue size, signal handling — can be checked before the model is available. + +use std::{io, path::PathBuf}; + +/// The stand-in engine: it holds nothing and answers with the request it was given. +#[derive(Debug)] +pub struct Passthrough; + +impl Passthrough { + /// Builds the engine on the worker thread, as a real engine's loader would. + pub fn load(_checkpoint: PathBuf) -> io::Result { + Ok(Self) + } +} diff --git a/src/frontend/src/worker.rs b/src/frontend/src/worker.rs new file mode 100644 index 0000000..6a85bfb --- /dev/null +++ b/src/frontend/src/worker.rs @@ -0,0 +1,473 @@ +//! One thread owns the engine, and a bounded queue owns admission. +//! +//! This is the half of the native path that does not depend on a model: the transport +//! hands over request bytes, the worker runs them one at a time, and everything a +//! deployment has to get right lives here — readiness that follows warmup, a queue that +//! refuses instead of growing, a request budget that can expire before work starts, and a +//! shutdown that stops admission without cutting a running inference. +//! +//! A model supplies three things and nothing else: +//! +//! ```no_run +//! # use std::{io, path::PathBuf}; +//! # use omni_jev::{engine::Answer, worker}; +//! # struct Model; +//! # impl Model { +//! # fn load(_: PathBuf) -> io::Result { Ok(Model) } +//! # fn warmup(&mut self) -> io::Result<()> { Ok(()) } +//! # fn infer(&mut self, _: &[u8]) -> Result, String> { Ok(Vec::new()) } +//! # } +//! let running = worker::spawn( +//! PathBuf::from("/checkpoints/laya"), +//! worker::Options::default(), +//! // 1. Build the engine. Runs on the worker thread, so a runtime that binds to its +//! // creating thread is created by the thread that keeps it. +//! Model::load, +//! // 2. Pay for compilation and allocation here, so readiness means something. +//! |model| model.warmup(), +//! // 3. Answer one request. The value becomes the response body. +//! |request| match request.engine.infer(request.body) { +//! Ok(body) => Ok(Answer { body, ..request.answer() }), +//! Err(detail) => Err(worker::Failure::InvalidRequest(detail)), +//! }, +//! ) +//! .expect("engine failed to load"); +//! ``` + +use std::{ + io, + panic::{AssertUnwindSafe, catch_unwind}, + path::PathBuf, + sync::{ + Arc, Mutex, MutexGuard, + atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}, + }, + thread, + time::{Duration, Instant}, +}; + +use tokio::sync::{mpsc, oneshot}; + +use crate::engine::{Answer, Engine, EngineError, Readiness, Reply, Report}; + +/// How the worker was asked to run. +#[derive(Clone, Debug)] +pub struct Options { + /// Requests accepted and not yet completed, running or queued. [`Engine::submit`] + /// refuses rather than waiting once this many are outstanding, which is what turns + /// overload into a fast `503` instead of an unbounded queue. The queue itself is only a + /// buffer for these jobs, so it is sized the same and enforces nothing on its own. + pub queue_capacity: usize, + /// How long a shutdown may spend waiting for accepted work. It bounds the wait, it does + /// not cancel anything: a kernel already submitted runs to completion. + pub drain: Duration, +} + +impl Default for Options { + fn default() -> Self { + Self { + queue_capacity: 32, + drain: Duration::from_secs(10), + } + } +} + +/// Why a request could not be executed. The worker decides between a caller's mistake and +/// its own failure; the transport decides which status code that becomes. +#[derive(Debug)] +pub enum Failure { + /// The body is not a request this engine accepts. The worker keeps serving. + InvalidRequest(String), + /// The engine failed at work it accepted, so it is retired: see [`run`]. + Inference, +} + +/// One request the worker has taken off the queue, with the state a model's own code +/// cannot see: its budget and whether anyone is still waiting for it. +pub struct Request<'a, E> { + pub engine: &'a mut E, + pub body: &'a [u8], + /// The caller's budget, covering upload and queueing as well as this inference. + pub deadline: Instant, + /// False once the client has gone, which happens when the transport gave up on the + /// deadline or the connection dropped. Work not yet submitted should be dropped; work + /// already submitted is never cancelled. + pub waiting: bool, + /// Queue state measured when this request was dequeued, for the answer to carry back. + depth: usize, +} + +impl Request<'_, E> { + /// An answer carrying the queue depth this request was dequeued at. Fill in `body`. + pub fn answer(&self) -> Answer { + Answer { + body: Vec::new(), + depth: self.depth, + } + } +} + +struct Job { + body: Vec, + deadline: Instant, + reply: oneshot::Sender>, +} + +/// The engine side of a running worker: what [`crate::engine::app`] talks to. +pub struct Handle { + /// Taken by shutdown to close the queue. `None` means admission is over. + sender: Mutex>>, + /// Serializes admission against shutdown, so a request cannot be accepted after the + /// queue has been closed. Held only for the length of a non-blocking send. + admission: Mutex<()>, + ready: AtomicBool, + stopped: AtomicBool, + depth: AtomicUsize, + /// The outstanding-request limit. Fixed for the life of the worker. + capacity: usize, + rejected: AtomicU64, + failure: Mutex>, +} + +/// A started worker: the engine handle and the thread that owns the engine. +pub struct Running { + /// The handle to give to [`crate::engine::app`]. + pub handle: Arc, + /// `Some` until [`Running::run_until_drained`] takes it to join. + worker: Option>, + drain: Duration, +} + +impl Running { + /// Closes the queue without waiting for accepted work. Idempotent, and called for you + /// by [`Running::run_until_drained`] and by this type's drop. + pub fn stop_accepting(&self) { + self.handle.stop_accepting(); + } + + /// How long a caller is willing to spend draining, so a shutdown sequence can compute + /// one deadline for all of its steps instead of giving each step its own. + pub fn drain_budget(&self) -> Duration { + self.drain + } + + /// Waits for accepted work to finish and for the worker thread to exit, no later than + /// `deadline`. + /// + /// The deadline is absolute and passed in rather than measured from here, because the + /// budget belongs to the shutdown sequence: a caller that has already spent part of it + /// stopping admission and closing the listener must not get a fresh budget for the + /// drain. Returns false if the deadline passed first, which is the honest answer when a + /// kernel or a device call is wedged — nothing here can interrupt it. + pub fn run_until_drained(self, deadline: Instant) -> bool { + // Taking a field out of a type that has a `Drop` impl needs the drop taken over + // first; the sequence below is what that `Drop` would have done anyway. + let mut this = std::mem::ManuallyDrop::new(self); + // The take cannot be empty: this is the only place the thread handle is moved out. + let worker = this + .worker + .take() + .expect("worker thread handle was already taken"); + this.handle.stop_accepting(); + while !worker.is_finished() && Instant::now() < deadline { + thread::sleep(Duration::from_millis(5)); + } + if !worker.is_finished() { + eprintln!("omni-jev: engine did not drain within {:?}", this.drain); + return false; + } + match worker.join() { + Ok(()) => true, + Err(_) => { + eprintln!("omni-jev: engine thread panicked"); + false + } + } + } +} + +impl Drop for Running { + fn drop(&mut self) { + // A caller that never reached shutdown must not leave the engine accepting work. + self.handle.stop_accepting(); + } +} + +impl Engine for Handle { + fn readiness(&self) -> Readiness { + if self.ready.load(Ordering::Acquire) { + Readiness::Ready + } else if lock(&self.failure).is_some() { + Readiness::Failed + } else { + Readiness::Starting + } + } + + fn submit(&self, body: Vec, deadline: Instant) -> Result { + // One lock orders admission against shutdown: without it a request can pass the + // readiness check and then be accepted after the queue was closed, leaving it + // accepted with no worker left to answer it. + let admission = lock(&self.admission); + if !self.ready.load(Ordering::Acquire) || self.stopped.load(Ordering::Acquire) { + return Err(EngineError::Unavailable); + } + let sender = lock(&self.sender); + let Some(sender) = sender.as_ref() else { + return Err(EngineError::Unavailable); + }; + // The chain is what makes the check a decision rather than a read: `depth` may only + // increase while the admission lock is held, so the count that was compared is still + // the count when the job is handed over. + let reserved = self.depth.fetch_add(1, Ordering::AcqRel) < self.capacity; + let refusal = if !reserved { + Some(EngineError::Busy) + } else { + let (reply, receiver) = oneshot::channel(); + match sender.try_send(Job { + body, + deadline, + reply, + }) { + // The worker is free to dequeue, answer and release immediately: it never + // takes the admission lock. Counting after the send would therefore release a + // slot that had not been taken yet, so the slot is taken before the handover. + Ok(()) => { + drop(admission); + return Ok(receiver); + } + Err(mpsc::error::TrySendError::Full(_)) => Some(EngineError::Busy), + Err(mpsc::error::TrySendError::Closed(_)) => Some(EngineError::Unavailable), + } + }; + // Nothing was accepted, so the reservation goes back; the refusal is then reported + // as such rather than left counted against the queue. + self.depth.fetch_sub(1, Ordering::AcqRel); + if let Some(EngineError::Busy) = refusal { + self.rejected.fetch_add(1, Ordering::Relaxed); + } + let refusal = refusal.expect("a refusal was decided above"); + drop(admission); + Err(refusal) + } + + fn report(&self) -> Option { + Some(Report { + depth: self.depth.load(Ordering::Acquire), + capacity: self.capacity, + rejected: self.rejected.load(Ordering::Relaxed), + }) + } + + fn failure(&self) -> Option { + lock(&self.failure).clone() + } +} + +impl Handle { + /// Closes the queue and clears readiness, without waiting for accepted work. + pub fn stop_accepting(&self) { + let admission = lock(&self.admission); + self.stopped.store(true, Ordering::Release); + self.ready.store(false, Ordering::Release); + // Dropped while the admission lock is held, so no request can slip a send between + // the close and the check in `submit`. + let sender = lock(&self.sender).take(); + drop(sender); + drop(admission); + } + + fn fail(&self, detail: String) { + eprintln!("omni-jev: engine failed: {detail}"); + *lock(&self.failure) = Some(detail); + self.ready.store(false, Ordering::Release); + } +} + +/// Starts the worker thread and returns once the engine has loaded and warmed up. +/// +/// `load` runs on the worker thread, which is the thread that keeps the engine, so a +/// runtime that binds to its creating thread is created in the right place. `warmup` runs +/// before readiness is reported, so the first request after `/health` is ready is not the +/// one that pays for compilation. A load failure returns `Err` and leaves no thread behind; +/// a failure after startup marks the handle [`Readiness::Failed`]. +pub fn spawn( + path: PathBuf, + options: Options, + load: L, + warmup: W, + process: P, +) -> io::Result +where + L: FnOnce(PathBuf) -> io::Result + Send + 'static, + W: FnOnce(&mut E) -> io::Result<()> + Send + 'static, + P: FnMut(Request<'_, E>) -> Result + Send + 'static, + E: Send + 'static, +{ + if options.queue_capacity == 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "queue capacity must be at least 1", + )); + } + + let handle = Arc::new(Handle { + sender: Mutex::new(None), + admission: Mutex::new(()), + ready: AtomicBool::new(false), + stopped: AtomicBool::new(false), + depth: AtomicUsize::new(0), + capacity: options.queue_capacity, + rejected: AtomicU64::new(0), + failure: Mutex::new(None), + }); + let (jobs, queue) = mpsc::channel::(options.queue_capacity); + let (started, loading) = oneshot::channel::>(); + + let worker = { + let handle = Arc::clone(&handle); + thread::Builder::new() + .name("omni-engine".into()) + .spawn(move || { + let outcome = catch_unwind(AssertUnwindSafe(|| { + load(path).and_then(|mut engine| warmup(&mut engine).map(|()| engine)) + })); + let engine = match outcome { + Ok(Ok(engine)) => engine, + Ok(Err(error)) => { + handle.fail(error.to_string()); + let _ = started.send(Err(error)); + return; + } + Err(panic) => { + let error = io::Error::other(format!( + "engine panicked while loading: {}", + panic_text(&panic) + )); + handle.fail(error.to_string()); + let _ = started.send(Err(error)); + return; + } + }; + // The queue is open before readiness is set, so a request that observes + // Ready always finds somewhere to go. + *lock(&handle.sender) = Some(jobs); + *lock(&handle.failure) = None; + handle.ready.store(true, Ordering::Release); + if started.send(Ok(())).is_err() { + return; + } + run(engine, queue, &handle, process); + // Reached on shutdown or on engine failure. Either way the process must + // stop reporting itself ready. + handle.ready.store(false, Ordering::Release); + *lock(&handle.sender) = None; + })? + }; + + // Waiting for the load result has to happen off the caller's thread. A caller that is + // already inside an async runtime cannot block, and blocking there starves the runtime + // that will later serve requests — which is how this is reached in practice, since the + // process's `main` is an async entry point. + let startup: io::Result> = thread::Builder::new() + .name("omni-engine-startup".into()) + .spawn(move || loading.blocking_recv())? + .join() + .map(|outcome| { + outcome.map_err(|_| io::Error::other("engine thread stopped during startup")) + }) + .unwrap_or_else(|_| Err(io::Error::other("engine thread stopped during startup"))); + + match startup { + Ok(Ok(())) => Ok(Running { + handle, + worker: Some(worker), + drain: options.drain, + }), + // The engine reported why it could not load, which is the message a caller wants. + Ok(Err(error)) | Err(error) => { + let _ = worker.join(); + Err(error) + } + } +} + +/// The worker loop: one job at a time, in arrival order, until the queue closes. +fn run(mut engine: E, mut queue: mpsc::Receiver, handle: &Handle, mut process: P) +where + P: FnMut(Request<'_, E>) -> Result, +{ + while let Some(job) = queue.blocking_recv() { + // The depth slot is released on every path out of this iteration, including the + // skip and panic paths, so a refused request cannot leave depth permanently high. + let slot = Slot(&handle.depth); + if job.reply.is_closed() || Instant::now() >= job.deadline { + // A cancelled or expired request never reaches the engine, which is the point + // of carrying a deadline this far: no GPU work for an answer nobody will read. + handle.rejected.fetch_add(1, Ordering::Relaxed); + continue; + } + let mut stop = false; + let body = &job.body; + let outcome = catch_unwind(AssertUnwindSafe(|| { + process(Request { + engine: &mut engine, + body, + deadline: job.deadline, + waiting: !job.reply.is_closed(), + depth: handle.depth.load(Ordering::Acquire), + }) + })); + let response = match outcome { + Ok(Ok(answer)) => Ok(answer), + // Bad input, not a broken engine. The worker keeps serving, so one caller's + // malformed body cannot take the deployment down. + Ok(Err(Failure::InvalidRequest(detail))) => Err(EngineError::InvalidRequest(detail)), + Ok(Err(Failure::Inference)) => { + // The engine failed at work it accepted: a device error, a lost context, an + // allocation that did not come back. The worker retires rather than offering + // the next caller an engine it can no longer vouch for, and `/health` carries + // the reason until a restart. + handle.fail("inference failed; the engine was retired".into()); + stop = true; + Err(EngineError::InferenceFailed) + } + Err(panic) => { + // A panic leaves the engine's state unknown for the same reason and is + // reported the same way. The payload is kept because it says more than + // "inference failed" does. + handle.fail(format!("engine panicked: {}", panic_text(&panic))); + stop = true; + Err(EngineError::InferenceFailed) + } + }; + let _ = job.reply.send(response); + // Released before the next dequeue, so a fresh answer never reports this request. + drop(slot); + if stop { + break; + } + } + handle.ready.store(false, Ordering::Release); +} + +/// Releases one unit of queue depth when it goes out of scope. +struct Slot<'a>(&'a AtomicUsize); + +impl Drop for Slot<'_> { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::AcqRel); + } +} + +fn panic_text(panic: &Box) -> &str { + panic + .downcast_ref::<&str>() + .copied() + .or_else(|| panic.downcast_ref::().map(String::as_str)) + .unwrap_or("non-string panic payload") +} + +fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex.lock().unwrap_or_else(|error| error.into_inner()) +} diff --git a/src/frontend/tests/engine_service.rs b/src/frontend/tests/engine_service.rs new file mode 100644 index 0000000..10dcda8 --- /dev/null +++ b/src/frontend/tests/engine_service.rs @@ -0,0 +1,628 @@ +//! End-to-end tests over real sockets for the in-process path: client -> omni-jev -> +//! a fake engine. The engine is a stand-in, so none of this needs a model or a GPU. +//! +//! The last test runs the compiled binary with `OMNI_SYSTEMONE_ENGINE=passthrough`, which +//! is the only way to cover the whole lifecycle: readiness, signals and process exit. + +use std::{ + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::{Duration, Instant}, +}; + +use axum::Router; +use omni_jev::engine::{ + self, Answer, Engine, EngineError, Readiness, Reply, Report, ServiceConfig, +}; +use tokio::{net::TcpListener, sync::oneshot}; + +const REQUEST: &str = r#"{"model":"english","state":"refund please","questions":{"q":{"type":"noul","instructions":"Ask?"}}}"#; + +/// How the fake engine answers. Every 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 with a full queue. + Full, + /// Echoes its own error variants. + Invalid, + Inference, +} + +struct Fake { + mode: Mode, + depth: usize, + capacity: usize, + rejected: u64, + /// Number of requests that reached `submit`. + submitted: Arc, +} + +impl Default for Fake { + fn default() -> Self { + Self { + mode: Mode::Ready, + depth: 0, + capacity: 4, + rejected: 0, + submitted: Arc::new(AtomicUsize::new(0)), + } + } +} + +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) -> Result { + self.submitted.fetch_add(1, Ordering::AcqRel); + match self.mode { + Mode::Full => Err(EngineError::Busy), + Mode::Invalid => Err(EngineError::InvalidRequest("bad \"envelope\"\n".into())), + Mode::Inference => Err(EngineError::InferenceFailed), + Mode::Silent => { + // Keeping the sender alive means the request neither completes nor fails: + // only the deadline can end it. + let (reply, receiver) = oneshot::channel(); + std::mem::forget(reply); + Ok(receiver) + } + _ => { + let answer = Answer { + body, + depth: self.depth, + }; + // Answered from another task, because a real engine never answers inside + // `submit` and the transport must not depend on getting one immediately. + let (reply, receiver) = oneshot::channel(); + tokio::spawn(async move { + let _ = reply.send(Ok(answer)); + }); + Ok(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()) + } +} + +fn client() -> reqwest::Client { + reqwest::Client::builder() + .no_proxy() + .redirect(reqwest::redirect::Policy::none()) + .timeout(Duration::from_secs(10)) + .build() + .unwrap() +} + +/// Starts the native service and returns its base URL. +async fn start(fake: Arc, config: ServiceConfig) -> String { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + let app = engine::app(fake, config); + tokio::spawn(async move { + engine::serve(listener, app, std::future::pending()) + .await + .unwrap(); + }); + url +} + +async fn post(url: &str, body: &str) -> reqwest::Response { + client() + .post(format!("{url}/v1/systemone")) + .body(body.to_owned()) + .send() + .await + .unwrap() +} + +/// The body a response carried. Asserting the whole JSON text rather than a parsed field +/// pins the wire format itself, and keeps the transport from depending on a parser it +/// deliberately does not have. +async fn body(response: reqwest::Response) -> String { + response.text().await.unwrap() +} + +#[tokio::test] +async fn readiness_and_failure_are_distinguishable() { + 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 fake = Arc::new(Fake { + mode, + ..Default::default() + }); + let url = start(fake.clone(), ServiceConfig::default()).await; + + let health = client().get(format!("{url}/health")).send().await.unwrap(); + assert_eq!(health.status().as_u16(), status, "mode {mode:?}"); + // A ready engine describes its queue; a failed one says why. A starting one claims + // neither, because a half-loaded engine knows nothing worth reporting. + assert_eq!(body(health).await, expected, "mode {mode:?}"); + + // Inference is refused before the engine is asked, so an unready engine never sees + // the request at all. + let response = post(&url, REQUEST).await; + assert_eq!(response.status().as_u16(), status, "mode {mode:?}"); + if mode != Mode::Ready { + assert_eq!(fake.submitted.load(Ordering::Acquire), 0); + } + } +} + +#[tokio::test] +async fn a_valid_request_keeps_its_bytes() { + let fake = Arc::new(Fake::default()); + let url = start(fake, 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_without_touching_the_engine() { + let fake = Arc::new(Fake::default()); + let url = start( + fake.clone(), + ServiceConfig { + max_body: 32, + ..Default::default() + }, + ) + .await; + + let response = post(&url, &"x".repeat(64)).await; + assert_eq!(response.status(), 413); + assert_eq!(fake.submitted.load(Ordering::Acquire), 0); +} + +#[tokio::test] +async fn a_full_queue_is_a_503_and_a_lost_slot() { + let fake = Arc::new(Fake { + mode: Mode::Full, + rejected: 7, + ..Default::default() + }); + let url = start(fake, ServiceConfig::default()).await; + + let response = post(&url, REQUEST).await; + assert_eq!(response.status(), 503); + assert_eq!(body(response).await, r#"{"error":"inference queue full"}"#); +} + +#[tokio::test] +async fn a_silent_engine_is_cut_off_by_the_deadline() { + let fake = Arc::new(Fake { + mode: Mode::Silent, + ..Default::default() + }); + let url = start( + fake, + ServiceConfig { + timeout: Duration::from_millis(150), + ..Default::default() + }, + ) + .await; + + let started = Instant::now(); + let response = post(&url, REQUEST).await; + assert_eq!(response.status(), 504); + assert!( + started.elapsed() < Duration::from_secs(5), + "the deadline, not the client, must end the wait" + ); + assert_eq!(body(response).await, r#"{"error":"inference timed out"}"#); +} + +#[tokio::test] +async fn engine_errors_map_to_one_status_each() { + for (mode, status, message) in [ + ( + Mode::Invalid, + 400, + "{\"error\":\"bad \\\"envelope\\\"\\n\"}", + ), + (Mode::Inference, 500, r#"{"error":"inference failed"}"#), + ] { + let fake = Arc::new(Fake { + mode, + ..Default::default() + }); + let url = start(fake, ServiceConfig::default()).await; + + let response = post(&url, REQUEST).await; + assert_eq!(response.status().as_u16(), status, "mode {mode:?}"); + // The engine'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!(body(response).await, message); + } +} + +#[tokio::test] +async fn the_answer_carries_the_depth_the_engine_reported() { + // The engine, not the transport, measures the queue, so the header is whatever the + // answer carries. A depth of zero is included rather than treated as absent. + let fake = Arc::new(Fake { + depth: 3, + ..Default::default() + }); + let url = start(fake, ServiceConfig::default()).await; + + let response = post(&url, REQUEST).await; + assert_eq!(response.status(), 200); + assert_eq!( + response.headers()["x-queue-depth"], + "3", + "backpressure must be visible before it becomes a refusal" + ); +} + +/// The forwarding path is unchanged, but both modes now share one serve loop, so this pins +/// that the shared loop binds and keeps serving rather than returning early. +#[tokio::test] +async fn engine_serve_keeps_serving_until_shutdown() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let app = Router::new(); + let handle = tokio::spawn(async move { + engine::serve(listener, app, std::future::pending()) + .await + .unwrap(); + }); + tokio::time::sleep(Duration::from_millis(20)).await; + assert!(!handle.is_finished(), "serve returned without a shutdown"); + + // And it stops when the shutdown future resolves. + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + engine::serve(listener, Router::new(), async {}) + .await + .unwrap(); + handle.abort(); +} + +/// Serving is not on a timer. A budget measured at startup would retire a healthy process +/// once it elapsed even though nothing asked it to stop — the drain budget belongs to +/// shutdown, and there is no shutdown until there is a signal. +#[cfg(unix)] +#[tokio::test] +async fn the_service_stays_up_beyond_the_drain_budget_without_a_signal() { + use tokio::io::AsyncBufReadExt; + + let mut child = tokio::process::Command::new(env!("CARGO_BIN_EXE_omni-jev")) + .env("OMNI_SYSTEMONE_ENGINE", "passthrough") + .env("OMNI_JEV_BIND", "127.0.0.1:0") + // The whole point: far shorter than the test waits below. + .env("OMNI_SYSTEMONE_DRAIN_MS", "200") + .stderr(std::process::Stdio::piped()) + .kill_on_drop(true) + .spawn() + .unwrap(); + + let mut stderr = tokio::io::BufReader::new(child.stderr.take().unwrap()); + let mut line = String::new(); + tokio::time::timeout(Duration::from_secs(10), stderr.read_line(&mut line)) + .await + .expect("binary did not start") + .unwrap(); + let address = line + .trim() + .strip_prefix("omni-jev listening on ") + .unwrap_or_else(|| panic!("unexpected startup line {line:?}")) + .split(' ') + .next() + .unwrap() + .to_owned(); + + // No signal is sent at any point in this test. + tokio::time::sleep(Duration::from_millis(800)).await; + assert!( + child.try_wait().unwrap().is_none(), + "the process exited on its own {DRAIN_MS} after starting", + DRAIN_MS = 200 + ); + let response = post(&format!("http://{address}"), REQUEST).await; + assert_eq!( + response.status(), + 200, + "the service stopped serving with no signal sent" + ); +} + +/// An unknown engine name is a configuration error. Falling through to forwarding would send +/// inference somewhere the operator did not ask for, and the README says as much. +#[cfg(unix)] +#[tokio::test] +async fn an_unknown_engine_name_is_a_startup_error() { + // Bounded, because the failure this guards against is the process *starting*: an + // unbounded wait would hang the suite instead of reporting a regression. + let output = tokio::time::timeout( + Duration::from_secs(10), + tokio::process::Command::new(env!("CARGO_BIN_EXE_omni-jev")) + .env("OMNI_SYSTEMONE_ENGINE", "typo-engine") + .env("OMNI_JEV_BIND", "127.0.0.1:0") + .kill_on_drop(true) + .output(), + ) + .await + .expect("an unknown engine name started a server instead of failing") + .unwrap(); + assert!( + !output.status.success(), + "an unknown engine name started the process anyway" + ); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!( + stderr.contains("typo-engine"), + "the error does not name the offending value: {stderr}" + ); +} + +/// SIGTERM bounds the *whole* shutdown, not just the part after the listener stopped. A +/// request whose body never arrives holds a handler open for its entire request budget, and +/// the process must still exit within the drain budget rather than waiting that budget out. +#[cfg(unix)] +#[tokio::test] +async fn shutdown_is_bounded_by_the_drain_budget_not_the_request_budget() { + use tokio::{io::AsyncBufReadExt, io::AsyncWriteExt}; + + let mut child = tokio::process::Command::new(env!("CARGO_BIN_EXE_omni-jev")) + .env("OMNI_SYSTEMONE_ENGINE", "passthrough") + .env("OMNI_JEV_BIND", "127.0.0.1:0") + .env("OMNI_SYSTEMONE_DRAIN_MS", "200") + // Far longer than the drain budget: without one budget for the whole shutdown, this + // is how long the process would stay alive after the signal. + .env("OMNI_SYSTEMONE_TIMEOUT_MS", "30000") + .stderr(std::process::Stdio::piped()) + .kill_on_drop(true) + .spawn() + .unwrap(); + + let mut stderr = tokio::io::BufReader::new(child.stderr.take().unwrap()); + let mut line = String::new(); + tokio::time::timeout(Duration::from_secs(10), stderr.read_line(&mut line)) + .await + .expect("binary did not start") + .unwrap(); + let address = line + .trim() + .strip_prefix("omni-jev listening on ") + .unwrap_or_else(|| panic!("unexpected startup line {line:?}")) + .split(' ') + .next() + .unwrap() + .to_owned(); + + // A request whose body never arrives holds a handler for its whole request budget, which + // is what the drain budget has to bound instead. + let mut stalled = tokio::net::TcpStream::connect(&address).await.unwrap(); + stalled + .write_all( + format!( + "POST /v1/systemone HTTP/1.1\r\nHost: {address}\r\nContent-Length: 100000\r\n\r\n" + ) + .as_bytes(), + ) + .await + .unwrap(); + stalled.flush().await.unwrap(); + // And then it stays open, deliberately. The transport waits for open connections, so this + // is the case where the budget has to be enforced rather than awaited: an earlier version + // of this test hung up first and so could not see the difference. + let pid = child.id().unwrap().to_string(); + let started = Instant::now(); + assert!( + std::process::Command::new("kill") + .args(["-TERM", &pid]) + .status() + .unwrap() + .success() + ); + let exit = tokio::time::timeout(Duration::from_secs(10), child.wait()) + .await + .expect("binary did not exit after SIGTERM") + .unwrap(); + let elapsed = started.elapsed(); + assert!(exit.success(), "exit status {exit}"); + // The budget has to bound the whole sequence, and the transport cannot be asked to stop: + // it waits on the connection this test deliberately keeps open. Without a bound here, an + // earlier version waited out the handler's entire request budget — 2.6 s of a 200 ms + // budget, measured. With the close window clamped to the deadline rather than added to it, + // the exit sits at ~240 ms; the threshold is tight enough that a version which adds its own + // period instead stays red. + assert!( + elapsed < Duration::from_millis(450), + "took {elapsed:?} to exit with a 200ms drain budget and a 30000ms request budget: the \ + whole sequence has to be bounded by the drain budget" + ); + drop(stalled); +} + +/// Runs the compiled binary in native mode: the full lifecycle in one process, including +/// the parts no in-process test can cover — a real worker thread, signal handling and a +/// clean exit. +/// A drain budget much larger than the shutdown needs must not become the shutdown's length. +/// Exit is bounded by the deadline taken at the signal, not by a period measured after the +/// drain: an earlier version drained instantly and then added its own grace to the budget, so +/// the process lived `budget + grace` however little work was left. With a budget this size, +/// that additive version takes ~5.4 s where the bounded one takes under 300 ms. +#[cfg(unix)] +#[tokio::test] +async fn shutdown_does_not_wait_out_a_budget_it_does_not_need() { + use tokio::{io::AsyncBufReadExt, io::AsyncWriteExt}; + + let mut child = tokio::process::Command::new(env!("CARGO_BIN_EXE_omni-jev")) + .env("OMNI_SYSTEMONE_ENGINE", "passthrough") + .env("OMNI_JEV_BIND", "127.0.0.1:0") + .env("OMNI_SYSTEMONE_DRAIN_MS", "5000") + .env("OMNI_SYSTEMONE_TIMEOUT_MS", "30000") + .stderr(std::process::Stdio::piped()) + .kill_on_drop(true) + .spawn() + .unwrap(); + + let mut stderr = tokio::io::BufReader::new(child.stderr.take().unwrap()); + let mut line = String::new(); + tokio::time::timeout(Duration::from_secs(10), stderr.read_line(&mut line)) + .await + .expect("binary did not start") + .unwrap(); + let address = line + .trim() + .strip_prefix("omni-jev listening on ") + .unwrap_or_else(|| panic!("unexpected startup line {line:?}")) + .split(' ') + .next() + .unwrap() + .to_owned(); + + // One answered request, and the client stays connected: an idle keep-alive connection is + // still an open one, so the transport has something to wait for. + let mut held = tokio::net::TcpStream::connect(&address).await.unwrap(); + held.write_all(format!("GET /health HTTP/1.1\r\nHost: {address}\r\n\r\n").as_bytes()) + .await + .unwrap(); + held.flush().await.unwrap(); + tokio::time::sleep(Duration::from_millis(100)).await; + + let pid = child.id().unwrap().to_string(); + let started = Instant::now(); + assert!( + std::process::Command::new("kill") + .args(["-TERM", &pid]) + .status() + .unwrap() + .success() + ); + let exit = tokio::time::timeout(Duration::from_secs(20), child.wait()) + .await + .expect("binary did not exit after SIGTERM") + .unwrap(); + let elapsed = started.elapsed(); + assert!(exit.success(), "exit status {exit}"); + assert!( + elapsed < Duration::from_millis(1500), + "took {elapsed:?} to exit with a 5000ms drain budget: the budget is an upper bound, \ + not a duration to wait out" + ); + drop(held); +} + +#[cfg(unix)] +#[tokio::test] +async fn binary_serves_a_linked_engine_and_drains_on_sigterm() { + use tokio::io::AsyncBufReadExt; + + let mut child = tokio::process::Command::new(env!("CARGO_BIN_EXE_omni-jev")) + .env("OMNI_SYSTEMONE_ENGINE", "passthrough") + .env("OMNI_JEV_BIND", "127.0.0.1:0") + .env("OMNI_SYSTEMONE_QUEUE", "4") + .env("OMNI_SYSTEMONE_MAX_BODY", "256") + .env("OMNI_SYSTEMONE_TIMEOUT_MS", "5000") + .stderr(std::process::Stdio::piped()) + .kill_on_drop(true) + .spawn() + .unwrap(); + + let mut stderr = tokio::io::BufReader::new(child.stderr.take().unwrap()); + let mut line = String::new(); + tokio::time::timeout(Duration::from_secs(10), stderr.read_line(&mut line)) + .await + .expect("binary did not start") + .unwrap(); + let address = line + .trim() + .strip_prefix("omni-jev listening on ") + .unwrap_or_else(|| panic!("unexpected startup line {line:?}")) + .split(' ') + .next() + .unwrap() + .to_owned(); + let url = format!("http://{address}"); + let client = client(); + + // The line above is printed only after loading and warmup, so readiness is already + // true here: a supervisor that reads the port from stderr never races startup. + let health = client.get(format!("{url}/health")).send().await.unwrap(); + assert_eq!(health.status(), 200); + assert_eq!( + body(health).await, + r#"{"status":"ok","depth":0,"capacity":4,"rejected":0}"# + ); + + let response = post(&url, REQUEST).await; + assert_eq!(response.status(), 200); + assert_eq!( + response.headers()["x-queue-depth"], + "1", + "the request being answered is still outstanding when it is dequeued" + ); + assert_eq!(body(response).await, REQUEST); + + // The configured body limit reaches the transport, not just the config. + assert_eq!(post(&url, &"x".repeat(1024)).await.status(), 413); + // And a refused request leaves the engine serving. + assert_eq!(post(&url, REQUEST).await.status(), 200); + + let pid = child.id().unwrap().to_string(); + assert!( + std::process::Command::new("kill") + .args(["-TERM", &pid]) + .status() + .unwrap() + .success() + ); + let exit = match tokio::time::timeout(Duration::from_secs(10), child.wait()).await { + Ok(exit) => exit.unwrap(), + Err(_) => { + let alive = std::process::Command::new("kill") + .args(["-0", &pid]) + .status() + .map(|s| s.success()) + .unwrap_or(false); + let mut rest = String::new(); + let _ = + tokio::time::timeout(Duration::from_millis(200), stderr.read_line(&mut rest)).await; + panic!("binary did not exit after SIGTERM (alive={alive}, stderr={rest:?})"); + } + }; + assert!(exit.success(), "exit status {exit}"); +} diff --git a/src/frontend/tests/worker_lifecycle.rs b/src/frontend/tests/worker_lifecycle.rs new file mode 100644 index 0000000..85d75b3 --- /dev/null +++ b/src/frontend/tests/worker_lifecycle.rs @@ -0,0 +1,613 @@ +//! Worker lifecycle tests: readiness, queue admission, request budget, panic isolation +//! and shutdown. All of it runs against a fake engine, so it needs no model and no GPU — +//! which is the point of keeping this half free of model code. + +use std::{ + io, + path::PathBuf, + sync::{ + Arc, Mutex, + atomic::{AtomicBool, AtomicUsize, Ordering}, + mpsc, + }, + thread, + time::{Duration, Instant}, +}; + +use omni_jev::{ + engine::{Answer, Engine, EngineError, Readiness, Reply}, + worker::{self, Failure, Options, Request}, +}; +use tokio::runtime::Runtime; + +/// Waits for a readiness state to settle. The worker answers a request and only then +/// retires, so a client can see the failure answer before the engine reads as failed. +fn wait_for(handle: &worker::Handle, wanted: Readiness) -> Readiness { + let deadline = Instant::now() + Duration::from_secs(5); + loop { + let state = handle.readiness(); + if state == wanted || Instant::now() >= deadline { + return state; + } + thread::sleep(Duration::from_millis(5)); + } +} + +/// One gate per inference a test intends to hold: `open` is signalled by the test, `held` +/// is what the fake inference waits on. +type Gate = (mpsc::Sender<()>, mpsc::Receiver<()>); + +/// The deadline a caller computes from its drain budget when shutdown begins. +fn drain_now(budget: Duration) -> Instant { + Instant::now() + budget +} + +fn checkpoint() -> PathBuf { + PathBuf::from("/nonexistent/checkpoint") +} + +/// The engine side of a running worker, plus the runtime needed to await its replies. +struct Harness { + handle: Arc, + runtime: Runtime, + /// The budget the worker under test was started with, used to end its shutdown. + drain: Duration, +} + +impl Harness { + fn submit(&self, body: &[u8]) -> Result { + self.handle + .submit(body.to_vec(), Instant::now() + Duration::from_secs(10)) + } + + /// Submits and waits, which is what an HTTP request does. A submission that is refused — + /// because the queue is full, say — comes back as a value rather than panicking, so a + /// caller that expects back-pressure can assert on it. + fn try_reply(&self, body: &[u8]) -> Result, EngineError>, EngineError> { + let reply = self.submit(body)?; + Ok(self + .runtime + .block_on(reply) + .expect("worker dropped an accepted request") + .map(|answer| answer.body)) + } + + /// As [`Harness::try_reply`], for the requests a test expects to be served. + fn reply(&self, body: &[u8]) -> Result, EngineError> { + self.try_reply(body).expect("submission refused") + } +} + +/// A fake engine, standing in for the model. Loading is instant; inference blocks until the +/// test opens a gate, so queue state can be observed instead of raced. +struct Fake { + /// Set by warmup, so a test can prove readiness follows it and not merely loading. + warmed: Arc, + /// Requests that reached inference. + served: Arc, + /// One per in-flight inference. The receiver is kept, not dropped: a dropped receiver + /// would open the gate immediately, and the point is to hold an inference open. + gates: Arc>>, + /// Makes inference fail instead of answering. + fail: bool, + /// Makes inference panic instead of answering. + panic: bool, + /// Signalled when an inference starts, so a test can wait for the worker to be inside + /// one instead of sleeping and hoping. + entered: Arc>>>, +} + +impl Fake { + fn new() -> Self { + Self { + warmed: Arc::new(AtomicBool::new(false)), + served: Arc::new(AtomicUsize::new(0)), + gates: Arc::new(Mutex::new(Vec::new())), + fail: false, + panic: false, + entered: Arc::new(Mutex::new(None)), + } + } + + /// Issues a gate for one inference the test intends to hold open. + fn expect_in_flight(&self) { + let pair = mpsc::channel(); + self.gates.lock().unwrap().push(pair); + } +} + +fn start( + fake: &Fake, + queue_capacity: usize, + drain: Duration, +) -> io::Result<(worker::Running, Harness)> { + let for_load = Arc::clone(&fake.warmed); + let for_warmup = Arc::clone(&fake.warmed); + let for_process = Arc::clone(&fake.served); + let gates = Arc::clone(&fake.gates); + let entered = Arc::clone(&fake.entered); + let (fail, panic) = (fake.fail, fake.panic); + let running = worker::spawn( + checkpoint(), + Options { + queue_capacity, + drain, + }, + // The engine value is a unit: these tests exercise the worker, not a model. + move |_path| { + let _ = &for_load; + Ok(()) + }, + move |_engine| { + for_warmup.store(true, Ordering::Release); + Ok(()) + }, + move |request: Request<'_, ()>| { + process(request, &for_process, &gates, &entered, fail, panic) + }, + )?; + let harness = Harness { + handle: Arc::clone(&running.handle), + runtime: Runtime::new().unwrap(), + drain, + }; + Ok((running, harness)) +} + +#[allow(clippy::too_many_arguments)] +fn process( + request: Request<'_, ()>, + served: &AtomicUsize, + gates: &Mutex>, + entered: &Mutex>>, + fail: bool, + panic: bool, +) -> Result { + served.fetch_add(1, Ordering::AcqRel); + if let Some(signal) = entered.lock().unwrap().take() { + let _ = signal.send(()); + } + if panic { + panic!("engine panicked on purpose"); + } + if fail { + return Err(Failure::Inference); + } + // Blocking is deliberate: a real inference occupies the worker for its duration, and + // these tests need one that stays occupied. Gates are consumed in the order the test + // declared them, which is the order the requests arrive in. + if !gates.lock().unwrap().is_empty() { + let (open, held) = gates.lock().unwrap().remove(0); + // Waits for the test to open this gate. Tests that hold an inference are explicit + // about it by calling `expect_in_flight`; one that does not is not held. + let _ = held.recv_timeout(Duration::from_secs(10)); + let _ = open.send(()); + } + Ok(Answer { + body: request.body.to_vec(), + ..request.answer() + }) +} + +#[test] +fn spawn_returns_only_after_loading_and_warmup() { + let fake = Fake::new(); + let (running, harness) = start(&fake, 1, Duration::from_secs(1)).unwrap(); + + // Both happened on the worker thread before `spawn` returned, so a caller holding a + // handle cannot observe a half-loaded engine. + assert!(fake.warmed.load(Ordering::Acquire)); + assert_eq!(harness.handle.readiness(), Readiness::Ready); + assert_eq!(harness.reply(b"hello").unwrap(), b"hello"); + + assert!(running.run_until_drained(drain_now(harness.drain))); + assert_eq!(harness.handle.readiness(), Readiness::Starting); +} + +#[test] +fn startup_failure_is_reported() { + let error = worker::spawn( + checkpoint(), + Options { + queue_capacity: 1, + drain: Duration::from_secs(1), + }, + |_path| Err(io::Error::other("no such checkpoint")), + |_engine| Ok(()), + |request: Request<'_, ()>| { + Ok(Answer { + body: Vec::new(), + ..request.answer() + }) + }, + ) + .err() + .expect("a failed load must not return a running worker"); + assert!(error.to_string().contains("no such checkpoint"), "{error}"); +} + +#[test] +fn queue_capacity_refuses_rather_than_grows() { + let fake = Fake::new(); + let (running, harness, blocked) = Blocked::wait(&fake, 2); + // The first request is held inside inference, so the depth below cannot drift: with the + // channel empty and this thread driving every submission, the two numbers are fixed. + let _first = harness.submit(b"first").expect("first refused"); + blocked.entered(); + let report = harness + .handle + .report() + .expect("the worker reports its queue"); + assert_eq!(report.depth, 1, "the running request counts"); + + // The second takes the remaining slot; the third has nowhere to go. An earlier version of + // this test submitted without waiting for the first to start, so whether the worker had + // already dequeued the second was a race, and it failed about once in eight runs. + let second = harness.submit(b"second").expect("second refused"); + match harness.submit(b"third") { + Err(EngineError::Busy) => {} + Err(other) => panic!("expected Busy, got {other:?}"), + Ok(_) => panic!("expected Busy, got a reply channel"), + } + let report = harness + .handle + .report() + .expect("the worker reports its queue"); + assert_eq!(report.depth, 2, "a refused request must not count"); + assert_eq!(report.capacity, 2); + assert_eq!(report.rejected, 1); + + // A refusal is not a lost slot: both accepted requests are still served, and once they + // have released their slots the same queue admits the next request. The first is answered + // by the blocked inference; the second is answered by the fake the moment it runs. + blocked.release(); + drop(second); + let deadline = Instant::now() + Duration::from_secs(5); + while harness.handle.report().unwrap().depth > 0 && Instant::now() < deadline { + thread::sleep(Duration::from_millis(5)); + } + assert_eq!( + harness.handle.report().unwrap().depth, + 0, + "the queue did not drain" + ); + assert_eq!(harness.reply(b"fifth").unwrap(), b"fifth"); + assert!(running.run_until_drained(drain_now(harness.drain))); +} + +/// Starts a worker whose inference spins on `hold` and then answers, taking no locks, so a +/// test can keep requests outstanding without ever holding a mutex the worker wants. +fn start_with_hold( + fake: &Fake, + queue_capacity: usize, + drain: Duration, + hold: &Arc, +) -> (worker::Running, Harness) { + let gate = Arc::clone(hold); + let running = worker::spawn( + checkpoint(), + Options { + queue_capacity, + drain, + }, + |_path| Ok(()), + |_engine| Ok(()), + move |request: Request<'_, ()>| { + let spin = Instant::now() + Duration::from_micros(20); + while gate.load(Ordering::Acquire) && Instant::now() < spin { + thread::yield_now(); + } + Ok(Answer { + body: request.body.to_vec(), + ..request.answer() + }) + }, + ) + .unwrap(); + let harness = Harness { + handle: Arc::clone(&running.handle), + runtime: Runtime::new().unwrap(), + drain, + }; + let _ = fake; + (running, harness) +} + +/// A worker held on one request, so a test can decide exactly what is queued behind it. +struct Blocked { + open: mpsc::Sender<()>, + entered: mpsc::Receiver<()>, +} + +impl Blocked { + /// Starts a worker whose next inference will block until [`Blocked::release`], and + /// signals when that inference has actually started. The caller submits the request that + /// takes the gate, so the depth it has to fit into is the caller's decision. + fn wait(fake: &Fake, queue_capacity: usize) -> (worker::Running, Harness, Self) { + let (open, held) = mpsc::channel(); + fake.gates.lock().unwrap().push((open.clone(), held)); + let (entered, started) = mpsc::channel(); + *fake.entered.lock().unwrap() = Some(entered); + let (ready, harness) = start(fake, queue_capacity, Duration::from_secs(2)).unwrap(); + let blocked = Self { + open, + entered: started, + }; + (ready, harness, blocked) + } + + fn entered(&self) { + self.entered + .recv_timeout(Duration::from_secs(5)) + .expect("the held inference never started"); + } + + fn release(&self) { + let _ = self.open.send(()); + } +} + +/// The configured limit covers requests that are *outstanding*, not only requests that are +/// waiting in the channel. The worker frees its channel slot the moment it dequeues a job, +/// so admission has to count the running one too — otherwise capacity 1 serves two at once +/// while the README promises a `503` for the second. +#[test] +fn a_running_request_counts_against_capacity() { + let fake = Fake::new(); + let (running, harness, blocked) = Blocked::wait(&fake, 1); + // The first request is inside inference, so the channel is empty and only the + // outstanding count can refuse the next one. + let _first = harness.submit(b"first").expect("the first request fits"); + blocked.entered(); + + match harness.submit(b"second") { + Err(EngineError::Busy) => {} + Err(other) => panic!("expected Busy while one request is running, got {other:?}"), + Ok(_) => panic!("a second request was accepted while one was already running"), + } + let report = harness + .handle + .report() + .expect("the worker reports its queue"); + assert_eq!(report.depth, 1, "a refused request must not be counted"); + assert_eq!(report.capacity, 1); + assert_eq!(report.rejected, 1); + + // The refusal is not a lost slot: admission resumes once the running request has + // finished. The slot is released as the worker sends its reply, so the next request can + // be refused for a moment after the reply arrives — retry rather than assume ordering. + blocked.release(); + let deadline = Instant::now() + Duration::from_secs(10); + loop { + match harness.try_reply(b"third") { + Ok(Ok(body)) => { + assert_eq!(body, b"third"); + break; + } + // Refused, not lost: retry until the finished request has released its slot. + Err(EngineError::Busy) if Instant::now() < deadline => thread::yield_now(), + other => panic!("admission never resumed: {other:?}"), + } + } + assert_eq!(harness.handle.report().unwrap().depth, 0); + assert!(running.run_until_drained(drain_now(harness.drain))); +} + +#[test] +fn cancelled_and_expired_requests_never_reach_the_engine() { + let fake = Fake::new(); + let (running, harness, blocked) = Blocked::wait(&fake, 4); + // The worker is parked inside inference, so anything submitted now is certainly queued + // rather than already picked up. Without this handshake the test races the worker. + let _blocker = harness.submit(b"blocker").unwrap(); + blocked.entered(); + + // A client that gave up while its request waited. + let abandoned = harness.submit(b"abandoned").unwrap(); + drop(abandoned); + + // A request whose budget expires while it waits. The wait below is what makes it + // expire: the worker is held inside inference, which is exactly the queueing delay the + // budget covers. + let late = harness + .handle + .submit(b"late".to_vec(), Instant::now() + Duration::from_millis(20)) + .unwrap(); + thread::sleep(Duration::from_millis(80)); + + assert!( + fake.served.load(Ordering::Acquire) == 1, + "only the blocker may have started" + ); + blocked.release(); + assert_eq!(harness.reply(b"served").unwrap(), b"served"); + + assert_eq!( + fake.served.load(Ordering::Acquire), + 2, + "only the blocker and the later request may run" + ); + assert_eq!( + harness.handle.report().unwrap().rejected, + 2, + "both skipped requests are reported as work the worker refused to run" + ); + assert_eq!( + harness.handle.report().unwrap().depth, + 0, + "a skipped request must release its queue slot" + ); + assert!( + harness.runtime.block_on(late).is_err(), + "the expired request must have been dropped, not answered" + ); + + assert!(running.run_until_drained(drain_now(harness.drain))); +} + +#[test] +fn a_failed_inference_retires_the_worker() { + let mut fake = Fake::new(); + fake.fail = true; + let (running, harness) = start(&fake, 1, Duration::from_secs(2)).unwrap(); + + let answer = harness.reply(b"body"); + assert!( + matches!(answer, Err(EngineError::InferenceFailed)), + "{answer:?}" + ); + + // An engine that could not run one request is not offered the next one. + assert_eq!( + wait_for(&harness.handle, Readiness::Failed), + Readiness::Failed + ); + assert!(matches!( + harness.submit(b"body"), + Err(EngineError::Unavailable) + )); + assert!(running.run_until_drained(drain_now(harness.drain))); +} + +#[test] +fn a_panicking_engine_answers_then_retires() { + let mut fake = Fake::new(); + fake.panic = true; + let (running, harness) = start(&fake, 1, Duration::from_secs(2)).unwrap(); + + let answer = harness.reply(b"body"); + + // The request is answered rather than dropped, so no client is left waiting, and the + // engine is retired because a panic leaves its state unknown. + assert!( + matches!(answer, Err(EngineError::InferenceFailed)), + "{answer:?}" + ); + assert_eq!( + wait_for(&harness.handle, Readiness::Failed), + Readiness::Failed + ); + let reason = harness + .handle + .failure() + .expect("a failed engine explains itself"); + assert!(reason.contains("on purpose"), "{reason}"); + assert!(running.run_until_drained(drain_now(harness.drain))); +} + +#[test] +fn drain_reports_when_it_cannot_finish() { + let fake = Fake::new(); + let (running, harness) = start(&fake, 4, Duration::from_millis(150)).unwrap(); + // A gate nobody opens stands in for a wedged kernel: nothing here can interrupt it. + // The second request queues behind it, so drain has both kinds of work to wait for. + fake.expect_in_flight(); + let _stuck = harness.submit(b"stuck").unwrap(); + let _queued = harness.submit(b"queued").unwrap(); + assert_eq!( + wait_for(&harness.handle, Readiness::Ready), + Readiness::Ready + ); + // The deadline is measured from here, as a caller's would be when shutdown begins. + let start = Instant::now(); + assert!( + !running.run_until_drained(drain_now(harness.drain)), + "drain must not claim success while work is outstanding" + ); + assert!( + start.elapsed() < Duration::from_secs(2), + "drain must respect its budget, took {:?}", + start.elapsed() + ); +} + +#[test] +fn a_drained_worker_refuses_further_work() { + let fake = Fake::new(); + let (running, harness) = start(&fake, 1, Duration::from_secs(2)).unwrap(); + assert_eq!(harness.handle.readiness(), Readiness::Ready); + assert!(running.run_until_drained(drain_now(harness.drain))); + + // The worker thread has exited, so nothing can answer even though the handle is alive. + // An engine that was never ready would also read as `Starting`, so check that this one + // did become ready before it was drained. + assert!(matches!( + harness.submit(b"body"), + Err(EngineError::Unavailable) + )); + assert_eq!(harness.handle.readiness(), Readiness::Starting); + assert!(fake.warmed.load(Ordering::Acquire)); +} + +#[test] +fn capacity_zero_is_rejected_at_startup() { + let error = worker::spawn( + checkpoint(), + Options { + queue_capacity: 0, + drain: Duration::from_secs(1), + }, + |_path| Ok(()), + |_engine| Ok(()), + |request: Request<'_, ()>| { + Ok(Answer { + body: Vec::new(), + ..request.answer() + }) + }, + ) + .err() + .expect("a zero-length queue cannot admit anything"); + assert!(error.to_string().contains("at least 1"), "{error}"); +} + +/// Queue depth is reserved before the job is published, so a consumer that wins the race to +/// dequeue cannot release a slot that was never taken. The bug this pins: incrementing after +/// the send lets the worker finish and decrement first, which reads as a wrong depth and +/// wraps the counter, leaving `/health` reporting a huge number. +#[test] +fn depth_is_reserved_before_the_job_is_published() { + let fake = Fake::new(); + // Each inference spins briefly. That is the window in which the worker and the caller + // race hardest, and it also keeps jobs outstanding long enough for the count below to + // mean something. + let hold = Arc::new(AtomicBool::new(true)); + let (running, harness) = start_with_hold(&fake, 64, Duration::from_secs(2), &hold); + + // Every reply is held rather than awaited, so nothing accepted can have completed and + // the accepted count is exactly what the reported depth has to be. + let deadline = Instant::now() + Duration::from_secs(2); + let mut replies = Vec::new(); + while replies.len() < 20_000 && Instant::now() < deadline { + if let Ok(reply) = harness.submit(b"storm") { + replies.push(reply); + let depth = harness.handle.report().unwrap().depth; + // A slot cannot be counted that was never taken, so depth can never exceed the + // requests this caller has accepted. Counting after the send broke it in the + // other direction: the released-before-reserved decrement wraps a `usize`, so + // the failure shows up here as a huge number. + assert!( + depth <= replies.len(), + "depth {depth} exceeded the {} requests accepted", + replies.len() + ); + } + } + assert!(replies.len() > 100, "the storm never got going"); + let accepted = replies.len(); + + // Stop holding inference, then let everything finish: depth has to return to zero + // rather than staying wrapped at a huge number. + hold.store(false, Ordering::Release); + drop(replies); + let deadline = Instant::now() + Duration::from_secs(20); + while harness.handle.report().unwrap().depth > 0 && Instant::now() < deadline { + thread::sleep(Duration::from_millis(5)); + } + assert_eq!( + harness.handle.report().unwrap().depth, + 0, + "the queue did not drain after {accepted} requests" + ); + assert!(running.run_until_drained(drain_now(harness.drain))); +}