feat(devtools): stage_pipeline --price-enumerated prices each enumerated memo choice and flags any cheaper than the selected plan

crates/devtools/src/bin/stage_pipeline.rs

@@ -135,7 +135,17 @@
 // the last step. Each choice has its `id` (`E1`, ...), `choice`, `dag`
 // (`memo::enumeration::compose`; null with `error` when it does not
 // compose), `budget_share_1` (every chosen candidate sized for the whole
-// budget) and `selected` (it is `selected_choice`). `priced` is null.
+// budget) and `selected` (it is `selected_choice`). With
+// `--price-enumerated <n>` (off by default), each of the first `n` listed
+// choices is also run through Stage 2 and Stage 3 as the selected plan is
+// (`memo_pricing::price_choices`, under the run's models): `valid`,
+// `rejection` (a Stage 3 rejection cause, or `compose_failed`,
+// `invalid_structure`, `stage2_failed`, `stage3_failed`, `rejected`) with
+// `rejection_reason`,
+// `cost_per_second` (its cheapest valid physical plan's) and
+// `cheaper_than_selected` (valid and cheaper than the selected plan); and
+// `priced` is `{count, cheaper_than_selected, selected_cost_per_second}`.
+// Without it `priced` is null.
 // `workload.data_workload` is the `DataWorkload` the run planned with, as
 // declared (`--workload`, `--example`) or built from `--sample-interval-ms`
 // (`--promql`); each query also carries its `recurrence` and
@@ -245,6 +255,7 @@ use std::collections::{BTreeSet, HashSet};
 use std::rc::Rc;
 
 use asap_devtools::candidate_labels::{candidate_labels, computation_label, query_of_each_target};
+use asap_devtools::memo_pricing::{price_choices, SelectedBaseline};
 use asap_devtools::planner_layering_examples::{
     example1, example2, pattern_a, pattern_b, PatternARecurrence, SHARED_ROWS_PER_SECOND,
 };
@@ -283,7 +294,7 @@ const USAGE: &str =
 [--cost-model usage-over-time|aws-fargate-us-east-1|provisioned-capacity|cluster-capacity [--cluster-memory-bytes <n> \
 --cluster-disk-bytes <n> [--cluster-cpu-cores <n>] --cluster-objective lowest-latency|least-memory|pareto-front]] \
 [--plan-selection-settings <file>] [--enforce-accuracy-guarantees] [--deployment-capabilities <json>] [--exhaustive [--max-combinations <n>] | --search candidate-connections] [--milp-time-limit-s <f64>] [--max-displayed <n>] \
-[--no-memo] [--max-memo-candidates <n>] [--memo-enumeration-limit <n>] [--max-enumerated <n>] --out <file>\n\
+[--no-memo] [--max-memo-candidates <n>] [--memo-enumeration-limit <n>] [--max-enumerated <n>] [--price-enumerated <n>] --out <file>\n\
 --search candidate-connections is the default search.";
 
 /// Default evaluator-assignment budget for an exhaustive archive run.
