diff --git a/.env.example b/.env.example index be039bf..5dd2e43 100644 --- a/.env.example +++ b/.env.example @@ -1,18 +1,6 @@ -# Redis Streams CONTINUUM_REDIS_URL=redis://localhost:6379 CONTINUUM_FRAME_STREAM=stream:camera_frames CONTINUUM_CONSUMER_GROUP=ai-worker CONTINUUM_CONSUMER_NAME=ai-worker-1 -CONTINUUM_STREAM_BLOCK_MS=5000 -CONTINUUM_STREAM_BATCH_SIZE=16 - -# Postgres (fast-path alert writes; same database as continuum-backend) -DATABASE_URL=postgres://continuum:continuum@localhost:5432/continuum - -# Inference -CONTINUUM_MODEL_PATH=./models/failure-detector.onnx -CONTINUUM_MODEL_INPUT_SIZE=640 -CONTINUUM_CONFIDENCE_THRESHOLD=0.75 -CONTINUUM_INTRA_OP_THREADS=4 RUST_LOG=info,continuum_ai_worker=debug diff --git a/Cargo.toml b/Cargo.toml index 38455b2..5cdceb4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,7 +2,7 @@ name = "continuum-ai-worker" version = "0.1.0" edition = "2021" -description = "Continuum AI worker — consumes camera frame batches from Redis Streams and runs ONNX Runtime print-failure detection" +description = "Continuum AI worker — consumes camera frame batches from Redis Streams and runs print-failure detection" license = "UNLICENSED" [[bin]] @@ -12,21 +12,13 @@ path = "src/main.rs" [dependencies] tokio = { version = "1.40", features = ["full"] } redis = { version = "0.27", features = ["tokio-comp", "streams"] } -ort = { version = "2.0.0-rc.9", features = ["ndarray", "download-binaries"] } -ndarray = "0.17" -image = "0.25" serde = { version = "1", features = ["derive"] } serde_json = "1" tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] } anyhow = "1" -thiserror = "1" -base64 = "0.22" -sqlx = { version = "0.8", features = ["runtime-tokio", "postgres", "chrono", "uuid"] } -uuid = { version = "1", features = ["v4", "serde"] } -chrono = { version = "0.4", features = ["serde"] } dotenvy = "0.15" -[profile.release] -opt-level = 3 -lto = true +# Real inference (ONNX Runtime via the `ort` crate) and the Postgres audit +# write (via `sqlx`) go here once src/inference and src/reporter grow past +# their stubs — see those modules' doc comments. diff --git a/README.md b/README.md index ec1760e..37f3e77 100644 --- a/README.md +++ b/README.md @@ -2,9 +2,18 @@ Computer-vision failure detection worker for the Continuum print farm platform. Consumes camera frame batches pushed onto the Redis Stream -`stream:camera_frames` (by `continuum-backend`'s `src/queue/`), runs them -through an ONNX Runtime model, and fast-paths high-confidence failure alerts -back to Redis and Postgres. +`stream:camera_frames` (by `continuum-backend`'s `src/queue/`) and reports +failures. + +## This is a learning-stage boilerplate + +Real ONNX Runtime inference (the `ort` crate) and the Postgres audit write +(`sqlx`) are **not** implemented yet — `src/inference/mod.rs` always +"detects" one fixed example failure, and `src/reporter/mod.rs` only prints +where the Postgres write would go. The Redis Streams consumer-group loop in +`main.rs` and the Redis pub/sub alert in `reporter/` are real. This lets the +whole pipeline (read a frame, "detect", report, ack) run and be understood +before adding a real model or a SQL database on top. ## Pipeline @@ -13,24 +22,24 @@ continuum-proxy --(frames)--> continuum-backend --(XADD)--> stream:camera_frames | XREADGROUP (this worker) v - src/inference (ort session) + inference::detect (stub for now) v - src/reporter (Redis pub/sub + Postgres INSERT) -``` - -## Structure - -``` -src/ - main.rs Consumer-group loop: XREADGROUP -> decode -> infer -> report -> XACK - inference/ Thread-safe ONNX Runtime session wrapper + pre/post-processing - reporter/ Alert fan-out: Redis pub/sub for realtime UI, Postgres for audit trail + reporter::report (Redis pub/sub — real) ``` ## Getting started ```bash cp .env.example .env -# place an ONNX failure-detection model at ./models/failure-detector.onnx cargo run ``` + +## Structure + +``` +src/ + main.rs Consumer-group loop: XREADGROUP -> detect -> report -> XACK + config.rs Loads settings from environment variables + inference/ Detection stub — real ONNX Runtime inference goes here later + reporter/ Redis pub/sub alert (real) + a note for the Postgres write (later) +``` diff --git a/models/.gitkeep b/models/.gitkeep deleted file mode 100644 index e69de29..0000000 diff --git a/src/config.rs b/src/config.rs index 3050b1d..f72b38f 100644 --- a/src/config.rs +++ b/src/config.rs @@ -6,13 +6,6 @@ pub struct Config { pub frame_stream: String, pub consumer_group: String, pub consumer_name: String, - pub stream_block_ms: usize, - pub stream_batch_size: usize, - pub database_url: String, - pub model_path: String, - pub model_input_size: u32, - pub confidence_threshold: f32, - pub intra_op_threads: usize, } impl Config { @@ -22,21 +15,10 @@ impl Config { frame_stream: env_or("CONTINUUM_FRAME_STREAM", "stream:camera_frames"), consumer_group: env_or("CONTINUUM_CONSUMER_GROUP", "ai-worker"), consumer_name: env_or("CONTINUUM_CONSUMER_NAME", "ai-worker-1"), - stream_block_ms: env_or("CONTINUUM_STREAM_BLOCK_MS", "5000").parse()?, - stream_batch_size: env_or("CONTINUUM_STREAM_BATCH_SIZE", "16").parse()?, - database_url: require("DATABASE_URL")?, - model_path: env_or("CONTINUUM_MODEL_PATH", "./models/failure-detector.onnx"), - model_input_size: env_or("CONTINUUM_MODEL_INPUT_SIZE", "640").parse()?, - confidence_threshold: env_or("CONTINUUM_CONFIDENCE_THRESHOLD", "0.75").parse()?, - intra_op_threads: env_or("CONTINUUM_INTRA_OP_THREADS", "4").parse()?, }) } } -fn require(key: &str) -> anyhow::Result { - env::var(key).map_err(|_| anyhow::anyhow!("missing required env var {key}")) -} - fn env_or(key: &str, default: &str) -> String { env::var(key).unwrap_or_else(|_| default.to_string()) } diff --git a/src/inference/detector.rs b/src/inference/detector.rs deleted file mode 100644 index 4fa25fc..0000000 --- a/src/inference/detector.rs +++ /dev/null @@ -1,88 +0,0 @@ -use std::sync::Arc; - -use image::DynamicImage; -use ort::session::builder::GraphOptimizationLevel; -use ort::session::Session; -use ort::value::Value; -use serde::Serialize; -use tokio::sync::Mutex; - -use super::preprocess::to_input_tensor; - -#[derive(Debug, Clone, Serialize)] -pub struct Detection { - pub class_id: u32, - pub label: String, - pub confidence: f32, - /// Normalized [0, 1] bounding box in (x_min, y_min, x_max, y_max). - pub bbox: [f32; 4], -} - -const CLASS_LABELS: &[&str] = &["spaghetti", "warping", "layer_shift", "bed_adhesion_failure"]; - -/// Thread-safe wrapper around an ONNX Runtime session. `ort::Session` is not -/// `Sync` on its own for concurrent `run()` calls, so callers share this -/// behind a `Mutex` and an `Arc` when fanning out across worker tasks. -pub struct FailureDetector { - session: Mutex, - input_size: u32, - confidence_threshold: f32, -} - -impl FailureDetector { - pub fn load(model_path: &str, input_size: u32, confidence_threshold: f32, intra_threads: usize) -> anyhow::Result> { - // ort's builder error types carry raw FFI pointers that aren't - // Send/Sync, so they can't flow through `?` into `anyhow::Result` - // directly (anyhow requires Send + Sync + 'static). Stringify at - // the boundary instead. - let session = Session::builder() - .map_err(|e| anyhow::anyhow!("failed to create ONNX Runtime session builder: {e}"))? - .with_optimization_level(GraphOptimizationLevel::Level3) - .map_err(|e| anyhow::anyhow!("failed to set graph optimization level: {e}"))? - .with_intra_threads(intra_threads) - .map_err(|e| anyhow::anyhow!("failed to set intra-op thread count: {e}"))? - .commit_from_file(model_path) - .map_err(|e| anyhow::anyhow!("failed to load ONNX model at {model_path}: {e}"))?; - - Ok(Arc::new(Self { - session: Mutex::new(session), - input_size, - confidence_threshold, - })) - } - - /// Runs inference on a single decoded frame and returns detections whose - /// confidence clears `confidence_threshold`. - pub async fn detect(&self, frame: &DynamicImage) -> anyhow::Result> { - let tensor = to_input_tensor(frame, self.input_size); - let input = Value::from_array(tensor)?; - - let mut session = self.session.lock().await; - let outputs = session.run(ort::inputs!["images" => input])?; - - // Exported with NMS baked into the graph: output0 is [1, N, 6] - // rows of (x1, y1, x2, y2, confidence, class_id) in input-pixel space. - let (shape, data) = outputs["output0"].try_extract_tensor::()?; - let num_boxes = shape[1] as usize; - let size = self.input_size as f32; - - let mut detections = Vec::new(); - for i in 0..num_boxes { - let row = &data[i * 6..i * 6 + 6]; - let confidence = row[4]; - if confidence < self.confidence_threshold { - continue; - } - - let class_id = row[5] as u32; - detections.push(Detection { - class_id, - label: CLASS_LABELS.get(class_id as usize).copied().unwrap_or("unknown").to_string(), - confidence, - bbox: [row[0] / size, row[1] / size, row[2] / size, row[3] / size], - }); - } - - Ok(detections) - } -} diff --git a/src/inference/mod.rs b/src/inference/mod.rs index c97aae8..f8a7913 100644 --- a/src/inference/mod.rs +++ b/src/inference/mod.rs @@ -1,4 +1,19 @@ -mod detector; -mod preprocess; +use serde::Serialize; -pub use detector::{Detection, FailureDetector}; +#[derive(Debug, Clone, Serialize)] +pub struct Detection { + pub label: String, + pub confidence: f32, +} + +/// Stand-in for real ONNX Runtime inference (the `ort` crate, using a model +/// trained on failure images). Always reports one fixed "failure" so the +/// rest of the pipeline — reporting, acking the stream entry — has +/// something to do and can be tested without a real model file. +/// +/// Swap this out once you're comfortable with the pipeline around it: +/// decode `frame_bytes` with the `image` crate, run it through an `ort` +/// session, and turn the model's output into real `Detection`s. +pub fn detect(_frame_bytes: &[u8]) -> Vec { + vec![Detection { label: "spaghetti".to_string(), confidence: 0.91 }] +} diff --git a/src/inference/preprocess.rs b/src/inference/preprocess.rs deleted file mode 100644 index 124b1fb..0000000 --- a/src/inference/preprocess.rs +++ /dev/null @@ -1,20 +0,0 @@ -use image::DynamicImage; -use ndarray::Array4; - -/// Resizes an RGB frame to a square `size x size` input tensor in CHW layout, -/// normalized to [0, 1] — the standard preprocessing for YOLO-family ONNX -/// export graphs. -pub fn to_input_tensor(image: &DynamicImage, size: u32) -> Array4 { - let resized = image.resize_exact(size, size, image::imageops::FilterType::Triangle); - let rgb = resized.to_rgb8(); - - let mut tensor = Array4::::zeros((1, 3, size as usize, size as usize)); - for (x, y, pixel) in rgb.enumerate_pixels() { - let [r, g, b] = pixel.0; - tensor[[0, 0, y as usize, x as usize]] = r as f32 / 255.0; - tensor[[0, 1, y as usize, x as usize]] = g as f32 / 255.0; - tensor[[0, 2, y as usize, x as usize]] = b as f32 / 255.0; - } - - tensor -} diff --git a/src/main.rs b/src/main.rs index 71d8896..8a43ffc 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,138 +2,59 @@ mod config; mod inference; mod reporter; -use std::collections::HashMap; - use redis::streams::{StreamKey, StreamReadOptions, StreamReadReply}; use redis::AsyncCommands; -use tracing::{error, info, warn}; use tracing_subscriber::EnvFilter; use crate::config::Config; -use crate::inference::FailureDetector; use crate::reporter::Reporter; #[tokio::main] async fn main() -> anyhow::Result<()> { tracing_subscriber::fmt() .with_env_filter(EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"))) - .json() .init(); - dotenvy_load(); + dotenvy::dotenv().ok(); let config = Config::from_env()?; - info!(model = %config.model_path, "loading ONNX Runtime failure-detection model"); - let detector = FailureDetector::load( - &config.model_path, - config.model_input_size, - config.confidence_threshold, - config.intra_op_threads, - )?; - - let mut reporter = Reporter::connect(&config.redis_url, &config.database_url).await?; - let client = redis::Client::open(config.redis_url.clone())?; let mut conn = client.get_multiplexed_tokio_connection().await?; + let mut reporter = Reporter::connect(&config.redis_url).await?; - ensure_consumer_group(&mut conn, &config).await?; - - info!( - stream = %config.frame_stream, - group = %config.consumer_group, - consumer = %config.consumer_name, - "starting frame consumer loop" - ); - - let read_opts = StreamReadOptions::default() - .group(&config.consumer_group, &config.consumer_name) - .block(config.stream_block_ms) - .count(config.stream_batch_size); - - loop { - let reply: StreamReadReply = match conn - .xread_options(&[&config.frame_stream], &[">"], &read_opts) - .await - { - Ok(reply) => reply, - Err(err) => { - error!(?err, "XREADGROUP failed, backing off"); - tokio::time::sleep(std::time::Duration::from_secs(2)).await; - continue; - } - }; - - for StreamKey { key, ids } in reply.keys { - for entry in ids { - let fields = entry_fields(&entry.map); - match process_entry(&detector, &mut reporter, &fields).await { - Ok(()) => {} - Err(err) => warn!(?err, id = %entry.id, "failed to process frame entry"), - } - - let _: redis::RedisResult = conn - .xack(&key, &config.consumer_group, &[entry.id.as_str()]) - .await; - } - } - } -} - -async fn ensure_consumer_group( - conn: &mut redis::aio::MultiplexedConnection, - config: &Config, -) -> anyhow::Result<()> { - let result: redis::RedisResult = conn + // Consumer groups have to exist before you can read from them. This + // fails with "BUSYGROUP" if it already exists — that's fine, ignore it. + let _: redis::RedisResult = conn .xgroup_create_mkstream(&config.frame_stream, &config.consumer_group, "$") .await; - if let Err(err) = result { - // BUSYGROUP means the group already exists — fine, everything else is real. - if !err.to_string().contains("BUSYGROUP") { - return Err(err.into()); + tracing::info!(stream = %config.frame_stream, group = %config.consumer_group, "waiting for camera frames"); + + let read_opts = StreamReadOptions::default() + .group(&config.consumer_group, &config.consumer_name) + .block(5000); + + loop { + let reply: StreamReadReply = conn.xread_options(&[&config.frame_stream], &[">"], &read_opts).await?; + + for StreamKey { key, ids } in reply.keys { + for entry in ids { + let printer_id = printer_id_field(&entry.map); + + let detections = inference::detect(&[]); + if !detections.is_empty() { + reporter.report(&printer_id, &detections).await?; + } + + let _: redis::RedisResult = conn.xack(&key, &config.consumer_group, &[entry.id.as_str()]).await; + } } } - - Ok(()) } -fn entry_fields(map: &HashMap) -> HashMap { - map.iter() - .filter_map(|(k, v)| match v { - redis::Value::BulkString(bytes) => { - Some((k.clone(), String::from_utf8_lossy(bytes).to_string())) - } - _ => None, - }) - .collect() -} - -async fn process_entry( - detector: &FailureDetector, - reporter: &mut Reporter, - fields: &HashMap, -) -> anyhow::Result<()> { - let printer_id = fields - .get("printer_id") - .ok_or_else(|| anyhow::anyhow!("frame entry missing printer_id field"))?; - let job_id = fields.get("job_id").map(String::as_str); - let frame_b64 = fields - .get("frame") - .ok_or_else(|| anyhow::anyhow!("frame entry missing frame field"))?; - - let bytes = base64::Engine::decode(&base64::engine::general_purpose::STANDARD, frame_b64)?; - let image = image::load_from_memory(&bytes)?; - - let detections = detector.detect(&image).await?; - if !detections.is_empty() { - reporter.report_failure(printer_id, job_id, &detections).await?; +fn printer_id_field(fields: &std::collections::HashMap) -> String { + match fields.get("printer_id") { + Some(redis::Value::BulkString(bytes)) => String::from_utf8_lossy(bytes).to_string(), + _ => "unknown".to_string(), } - - Ok(()) -} - -/// Loads `.env` if present; a missing file is fine (env vars may already be -/// set by the process supervisor), so the error is deliberately ignored. -fn dotenvy_load() { - let _ = dotenvy::dotenv(); } diff --git a/src/reporter/mod.rs b/src/reporter/mod.rs index b90b11a..8736e0a 100644 --- a/src/reporter/mod.rs +++ b/src/reporter/mod.rs @@ -1,42 +1,40 @@ -mod postgres; -mod redis_pub; - -use crate::inference::Detection; -use sqlx::PgPool; +use redis::AsyncCommands; use tracing::info; +use crate::inference::Detection; + +const ALERT_CHANNEL: &str = "channel:print_alerts"; + +/// Fast-paths a failure alert out to whatever's listening. pub struct Reporter { redis: redis::aio::MultiplexedConnection, - pg: PgPool, } impl Reporter { - pub async fn connect(redis_url: &str, database_url: &str) -> anyhow::Result { - let redis_client = redis::Client::open(redis_url)?; - let redis = redis_client.get_multiplexed_tokio_connection().await?; - let pg = PgPool::connect(database_url).await?; - - Ok(Self { redis, pg }) + pub async fn connect(redis_url: &str) -> anyhow::Result { + let client = redis::Client::open(redis_url)?; + let redis = client.get_multiplexed_tokio_connection().await?; + Ok(Self { redis }) } - /// Fast-paths a failure alert: publishes to Redis immediately for any - /// live UI subscribers, then persists to Postgres for the audit trail - /// and downstream farm-manager notifications. - pub async fn report_failure( - &mut self, - printer_id: &str, - job_id: Option<&str>, - detections: &[Detection], - ) -> anyhow::Result<()> { - let best = detections - .iter() - .max_by(|a, b| a.confidence.total_cmp(&b.confidence)) - .expect("report_failure called with no detections"); - + /// Publishes to Redis (real — this is what `continuum-backend`'s + /// WebSocket fan-out subscribes to, for live operator-console alerts). + /// + /// A real version would also `INSERT` into Postgres here via `sqlx` for + /// the permanent audit trail — that's left as a `println!` for now; + /// come back to it once async SQL feels comfortable. + pub async fn report(&mut self, printer_id: &str, detections: &[Detection]) -> anyhow::Result<()> { + let best = &detections[0]; info!(printer_id, label = %best.label, confidence = best.confidence, "reporting print failure"); - redis_pub::publish_alert(&mut self.redis, printer_id, best).await?; - postgres::insert_alert(&self.pg, printer_id, job_id, detections).await?; + let payload = serde_json::json!({ + "printerId": printer_id, + "label": best.label, + "confidence": best.confidence, + }); + self.redis.publish::<_, _, ()>(ALERT_CHANNEL, payload.to_string()).await?; + + println!("(would also INSERT this alert into Postgres for the audit trail)"); Ok(()) } diff --git a/src/reporter/postgres.rs b/src/reporter/postgres.rs deleted file mode 100644 index 4caae22..0000000 --- a/src/reporter/postgres.rs +++ /dev/null @@ -1,31 +0,0 @@ -use sqlx::PgPool; -use uuid::Uuid; - -use crate::inference::Detection; - -/// Persists the full detection set for a failure event to the shared -/// Postgres database (the same instance `continuum-backend` writes to via -/// Drizzle), for audit history and the farm-manager review queue. -pub async fn insert_alert( - pool: &PgPool, - printer_id: &str, - job_id: Option<&str>, - detections: &[Detection], -) -> anyhow::Result<()> { - let detections_json = serde_json::to_value(detections)?; - - sqlx::query( - r#" - INSERT INTO print_failure_alerts (id, printer_id, print_job_id, detections, created_at) - VALUES ($1, $2, $3, $4, now()) - "#, - ) - .bind(Uuid::new_v4()) - .bind(printer_id) - .bind(job_id) - .bind(detections_json) - .execute(pool) - .await?; - - Ok(()) -} diff --git a/src/reporter/redis_pub.rs b/src/reporter/redis_pub.rs deleted file mode 100644 index 98da27d..0000000 --- a/src/reporter/redis_pub.rs +++ /dev/null @@ -1,26 +0,0 @@ -use redis::AsyncCommands; -use serde_json::json; - -use crate::inference::Detection; - -const ALERT_CHANNEL: &str = "channel:print_alerts"; - -/// Publishes a low-latency alert for anything subscribed to -/// `channel:print_alerts` — primarily `continuum-backend`'s WebSocket -/// fan-out to connected operator consoles. -pub async fn publish_alert( - conn: &mut redis::aio::MultiplexedConnection, - printer_id: &str, - detection: &Detection, -) -> anyhow::Result<()> { - let payload = json!({ - "type": "printer.failure_detected", - "printerId": printer_id, - "label": detection.label, - "confidence": detection.confidence, - "bbox": detection.bbox, - }); - - conn.publish::<_, _, ()>(ALERT_CHANNEL, payload.to_string()).await?; - Ok(()) -}