feat(planner): build rqe-optimizer Raqes and facts from the workload config

asap-planner-rs/src/bin/optimizer_cli.rs

@@ -7,10 +7,12 @@
 use std::path::PathBuf;
 
 use asap_planner::optimizer::{
-    load_optional_selected_atomic_cost_table, run_greedy_pipeline, AtomicCostTable, LabelSetFacts,
+    build_milp_workload, load_optional_selected_atomic_cost_table, load_workload_facts,
+    run_greedy_pipeline, solve_milp, AtomicCostTable, LabelSetFacts,
 };
 use asap_planner::ControllerConfig;
 use clap::Parser;
+use rqe_optimizer::milp::Objective;
 
 #[derive(Parser, Debug)]
 #[command(
@@ -27,26 +29,53 @@ struct Args {
     #[arg(long = "data-ingestion-interval-ms", value_parser = clap::value_parser!(u64).range(1..))]
     data_ingestion_interval_ms: u64,
 
-    /// YAML label-set facts: `series_count` per (metric, spatial filter) and
-    /// `cardinality` per (metric, spatial filter, grouping labels).
-    #[arg(long = "label-set-facts")]
-    label_set_facts: PathBuf,
-
-    /// Path to the versioned atomic-cost document sketch-bench's `atomic-costs`
-    /// subcommand exports. Requires --atomic-cost-workload to select exactly
-    /// one measured workload profile. Omitted: every
-    /// benchmarked-family candidate (CMS/HLL/KLL) is dropped, since there is
-    /// no data to cost it at — only trivial accumulators and EXACT remain
-    /// selectable.
-    #[arg(long = "atomic-costs")]
+    /// Greedy only. YAML label-set facts: `series_count` per (metric, spatial
+    /// filter) and `cardinality` per (metric, spatial filter, grouping labels).
+    #[arg(long = "label-set-facts", required_unless_present = "milp")]
+    label_set_facts: Option<PathBuf>,
+
+    /// Greedy: the versioned atomic-cost document sketch-bench's
+    /// `atomic-costs` subcommand exports; requires --atomic-cost-workload.
+    /// Omitted: every benchmarked-family candidate (CMS/HLL/KLL) is dropped,
+    /// leaving only trivial accumulators and EXACT.
+    /// MILP: the flat cost table `export_rqe_optimizer_costs.sh` writes
+    /// (`rqe_atomic_costs.json`); required.
+    #[arg(long = "atomic-costs", required_if_eq("milp", "true"))]
     atomic_costs: Option<PathBuf>,
 
-    /// JSON `profiles[].workload` value copied from the sketch-bench atomic-cost
-    /// document. This makes the empirical workload profile explicit and avoids
-    /// mixing costs from different traces or time windows.
-    #[arg(long = "atomic-cost-workload", requires = "atomic_costs")]
+    /// Greedy only. JSON `profiles[].workload` value copied from the
+    /// sketch-bench atomic-cost document. This makes the empirical workload
+    /// profile explicit and avoids mixing costs from different traces or time
+    /// windows.
+    #[arg(
+        long = "atomic-cost-workload",
+        requires = "atomic_costs",
+        conflicts_with = "milp"
+    )]
     atomic_cost_workload: Option<PathBuf>,
 
+    /// Plan with sketch-bench's rqe-optimizer MILP and print the plan; writes
+    /// no configs yet.
+    #[arg(long)]
+    milp: bool,
+
+    /// MILP only. YAML workload facts: per metric, `cardinality` per label
+    /// set, including the set of all its labels (the series count).
+    #[arg(
+        long = "workload-facts",
+        required_if_eq("milp", "true"),
+        requires = "milp"
+    )]
+    workload_facts: Option<PathBuf>,
+
+    /// MILP only. Objective weight on CPU-sec/sec. Default: rqe-optimizer's.
+    #[arg(long = "w-cpu", requires = "milp")]
+    w_cpu: Option<f64>,
+
+    /// MILP only. Objective weight on memory GiB. Default: rqe-optimizer's.
+    #[arg(long = "w-mem", requires = "milp")]
+    w_mem: Option<f64>,
+
     #[arg(short, long, action = clap::ArgAction::Count)]
     verbose: u8,
 }
