Rebuild conclusion stage summaries from run events

This commit is contained in:
Bryan Helmkamp 2026-09-05 11:44:23 -04:00
parent efa0d16588
commit a7ed03e175
No known key found for this signature in database
5 changed files with 442 additions and 275 deletions

View file

@ -19,7 +19,7 @@ use fabro_types::{
RunStatus, RunTimestamps, SandboxProviderKind, StageCompletion, StageHandler, StageId,
StageInferenceProjection, StageModelUsage, StageOutcome, StageProjection, StageState,
StartRecord, SubAgentProjection, SubAgentStatus, TodoListKind, TodoListProjection,
TodoProjection, WorkflowRef, first_event_seq, timing,
TodoProjection, WorkflowRef, billing_rollup, first_event_seq, timing,
};
use fabro_util::error::render_compact_with_causes;
@ -247,7 +247,7 @@ impl RunProjectionReducer for RunProjection {
ts,
)?;
self.pending_control = None;
self.conclusion = Some(conclusion_from_completed(props, ts)?);
self.conclusion = Some(conclusion_from_completed(self, props, ts)?);
self.pending_interviews.clear();
}
EventBody::RunFailed(props) => {
@ -258,7 +258,7 @@ impl RunProjectionReducer for RunProjection {
ts,
)?;
self.pending_control = None;
self.conclusion = Some(conclusion_from_failed(props, ts));
self.conclusion = Some(conclusion_from_failed(self, props, ts));
self.pending_interviews.clear();
finalize_unfinished_stages_after_run_failed(self, props, ts);
}
@ -1560,9 +1560,12 @@ fn diff_from_checkpoint_props(props: &CheckpointCompletedProps) -> RunDiff {
}
fn conclusion_from_completed(
projection: &RunProjection,
props: &RunCompletedProps,
timestamp: DateTime<Utc>,
) -> Result<Conclusion> {
let (stages, total_retries) = billing_rollup::billing_rollup_from_projection(projection, None)
.conclusion_stages(projection);
Ok(Conclusion {
timestamp,
status: StageOutcome::from_str(&props.status)
@ -1570,9 +1573,9 @@ fn conclusion_from_completed(
timing: props.timing,
failure: None,
final_git_commit_sha: props.final_git_commit_sha.clone(),
stages: Vec::new(),
stages,
billing: props.billing.clone(),
total_retries: 0,
total_retries,
diff: RunDiff {
patch: props.final_patch.clone(),
summary: props.diff_summary,
@ -1580,7 +1583,13 @@ fn conclusion_from_completed(
})
}
fn conclusion_from_failed(props: &RunFailedProps, timestamp: DateTime<Utc>) -> Conclusion {
fn conclusion_from_failed(
projection: &RunProjection,
props: &RunFailedProps,
timestamp: DateTime<Utc>,
) -> Conclusion {
let (stages, total_retries) = billing_rollup::billing_rollup_from_projection(projection, None)
.conclusion_stages(projection);
Conclusion {
timestamp,
status: StageOutcome::Failed {
@ -1589,9 +1598,9 @@ fn conclusion_from_failed(props: &RunFailedProps, timestamp: DateTime<Utc>) -> C
timing: props.timing,
failure: Some(props.failure.clone()),
final_git_commit_sha: props.final_git_commit_sha.clone(),
stages: Vec::new(),
stages,
billing: props.billing.clone(),
total_retries: 0,
total_retries,
diff: RunDiff {
patch: props.final_patch.clone(),
summary: props.diff_summary,
@ -4416,6 +4425,170 @@ mod tests {
);
}
#[test]
fn terminal_conclusion_replays_stage_summaries_without_metadata() {
for terminal_name in ["run.completed", "run.failed"] {
let mut settings = WorkflowSettings::default();
settings.run.meta_branch.enabled = false;
let mut events = vec![
test_raw_event(
1,
"run.created",
&json!({
"settings": settings,
"graph": { "name": "test", "nodes": {}, "edges": [], "attrs": {} },
"labels": {},
"provenance": test_support::test_run_provenance()
}),
None,
),
test_raw_event(
2,
"run.runnable",
&json!({ "source": "start_requested" }),
None,
),
test_raw_event(3, "run.starting", &json!({}), None),
test_raw_event(4, "run.running", &json!({}), None),
];
// Two executions of zebra share one conclusion row. First-event
// order differs from checkpoint order, and skipped has no completion.
for (seq, node, visit, millis, tokens) in [
(5, "zebra", 1, 1200, 100),
(6, "apple", 1, 300, 20),
(7, "zebra", 2, 800, 200),
] {
let mut props = completed_props(millis, StageOutcome::Succeeded);
props.billing = Some(test_usage("test-model", tokens, 10));
events.push(test_stage_event(
seq,
EventBody::StageCompleted(props),
StageId::new(node, visit),
));
}
events.push(test_raw_event(8, "checkpoint.completed", &json!({
"status": "succeeded",
"current_node": "zebra",
"completed_nodes": ["apple", "zebra", "zebra"],
"node_retries": { "zebra": 3, "apple": 1 },
"node_outcomes": {
"apple": Outcome::<Option<BilledModelUsage>>::success(),
"zebra": Outcome::<Option<BilledModelUsage>>::success(),
"skipped": Outcome::<Option<BilledModelUsage>>::skipped("condition was false")
},
"context_values": {},
"node_visits": { "zebra": 2, "apple": 1, "skipped": 1 },
"git_commit_sha": "checkpoint-sha"
}), Some("zebra")));
let terminal_billing = usage_counts(&test_usage("test-model", 320, 30));
let terminal_props = if terminal_name == "run.completed" {
json!({
"status": "succeeded", "reason": "completed",
"timing": fabro_types::RunTiming::wall_only(9000),
"artifact_count": 0, "billing": terminal_billing,
"final_git_commit_sha": "final-sha", "final_patch": "final patch"
})
} else {
let mut props = run_failed_props(FailureReason::WorkflowError);
props.timing = fabro_types::RunTiming::wall_only(9000);
props.billing = Some(terminal_billing.clone());
props.final_git_commit_sha = Some("final-sha".to_string());
props.final_patch = Some("final patch".to_string());
serde_json::to_value(props).unwrap()
};
events.push(test_raw_event(9, terminal_name, &terminal_props, None));
for event in &mut events {
event.event.ts = test_dt("2026-04-07T12:00:00Z")
+ chrono::Duration::seconds(i64::from(event.seq));
}
// Cross the persisted wire boundary before both incremental and full replay.
let events: Vec<EventEnvelope> =
serde_json::from_slice(&serde_json::to_vec(&events).unwrap()).unwrap();
let mut live = RunProjection::apply_events(&events[..1]).unwrap();
for event in &events[1..] {
live.apply_event(event).unwrap();
}
let replayed = RunProjection::apply_events(&events).unwrap();
let conclusion = replayed.conclusion.as_ref().unwrap();
assert_eq!(
serde_json::to_value(&live.conclusion).unwrap(),
serde_json::to_value(conclusion).unwrap(),
);
assert_eq!(conclusion.timestamp, events.last().unwrap().event.ts);
assert_eq!(conclusion.timing.wall_time_ms, 9000);
assert_eq!(conclusion.billing, Some(terminal_billing));
assert_eq!(
conclusion.final_git_commit_sha.as_deref(),
Some("final-sha")
);
assert_eq!(conclusion.diff.patch.as_deref(), Some("final patch"));
assert_eq!(conclusion.failure.is_some(), terminal_name == "run.failed");
insta::allow_duplicates! {
insta::assert_snapshot!(serde_json::to_string_pretty(&json!({
"stages": conclusion.stages,
"total_retries": conclusion.total_retries,
})).unwrap(), @r###"
{
"stages": [
{
"stage_id": "zebra",
"stage_label": "zebra",
"timing": {
"wall_time_ms": 2000,
"inference_time_ms": 0,
"tool_time_ms": 0,
"active_time_ms": 0
},
"billing_usd_micros": 320,
"retries": 2
},
{
"stage_id": "apple",
"stage_label": "apple",
"timing": {
"wall_time_ms": 300,
"inference_time_ms": 0,
"tool_time_ms": 0,
"active_time_ms": 0
},
"billing_usd_micros": 30,
"retries": 0
},
{
"stage_id": "skipped",
"stage_label": "skipped",
"timing": {
"wall_time_ms": 0,
"inference_time_ms": 0,
"tool_time_ms": 0,
"active_time_ms": 0
},
"retries": 0
}
],
"total_retries": 2
}
"###);
}
}
}
#[test]
fn terminal_conclusion_without_checkpoint_has_no_stage_summaries() {
let mut state = running_projection();
let terminal = test_event(
4,
EventBody::RunFailed(run_failed_props(FailureReason::WorkflowError)),
None,
);
state.apply_event(&terminal).unwrap();
let conclusion = state.conclusion.unwrap();
assert!(conclusion.stages.is_empty());
assert_eq!(conclusion.total_retries, 0);
assert_eq!(conclusion.timestamp, terminal.event.ts);
}
#[test]
fn run_failed_with_final_patch_populates_projection() {
let mut state = running_projection();

View file

@ -1,128 +1,7 @@
use std::collections::HashMap;
use fabro_model::Catalog;
use fabro_types::{BilledTokenCounts, ModelRef, RunProjection, RunTiming, StageTiming};
#[derive(Debug, Clone, PartialEq)]
pub struct ProjectionBillingStage {
pub node_id: String,
pub billing: BilledTokenCounts,
/// Per-node timing summed across every visit of that node within this
/// projection. `wall_time_ms`, `inference_time_ms`, `tool_time_ms`, and
/// `active_time_ms` are all summed in lockstep.
pub timing: StageTiming,
pub model: Option<ModelRef>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProjectionBillingByModel {
pub model: ModelRef,
pub stages: i64,
pub billing: BilledTokenCounts,
}
#[derive(Debug, Clone, Default, PartialEq)]
pub struct ProjectionBillingRollup {
pub stages: Vec<ProjectionBillingStage>,
pub totals: BilledTokenCounts,
pub by_model: Vec<ProjectionBillingByModel>,
/// Run-level timing summed across every stage visit. `wall_time_ms` is
/// the sum of stage visit wall times (not the run clock duration).
pub timing: RunTiming,
pub billed_visit_count: usize,
}
impl ProjectionBillingRollup {
#[must_use]
pub fn billing_if_present(&self) -> Option<BilledTokenCounts> {
(self.billed_visit_count > 0).then(|| self.totals.clone())
}
}
#[must_use]
pub fn billing_rollup_from_projection(
projection: &RunProjection,
catalog: Option<&Catalog>,
) -> ProjectionBillingRollup {
let mut stage_indices = HashMap::<String, usize>::new();
let mut stages = Vec::<ProjectionBillingStage>::new();
let mut by_model = HashMap::<ModelRef, ProjectionBillingByModel>::new();
let mut totals = BilledTokenCounts::default();
let mut run_timing = RunTiming::default();
let mut billed_visit_count = 0_usize;
for (stage_id, stage) in projection.iter_stages() {
if projection.is_boundary_stage(stage_id.node_id()) {
continue;
}
let usage = stage.billed_usage(catalog);
let usage = usage.as_ref();
if stage.completion.is_none() && stage.timing.is_none() && usage.is_zero() {
continue;
}
let node_id = stage_id.node_id();
let index = *stage_indices.entry(node_id.to_string()).or_insert_with(|| {
let index = stages.len();
stages.push(ProjectionBillingStage {
node_id: node_id.to_string(),
billing: BilledTokenCounts::default(),
timing: StageTiming::default(),
model: None,
});
index
});
let row = &mut stages[index];
if let Some(timing) = stage.timing {
row.timing = row.timing.saturating_add(&timing);
run_timing = run_timing.saturating_add(&RunTiming::from(timing));
}
if !usage.is_zero() {
billed_visit_count += 1;
row.billing.add_counts(usage);
totals.add_counts(usage);
if let Some(model) = &stage.model {
row.model = Some(model.clone());
let model_entry =
by_model
.entry(model.clone())
.or_insert_with(|| ProjectionBillingByModel {
model: model.clone(),
stages: 0,
billing: BilledTokenCounts::default(),
});
model_entry.stages += 1;
model_entry.billing.add_counts(usage);
}
}
}
let mut by_model = by_model.into_values().collect::<Vec<_>>();
by_model.sort_by(|left, right| {
let left_provider = left.model.provider.to_string();
let right_provider = right.model.provider.to_string();
left_provider
.cmp(&right_provider)
.then_with(|| left.model.model_id.cmp(&right.model.model_id))
.then_with(|| {
left.model
.speed
.map(<&'static str>::from)
.cmp(&right.model.speed.map(<&'static str>::from))
})
});
ProjectionBillingRollup {
stages,
totals,
by_model,
timing: run_timing,
billed_visit_count,
}
}
pub use fabro_types::billing_rollup::{
ProjectionBillingByModel, ProjectionBillingRollup, ProjectionBillingStage,
billing_rollup_from_projection,
};
#[cfg(test)]
mod tests {

View file

@ -1,4 +1,3 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Instant;
@ -10,17 +9,17 @@ use fabro_util::error::collect_causes;
use fabro_util::time::elapsed_ms;
use super::types::{Concluded, Executed, FinalizeOptions, Finalized, PublishOutcome, Published};
use crate::billing_rollup;
use crate::error::{Error, run_failure_from_error, run_failure_from_outcome_failure};
use crate::event::{Event, RunNoticeCode, RunNoticeLevel};
use crate::outcome::{Outcome, StageOutcome};
use crate::records::{Checkpoint, Conclusion, StageSummary};
use crate::records::Conclusion;
use crate::run_metadata::{MetadataSnapshot, metadata_push_failure_is_transient};
use crate::run_options::RunOptions;
use crate::run_status::{FailureReason, RunStatus, SuccessReason};
use crate::runtime_store::RunStoreHandle;
use crate::sandbox_git::{git_diff_with_timeout, list_diff_numstat, summarize_diff_numstat};
use crate::services::RunServices;
use crate::{ProjectionBillingRollup, billing_rollup_from_projection};
pub fn classify_engine_result(
engine_result: &Result<Outcome, Error>,
@ -65,22 +64,8 @@ pub(crate) async fn build_conclusion_from_store(
final_git_commit_sha: Option<String>,
) -> Conclusion {
let projection = run_store.state().await.ok();
let projection_order = projection
.as_ref()
.map(stage_projection_order)
.unwrap_or_default();
let projection_billing = projection
.as_ref()
.map(|projection| billing_rollup_from_projection(projection, None))
.unwrap_or_default();
let checkpoint = projection
.as_ref()
.and_then(|state| state.current_checkpoint());
build_conclusion_from_parts(
checkpoint,
&projection_billing,
&projection_order,
build_conclusion_from_projection(
projection.as_ref(),
status,
failure,
run_wall_time_ms,
@ -88,113 +73,32 @@ pub(crate) async fn build_conclusion_from_store(
)
}
fn build_conclusion_from_parts(
checkpoint: Option<&Checkpoint>,
projection_billing: &ProjectionBillingRollup,
projection_order: &HashMap<String, u32>,
fn build_conclusion_from_projection(
projection: Option<&RunProjection>,
status: StageOutcome,
failure: Option<RunFailure>,
run_wall_time_ms: u64,
final_git_commit_sha: Option<String>,
) -> Conclusion {
// Looping workflows revisit nodes; `completed_nodes` accumulates duplicates
// while the other checkpoint maps are keyed by node_id. Dedupe to one row
// per node so the stages table matches the deduped billing total.
let (stages, total_retries) = if let Some(cp) = checkpoint {
let billing_by_node = projection_billing
.stages
.iter()
.map(|stage| (stage.node_id.as_str(), stage))
.collect::<HashMap<_, _>>();
let mut stage_rows = Vec::new();
let mut seen = std::collections::HashSet::new();
let mut retries_sum: u32 = 0;
let mut stage_order = Vec::new();
for (original_checkpoint_order, node_id) in cp.completed_nodes.iter().enumerate() {
if !seen.insert(node_id.as_str()) {
continue;
}
stage_order.push((original_checkpoint_order, node_id.as_str()));
}
let mut extra_node_outcomes = cp
.node_outcomes
.keys()
.filter(|node_id| !seen.contains(node_id.as_str()))
.map(String::as_str)
.collect::<Vec<_>>();
extra_node_outcomes.sort_unstable();
let extra_offset = stage_order.len();
for (extra_index, node_id) in extra_node_outcomes.into_iter().enumerate() {
seen.insert(node_id);
stage_order.push((extra_offset + extra_index, node_id));
}
for (original_checkpoint_order, node_id) in stage_order {
let retries = cp
.node_retries
.get(node_id)
.copied()
.unwrap_or(1)
.saturating_sub(1);
retries_sum += retries;
let billing = billing_by_node.get(node_id);
let summary = StageSummary {
stage_id: node_id.to_string(),
stage_label: node_id.to_string(),
timing: billing
.map_or_else(fabro_types::StageTiming::default, |stage| stage.timing),
billing_usd_micros: billing.and_then(|stage| stage.billing.total_usd_micros),
retries,
};
stage_rows.push((
projection_order.get(node_id).copied().unwrap_or(u32::MAX),
original_checkpoint_order,
summary,
));
}
stage_rows.sort_by(|left, right| {
left.0
.cmp(&right.0)
.then_with(|| left.1.cmp(&right.1))
.then_with(|| left.2.stage_id.cmp(&right.2.stage_id))
});
let stages = stage_rows
.into_iter()
.map(|(_, _, summary)| summary)
.collect();
(stages, retries_sum)
} else {
(vec![], 0)
};
let billing = projection
.map(|projection| billing_rollup::billing_rollup_from_projection(projection, None))
.unwrap_or_default();
let (stages, total_retries) = projection
.map(|projection| billing.conclusion_stages(projection))
.unwrap_or_default();
Conclusion {
timestamp: chrono::Utc::now(),
status,
timing: projection_billing.timing.with_wall_time(run_wall_time_ms),
timing: billing.timing.with_wall_time(run_wall_time_ms),
failure,
final_git_commit_sha,
stages,
billing: projection_billing.billing_if_present(),
billing: billing.billing_if_present(),
total_retries,
diff: fabro_types::RunDiff::default(),
}
}
fn stage_projection_order(state: &RunProjection) -> HashMap<String, u32> {
let mut order = HashMap::new();
for (stage_id, stage) in state.iter_stages() {
order
.entry(stage_id.node_id().to_string())
.and_modify(|first_seq: &mut u32| {
*first_seq = (*first_seq).min(stage.first_event_seq.get());
})
.or_insert_with(|| stage.first_event_seq.get());
}
order
}
/// `conclusion` is injected because the terminal event hasn't been emitted
/// yet — the run store's `projection.conclusion` is still `None` at this point.
pub async fn write_finalize_commit(
@ -439,7 +343,7 @@ async fn compute_final_patch(
#[cfg(any(test, feature = "test-support"))]
pub(crate) fn billing_from_projection(projection: &RunProjection) -> Option<BilledTokenCounts> {
billing_rollup_from_projection(projection, None).billing_if_present()
billing_rollup::billing_rollup_from_projection(projection, None).billing_if_present()
}
pub(crate) fn build_terminal_event(
@ -543,21 +447,8 @@ pub async fn conclude(executed: Executed, options: &FinalizeOptions) -> Result<C
.filter(|envelope| matches!(envelope.event.body, EventBody::ArtifactCaptured(_)))
.count();
let projection = services.run_store.state().await.ok();
let projection_order = projection
.as_ref()
.map(stage_projection_order)
.unwrap_or_default();
let projection_billing = projection
.as_ref()
.map(|projection| billing_rollup_from_projection(projection, None))
.unwrap_or_default();
let checkpoint = projection
.as_ref()
.and_then(|state| state.current_checkpoint());
let mut conclusion = build_conclusion_from_parts(
checkpoint,
&projection_billing,
&projection_order,
let mut conclusion = build_conclusion_from_projection(
projection.as_ref(),
final_status,
failure_reason,
wall_time_ms,
@ -703,6 +594,7 @@ mod tests {
use crate::context::Context;
use crate::error::ErrorStage;
use crate::event::{Emitter, StoreProgressLogger, append_event};
use crate::records::Checkpoint;
use crate::run_metadata::{RunMetadataRuntime, RunMetadataWriterHandle};
use crate::run_options::{GitCheckpointOptions, RunOptions};
use crate::runtime_store::{RunStoreBackend, RunStoreHandle};
@ -956,7 +848,6 @@ mod tests {
let mut projection = test_projection();
projection.stage_entry("zebra", 1, first_event_seq(1));
projection.stage_entry("apple", 1, first_event_seq(2));
let projection_order = stage_projection_order(&projection);
let checkpoint = checkpoint_with(
vec!["apple", "zebra"],
HashMap::from([
@ -965,10 +856,13 @@ mod tests {
]),
);
let conclusion = build_conclusion_from_parts(
Some(&checkpoint),
&ProjectionBillingRollup::default(),
&projection_order,
projection.checkpoints.push(fabro_types::CheckpointRecord {
seq: 10,
checkpoint,
diff: fabro_types::RunDiff::default(),
});
let conclusion = build_conclusion_from_projection(
Some(&projection),
StageOutcome::Succeeded,
None,
10,
@ -988,7 +882,6 @@ mod tests {
let mut projection = test_projection();
projection.stage_entry("skipped", 1, first_event_seq(4));
projection.stage_entry("finished", 1, first_event_seq(5));
let projection_order = stage_projection_order(&projection);
let checkpoint = checkpoint_with(
vec!["finished"],
HashMap::from([
@ -1000,10 +893,13 @@ mod tests {
]),
);
let conclusion = build_conclusion_from_parts(
Some(&checkpoint),
&ProjectionBillingRollup::default(),
&projection_order,
projection.checkpoints.push(fabro_types::CheckpointRecord {
seq: 10,
checkpoint,
diff: fabro_types::RunDiff::default(),
});
let conclusion = build_conclusion_from_projection(
Some(&projection),
StageOutcome::Succeeded,
None,
10,
@ -1047,8 +943,6 @@ mod tests {
timestamp: chrono::Utc::now(),
});
let projection_order = stage_projection_order(&projection);
let projection_billing = billing_rollup_from_projection(&projection, None);
let mut latest_outcome = Outcome::success();
latest_outcome.usage = Some(success_usage);
latest_outcome.timing = Some(fabro_types::StageTiming::wall_only(800));
@ -1058,10 +952,13 @@ mod tests {
);
checkpoint.node_retries.insert("verify".to_string(), 2);
let conclusion = build_conclusion_from_parts(
Some(&checkpoint),
&projection_billing,
&projection_order,
projection.checkpoints.push(fabro_types::CheckpointRecord {
seq: 10,
checkpoint,
diff: fabro_types::RunDiff::default(),
});
let conclusion = build_conclusion_from_projection(
Some(&projection),
StageOutcome::Succeeded,
None,
10,

View file

@ -0,0 +1,217 @@
use std::collections::HashMap;
use fabro_model::Catalog;
use crate::{BilledTokenCounts, ModelRef, RunProjection, RunTiming, StageSummary, StageTiming};
#[derive(Debug, Clone, PartialEq)]
pub struct ProjectionBillingStage {
pub node_id: String,
pub billing: BilledTokenCounts,
/// Per-node timing summed across every visit of that node within this
/// projection. `wall_time_ms`, `inference_time_ms`, `tool_time_ms`, and
/// `active_time_ms` are all summed in lockstep.
pub timing: StageTiming,
pub model: Option<ModelRef>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProjectionBillingByModel {
pub model: ModelRef,
pub stages: i64,
pub billing: BilledTokenCounts,
}
#[derive(Debug, Clone, Default, PartialEq)]
pub struct ProjectionBillingRollup {
pub stages: Vec<ProjectionBillingStage>,
pub totals: BilledTokenCounts,
pub by_model: Vec<ProjectionBillingByModel>,
/// Run-level timing summed across every stage visit. `wall_time_ms` is
/// the sum of stage visit wall times (not the run clock duration).
pub timing: RunTiming,
pub billed_visit_count: usize,
}
impl ProjectionBillingRollup {
#[must_use]
pub fn billing_if_present(&self) -> Option<BilledTokenCounts> {
(self.billed_visit_count > 0).then(|| self.totals.clone())
}
/// Reconstruct the conclusion's per-node summaries from checkpoint and
/// stage events. Repeated visits share one row, ordered by the node's
/// first stage event.
#[must_use]
pub fn conclusion_stages(&self, projection: &RunProjection) -> (Vec<StageSummary>, u32) {
let projection_order = stage_projection_order(projection);
// Looping workflows revisit nodes; `completed_nodes` accumulates duplicates
// while the other checkpoint maps are keyed by node_id. Dedupe to one row
// per node so the stages table matches the deduped billing total.
if let Some(cp) = projection.current_checkpoint() {
let billing_by_node = self
.stages
.iter()
.map(|stage| (stage.node_id.as_str(), stage))
.collect::<HashMap<_, _>>();
let mut stage_rows = Vec::new();
let mut seen = std::collections::HashSet::new();
let mut retries_sum: u32 = 0;
let mut stage_order = Vec::new();
for (original_checkpoint_order, node_id) in cp.completed_nodes.iter().enumerate() {
if !seen.insert(node_id.as_str()) {
continue;
}
stage_order.push((original_checkpoint_order, node_id.as_str()));
}
let mut extra_node_outcomes = cp
.node_outcomes
.keys()
.filter(|node_id| !seen.contains(node_id.as_str()))
.map(String::as_str)
.collect::<Vec<_>>();
extra_node_outcomes.sort_unstable();
let extra_offset = stage_order.len();
for (extra_index, node_id) in extra_node_outcomes.into_iter().enumerate() {
seen.insert(node_id);
stage_order.push((extra_offset + extra_index, node_id));
}
for (original_checkpoint_order, node_id) in stage_order {
let retries = cp
.node_retries
.get(node_id)
.copied()
.unwrap_or(1)
.saturating_sub(1);
retries_sum += retries;
let billing = billing_by_node.get(node_id);
let summary = StageSummary {
stage_id: node_id.to_string(),
stage_label: node_id.to_string(),
timing: billing.map_or_else(StageTiming::default, |stage| stage.timing),
billing_usd_micros: billing.and_then(|stage| stage.billing.total_usd_micros),
retries,
};
stage_rows.push((
projection_order.get(node_id).copied().unwrap_or(u32::MAX),
original_checkpoint_order,
summary,
));
}
stage_rows.sort_by(|left, right| {
left.0
.cmp(&right.0)
.then_with(|| left.1.cmp(&right.1))
.then_with(|| left.2.stage_id.cmp(&right.2.stage_id))
});
let stages = stage_rows
.into_iter()
.map(|(_, _, summary)| summary)
.collect();
(stages, retries_sum)
} else {
(vec![], 0)
}
}
}
#[must_use]
pub fn billing_rollup_from_projection(
projection: &RunProjection,
catalog: Option<&Catalog>,
) -> ProjectionBillingRollup {
let mut stage_indices = HashMap::<String, usize>::new();
let mut stages = Vec::<ProjectionBillingStage>::new();
let mut by_model = HashMap::<ModelRef, ProjectionBillingByModel>::new();
let mut totals = BilledTokenCounts::default();
let mut run_timing = RunTiming::default();
let mut billed_visit_count = 0_usize;
for (stage_id, stage) in projection.iter_stages() {
if projection.is_boundary_stage(stage_id.node_id()) {
continue;
}
let usage = stage.billed_usage(catalog);
let usage = usage.as_ref();
if stage.completion.is_none() && stage.timing.is_none() && usage.is_zero() {
continue;
}
let node_id = stage_id.node_id();
let index = *stage_indices.entry(node_id.to_string()).or_insert_with(|| {
let index = stages.len();
stages.push(ProjectionBillingStage {
node_id: node_id.to_string(),
billing: BilledTokenCounts::default(),
timing: StageTiming::default(),
model: None,
});
index
});
let row = &mut stages[index];
if let Some(timing) = stage.timing {
row.timing = row.timing.saturating_add(&timing);
run_timing = run_timing.saturating_add(&RunTiming::from(timing));
}
if !usage.is_zero() {
billed_visit_count += 1;
row.billing.add_counts(usage);
totals.add_counts(usage);
if let Some(model) = &stage.model {
row.model = Some(model.clone());
let model_entry =
by_model
.entry(model.clone())
.or_insert_with(|| ProjectionBillingByModel {
model: model.clone(),
stages: 0,
billing: BilledTokenCounts::default(),
});
model_entry.stages += 1;
model_entry.billing.add_counts(usage);
}
}
}
let mut by_model = by_model.into_values().collect::<Vec<_>>();
by_model.sort_by(|left, right| {
let left_provider = left.model.provider.to_string();
let right_provider = right.model.provider.to_string();
left_provider
.cmp(&right_provider)
.then_with(|| left.model.model_id.cmp(&right.model.model_id))
.then_with(|| {
left.model
.speed
.map(<&'static str>::from)
.cmp(&right.model.speed.map(<&'static str>::from))
})
});
ProjectionBillingRollup {
stages,
totals,
by_model,
timing: run_timing,
billed_visit_count,
}
}
fn stage_projection_order(state: &RunProjection) -> HashMap<String, u32> {
let mut order = HashMap::new();
for (stage_id, stage) in state.iter_stages() {
order
.entry(stage_id.node_id().to_string())
.and_modify(|first_seq: &mut u32| {
*first_seq = (*first_seq).min(stage.first_event_seq.get());
})
.or_insert_with(|| stage.first_event_seq.get());
}
order
}

View file

@ -3,6 +3,7 @@ extern crate self as fabro_types;
pub mod artifact;
pub mod auth;
pub mod billing;
pub mod billing_rollup;
pub mod blob_hash;
pub mod blob_ref;
pub mod checkpoint;