use crate::caps::Capabilities;
use crate::handles::Handles;
use crate::infer::{InferBackend, InferRequest, InferResult};
use cuttlefish_abi::{error_codes, Command, Envelope, Event, JobError, JobStatus, Usage};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Instant;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use wasmtime::{Engine, Instance, Linker, Memory, Module, Store, TypedFunc};
const RENDER_WIDTH: u16 = 1024;
#[derive(Debug, Clone)]
pub enum JobEvent {
Token(String),
Progress(serde_json::Value),
}
pub type Alternates =
std::collections::HashMap<cuttlefish_core::spec::ModelRef, Arc<dyn InferBackend>>;
pub fn alternate_models_of(
spec: &cuttlefish_core::spec::Spec,
) -> Vec<cuttlefish_core::spec::ModelRef> {
use cuttlefish_core::graph::{AcceptCheck, Rung};
let mut seen: Vec<cuttlefish_core::spec::ModelRef> = Vec::new();
let mut push = |model: cuttlefish_core::spec::ModelRef| {
if model != spec.model && !seen.contains(&model) {
seen.push(model);
}
};
for (_, node) in &spec.nodes.nodes {
for rung in &node.on_fail {
if let Rung::Reroute(model) = rung {
push(model.clone());
}
}
for check in &node.accept {
if let AcceptCheck::Judge {
model: Some(model), ..
} = check
{
push(model.clone());
}
}
}
seen
}
pub struct JobSpec {
pub nodes: Vec<crate::dag::CheckedNode>,
pub exclusive_to: std::collections::HashMap<String, crate::dag::BranchExclusivity>,
pub input: serde_json::Value,
pub caps: Capabilities,
pub embedder: Option<Arc<dyn InferBackend>>,
pub warehouse: Option<WarehousePlan>,
pub alternates: Alternates,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Abi {
W32,
W64,
}
impl Abi {
fn ptr_size(self) -> usize {
match self {
Abi::W32 => 4,
Abi::W64 => 8,
}
}
}
struct Guest {
store: Store<()>,
memory: Memory,
abi: Abi,
alloc: TypedFunc<u32, u32>,
init: TypedFunc<(u32, u32), u32>,
step: TypedFunc<(u32, u32), u32>,
on_token: Option<TypedFunc<(u32, u32), i32>>,
}
impl Guest {
fn new(
engine: &Engine,
cache: &crate::module_cache::ModuleCache,
module_bytes: &[u8],
) -> anyhow::Result<Self> {
let module = cache.compile(engine, module_bytes)?;
let linker: Linker<()> = Linker::new(engine);
let mut store = Store::new(engine, ());
let instance: Instance = linker.instantiate(&mut store, &module)?;
let memory = instance
.get_memory(&mut store, "memory")
.ok_or_else(|| anyhow::anyhow!("guest exports no memory"))?;
let abi = if memory.ty(&store).is_64() {
Abi::W64
} else {
Abi::W32
};
if abi == Abi::W64 {
anyhow::bail!("guest uses 64-bit memory; only 32-bit guests are supported");
}
Ok(Self {
alloc: instance.get_typed_func(&mut store, "cf_alloc")?,
init: instance.get_typed_func(&mut store, "cf_init")?,
step: instance.get_typed_func(&mut store, "cf_step")?,
on_token: instance.get_typed_func(&mut store, "cf_on_token").ok(),
memory,
abi,
store,
})
}
fn write(&mut self, bytes: &[u8]) -> anyhow::Result<(u32, u32)> {
let len = bytes.len() as u32;
let ptr = self.alloc.call(&mut self.store, len)?;
self.memory.write(&mut self.store, ptr as usize, bytes)?;
Ok((ptr, len))
}
fn read_desc(&mut self, desc_ptr: u32) -> anyhow::Result<Vec<u8>> {
let w = self.abi.ptr_size();
let mut desc = vec![0u8; 2 * w];
self.memory
.read(&mut self.store, desc_ptr as usize, &mut desc)?;
let field = |bytes: &[u8]| -> u64 {
match w {
4 => u32::from_le_bytes(bytes.try_into().expect("4 bytes")) as u64,
_ => u64::from_le_bytes(bytes.try_into().expect("8 bytes")),
}
};
let ptr = field(&desc[..w]) as usize;
let len = field(&desc[w..]) as usize;
let mut buf = vec![0u8; len];
self.memory.read(&mut self.store, ptr, &mut buf)?;
Ok(buf)
}
fn call_init(&mut self, input: &serde_json::Value) -> anyhow::Result<Command> {
let bytes = serde_json::to_vec(input)?;
let (ptr, len) = self.write(&bytes)?;
let desc = self.init.call(&mut self.store, (ptr, len))?;
Ok(serde_json::from_slice(&self.read_desc(desc)?)?)
}
fn call_step(&mut self, event: &Event) -> anyhow::Result<Command> {
let bytes = serde_json::to_vec(event)?;
let (ptr, len) = self.write(&bytes)?;
let desc = self.step.call(&mut self.store, (ptr, len))?;
Ok(serde_json::from_slice(&self.read_desc(desc)?)?)
}
fn call_on_token(&mut self, token: &str) -> anyhow::Result<bool> {
let Some(f) = self.on_token.clone() else {
return Ok(true);
};
let (ptr, len) = self.write(token.as_bytes())?;
Ok(f.call(&mut self.store, (ptr, len))? == 0)
}
}
pub fn read_signature(
engine: &Engine,
module_bytes: &[u8],
) -> anyhow::Result<cuttlefish_abi::Signature> {
let permissive = cuttlefish_abi::Signature {
input: cuttlefish_abi::Ty::Json,
output: cuttlefish_abi::Ty::Json,
};
let module = Module::new(engine, module_bytes)?;
let linker: Linker<()> = Linker::new(engine);
let mut store = Store::new(engine, ());
let instance = linker.instantiate(&mut store, &module)?;
let Ok(signature) = instance.get_typed_func::<(), u32>(&mut store, "cf_signature") else {
return Ok(permissive);
};
let Some(memory) = instance.get_memory(&mut store, "memory") else {
return Ok(permissive);
};
let desc_ptr = signature.call(&mut store, ())? as usize;
let mut desc = [0u8; 8];
memory.read(&mut store, desc_ptr, &mut desc)?;
let ptr = u32::from_le_bytes(desc[..4].try_into().expect("4 bytes")) as usize;
let len = u32::from_le_bytes(desc[4..].try_into().expect("4 bytes")) as usize;
let mut buf = vec![0u8; len];
memory.read(&mut store, ptr, &mut buf)?;
Ok(serde_json::from_slice(&buf)?)
}
fn fail(code: &str, message: impl Into<String>, usage: Usage) -> Envelope {
Envelope {
status: JobStatus::Failed,
result: None,
error: Some(JobError {
code: code.into(),
message: message.into(),
}),
usage,
}
}
fn references_any(
expr: &cuttlefish_core::graph::InputExpr,
skipped: &std::collections::HashSet<String>,
) -> bool {
use cuttlefish_core::graph::InputExpr;
match expr {
InputExpr::FromNode(n) => skipped.contains(n),
InputExpr::Record(fields) => fields.values().any(|e| references_any(e, skipped)),
InputExpr::List(items) => items.iter().any(|e| references_any(e, skipped)),
}
}
fn evaluate_input(
expr: &cuttlefish_core::graph::InputExpr,
outputs: &std::collections::HashMap<String, serde_json::Value>,
) -> serde_json::Value {
use cuttlefish_core::graph::InputExpr;
match expr {
InputExpr::FromNode(n) => outputs.get(n).cloned().unwrap_or(serde_json::Value::Null),
InputExpr::Record(fields) => {
let mut map = serde_json::Map::new();
for (k, v) in fields {
map.insert(k.clone(), evaluate_input(v, outputs));
}
serde_json::Value::Object(map)
}
InputExpr::List(items) => {
serde_json::Value::Array(items.iter().map(|e| evaluate_input(e, outputs)).collect())
}
}
}
fn cancelled(usage: Usage, message: &str) -> Envelope {
Envelope {
status: JobStatus::Cancelled,
result: None,
error: Some(JobError {
code: error_codes::CANCELLED.into(),
message: message.into(),
}),
usage,
}
}
#[derive(Debug, Clone)]
pub struct WarehousePlan {
pub root: std::path::PathBuf,
pub spec_name: String,
pub model: String,
pub embedding_model: Option<String>,
}
#[derive(Debug, Clone)]
pub struct WarehouseNode {
name: String,
is_fanout: bool,
item_output: Option<cuttlefish_abi::Ty>,
output_ty: cuttlefish_abi::Ty,
}
impl WarehouseNode {
fn of(node: &crate::dag::CheckedNode) -> Self {
Self {
name: node.name.clone(),
is_fanout: node.over.is_some(),
item_output: node.item_output.clone(),
output_ty: node.signature.output.clone(),
}
}
}
fn write_warehouse(
plan: &WarehousePlan,
nodes: &[WarehouseNode],
ledger: &crate::ledger::Ledger,
envelope: &Envelope,
) {
use crate::warehouse as wh;
let job_id = ledger
.job_dir()
.file_name()
.map(|n| n.to_string_lossy().into_owned())
.unwrap_or_else(|| "unknown".into());
let lineage = wh::Lineage {
job_id: job_id.clone(),
spec_name: plan.spec_name.clone(),
spec_fingerprint: ledger.graph_fingerprint().unwrap_or_default(),
model: plan.model.clone(),
embedding_model: plan.embedding_model.clone(),
cuttlefish_version: env!("CARGO_PKG_VERSION").to_string(),
};
let (mut bronze, mut silver, mut gold) = (
std::collections::BTreeMap::new(),
std::collections::BTreeMap::new(),
std::collections::BTreeMap::new(),
);
for node in nodes {
if !node.is_fanout {
continue;
}
let rows: Vec<wh::Row> = match ledger.concluded_rows(&node.name) {
Ok(rows) => rows
.into_iter()
.map(|r| wh::Row {
node: node.name.clone(),
item: r.item,
status: r.status,
concluded_at: r.concluded_at,
source_input: r.source_input,
output: r.output,
error: r.error,
})
.collect(),
Err(e) => {
eprintln!("warning: warehouse: reading node `{}`: {e}", node.name);
continue;
}
};
if rows.is_empty() {
continue;
}
match wh::bronze_batch(&rows, &lineage) {
Ok(batch) => {
let path = plan
.root
.join("bronze")
.join(format!("{}.parquet", node.name));
match wh::write_parquet(&path, &batch) {
Ok(()) => {
bronze.insert(
node.name.clone(),
wh::Layer::Written(wh::entry_for(&plan.root, &path, &batch)),
);
}
Err(e) => eprintln!("warning: warehouse: writing bronze: {e}"),
}
}
Err(e) => eprintln!("warning: warehouse: building bronze: {e}"),
}
let item_output = node.item_output.clone().unwrap_or(cuttlefish_abi::Ty::Json);
match wh::silver_batch(&rows, &lineage, &item_output) {
Ok(Some(batch)) => {
let path = plan
.root
.join("silver")
.join(format!("{}.parquet", node.name));
match wh::write_parquet(&path, &batch) {
Ok(()) => {
silver.insert(
node.name.clone(),
wh::Layer::Written(wh::entry_for(&plan.root, &path, &batch)),
);
}
Err(e) => eprintln!("warning: warehouse: writing silver: {e}"),
}
}
Ok(None) => {
silver.insert(
node.name.clone(),
wh::Layer::Skipped {
skipped: format!(
"node `{}` declares its per-item output as `{}`, which names no \
fields to validate against — declare a record of `text`, `number` \
or `bool` fields to get a typed silver table",
node.name,
item_output.describe()
),
},
);
}
Err(e) => eprintln!("warning: warehouse: building silver: {e}"),
}
}
if let (Some(value), Some(last)) = (envelope.result.as_ref(), nodes.last()) {
let row = wh::Row {
node: last.name.clone(),
item: 0,
status: "completed".into(),
concluded_at: wh::now_rfc3339(),
source_input: None,
output: Some(value.clone()),
error: None,
};
let ty = last.output_ty.clone();
match wh::silver_batch(std::slice::from_ref(&row), &lineage, &ty) {
Ok(Some(batch)) => {
let path = plan
.root
.join("gold")
.join(format!("{}.parquet", last.name));
match wh::write_parquet(&path, &batch) {
Ok(()) => {
gold.insert(
last.name.clone(),
wh::Layer::Written(wh::entry_for(&plan.root, &path, &batch)),
);
}
Err(e) => eprintln!("warning: warehouse: writing gold: {e}"),
}
}
Ok(None) => {
gold.insert(
last.name.clone(),
wh::Layer::Skipped {
skipped: format!(
"node `{}` declares its output as `{}`, which names no fields to \
validate against",
last.name,
ty.describe()
),
},
);
}
Err(e) => eprintln!("warning: warehouse: building gold: {e}"),
}
}
let manifest = wh::Manifest {
job_id,
spec_name: plan.spec_name.clone(),
spec_fingerprint: lineage.spec_fingerprint.clone(),
model: plan.model.clone(),
embedding_model: plan.embedding_model.clone(),
cuttlefish_version: lineage.cuttlefish_version.clone(),
written_at: wh::now_rfc3339(),
bronze,
silver,
gold,
};
if let Err(e) = wh::write_manifest(&plan.root, &manifest) {
eprintln!("warning: warehouse: writing the manifest: {e}");
}
}
pub async fn run_job(
engine: Arc<Engine>,
backend: Arc<dyn InferBackend>,
job: JobSpec,
events: mpsc::Sender<JobEvent>,
cancel: CancellationToken,
ledger: &crate::ledger::Ledger,
cache: &crate::module_cache::ModuleCache,
) -> Envelope {
let started = Instant::now();
let mut usage = Usage {
model: backend.model_name(),
..Usage::default()
};
let mut handles = Handles::default();
let job_warehouse = job.warehouse.clone();
let warehouse_nodes: Vec<WarehouseNode> = match &job_warehouse {
Some(_) => job.nodes.iter().map(WarehouseNode::of).collect(),
None => Vec::new(),
};
let envelope = 'run: {
if job.nodes.is_empty() {
usage.duration_ms = started.elapsed().as_millis() as u64;
break 'run fail(
error_codes::SCHEMA_VALIDATION_FAILED,
"this job has no nodes to run",
usage,
);
}
let checks: Vec<crate::accept::CompiledChecks> = match job
.nodes
.iter()
.map(|n| {
crate::accept::CompiledChecks::compile(&n.accept)
.map_err(|e| format!("node `{}`: {e}", n.name))
})
.collect()
{
Ok(c) => c,
Err(message) => {
usage.duration_ms = started.elapsed().as_millis() as u64;
break 'run fail(error_codes::SCHEMA_VALIDATION_FAILED, message, usage);
}
};
let total = job.nodes.len();
let mut outputs: std::collections::HashMap<String, serde_json::Value> =
std::collections::HashMap::new();
let mut skipped: std::collections::HashSet<String> = std::collections::HashSet::new();
let mut route_taken: std::collections::HashMap<String, String> =
std::collections::HashMap::new();
for (index, node) in job.nodes.iter().enumerate() {
match ledger.is_skipped(&node.name) {
Ok(true) => {
skipped.insert(node.name.clone());
continue;
}
Ok(false) => {}
Err(e) => {
usage.duration_ms = started.elapsed().as_millis() as u64;
break 'run fail(
error_codes::SCHEMA_VALIDATION_FAILED,
format!("reading ledger skip state for node `{}`: {e}", node.name),
usage,
);
}
}
match ledger.get_completed(&node.name) {
Ok(Some(cached)) => {
if let Some(route) = cached.get("route").and_then(|v| v.as_str()) {
route_taken.insert(node.name.clone(), route.to_string());
}
outputs.insert(node.name.clone(), cached);
continue;
}
Ok(None) => {}
Err(e) => {
usage.duration_ms = started.elapsed().as_millis() as u64;
break 'run fail(
error_codes::SCHEMA_VALIDATION_FAILED,
format!("reading ledger checkpoint for node `{}`: {e}", node.name),
usage,
);
}
}
if total > 1 {
let _ = events
.send(JobEvent::Progress(serde_json::json!({
"stage": index + 1,
"of": total,
"node": node.name,
})))
.await;
}
if let Some(ex) = job.exclusive_to.get(&node.name) {
if let Some(taken) = route_taken.get(&ex.decision) {
if taken != &ex.label {
skipped.insert(node.name.clone());
if let Err(e) = ledger.write_skipped(&node.name) {
usage.duration_ms = started.elapsed().as_millis() as u64;
break 'run fail(
error_codes::SCHEMA_VALIDATION_FAILED,
format!("recording skip for node `{}` in ledger: {e}", node.name),
usage,
);
}
continue;
}
}
}
if let Some(expr) = &node.input {
if references_any(expr, &skipped) {
skipped.insert(node.name.clone());
if let Err(e) = ledger.write_skipped(&node.name) {
usage.duration_ms = started.elapsed().as_millis() as u64;
break 'run fail(
error_codes::SCHEMA_VALIDATION_FAILED,
format!("recording skip for node `{}` in ledger: {e}", node.name),
usage,
);
}
continue;
}
}
let node_input = match &node.input {
None => job.input.clone(),
Some(expr) => evaluate_input(expr, &outputs),
};
if node.over.is_some() {
let collected = match run_fanout_node(
&engine,
cache,
&backend,
&job.alternates,
&checks[index],
node,
&job.caps,
job.embedder.as_ref(),
&mut handles,
&events,
&cancel,
&mut usage,
started,
index,
total,
ledger,
)
.await
{
Ok(v) => v,
Err(envelope) => break 'run envelope,
};
if let Err(e) = ledger.write_completed(&node.name, &collected) {
usage.duration_ms = started.elapsed().as_millis() as u64;
break 'run fail(
error_codes::SCHEMA_VALIDATION_FAILED,
format!(
"recording checkpoint for node `{}` in ledger: {e}",
node.name
),
usage,
);
}
outputs.insert(node.name.clone(), collected);
continue;
}
let job_dir = ledger.job_dir();
let ladder = Ladder {
engine: &engine,
embedder: job.embedder.as_ref(),
job_dir,
cache,
default_backend: &backend,
alternates: &job.alternates,
checks: &checks[index],
expected: &node.signature.output,
on_fail: &node.on_fail,
repeat_until: node.repeat_until.as_deref(),
max_iterations: node.max_iterations,
caps: &job.caps,
events: &events,
cancel: &cancel,
started,
index,
};
let result = match ladder
.run(
&node.module_bytes,
node.script.as_deref(),
node_input.clone(),
&mut handles,
&mut usage,
)
.await
{
Ok(value) => value,
Err(LadderError::Fatal(envelope)) => break 'run *envelope,
Err(LadderError::Exhausted {
reason,
escalated,
envelope,
}) => {
if escalated {
if let Err(e) =
ledger.write_escalated(&node.name, None, &reason, Some(&node_input))
{
eprintln!("recording escalation for node `{}`: {e}", node.name);
}
}
usage.duration_ms = started.elapsed().as_millis() as u64;
break 'run match envelope {
Some(e) => *e,
None => fail(
error_codes::SCHEMA_VALIDATION_FAILED,
format!("node `{}`: {reason}", node.name),
usage,
),
};
}
};
let valid_labels: Vec<&str> = job
.exclusive_to
.values()
.filter(|ex| ex.decision == node.name)
.map(|ex| ex.label.as_str())
.collect();
if !valid_labels.is_empty() {
match result.get("route").and_then(|v| v.as_str()) {
Some(route) if valid_labels.contains(&route) => {
route_taken.insert(node.name.clone(), route.to_string());
}
Some(route) => {
usage.duration_ms = started.elapsed().as_millis() as u64;
break 'run fail(
error_codes::SCHEMA_VALIDATION_FAILED,
format!(
"node `{}` produced route \"{route}\", which doesn't match any \
declared label for this branches decision",
node.name
),
usage,
);
}
None => {
usage.duration_ms = started.elapsed().as_millis() as u64;
break 'run fail(
error_codes::SCHEMA_VALIDATION_FAILED,
format!(
"node `{}` is a branches decision but its output has no `route` field",
node.name
),
usage,
);
}
}
} else if let Some(route) = result.get("route").and_then(|v| v.as_str()) {
route_taken.insert(node.name.clone(), route.to_string());
}
if let Err(e) = ledger.write_completed(&node.name, &result) {
usage.duration_ms = started.elapsed().as_millis() as u64;
break 'run fail(
error_codes::SCHEMA_VALIDATION_FAILED,
format!(
"recording checkpoint for node `{}` in ledger: {e}",
node.name
),
usage,
);
}
outputs.insert(node.name.clone(), result);
}
usage.duration_ms = started.elapsed().as_millis() as u64;
let result = job
.nodes
.iter()
.rev()
.find_map(|n| outputs.get(&n.name).cloned());
match result {
Some(value) => Envelope {
status: JobStatus::Completed,
result: Some(value),
error: None,
usage,
},
None => fail(
error_codes::SCHEMA_VALIDATION_FAILED,
"every node in this job was skipped; there is no result",
usage,
),
}
};
let ledger_status = match envelope.status {
JobStatus::Completed => "completed",
JobStatus::Failed => "failed",
JobStatus::Cancelled => "cancelled",
_ => "running", };
if let Err(e) = ledger.finish(ledger_status) {
eprintln!("warning: failed to record job {ledger_status} status in ledger: {e}");
}
if let Some(plan) = &job_warehouse {
write_warehouse(plan, &warehouse_nodes, ledger, &envelope);
}
envelope
}
#[allow(clippy::too_many_arguments)]
async fn run_fanout_node(
engine: &Engine,
cache: &crate::module_cache::ModuleCache,
backend: &Arc<dyn InferBackend>,
alternates: &Alternates,
checks: &crate::accept::CompiledChecks,
node: &crate::dag::CheckedNode,
caps: &Capabilities,
embedder: Option<&Arc<dyn InferBackend>>,
handles: &mut Handles,
events: &mpsc::Sender<JobEvent>,
cancel: &CancellationToken,
usage: &mut Usage,
started: Instant,
index: usize,
total: usize,
ledger: &crate::ledger::Ledger,
) -> Result<serde_json::Value, Envelope> {
use sha2::{Digest, Sha256};
let manifest_path = node
.over
.as_ref()
.expect("run_fanout_node is only called for a node with `over`");
let node_name = node.name.as_str();
let item_output = node
.item_output
.as_ref()
.unwrap_or(&node.signature.output)
.clone();
let bail = |message: String, usage: &mut Usage| -> Envelope {
usage.duration_ms = started.elapsed().as_millis() as u64;
fail(
error_codes::SCHEMA_VALIDATION_FAILED,
message,
usage.clone(),
)
};
let bytes = std::fs::read(manifest_path).map_err(|e| {
bail(
format!(
"node `{node_name}`: reading fan-out manifest {}: {e}",
manifest_path.display()
),
usage,
)
})?;
let text = String::from_utf8(bytes.clone()).map_err(|e| {
bail(
format!(
"node `{node_name}`: fan-out manifest {} is not valid UTF-8: {e}",
manifest_path.display()
),
usage,
)
})?;
let mut items: Vec<serde_json::Value> = Vec::new();
for (i, line) in text.lines().enumerate() {
if line.trim().is_empty() {
continue;
}
match serde_json::from_str(line) {
Ok(v) => items.push(v),
Err(e) => {
return Err(bail(
format!(
"node `{node_name}`: fan-out manifest {} line {} is not valid JSON: {e}",
manifest_path.display(),
i + 1
),
usage,
))
}
}
}
if items.is_empty() {
return Err(bail(
format!(
"node `{node_name}`: fan-out manifest {} is empty — zero items almost always \
means the step that produced it failed, and reducing over nothing would \
silently look like success",
manifest_path.display()
),
usage,
));
}
let digest = crate::hex::encode(Sha256::digest(&bytes));
match ledger.check_or_record_manifest(node_name, &digest, items.len()) {
Ok(Ok(())) => {}
Ok(Err(previous)) => {
return Err(bail(
format!(
"node `{node_name}`: fan-out manifest {} has changed since this job first \
ran (was {previous}, now {digest}) — recorded item indices no longer refer \
to the same inputs, so resuming would pair results with the wrong items; \
re-submit the job instead",
manifest_path.display()
),
usage,
))
}
Err(e) => {
return Err(bail(
format!("node `{node_name}`: recording fan-out manifest digest: {e}"),
usage,
))
}
}
let mut succeeded = 0usize;
let mut failed = 0usize;
for (item_index, item_input) in items.iter().enumerate() {
if cancel.is_cancelled() {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(Envelope {
status: JobStatus::Cancelled,
result: None,
error: None,
usage: usage.clone(),
});
}
match ledger.item_concluded(node_name, item_index) {
Ok(true) => {
if ledger
.get_item_completed(node_name, item_index)
.unwrap_or(None)
.is_some()
{
succeeded += 1;
} else {
failed += 1;
}
continue;
}
Ok(false) => {}
Err(e) => {
return Err(bail(
format!("node `{node_name}`: reading item {item_index} from ledger: {e}"),
usage,
))
}
}
let _ = events
.send(JobEvent::Progress(serde_json::json!({
"stage": index + 1,
"of": total,
"node": node_name,
"item": item_index,
"items": items.len(),
"succeeded": succeeded,
"failed": failed,
})))
.await;
if !node.signature.input.matches_value(item_input) {
failed += 1;
let message = format!(
"item {item_index} does not match the block's declared input `{}`",
node.signature.input
);
if let Err(e) =
ledger.write_item_failed(node_name, item_index, &message, Some(item_input))
{
return Err(bail(
format!("node `{node_name}`: recording item {item_index} failure: {e}"),
usage,
));
}
continue;
}
let job_dir = ledger.job_dir();
let ladder = Ladder {
engine,
embedder,
job_dir,
cache,
default_backend: backend,
alternates,
checks,
expected: &item_output,
on_fail: &node.on_fail,
repeat_until: None,
max_iterations: None,
caps,
events,
cancel,
started,
index,
};
match ladder
.run(
&node.module_bytes,
node.script.as_deref(),
item_input.clone(),
handles,
usage,
)
.await
{
Ok(value) => {
succeeded += 1;
if let Err(e) =
ledger.write_item_completed(node_name, item_index, &value, Some(item_input))
{
return Err(bail(
format!("node `{node_name}`: recording item {item_index} result: {e}"),
usage,
));
}
}
Err(LadderError::Fatal(envelope)) => return Err(*envelope),
Err(LadderError::Exhausted {
reason, escalated, ..
}) => {
failed += 1;
let message = format!("item {item_index} {reason}");
let recorded = if escalated {
ledger.write_escalated(node_name, Some(item_index), &message, Some(item_input))
} else {
ledger.write_item_failed(node_name, item_index, &message, Some(item_input))
};
if let Err(e) = recorded {
return Err(bail(
format!("node `{node_name}`: recording item {item_index} failure: {e}"),
usage,
));
}
}
}
}
if succeeded == 0 {
let first_error = ledger
.concluded_items(node_name)
.ok()
.and_then(|items| items.into_iter().find_map(|(_, _, err)| err))
.unwrap_or_else(|| "no error recorded".to_string());
return Err(bail(
format!(
"node `{node_name}`: all {failed} fan-out item(s) failed — there is nothing for \
a downstream node to reduce over. First failure: {first_error}"
),
usage,
));
}
let results_dir = ledger.job_dir().join("results");
std::fs::create_dir_all(&results_dir).map_err(|e| {
bail(
format!(
"node `{node_name}`: creating {}: {e}",
results_dir.display()
),
usage,
)
})?;
let results_path = results_dir.join(format!("{node_name}.results.jsonl"));
let failures_path = results_dir.join(format!("{node_name}.failures.jsonl"));
let concluded = ledger.concluded_items(node_name).map_err(|e| {
bail(
format!("node `{node_name}`: reading concluded items from ledger: {e}"),
usage,
)
})?;
let (mut results_out, mut failures_out) = (String::new(), String::new());
for (item_index, output, error) in concluded {
match (output, error) {
(Some(value), _) => results_out.push_str(&format!(
"{}\n",
serde_json::json!({"item": item_index, "result": value})
)),
(None, Some(message)) => failures_out.push_str(&format!(
"{}\n",
serde_json::json!({"item": item_index, "error": message})
)),
(None, None) => {}
}
}
for (path, contents) in [
(&results_path, &results_out),
(&failures_path, &failures_out),
] {
std::fs::write(path, contents).map_err(|e| {
bail(
format!("node `{node_name}`: writing {}: {e}", path.display()),
usage,
)
})?;
}
Ok(serde_json::json!({
"results_path": results_path.to_string_lossy(),
"failures_path": failures_path.to_string_lossy(),
"succeeded": succeeded,
"failed": failed,
}))
}
enum LadderError {
Exhausted {
reason: String,
escalated: bool,
envelope: Option<Box<Envelope>>,
},
Fatal(Box<Envelope>),
}
struct Ladder<'a> {
engine: &'a Engine,
embedder: Option<&'a Arc<dyn InferBackend>>,
job_dir: &'a std::path::Path,
cache: &'a crate::module_cache::ModuleCache,
default_backend: &'a Arc<dyn InferBackend>,
alternates: &'a Alternates,
checks: &'a crate::accept::CompiledChecks,
expected: &'a cuttlefish_abi::Ty,
on_fail: &'a [cuttlefish_core::graph::Rung],
repeat_until: Option<&'a str>,
max_iterations: Option<u32>,
caps: &'a Capabilities,
events: &'a mpsc::Sender<JobEvent>,
cancel: &'a CancellationToken,
started: Instant,
index: usize,
}
impl Ladder<'_> {
#[allow(clippy::too_many_arguments)]
async fn run(
&self,
module_bytes: &[u8],
script: Option<&str>,
input: serde_json::Value,
handles: &mut Handles,
usage: &mut Usage,
) -> Result<serde_json::Value, LadderError> {
use cuttlefish_core::graph::Rung;
let mut backend = self.default_backend.clone();
let mut rung = 0usize;
let mut retries_left = 0u32;
'ladder: loop {
let (reason, envelope) = match self
.attempt(
module_bytes,
script,
input.clone(),
&backend,
handles,
usage,
)
.await
{
Ok(value) => return Ok(value),
Err(LadderError::Fatal(e)) => return Err(LadderError::Fatal(e)),
Err(LadderError::Exhausted {
reason, envelope, ..
}) => (reason, envelope),
};
loop {
if retries_left > 0 {
retries_left -= 1;
continue 'ladder;
}
match self.on_fail.get(rung) {
None => {
return Err(LadderError::Exhausted {
reason,
escalated: false,
envelope,
})
}
Some(Rung::Retry(n)) => {
rung += 1;
retries_left = *n;
}
Some(Rung::Reroute(model)) => {
rung += 1;
match self.alternates.get(model) {
Some(b) => {
backend = b.clone();
continue 'ladder;
}
None => {
return Err(LadderError::Exhausted {
reason: format!(
"reroute names model `{model}`, which was not resolved \
at startup (after: {reason})"
),
escalated: false,
envelope: None,
});
}
}
}
Some(Rung::Escalate) => {
return Err(LadderError::Exhausted {
reason,
escalated: true,
envelope,
})
}
}
}
}
}
#[allow(clippy::too_many_arguments)]
async fn attempt(
&self,
module_bytes: &[u8],
script: Option<&str>,
input: serde_json::Value,
backend: &Arc<dyn InferBackend>,
handles: &mut Handles,
usage: &mut Usage,
) -> Result<serde_json::Value, LadderError> {
let rejected = |reason: String| LadderError::Exhausted {
reason,
escalated: false,
envelope: None,
};
let mut current = input.clone();
let mut iterations: u32 = 0;
let value = loop {
let produced = match run_stage(
self.engine,
self.cache,
backend,
module_bytes,
current.clone(),
script,
self.caps,
self.embedder,
self.job_dir,
handles,
self.events,
self.cancel,
usage,
self.started,
self.index,
)
.await
{
Ok(v) => v,
Err(envelope) if envelope.status == JobStatus::Cancelled => {
return Err(LadderError::Fatal(Box::new(envelope)))
}
Err(envelope) => {
let reason = envelope
.error
.as_ref()
.map(|e| format!("{}: {}", e.code, e.message))
.unwrap_or_else(|| "failed with no error detail".to_string());
return Err(LadderError::Exhausted {
reason,
escalated: false,
envelope: Some(Box::new(envelope)),
});
}
};
let Some(field) = self.repeat_until else {
break produced;
};
iterations += 1;
if produced.get(field).and_then(|v| v.as_str()) == Some("done") {
break produced;
}
let max = self
.max_iterations
.expect("repeat_until requires max_iterations, enforced at parse time");
if iterations >= max {
return Err(rejected(format!(
"did not reach repeat_until=\"done\" within max_iterations={max}"
)));
}
current = produced;
};
if !self.expected.matches_value(&value) {
return Err(rejected(format!(
"produced {value}, which doesn't match the declared output `{}`",
self.expected
)));
}
if let Err(why) = self.checks.check_schemas(&value) {
return Err(rejected(why));
}
match self
.checks
.run_judges(&input, &value, self.default_backend, self.alternates)
.await
{
crate::accept::JudgeVerdict::Accepted => Ok(value),
crate::accept::JudgeVerdict::Rejected(why) => Err(rejected(format!("judge: {why}"))),
crate::accept::JudgeVerdict::Unusable(why) => {
Err(rejected(format!("judge gave no usable verdict: {why}")))
}
}
}
}
#[allow(clippy::too_many_arguments)]
async fn run_stage(
engine: &Engine,
cache: &crate::module_cache::ModuleCache,
backend: &Arc<dyn InferBackend>,
module_bytes: &[u8],
input: serde_json::Value,
script: Option<&str>,
caps: &Capabilities,
embedder: Option<&Arc<dyn InferBackend>>,
job_dir: &std::path::Path,
handles: &mut Handles,
events: &mpsc::Sender<JobEvent>,
cancel: &CancellationToken,
usage: &mut Usage,
started: Instant,
stage_index: usize,
) -> Result<serde_json::Value, Envelope> {
let blame = |message: String| -> String {
if stage_index == 0 {
message
} else {
format!("block {} of the pipeline: {message}", stage_index + 1)
}
};
let mut doc_paths: std::collections::HashMap<u32, std::path::PathBuf> =
std::collections::HashMap::new();
let mut doc_texts: std::collections::HashMap<u32, std::sync::Arc<String>> =
std::collections::HashMap::new();
let mut doc_pages: std::collections::HashMap<u32, u32> = std::collections::HashMap::new();
let mut guest = match Guest::new(engine, cache, module_bytes) {
Ok(g) => g,
Err(e) => {
return Err(fail(
error_codes::WASM_TRAP,
blame(e.to_string()),
usage.clone(),
))
}
};
let input = match script {
Some(script) => serde_json::json!({
"__cuttlefish_script": script,
"input": input,
}),
None => input,
};
let mut command = match guest.call_init(&input) {
Ok(c) => c,
Err(e) => {
return Err(fail(
error_codes::WASM_TRAP,
blame(e.to_string()),
usage.clone(),
))
}
};
loop {
if cancel.is_cancelled() {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(cancelled(usage.clone(), "job cancelled"));
}
let event = match command {
Command::Done { result } => return Ok(result),
Command::Fail { code, message } => {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(&code, blame(message), usage.clone()));
}
Command::Emit { progress } => {
let _ = events.send(JobEvent::Progress(progress)).await;
Event::Emitted
}
Command::Embed { texts } => {
let Some(backend) = embedder else {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(
error_codes::UNSUPPORTED,
"this spec declares no `embedding_model`, so `embed` has nothing to \
call. Add one, e.g. `embedding_model = Ollama \"nomic-embed-text\";` \
— it is deliberately separate from `model`, since a chat model \
cannot produce embeddings."
.to_string(),
usage.clone(),
));
};
match backend.embed(&texts).await {
Ok(vectors) => Event::Embedded { vectors },
Err(e) => {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(error_codes::UNSUPPORTED, e.to_string(), usage.clone()));
}
}
}
Command::Fetch { url } => {
if !caps.allows_fetch(&url) {
usage.duration_ms = started.elapsed().as_millis() as u64;
let granted = if caps.fetch_prefixes().is_empty() {
"this spec grants no `Fetch` capability at all".to_string()
} else {
format!("granted prefixes: {}", caps.fetch_prefixes().join(", "))
};
return Err(fail(
error_codes::CAPABILITY_DENIED,
format!(
"fetch not permitted: {url}\nAdd `Fetch \"<prefix>\"` to the \
spec's `capabilities`. {granted}"
),
usage.clone(),
));
}
match crate::fetch::fetch_to_file(&url, job_dir).await {
Ok(path) => match handles.open(&path) {
Ok((handle, len, kind)) => Event::Opened { handle, len, kind },
Err(e) => {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(
error_codes::UNSUPPORTED,
format!("opening the fetched copy of {url}: {e}"),
usage.clone(),
));
}
},
Err(e) => {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(error_codes::UNSUPPORTED, e.to_string(), usage.clone()));
}
}
}
Command::Open { path } => {
let p = std::path::PathBuf::from(&path);
if let Some(denial) = caps.read_denial(&p) {
usage.duration_ms = started.elapsed().as_millis() as u64;
let (code, message) = match denial {
crate::caps::ReadDenial::Missing => (
error_codes::NOT_FOUND,
format!(
"no such file: {path} — it is inside a granted root, so this is \
the path being wrong rather than the grant"
),
),
crate::caps::ReadDenial::NotGranted => (
error_codes::CAPABILITY_DENIED,
format!("read not permitted: {path}"),
),
};
return Err(fail(code, message, usage.clone()));
}
match handles.open(&p) {
Ok((handle, len, kind)) => {
let kind = match kind {
cuttlefish_abi::MediaKind::Document { .. } => {
match crate::documents::inspect(&p) {
Ok(info) => cuttlefish_abi::MediaKind::Document {
pages: info.pages,
has_text_layer: info.has_text_layer,
},
Err(_) => cuttlefish_abi::MediaKind::Binary,
}
}
other => other,
};
doc_paths.insert(handle, p.clone());
Event::Opened { handle, len, kind }
}
Err(e) => {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(
error_codes::CAPABILITY_DENIED,
e.to_string(),
usage.clone(),
));
}
}
}
Command::Slice {
handle,
offset,
len,
} => match handles.slice(handle, offset, len) {
Ok(w) => Event::Sliced {
text: w.text,
next_offset: w.next_offset,
},
Err(e) => {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(
error_codes::CAPABILITY_DENIED,
e.to_string(),
usage.clone(),
));
}
},
Command::SliceBytes {
handle,
offset,
len,
} => match handles.slice_bytes(handle, offset, len) {
Ok((bytes, next_offset)) => {
use base64::Engine;
Event::SlicedBytes {
bytes_base64: base64::engine::general_purpose::STANDARD.encode(&bytes),
next_offset,
}
}
Err(e) => {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(
error_codes::CAPABILITY_DENIED,
e.to_string(),
usage.clone(),
));
}
},
Command::PageText { handle, page } => {
let Some(path) = doc_paths.get(&handle).cloned() else {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(
error_codes::CAPABILITY_DENIED,
format!("no such handle: {handle}"),
usage.clone(),
));
};
let text = match doc_text(&mut doc_texts, handle, &path) {
Ok(t) => t,
Err(e) => {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(error_codes::UNSUPPORTED, e.to_string(), usage.clone()));
}
};
let page_tree_count = doc_page_count(&mut doc_pages, handle, &path);
match crate::documents::page_text_from(&text, page, page_tree_count) {
Ok(text) => Event::PageTexted { text },
Err(e) => {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(error_codes::UNSUPPORTED, e.to_string(), usage.clone()));
}
}
}
Command::PageTextOpt { handle, page } => {
let Some(path) = doc_paths.get(&handle).cloned() else {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(
error_codes::CAPABILITY_DENIED,
format!("no such handle: {handle}"),
usage.clone(),
));
};
match doc_text(&mut doc_texts, handle, &path) {
Err(e) => Event::PageTextAttempted {
text: None,
error: Some(e.to_string()),
},
Ok(text) => {
let page_tree_count = doc_page_count(&mut doc_pages, handle, &path);
match crate::documents::page_text_from(&text, page, page_tree_count) {
Ok(text) => Event::PageTextAttempted {
text: Some(text),
error: None,
},
Err(e) => Event::PageTextAttempted {
text: None,
error: Some(e.to_string()),
},
}
}
}
}
Command::DocumentText { handle } => {
let Some(path) = doc_paths.get(&handle).cloned() else {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(
error_codes::CAPABILITY_DENIED,
format!("no such handle: {handle}"),
usage.clone(),
));
};
match doc_text(&mut doc_texts, handle, &path) {
Ok(text) => Event::PageTexted {
text: text.as_ref().clone(),
},
Err(e) => {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(error_codes::UNSUPPORTED, e.to_string(), usage.clone()));
}
}
}
Command::PageImage { handle, page } => {
let Some(path) = doc_paths.get(&handle).cloned() else {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(
error_codes::CAPABILITY_DENIED,
format!("no such handle: {handle}"),
usage.clone(),
));
};
match crate::documents::render_page(&path, page, RENDER_WIDTH) {
Ok(png) => {
let (handle, len) = handles.insert_bytes(
png,
cuttlefish_abi::MediaKind::Image {
format: "png".into(),
},
);
Event::PageImaged { handle, len }
}
Err(e) => {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(error_codes::UNSUPPORTED, e.to_string(), usage.clone()));
}
}
}
Command::ImageOp { handle, op } => {
let bytes = match handles.read_all(handle) {
Ok(b) => b,
Err(e) => {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(
error_codes::CAPABILITY_DENIED,
e.to_string(),
usage.clone(),
));
}
};
match crate::images::apply(&bytes, &op) {
Ok(png) => {
let (handle, len) = handles.insert_bytes(
png,
cuttlefish_abi::MediaKind::Image {
format: "png".into(),
},
);
Event::PageImaged { handle, len }
}
Err(e) => {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(error_codes::UNSUPPORTED, e.to_string(), usage.clone()));
}
}
}
Command::Infer {
prompt,
max_tokens,
images,
} => {
if !images.is_empty() && !backend.supports_images() {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(
error_codes::UNSUPPORTED,
format!(
"this job supplied {} image(s), but the backend serving `{}` cannot \
accept them. Use a vision-capable model through the `ollama` \
provider, or change the block to send text only.",
images.len(),
backend.model_name()
),
usage.clone(),
));
}
let mut image_bytes = Vec::with_capacity(images.len());
for handle in &images {
match handles.read_all(*handle) {
Ok(bytes) => image_bytes.push(bytes),
Err(e) => {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(
error_codes::CAPABILITY_DENIED,
e.to_string(),
usage.clone(),
));
}
}
}
let (tx, mut rx) = mpsc::unbounded_channel::<String>();
let stop = Arc::new(AtomicBool::new(false));
let sink_stop = stop.clone();
let mut sink = move |t: &str| {
tx.send(t.to_string()).is_ok() && !sink_stop.load(Ordering::Relaxed)
};
let mut trap: Option<String> = None;
let outcome: Option<anyhow::Result<InferResult>> = {
let request = InferRequest {
prompt: &prompt,
max_tokens,
images: &image_bytes,
};
let infer = backend.infer(request, &mut sink);
tokio::pin!(infer);
loop {
tokio::select! {
biased;
_ = cancel.cancelled() => break None,
Some(tok) = rx.recv() => {
let _ = events.send(JobEvent::Token(tok.clone())).await;
match guest.call_on_token(&tok) {
Ok(true) => {}
Ok(false) => stop.store(true, Ordering::Relaxed),
Err(e) => {
trap = Some(e.to_string());
break None;
}
}
}
r = &mut infer => break Some(r),
}
}
};
if let Some(message) = trap {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(error_codes::WASM_TRAP, message, usage.clone()));
}
while let Ok(tok) = rx.try_recv() {
let _ = events.send(JobEvent::Token(tok)).await;
}
match outcome {
None => {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(cancelled(usage.clone(), "cancelled during inference"));
}
Some(Err(e)) => {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(
error_codes::MODEL_LOAD_FAILED,
e.to_string(),
usage.clone(),
));
}
Some(Ok(r)) => {
usage.tokens_in += r.tokens_in;
usage.tokens_out += r.tokens_out;
Event::InferDone {
text: r.text,
tokens_out: r.tokens_out,
}
}
}
}
};
command = match guest.call_step(&event) {
Ok(c) => c,
Err(e) => {
usage.duration_ms = started.elapsed().as_millis() as u64;
return Err(fail(
error_codes::WASM_TRAP,
blame(e.to_string()),
usage.clone(),
));
}
};
}
}
fn doc_text(
cache: &mut std::collections::HashMap<u32, std::sync::Arc<String>>,
handle: u32,
path: &std::path::Path,
) -> anyhow::Result<std::sync::Arc<String>> {
if let Some(hit) = cache.get(&handle) {
return Ok(hit.clone());
}
let text = std::sync::Arc::new(crate::documents::document_text(path)?);
cache.insert(handle, text.clone());
Ok(text)
}
fn doc_page_count(
cache: &mut std::collections::HashMap<u32, u32>,
handle: u32,
path: &std::path::Path,
) -> u32 {
if let Some(hit) = cache.get(&handle) {
return *hit;
}
let count = crate::documents::page_count(path).unwrap_or(0);
cache.insert(handle, count);
count
}