bb_compiler/
derive_wire_deadlines.rs

1//!
2//! Stamps each `wire.Send`'s static deadline from the (per_hop_budget_ns ×
3//! chain_depth) formula, defaulting to chain_depth = 1 when no
4//! `chain_depth` on every target-boundary `wire.Send`. For each
5//! such Send, computes the static deadline as
6//! `chain_depth * per_hop_budget_ns` and stamps it as a
7//! `deadline_ns: i64` attribute on the NodeProto. The existing
8//! [`mod@super::insert_async_deadlines`] pass then inserts a
9//! `DeadlineCheck` gate upstream so the deadline is enforced at
10//! runtime.
11//!
12//! The per-hop budget comes from the [`crate::Compiler`]'s
13//! `per_hop_budget_ns` field (default
14//! [`bb_ir::syscall_ids::DEFAULT_PER_HOP_BUDGET_NS`]). Build-time
15//! CI matrices that produce deployment bundles with different
16//! latency assumptions can override via
17//! `Compiler::with_per_hop_budget_ns(ns)`.
18//!
19//! Runtime override note: the runtime side
20//! ([`crate::node::config::NodeConfig::per_hop_budget_ns`]) is the
21//! source of truth at delivery time. When `chain_depth` metadata
22//! survives on the NodeProto, the engine can replace the static
23//! deadline with `chain_depth * NodeConfig.per_hop_budget_ns` at
24//! dispatch time. The compiler pass exists so single-Node
25//! deployments + tests get a sane baseline deadline without
26//! requiring runtime fixup; multi-Node deployments with bespoke
27//! latency profiles rely on the runtime override.
28
29use crate::error::CompileError;
30use crate::partition_by_wire_ops::WIRE_DOMAIN;
31use bb_ir::keys::CHAIN_DEPTH_KEY;
32use bb_ir::proto::onnx::{attribute_proto, AttributeProto, ModelProto};
33use bb_ir::syscall_ids::ATTR_DEADLINE_NS;
34
35const SEND_OP: &str = "Send";
36
37/// Walk every `wire.Send` NodeProto with `chain_depth` metadata and
38/// stamp a `deadline_ns: i64` attribute. Idempotent - re-runs
39/// overwrite the previous attribute in place.
40///
41/// Returns the count of stamps applied.
42pub fn derive_wire_deadlines(
43    model: &mut ModelProto,
44    per_hop_budget_ns: u64,
45) -> Result<usize, CompileError> {
46    let mut stamp_count = 0usize;
47    for func in model.functions.iter_mut() {
48        for node in func.node.iter_mut() {
49            if node.domain != WIRE_DOMAIN || node.op_type != SEND_OP {
50                continue;
51            }
52            // Missing chain_depth defaults to 1 (single-hop) — covers
53            // every wire.Send the compiler hasn't paired with a
54            // multi-hop chain. Multi-hop chains stamp explicit
55            // chain_depth metadata via a later pass; both branches feed
56            // the same `chain_depth * per_hop_budget_ns` formula.
57            let chain_depth = read_chain_depth(node).unwrap_or(1);
58            let deadline_ns = (chain_depth as i64).saturating_mul(per_hop_budget_ns as i64);
59            upsert_deadline(node, deadline_ns);
60            stamp_count += 1;
61        }
62    }
63    Ok(stamp_count)
64}
65
66fn read_chain_depth(node: &bb_ir::proto::onnx::NodeProto) -> Option<u64> {
67    node.metadata_props
68        .iter()
69        .find(|p| p.key == CHAIN_DEPTH_KEY)
70        .and_then(|p| p.value.parse().ok())
71}
72
73fn upsert_deadline(node: &mut bb_ir::proto::onnx::NodeProto, deadline_ns: i64) {
74    if let Some(existing) = node
75        .attribute
76        .iter_mut()
77        .find(|a| a.name == ATTR_DEADLINE_NS)
78    {
79        existing.i = deadline_ns;
80        return;
81    }
82    node.attribute.push(AttributeProto {
83        name: ATTR_DEADLINE_NS.to_string(),
84        i: deadline_ns,
85        r#type: attribute_proto::AttributeType::Int as i32,
86        ..Default::default()
87    });
88}
89
bb_compiler/derive_wire_deadlines.rs

bb_compiler/
derive_wire_deadlines.rs