@@ -382,6 +393,7 @@ fn run(args: Vec<String>) -> Result<(), Failure> {
         candidates_per_group: 64,
         enumeration_limit: 1000,
         enumerated: 200,
+        choices_to_price: 0,
     };
     let mut args = args.into_iter();
     while let Some(flag) = args.next() {
@@ -454,6 +466,9 @@ fn run(args: Vec<String>) -> Result<(), Failure> {
             "--max-enumerated" => {
                 memo_caps.enumerated = value()?.parse().map_err(|e| format!("{e}"))?
             }
+            "--price-enumerated" => {
+                memo_caps.choices_to_price = value()?.parse().map_err(|e| format!("{e}"))?
+            }
             "--search" => {
                 search_algorithm = Some(match value()?.as_str() {
                     "candidate-connections" => SearchAlgorithm::CandidateConnections,
@@ -576,6 +591,23 @@ fn run(args: Vec<String>) -> Result<(), Failure> {
         },
     };
     plan_selection.ignore_accuracy_guarantees = !enforce_accuracy_guarantees;
+    if memo_caps.choices_to_price > 0 {
+        if !write_memo {
+            return Err("--price-enumerated needs the memo (drop --no-memo)".into());
+        }
+        // The exhaustive search reports no calibration factor to price each
+        // choice on the selection's scale.
+        if exhaustive
+            && plan_selection.frozen_calibration_factor.is_none()
+            && !plan_selection.measured_operation_costs.operators.is_empty()
+        {
+            return Err(
+                "--price-enumerated with --exhaustive needs a frozen_calibration_factor \
+                        when measured operation costs are given"
+                    .into(),
+            );
+        }
+    }
     let from_files = if workload_paths.is_empty() {
         None
     } else if example.is_none() && queries.is_empty() {
@@ -743,7 +775,20 @@ fn run(args: Vec<String>) -> Result<(), Failure> {
         .ok()
         .and_then(|run| run.memo_steps.as_ref())
         .filter(|_| write_memo)
-        .map(|steps| memo_json(steps, &memo_caps));
+        .map(|steps| {
+            let pricing = Pricing {
+                stage0_roots: &roots,
+                entries: &entries,
+                data: &data,
+                models,
+                selected: result
+                    .as_ref()
+                    .ok()
+                    .and_then(|run| run.stage3_plan_selection.as_ref().ok())
+                    .map(SelectedBaseline::of),
+            };
+            memo_json(steps, &memo_caps, &pricing)
+        });
     let mut document = document_of(described, &roots, &entries, result, max_displayed)?;
     if let Some(memo) = memo {
         document["memo"] = memo;
@@ -1447,18 +1492,30 @@ fn query_batch(
 }
 
 /// How much of the memo the document writes (`--max-memo-candidates`,
-/// `--memo-enumeration-limit`, `--max-enumerated`).
+/// `--memo-enumeration-limit`, `--max-enumerated`) and how many listed
+/// choices it prices (`--price-enumerated`).
 struct MemoCaps {
     candidates_per_group: usize,
     enumeration_limit: usize,
     enumerated: usize,
+    choices_to_price: usize,
+}
+
+/// What pricing an enumerated choice needs: the run's workload, models and
+/// selection.
+struct Pricing<'a> {
+    stage0_roots: &'a [QueryRoot],
+    entries: &'a [QueryWorkloadEntry],
+    data: &'a DataWorkload,
+    models: PlanningModels<'a>,
+    selected: Option<SelectedBaseline>,
 }
 
 /// The document's `memo` section: the memo after each step, the selected
 /// plan's choice over the last step's memo, and the candidate DAGs each
 /// step's memo stands for (`memo::enumeration::enumerate`; each listed
 /// choice composed by `memo::enumeration::compose`).
-fn memo_json(steps: &MemoSteps, caps: &MemoCaps) -> Value {
+fn memo_json(steps: &MemoSteps, caps: &MemoCaps, pricing: &Pricing<'_>) -> Value {
     let named = [
         (
             "stage0",
@@ -1506,8 +1563,19 @@ fn memo_json(steps: &MemoSteps, caps: &MemoCaps) -> Value {
     let last = &steps.concrete_states;
     let (step, memo) = listed.unwrap_or(("concrete_states", last));
     let selected = steps.selected_choice.as_ref().ok();
-    let choices: Vec<_> = enumerate(memo, caps.enumerated)
-        .choices
+    let enumerated = enumerate(memo, caps.enumerated).choices;
+    let priced = (caps.choices_to_price > 0).then(|| {
+        price_choices(
+            memo,
+            &enumerated[..caps.choices_to_price.min(enumerated.len())],
+            pricing.stage0_roots,
+            pricing.entries,
+            pricing.data,
+            pricing.models,
+            pricing.selected,
+        )
+    });
+    let mut choices: Vec<_> = enumerated
         .iter()
         .enumerate()
         .map(|(index, choice)| {
@@ -1535,6 +1603,20 @@ fn memo_json(steps: &MemoSteps, caps: &MemoCaps) -> Value {
             entry
         })
         .collect();
+    for (entry, priced) in choices.iter_mut().zip(priced.iter().flatten()) {
+        entry["valid"] = json!(priced.valid);
+        entry["rejection"] = json!(priced.rejection);
+        entry["rejection_reason"] = json!(priced.rejection_reason);
+        entry["cost_per_second"] = json!(priced.cost_per_second);
+        entry["cheaper_than_selected"] = json!(priced.cheaper_than_selected);
+    }
+    let priced = priced.map_or(Value::Null, |priced| {
+        json!({
+            "count": priced.len(),
+            "cheaper_than_selected": priced.iter().filter(|p| p.cheaper_than_selected).count(),
+            "selected_cost_per_second": pricing.selected.and_then(|s| s.cost_per_second),
+        })
+    });
     let mut memo = json!({
         "steps": json_steps,
         "selected_choice": selected.map_or(Value::Null, choice_json),
@@ -1543,7 +1625,7 @@ fn memo_json(steps: &MemoSteps, caps: &MemoCaps) -> Value {
             "per_step": per_step,
             "limit": caps.enumeration_limit,
             "choices": choices,
-            "priced": null,
+            "priced": priced,
         },
     });
     if let Err(reason) = &steps.selected_choice {

crates/devtools/src/lib.rs

@@ -21,6 +21,7 @@
 //! missed reuse and sharing behind the `planner_fuzz` bin.
 
 pub mod candidate_labels;
+pub mod memo_pricing;
 pub mod planner_fuzz;
 pub mod planner_layering_examples;
 pub mod planning_stages;