mirror of
https://github.com/fabro-sh/fabro.git
synced 2026-09-07 08:27:12 +00:00
Rebuild conclusion stage summaries from run events
This commit is contained in:
parent
efa0d16588
commit
a7ed03e175
5 changed files with 442 additions and 275 deletions
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
217
lib/foundation/fabro-types/src/billing_rollup.rs
Normal file
217
lib/foundation/fabro-types/src/billing_rollup.rs
Normal 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
|
||||
}
|
||||
|
|
@ -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;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue