use std::collections::BTreeMap;
use std::path::Path;
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use crate::graph::{build_from_trace, ExecutionGraph};
use crate::nsight::{GpuEvidenceStatus, NsightEvidence};
use crate::trace::{analyze_health, parse_trace, TraceDocument, TraceHealth, TraceRunMeta};
pub const EVIDENCE_SCHEMA: &str = "candle-graph/evidence/1";
pub const COMPARISON_SCHEMA: &str = "candle-graph/comparison/1";
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct EvidencePacket {
pub schema: String,
pub provenance: TraceRunMeta,
pub health: TraceHealth,
pub findings: Vec<String>,
pub facts: Vec<EvidenceFact>,
pub gaps: Vec<String>,
pub graph: ExecutionGraph,
pub gpu: NsightEvidence,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub comparison: Option<Comparison>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct EvidenceFact {
pub code: String,
pub label: String,
pub value: f64,
pub unit: String,
pub source: String,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Comparison {
pub schema: String,
pub baseline_run_id: String,
pub candidate_run_id: String,
pub comparable: bool,
pub warnings: Vec<String>,
pub total_delta_ns: i128,
pub total_delta_percent: Option<f64>,
pub baseline_peak_bytes: u64,
pub candidate_peak_bytes: u64,
pub peak_delta_bytes: i128,
pub spans: Vec<SpanComparison>,
pub gradients: Vec<GradientComparison>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SpanComparison {
pub name: String,
pub baseline_count: usize,
pub candidate_count: usize,
pub baseline_total_ns: u64,
pub candidate_total_ns: u64,
pub baseline_mean_ns: u64,
pub candidate_mean_ns: u64,
pub delta_ns: i128,
pub delta_percent: Option<f64>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct GradientComparison {
pub parameter: String,
pub baseline_state: Option<String>,
pub candidate_state: Option<String>,
pub baseline_norm: Option<f64>,
pub candidate_norm: Option<f64>,
pub norm_delta: Option<f64>,
}
pub fn build_evidence(
trace: &Path,
baseline: Option<&Path>,
nsight_dir: Option<&Path>,
) -> Result<EvidencePacket> {
let doc = parse_trace(trace).with_context(|| format!("parse trace {}", trace.display()))?;
let comparison = baseline
.map(|path| {
let baseline =
parse_trace(path).with_context(|| format!("parse baseline {}", path.display()))?;
Ok::<_, anyhow::Error>(compare_documents(&baseline, &doc))
})
.transpose()?;
let measured = measured_span_ids(&doc);
let expected_semantic_keys = doc
.spans
.iter()
.filter(|span| measured.contains(span.id.as_str()))
.map(|span| span.name.clone())
.collect::<Vec<_>>();
EvidencePacket::from_document(
doc,
NsightEvidence::load_optional(nsight_dir, &expected_semantic_keys),
comparison,
)
}
impl EvidencePacket {
pub fn from_document(
doc: TraceDocument,
gpu: NsightEvidence,
comparison: Option<Comparison>,
) -> Result<Self> {
let health = analyze_health(&doc);
let graph = build_from_trace(&doc)?;
let mut findings = graph
.summary
.slowest_spans
.iter()
.take(5)
.map(|span| {
format!(
"{}: {:.2} ms self time",
span.name,
span.self_time_ns as f64 / 1_000_000.0
)
})
.collect::<Vec<_>>();
let mut facts = vec![EvidenceFact {
code: "measured_total".into(),
label: "Measured update total".into(),
value: graph.summary.total_ms,
unit: "ms".into(),
source: "measured_span".into(),
}];
facts.extend(
graph
.summary
.slowest_spans
.iter()
.take(8)
.map(|span| EvidenceFact {
code: "span_self_time".into(),
label: span.name.clone(),
value: span.self_time_ns as f64,
unit: "ns".into(),
source: span.id.clone(),
}),
);
if !graph.gradients.is_empty() {
let concerning = graph
.gradients
.iter()
.filter(|gradient| {
!matches!(gradient.state, crate::graph::GradientRecordState::Present)
})
.count();
findings.push(format!(
"{} gradient facts captured; {} require attention",
graph.gradients.len(),
concerning
));
}
if gpu.status == GpuEvidenceStatus::Available {
if let Some(kernel) = gpu.kernels.first() {
findings.push(format!(
"top GPU kernel `{}`: {:.2} ms total",
kernel.name,
kernel.total_ns as f64 / 1_000_000.0
));
}
}
let mut gaps = health
.gaps()
.map(|issue| issue.message.clone())
.collect::<Vec<_>>();
if gpu.status != GpuEvidenceStatus::Available {
gaps.push(
gpu.reason
.clone()
.unwrap_or_else(|| "GPU evidence is unavailable".into()),
);
}
Ok(Self {
schema: EVIDENCE_SCHEMA.into(),
provenance: doc.run,
health,
findings,
facts,
gaps,
graph,
gpu,
comparison,
})
}
pub fn markdown(&self) -> String {
let trust = if self.health.trusted {
"TRUSTED"
} else {
"UNTRUSTED"
};
let mut out = format!(
"# candle-graph evidence\n\n- Status: **{trust}**\n- Entrypoint: `{}`\n- Capture update: {} ({} warmup update{})\n- Device: `{}`\n- Timing mode: `{:?}`\n- Measured region device-synchronized: {}\n- Total: {:.2} ms\n\n## Findings\n\n",
self.provenance.entrypoint,
self.provenance.capture_step,
self.provenance.warmup_steps,
if self.provenance.warmup_steps == 1 { "" } else { "s" },
self.provenance.device,
self.provenance.timing_mode,
self.provenance.measured_region_device_synchronized,
self.graph.summary.total_ms,
);
push_list(
&mut out,
&self.findings,
"No trusted findings were derived.",
);
out.push_str("\n## Evidence gaps\n\n");
push_list(&mut out, &self.gaps, "No known gaps.");
out.push_str("\n## Coverage\n\n```json\n");
out.push_str(&serde_json::to_string_pretty(&self.health.coverage).unwrap_or_default());
out.push_str("\n```\n");
out.push_str("\n## Tensor checkpoints\n\n");
if self.graph.tensors.is_empty() {
out.push_str("No tensor checkpoints were captured.\n");
} else {
out.push_str("| Tensor | Shape | Dtype | Device | Storage | Requires grad |\n| --- | --- | --- | --- | ---: | --- |\n");
for tensor in &self.graph.tensors {
out.push_str(&format!(
"| `{}` | `{:?}` | `{}` | `{}` | {} B | {} |\n",
tensor.tensor_id,
tensor.shape,
tensor.dtype,
tensor.device,
tensor.storage_bytes,
tensor.requires_grad
));
}
}
out.push_str("\n## Gradient evidence\n\n");
out.push_str(&format!(
"{} parameter gradients captured; {} require attention.\n",
self.graph.gradients.len(),
self.graph
.gradients
.iter()
.filter(|gradient| !matches!(
gradient.state,
crate::graph::GradientRecordState::Present
))
.count()
));
if let Some(comparison) = &self.comparison {
out.push_str("\n## Baseline comparison\n\n");
match comparison.total_delta_percent {
Some(percent) => out.push_str(&format!("Total time changed by {percent:+.2}%.\n")),
None => out.push_str(
"Total-time percentage is unavailable because the baseline was zero.\n",
),
}
for warning in &comparison.warnings {
out.push_str(&format!("- Warning: {warning}\n"));
}
}
out
}
}
pub fn compare_documents(baseline: &TraceDocument, candidate: &TraceDocument) -> Comparison {
let mut warnings = Vec::new();
let mut comparable = analyze_health(baseline).trusted && analyze_health(candidate).trusted;
for (name, left, right) in [
(
"entrypoint",
baseline.run.entrypoint.as_str(),
candidate.run.entrypoint.as_str(),
),
(
"phase",
baseline.run.phase.as_str(),
candidate.run.phase.as_str(),
),
(
"device",
baseline.run.device.as_str(),
candidate.run.device.as_str(),
),
] {
if left != right {
comparable = false;
warnings.push(format!("{name} differs: `{left}` vs `{right}`"));
}
}
if baseline.run.timing_mode != candidate.run.timing_mode {
comparable = false;
warnings.push("timing mode differs".into());
}
if baseline.run.measured_region_device_synchronized
!= candidate.run.measured_region_device_synchronized
{
comparable = false;
warnings.push("measured-region device synchronization differs".into());
}
if baseline.run.warmup_steps != candidate.run.warmup_steps {
comparable = false;
warnings.push(format!(
"warmup differs: {} vs {}",
baseline.run.warmup_steps, candidate.run.warmup_steps
));
}
let descriptive = ["source_revision", "source_commit", "build_id"];
let baseline_conditions = baseline
.run
.tags
.iter()
.filter(|(key, _)| !descriptive.contains(&key.as_str()))
.collect::<BTreeMap<_, _>>();
let candidate_conditions = candidate
.run
.tags
.iter()
.filter(|(key, _)| !descriptive.contains(&key.as_str()))
.collect::<BTreeMap<_, _>>();
if baseline_conditions != candidate_conditions {
comparable = false;
warnings.push("workload tags differ; batch/model/precision conditions must match".into());
}
for key in descriptive {
if baseline.run.tags.get(key) != candidate.run.tags.get(key) {
warnings.push(format!(
"descriptive `{key}` differs, as expected for a code-change comparison"
));
}
}
let baseline_spans = aggregate_spans(baseline);
let candidate_spans = aggregate_spans(candidate);
let mut names = baseline_spans
.keys()
.chain(candidate_spans.keys())
.cloned()
.collect::<Vec<_>>();
names.sort();
names.dedup();
let mut spans = names
.into_iter()
.map(|name| {
let (baseline_count, baseline_total_ns) =
baseline_spans.get(&name).copied().unwrap_or_default();
let (candidate_count, candidate_total_ns) =
candidate_spans.get(&name).copied().unwrap_or_default();
let delta_ns = candidate_total_ns as i128 - baseline_total_ns as i128;
SpanComparison {
name,
baseline_count,
candidate_count,
baseline_total_ns,
candidate_total_ns,
baseline_mean_ns: mean(baseline_total_ns, baseline_count),
candidate_mean_ns: mean(candidate_total_ns, candidate_count),
delta_ns,
delta_percent: percent(baseline_total_ns, candidate_total_ns),
}
})
.collect::<Vec<_>>();
spans.sort_by_key(|span| std::cmp::Reverse(span.delta_ns.unsigned_abs()));
spans.truncate(50);
let baseline_total = root_total(baseline);
let candidate_total = root_total(candidate);
let baseline_peak = crate::trace::memory::analyze_memory(baseline)
.summary
.peak_bytes;
let candidate_peak = crate::trace::memory::analyze_memory(candidate)
.summary
.peak_bytes;
Comparison {
schema: COMPARISON_SCHEMA.into(),
baseline_run_id: baseline.run.run_id.clone(),
candidate_run_id: candidate.run.run_id.clone(),
comparable,
warnings,
total_delta_ns: candidate_total as i128 - baseline_total as i128,
total_delta_percent: percent(baseline_total, candidate_total),
baseline_peak_bytes: baseline_peak,
candidate_peak_bytes: candidate_peak,
peak_delta_bytes: candidate_peak as i128 - baseline_peak as i128,
spans,
gradients: compare_gradients(baseline, candidate),
}
}
fn aggregate_spans(doc: &TraceDocument) -> BTreeMap<String, (usize, u64)> {
let mut result = BTreeMap::new();
let measured = measured_span_ids(doc);
for span in doc
.spans
.iter()
.filter(|span| measured.contains(span.id.as_str()))
{
let mut path = vec![span.name.clone()];
let mut parent = span.parent_id.as_deref();
let mut seen = std::collections::HashSet::new();
while let Some(id) = parent {
if !seen.insert(id) {
break;
}
let Some(parent_span) = doc.spans.iter().find(|candidate| candidate.id == id) else {
break;
};
path.push(parent_span.name.clone());
parent = parent_span.parent_id.as_deref();
}
path.reverse();
let step = span
.step
.map(|step| format!("/{step:?}"))
.unwrap_or_default();
let key = format!("{} [{}]{step}", path.join("/"), span.kind);
let entry = result.entry(key).or_insert((0usize, 0u64));
entry.0 += 1;
entry.1 = entry.1.saturating_add(span.duration_ns);
}
result
}
fn measured_span_ids(doc: &TraceDocument) -> std::collections::HashSet<&str> {
let mut ids = doc
.spans
.iter()
.filter(|span| span.measured)
.map(|span| span.id.as_str())
.collect::<std::collections::HashSet<_>>();
loop {
let before = ids.len();
for span in &doc.spans {
if span
.parent_id
.as_deref()
.is_some_and(|parent| ids.contains(parent))
{
ids.insert(span.id.as_str());
}
}
if ids.len() == before {
return ids;
}
}
}
fn compare_gradients(
baseline: &TraceDocument,
candidate: &TraceDocument,
) -> Vec<GradientComparison> {
let baseline = baseline
.gradients
.iter()
.map(|item| (format!("{}/{}", item.root, item.key), item))
.collect::<BTreeMap<_, _>>();
let candidate = candidate
.gradients
.iter()
.map(|item| (format!("{}/{}", item.root, item.key), item))
.collect::<BTreeMap<_, _>>();
let mut keys = baseline
.keys()
.chain(candidate.keys())
.cloned()
.collect::<Vec<_>>();
keys.sort();
keys.dedup();
keys.into_iter()
.filter_map(|parameter| {
let left = baseline.get(¶meter).copied();
let right = candidate.get(¶meter).copied();
let changed = left.map(|x| (x.state, x.norm)) != right.map(|x| (x.state, x.norm));
changed.then(|| GradientComparison {
parameter,
baseline_state: left.map(|x| x.state.to_string()),
candidate_state: right.map(|x| x.state.to_string()),
baseline_norm: left.and_then(|x| x.norm),
candidate_norm: right.and_then(|x| x.norm),
norm_delta: left
.and_then(|x| x.norm)
.zip(right.and_then(|x| x.norm))
.map(|(a, b)| b - a),
})
})
.take(100)
.collect()
}
fn mean(total: u64, count: usize) -> u64 {
if count == 0 {
0
} else {
total / count as u64
}
}
fn root_total(doc: &TraceDocument) -> u64 {
doc.spans
.iter()
.filter(|span| span.measured)
.map(|span| span.duration_ns)
.sum()
}
fn percent(baseline: u64, candidate: u64) -> Option<f64> {
(baseline > 0).then(|| (candidate as f64 - baseline as f64) * 100.0 / baseline as f64)
}
fn push_list(out: &mut String, values: &[String], empty: &str) {
if values.is_empty() {
out.push_str(&format!("- {empty}\n"));
} else {
for value in values {
out.push_str(&format!("- {value}\n"));
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::phase::{ExecutionPhase, ExecutionStep};
use crate::trace::{
GradientEvent, GradientState, SpanKind, SpanRecord, TimingMode, TraceRunMeta, SCHEMA,
};
fn document(run_id: &str, measured_ns: u64, forward_ns: u64) -> TraceDocument {
TraceDocument {
schema: SCHEMA.into(),
run: TraceRunMeta {
run_id: run_id.into(),
correlation_id: format!("demo/{run_id}"),
entrypoint: "demo::update".into(),
phase: ExecutionPhase::Train,
timestamp: "2026-08-08T00:00:00Z".into(),
capture_step: 1,
warmup_steps: 0,
device: "cpu".into(),
measured_region_device_synchronized: false,
timing_mode: TimingMode::Host,
tags: [("physical_batch".into(), "2".into())].into(),
candle_version: None,
},
spans: vec![
span("session", None, false, measured_ns + 100, None),
span("update", Some("session"), true, measured_ns, None),
span(
"forward",
Some("update"),
false,
forward_ns,
Some(ExecutionStep::Forward),
),
span(
"backward",
Some("update"),
false,
measured_ns.saturating_sub(forward_ns + 10),
Some(ExecutionStep::Backward),
),
span(
"optimizer",
Some("update"),
false,
10,
Some(ExecutionStep::Optimizer),
),
],
ops: vec![],
tensors: vec![],
memory: vec![],
device_memory: vec![],
gradients: vec![GradientEvent {
event_id: "g".into(),
root: "vb".into(),
key: "weight".into(),
state: GradientState::Present,
norm: Some(forward_ns as f64),
}],
edges: vec![],
}
}
fn span(
id: &str,
parent: Option<&str>,
measured: bool,
duration_ns: u64,
step: Option<ExecutionStep>,
) -> SpanRecord {
SpanRecord {
id: id.into(),
parent_id: parent.map(str::to_string),
name: id.into(),
kind: SpanKind::Function,
measured,
start_ns: 0,
closed: true,
duration_ns,
step,
}
}
#[test]
fn compares_measured_region_and_semantic_paths() {
let comparison =
compare_documents(&document("base", 100, 60), &document("candidate", 80, 40));
assert!(comparison.comparable);
assert_eq!(comparison.total_delta_ns, -20);
assert_eq!(comparison.total_delta_percent, Some(-20.0));
assert!(comparison
.spans
.iter()
.any(|span| span.name.contains("session/update/forward")));
assert_eq!(comparison.gradients.len(), 1);
}
#[test]
fn rejects_different_measured_region_synchronization_contracts() {
let baseline = document("base", 100, 60);
let mut candidate = document("candidate", 80, 40);
candidate.run.measured_region_device_synchronized = true;
let comparison = compare_documents(&baseline, &candidate);
assert!(!comparison.comparable);
assert!(comparison
.warnings
.iter()
.any(|warning| warning.contains("measured-region device synchronization differs")));
}
#[test]
fn packet_markdown_exposes_trust_gaps_and_coverage() {
let packet = EvidencePacket::from_document(
document("candidate", 80, 40),
NsightEvidence::unavailable("nsys not installed"),
None,
)
.unwrap();
let markdown = packet.markdown();
assert!(markdown.contains("Status: **TRUSTED**"));
assert!(markdown.contains("Measured region device-synchronized: false"));
assert!(markdown.contains("nsys not installed"));
assert!(markdown.contains("optimizer_spans"));
assert!(!packet.facts.is_empty());
}
}