use bb_ir::proto::onnx::GraphProto;
use crate::error::CompileError;
use crate::partition_by_wire_ops::WIRE_DOMAIN;
const SYSCALL_DOMAIN: &str = "ai.bytesandbrains.syscall";
const SEND_OP: &str = "Send";
const RECV_OP: &str = "Recv";
const PEER_HEALTH_GATE_TX_OP: &str = "PeerHealthGateTx";
const BACKOFF_GATE_TX_OP: &str = "BackoffGateTx";
const DEDUP_GATE_RX_OP: &str = "DedupGateRx";
const PEER_HEALTH_GATE_RX_OP: &str = "PeerHealthGateRx";
const BACKOFF_GATE_RX_OP: &str = "BackoffGateRx";
const DEADLINE_CHECK_OP: &str = "DeadlineCheck";
const DEADLINE_NS_ATTR: &str = "deadline_ns";
const PEER_ATTR: &str = "peer";
pub fn validate_runtime_complete(sub_graph: &GraphProto) -> Result<(), CompileError> {
let nodes = &sub_graph.node;
let has_op = |op_type: &str, domain: &str| -> bool {
nodes
.iter()
.any(|n| n.op_type == op_type && n.domain == domain)
};
let has_peer_send = nodes.iter().any(|n| {
n.domain == WIRE_DOMAIN
&& n.op_type == SEND_OP
&& n.attribute.iter().any(|a| a.name == PEER_ATTR)
});
if has_peer_send {
if !has_op(PEER_HEALTH_GATE_TX_OP, SYSCALL_DOMAIN) {
return Err(CompileError::Internal {
detail: format!(
"validate_runtime_complete: partition `{}` has a peer-routed wire.Send but no PeerHealthGateTx",
sub_graph.name,
),
});
}
if !has_op(BACKOFF_GATE_TX_OP, SYSCALL_DOMAIN) {
return Err(CompileError::Internal {
detail: format!(
"validate_runtime_complete: partition `{}` has a peer-routed wire.Send but no BackoffGateTx",
sub_graph.name,
),
});
}
}
let has_peer_recv = nodes.iter().any(|n| {
n.domain == WIRE_DOMAIN
&& n.op_type == RECV_OP
&& n.attribute.iter().any(|a| a.name == PEER_ATTR)
});
if has_peer_recv {
if !has_op(DEDUP_GATE_RX_OP, SYSCALL_DOMAIN) {
return Err(CompileError::Internal {
detail: format!(
"validate_runtime_complete: partition `{}` has a peer-routed wire.Recv but no DedupGateRx",
sub_graph.name,
),
});
}
if !has_op(PEER_HEALTH_GATE_RX_OP, SYSCALL_DOMAIN) {
return Err(CompileError::Internal {
detail: format!(
"validate_runtime_complete: partition `{}` has a peer-routed wire.Recv but no PeerHealthGateRx",
sub_graph.name,
),
});
}
if !has_op(BACKOFF_GATE_RX_OP, SYSCALL_DOMAIN) {
return Err(CompileError::Internal {
detail: format!(
"validate_runtime_complete: partition `{}` has a peer-routed wire.Recv but no BackoffGateRx",
sub_graph.name,
),
});
}
}
let any_deadline = nodes
.iter()
.any(|n| n.attribute.iter().any(|a| a.name == DEADLINE_NS_ATTR));
if any_deadline && !has_op(DEADLINE_CHECK_OP, SYSCALL_DOMAIN) {
return Err(CompileError::Internal {
detail: format!(
"validate_runtime_complete: partition `{}` carries deadline_ns but no DeadlineCheck",
sub_graph.name,
),
});
}
for reg in crate::gate_contract::contracts() {
reg.contract.assert_inserted(sub_graph)?;
}
Ok(())
}