From 17f8e4b4879b6f23b56cf5e23957f65bcf8513a7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Iwo=20Strzebo=C5=84ski?= Date: Fri, 28 Aug 2026 16:26:53 +0000 Subject: [PATCH] Initial boilerplate scaffold for continuum-ai-worker --- .env.example | 18 +++++ .gitignore | 5 ++ Cargo.toml | 31 ++++++++ README.md | 36 +++++++++ models/.gitkeep | 0 src/config.rs | 42 ++++++++++ src/inference/detector.rs | 80 +++++++++++++++++++ src/inference/mod.rs | 4 + src/inference/preprocess.rs | 24 ++++++ src/main.rs | 149 ++++++++++++++++++++++++++++++++++++ src/reporter/mod.rs | 43 +++++++++++ src/reporter/postgres.rs | 31 ++++++++ src/reporter/redis_pub.rs | 26 +++++++ 13 files changed, 489 insertions(+) create mode 100644 .env.example create mode 100644 .gitignore create mode 100644 Cargo.toml create mode 100644 README.md create mode 100644 models/.gitkeep create mode 100644 src/config.rs create mode 100644 src/inference/detector.rs create mode 100644 src/inference/mod.rs create mode 100644 src/inference/preprocess.rs create mode 100644 src/main.rs create mode 100644 src/reporter/mod.rs create mode 100644 src/reporter/postgres.rs create mode 100644 src/reporter/redis_pub.rs diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..be039bf --- /dev/null +++ b/.env.example @@ -0,0 +1,18 @@ +# 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/.gitignore b/.gitignore new file mode 100644 index 0000000..6f77511 --- /dev/null +++ b/.gitignore @@ -0,0 +1,5 @@ +/target +Cargo.lock +*.onnx +.env +*.log diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..9601b49 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,31 @@ +[package] +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" +license = "UNLICENSED" + +[[bin]] +name = "continuum-ai-worker" +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.16" +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"] } + +[profile.release] +opt-level = 3 +lto = true diff --git a/README.md b/README.md new file mode 100644 index 0000000..ec1760e --- /dev/null +++ b/README.md @@ -0,0 +1,36 @@ +# continuum-ai-worker + +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. + +## Pipeline + +``` +continuum-proxy --(frames)--> continuum-backend --(XADD)--> stream:camera_frames + | + XREADGROUP (this worker) + v + src/inference (ort session) + 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 +``` + +## Getting started + +```bash +cp .env.example .env +# place an ONNX failure-detection model at ./models/failure-detector.onnx +cargo run +``` diff --git a/models/.gitkeep b/models/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..3050b1d --- /dev/null +++ b/src/config.rs @@ -0,0 +1,42 @@ +use std::env; + +#[derive(Debug, Clone)] +pub struct Config { + pub redis_url: String, + 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 { + pub fn from_env() -> anyhow::Result { + Ok(Self { + redis_url: env_or("CONTINUUM_REDIS_URL", "redis://localhost:6379"), + 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 new file mode 100644 index 0000000..cdcaa8f --- /dev/null +++ b/src/inference/detector.rs @@ -0,0 +1,80 @@ +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> { + let session = Session::builder()? + .with_optimization_level(GraphOptimizationLevel::Level3)? + .with_intra_threads(intra_threads)? + .commit_from_file(model_path)?; + + 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 new file mode 100644 index 0000000..c97aae8 --- /dev/null +++ b/src/inference/mod.rs @@ -0,0 +1,4 @@ +mod detector; +mod preprocess; + +pub use detector::{Detection, FailureDetector}; diff --git a/src/inference/preprocess.rs b/src/inference/preprocess.rs new file mode 100644 index 0000000..57f974c --- /dev/null +++ b/src/inference/preprocess.rs @@ -0,0 +1,24 @@ +use image::{DynamicImage, GenericImageView}; +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 +} + +pub fn frame_dimensions(image: &DynamicImage) -> (u32, u32) { + image.dimensions() +} diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..56ec210 --- /dev/null +++ b/src/main.rs @@ -0,0 +1,149 @@ +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(); + 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?; + + 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 + .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()); + } + } + + 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?; + } + + Ok(()) +} + +fn dotenvy_load() { + if let Ok(contents) = std::fs::read_to_string(".env") { + for line in contents.lines() { + let line = line.trim(); + if line.is_empty() || line.starts_with('#') { + continue; + } + if let Some((key, value)) = line.split_once('=') { + if std::env::var(key).is_err() { + std::env::set_var(key, value); + } + } + } + } +} diff --git a/src/reporter/mod.rs b/src/reporter/mod.rs new file mode 100644 index 0000000..b90b11a --- /dev/null +++ b/src/reporter/mod.rs @@ -0,0 +1,43 @@ +mod postgres; +mod redis_pub; + +use crate::inference::Detection; +use sqlx::PgPool; +use tracing::info; + +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 }) + } + + /// 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"); + + 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?; + + Ok(()) + } +} diff --git a/src/reporter/postgres.rs b/src/reporter/postgres.rs new file mode 100644 index 0000000..4caae22 --- /dev/null +++ b/src/reporter/postgres.rs @@ -0,0 +1,31 @@ +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 new file mode 100644 index 0000000..98da27d --- /dev/null +++ b/src/reporter/redis_pub.rs @@ -0,0 +1,26 @@ +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(()) +}