@@ -64,7 +93,14 @@ fn main() -> anyhow::Result<()> {
 
     let yaml_str = std::fs::read_to_string(&args.input_config)?;
     let config: ControllerConfig = serde_yaml::from_str(&yaml_str)?;
-    let facts = LabelSetFacts::from_path(&args.label_set_facts)?;
+    if args.milp {
+        return run_milp(&args, &config);
+    }
+    let facts = LabelSetFacts::from_path(
+        args.label_set_facts
+            .as_deref()
+            .expect("clap requires --label-set-facts without --milp"),
+    )?;
 
     let atomic_cost_table = match load_optional_selected_atomic_cost_table(
         args.atomic_costs.as_deref(),
@@ -108,3 +144,76 @@ fn main() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+fn run_milp(args: &Args, config: &ControllerConfig) -> anyhow::Result<()> {
+    let hints = config.metrics.as_deref().unwrap_or_default();
+    let facts = load_workload_facts(
+        args.workload_facts
+            .as_deref()
+            .expect("clap requires --workload-facts with --milp"),
+        hints,
+        args.data_ingestion_interval_ms,
+    )?;
+    let costs_path = args
+        .atomic_costs
+        .as_deref()
+        .expect("clap requires --atomic-costs with --milp");
+    let costs: AtomicCostTable = serde_json::from_str(&std::fs::read_to_string(costs_path)?)
+        .map_err(|e| anyhow::anyhow!("parsing cost table {}: {e}", costs_path.display()))?;
+    let Objective::AUCCost { w_cpu, w_mem } = Objective::default();
+    let objective = Objective::AUCCost {
+        w_cpu: args.w_cpu.unwrap_or(w_cpu),
+        w_mem: args.w_mem.unwrap_or(w_mem),
+    };
+    tracing::debug!(?objective, cost_rows = costs.len(), "milp: inputs loaded");
+
+    let workload = build_milp_workload(config, &facts, args.data_ingestion_interval_ms)?;
+    let (deployments, solution) = solve_milp(&workload, &facts, &costs, objective)?;
+
+    let mut active: Vec<usize> = solution.mapping.clone();
+    active.sort();
+    active.dedup();
+    println!("=== Deployments: {} ===", active.len());
+    for &d in &active {
+        let dep = &deployments[d];
+        println!(
+            "  [{d}] {:?} {} config={} metric={} grouping={:?} window={}ms slide={}ms",
+            dep.capability,
+            dep.config.sketch,
+            dep.config.sketch_config,
+            dep.metric,
+            dep.grouping_labels,
+            dep.window_ms,
+            dep.slide_ms,
+        );
+    }
+    println!("\n=== Raqes: {} ===", workload.raqes.len());
+    for ((raqe, &d), latency_ms) in workload
+        .raqes
+        .iter()
+        .zip(&solution.mapping)
+        .zip(&solution.plan_cost.query_latency_ms)
+    {
+        println!("  {} -> [{d}] latency={latency_ms:.3e}ms", raqe.id);
+    }
+    let cost = &solution.plan_cost;
+    println!(
+        "\nobjective={:.6e} cpu={:.6e} cpu-sec/sec memory={:.3} MB",
+        objective.value(cost),
+        cost.cpu_secs_per_sec(),
+        cost.memory_bytes() / 1e6,
+    );
+    for (phase, c) in [
+        ("ingest", &cost.ingest),
+        ("merge", &cost.merge),
+        ("query", &cost.query),
+        ("storage", &cost.storage),
+    ] {
+        println!(
+            "  {phase}: cpu={:.6e} cpu-sec/sec memory={:.3} MB",
+            c.cpu_secs_per_sec,
+            c.memory_bytes / 1e6
+        );
+    }
+    Ok(())
+}

asap-planner-rs/src/optimizer/atomic_costs.rs

@@ -1,14 +1,8 @@
 //! The atomic-cost table exported by sketch-bench (sketch-bench#30,
 //! `scripts/export_atomic_costs.sh`), and the (sketch_type, params) lookup
 //! that resolves a candidate's [`AtomicCosts`] from it.
-//!
-//! `AtomicCostEntry`/`AtomicCostTable` are a deliberate duplicate of
-//! sketch-bench's `aqpbm_core::atomic_costs` types, not a shared dependency —
-//! see ASAPQuery#524 and sketch-bench#30 for why. Keep the two in sync by
-//! hand; `atomic_cost_entry_deserializes_sketch_benchs_documented_shape`
-//! below is a canary for drift.
-
-use std::collections::{BTreeMap, HashMap};
+
+use std::collections::HashMap;
 use std::path::Path;
 
 use promql_utilities::query_logics::enums::AggregationType;
@@ -73,25 +67,7 @@ pub struct ExternalWorkload {
     pub timestamp_unit: String,
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
-#[serde(deny_unknown_fields)]
-pub struct AtomicCostEntry {
-    pub sketch: String,
-    pub sketch_config: Value,
-    pub mem_bytes_per_instance: f64,
-    pub insert_cpu_secs: f64,
-    pub merge_cpu_secs: f64,
-    pub query_cpu_secs: f64,
-    pub query_accuracy: BTreeMap<String, f64>,
-    /// The conditions sketch-bench measured the row under (items, keys and
-    /// value range per instance, merge operand size, distribution;
-    /// sketch-bench#147). Carried opaquely, like a synthetic workload
-    /// description; absent in tables written before it existed.
-    #[serde(default, skip_serializing_if = "Option::is_none")]
-    pub measured_at: Option<Value>,
-}
-
-pub type AtomicCostTable = Vec<AtomicCostEntry>;
+pub use rqe_optimizer::{AtomicCostEntry, AtomicCostTable};
 
 /// Parse a standalone JSON workload selector. The selector is the exact
 /// `profiles[].workload` value copied from the benchmark artifact, making the
@@ -472,6 +448,8 @@ fn valid_cost_entry(entry: &AtomicCostEntry) -> bool {
 
 #[cfg(test)]
 mod tests {
+    use std::collections::BTreeMap;
+
     use super::*;
 
     #[test]
@@ -578,21 +556,17 @@ mod tests {
         assert!(serde_json::from_str::<AtomicCostEntry>(json).is_err());
     }
 
-    /// sketch-bench#147 adds `measured_at`; it loads, and older rows without
-    /// it still do. Other unknown fields are still refused.
+    /// Rows with `measured_at` load, and older rows without it still do.
     #[test]
     fn atomic_cost_entry_accepts_optional_measured_at() {
         let base = r#""sketch":"kll-percall","sketch_config":null,"mem_bytes_per_instance":1.0,"insert_cpu_secs":1.0,"merge_cpu_secs":1.0,"query_cpu_secs":1.0,"query_accuracy":{}"#;
         let with = format!(
             r#"{{{base},"measured_at":{{"items_per_instance":1000000,"keys_per_instance":100000,"value_range":[1.0,100000.0],"merge_operand_items":62500,"distribution":{{"kind":"zipf","skewness":1.1,"population_size":100000,"seed":42}}}}}}"#
         );
         let entry: AtomicCostEntry = serde_json::from_str(&with).unwrap();
-        assert_eq!(entry.measured_at.unwrap()["items_per_instance"], 1_000_000);
+        assert_eq!(entry.measured_at.unwrap().items_per_instance, 1_000_000);
         let without: AtomicCostEntry = serde_json::from_str(&format!("{{{base}}}")).unwrap();
         assert!(without.measured_at.is_none());
-        assert!(
-            serde_json::from_str::<AtomicCostEntry>(&format!(r#"{{{base},"other":1}}"#)).is_err()
-        );
     }
 
     #[test]
@@ -617,6 +591,7 @@ mod tests {
             merge_cpu_secs: 4.5e-4,
             query_cpu_secs: 7.8e-8,
             query_accuracy: BTreeMap::new(),
+            merge_accuracy: BTreeMap::new(),
             measured_at: None,
         }
     }
@@ -640,6 +615,7 @@ mod tests {
             merge_cpu_secs: 4.0,
             query_cpu_secs: 8.0,
             query_accuracy: BTreeMap::new(),
+            merge_accuracy: BTreeMap::new(),
             measured_at: None,
         }
     }
@@ -652,15 +628,6 @@ mod tests {
         ])
     }
 
-    #[test]
-    fn atomic_cost_entry_deserializes_sketch_benchs_documented_shape() {
-        // Pinned against the current sketch-bench atomic-cost entry shape.
-        let json = r#"{"sketch":"cms-fastpath-vector2d","sketch_config":{"algorithm":"cms-fastpath-vector2d","params":{"cols":1024,"rows":3}},"mem_bytes_per_instance":12288.0,"insert_cpu_secs":8.484689139741214e-9,"merge_cpu_secs":0.00045364040539336466,"query_cpu_secs":7.799774697708031e-8,"query_accuracy":{"relative_error":0.01}}"#;
-        let entry: AtomicCostEntry = serde_json::from_str(json).expect("documented shape parses");
-        assert_eq!(entry.sketch, "cms-fastpath-vector2d");
-        assert_eq!(entry.mem_bytes_per_instance, 12288.0);
-    }
-
     #[test]
     fn cms_candidate_resolves_by_exact_key_regardless_of_value_type() {
         // ASAPQuery's grid stores depth/width as u64; sketch-bench's exported
@@ -817,6 +784,7 @@ mod tests {
             merge_cpu_secs: 2.76e-4,
             query_cpu_secs: 1.23e-4,
             query_accuracy: BTreeMap::new(),
+            merge_accuracy: BTreeMap::new(),
             measured_at: None,
         }];
         let hll_params = HashMap::from([("precision".to_string(), Value::from(14u64))]);
@@ -830,6 +798,7 @@ mod tests {
             merge_cpu_secs: 1.0e-3,
             query_cpu_secs: 1.6e-4,
             query_accuracy: BTreeMap::new(),
+            merge_accuracy: BTreeMap::new(),
             measured_at: None,
         }];
         let kll_params = HashMap::from([("K".to_string(), Value::from(200u64))]);

asap-planner-rs/src/optimizer/greedy.rs

@@ -237,6 +237,7 @@ mod tests {
             merge_cpu_secs: 0.0,
             query_cpu_secs: 0.0,
             query_accuracy: std::collections::BTreeMap::new(),
+            merge_accuracy: std::collections::BTreeMap::new(),
             measured_at: None,
         }];
         let aqe = make_aqe(Statistic::Topk, 60_000, 60_000, 1.0 / 60.0);
@@ -276,6 +277,7 @@ mod tests {
             merge_cpu_secs: 0.0,
             query_cpu_secs: 0.0,
             query_accuracy: std::collections::BTreeMap::new(),
+            merge_accuracy: std::collections::BTreeMap::new(),
             measured_at: None,
         }];
         let mut aqe = make_aqe(Statistic::Sum, 60_000, 60_000, 1.0 / 60.0);