feat(planner): emit streaming and inference configs from the MILP plan

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

@@ -8,8 +8,8 @@ use std::path::PathBuf;
 
 use asap_planner::optimizer::{
     build_milp_workload, load_flat_atomic_cost_table, load_optional_selected_atomic_cost_table,
-    load_workload_facts, run_greedy_pipeline, solve_milp, AtomicCostTable, LabelSetFacts,
-    LabelSetFactsError,
+    load_workload_facts, plan_to_planner_output, reject_avg_queries, run_greedy_pipeline,
+    solve_milp, AtomicCostTable, LabelSetFacts, LabelSetFactsError,
 };
 use asap_planner::ControllerConfig;
 use clap::Parser;
@@ -59,11 +59,24 @@ struct Args {
     )]
     atomic_cost_workload: Option<PathBuf>,
 
-    /// Plan with sketch-bench's rqe-optimizer MILP and print the plan; writes
-    /// no configs yet.
+    /// Plan with sketch-bench's rqe-optimizer MILP and print the plan.
     #[arg(long)]
     milp: bool,
 
+    /// MILP only. Write `streaming_config.yaml` and `inference_config.yaml`
+    /// for the plan here.
+    #[arg(long = "output-dir", requires = "milp")]
+    output_dir: Option<PathBuf>,
+
+    /// MILP only. Also plan with families the engine can't deploy; prints the
+    /// plan and writes no configs.
+    #[arg(
+        long = "allow-undeployable-families",
+        requires = "milp",
+        conflicts_with = "output_dir"
+    )]
+    allow_undeployable_families: bool,
+
     /// MILP only. YAML workload facts: per metric, `cardinality` per label
     /// set, including the set of all its labels (the series count).
     #[arg(
@@ -187,8 +200,18 @@ fn run_milp(args: &Args, config: &ControllerConfig) -> anyhow::Result<()> {
     let objective = Objective::AUCCost { w_cpu, w_mem };
     tracing::debug!(?objective, cost_rows = costs.len(), "milp: inputs loaded");
 
+    // Fail before solving when the plan would be written but can't be.
+    if args.output_dir.is_some() {
+        reject_avg_queries(config)?;
+    }
     let workload = build_milp_workload(config, &facts, args.data_ingestion_interval_ms)?;
-    let solution = solve_milp(&workload, &facts, &costs, objective)?;
+    let solution = solve_milp(
+        &workload,
+        &facts,
+        &costs,
+        objective,
+        args.allow_undeployable_families,
+    )?;
 
     println!("=== Deployments: {} ===", solution.deployments.len());
     for (d, planned) in solution.deployments.iter().enumerate() {
@@ -237,6 +260,18 @@ fn run_milp(args: &Args, config: &ControllerConfig) -> anyhow::Result<()> {
             c.memory_bytes / 1e6
         );
     }
+
+    if let Some(dir) = &args.output_dir {
+        let output = plan_to_planner_output(config, &workload, &solution)?;
+        // Serialize both before writing either, so a failure can't leave a
+        // new streaming config next to a stale inference config.
+        let streaming = output.to_streaming_yaml_string()?;
+        let inference = output.to_inference_yaml_string()?;
+        std::fs::create_dir_all(dir)?;
+        std::fs::write(dir.join("streaming_config.yaml"), streaming)?;
+        std::fs::write(dir.join("inference_config.yaml"), inference)?;
+        println!("\nwrote configs to {}", dir.display());
+    }
     Ok(())
 }
 

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

@@ -157,6 +157,16 @@ fn decompose_to_leaves(query: &str) -> Vec<String> {
     leaves
 }
 
+/// Whether any leaf of `query` is an `avg` that `decompose_to_leaves` rewrites.
+pub(crate) fn contains_avg(query: &str) -> bool {
+    match parse_binary_arms(query) {
+        Some((lhs, rhs)) => [lhs, rhs]
+            .into_iter()
+            .any(|arm| matches!(arm, BinaryArm::Query(q) if contains_avg(&q))),
+        None => rewrite_avg(query).len() > 1,
+    }
+}
+
 /// Rewrite an `avg` leaf into the sum and count leaves it is computed from:
 /// `avg by (l) (x)` → `sum by (l) (x)`, `count by (l) (x)`, and
 /// `avg_over_time(x[5m])` → `sum_over_time(x[5m])`, `count_over_time(x[5m])`.

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

@@ -46,14 +46,17 @@ pub enum MilpError {
 }
 
 /// The cheapest plan. Only families in sketch-bench's `DEPLOYABLE_FAMILIES`
-/// are candidates.
+/// are candidates unless `allow_undeployable_families`, whose plan can be
+/// studied but not deployed.
 pub fn solve_milp(
     workload: &MilpWorkload,
     facts: &WorkloadFacts,
     costs: &[AtomicCostEntry],
     objective: Objective,
+    allow_undeployable_families: bool,
 ) -> Result<MilpSolution, MilpError> {
-    let deployments = build_all_candidates(&workload.raqes, costs, facts, false);
+    let deployments =
+        build_all_candidates(&workload.raqes, costs, facts, allow_undeployable_families);
     tracing::debug!(
         candidates = deployments.len(),
         cost_rows = costs.len(),
@@ -251,13 +254,13 @@ fn accuracy_target(
 }
 
 #[cfg(test)]
-mod tests {
+pub(super) mod tests {
     use std::collections::BTreeMap;
 
     use super::*;
     use crate::optimizer::workload_facts::parse_workload_facts;
 
-    const SCRAPE_MS: u64 = 15_000;
+    pub(crate) const SCRAPE_MS: u64 = 15_000;
 
     const FACTS: &str = r#"
 metrics:
@@ -271,20 +274,20 @@ metrics:
         cardinality: 1
 "#;
 
-    fn config(groups: &str) -> ControllerConfig {
+    pub(crate) fn config(groups: &str) -> ControllerConfig {
         let yaml = format!(
             "query_groups:\n{groups}\nmetrics:\n  - metric: http_requests_total\n    labels: [instance, job]\n"
         );
         serde_yaml::from_str(&yaml).unwrap()
     }
 
-    fn group(query: &str, accuracy_sla: f64) -> String {
+    pub(crate) fn group(query: &str, accuracy_sla: f64) -> String {
         format!(
             "  - queries: [\"{query}\"]\n    repetition_delay_ms: 60000\n    controller_options: {{accuracy_sla: {accuracy_sla}}}\n"
         )
     }
 
-    fn facts(config: &ControllerConfig) -> WorkloadFacts {
+    pub(crate) fn facts(config: &ControllerConfig) -> WorkloadFacts {
         parse_workload_facts(FACTS, config.metrics.as_deref().unwrap(), SCRAPE_MS).unwrap()
     }
 
@@ -311,7 +314,7 @@ metrics:
         }
     }
 
-    fn costs() -> Vec<AtomicCostEntry> {
+    pub(crate) fn costs() -> Vec<AtomicCostEntry> {
         vec![
             cost("exact-sum", &[(RELATIVE_ERROR, 0.0)]),
             cost("kll-percall", &[(MAX_RANK_ERROR, 0.005)]),
@@ -469,7 +472,7 @@ metrics:
         let config = config(&(group(query, 0.99) + &group(query, 0.99)));
         let facts = facts(&config);
         let w = build_milp_workload(&config, &facts, SCRAPE_MS).unwrap();
-        let solution = solve_milp(&w, &facts, &costs(), Objective::default()).unwrap();
+        let solution = solve_milp(&w, &facts, &costs(), Objective::default(), false).unwrap();
         assert_eq!(solution.deployments.len(), 1);
         assert_eq!(solution.raqes[0].deployment, solution.raqes[1].deployment);
         assert_eq!(
@@ -487,7 +490,7 @@ metrics:
         ));
         let facts = facts(&config);
         let w = build_milp_workload(&config, &facts, SCRAPE_MS).unwrap();
-        let err = solve_milp(&w, &facts, &costs(), Objective::default()).unwrap_err();
+        let err = solve_milp(&w, &facts, &costs(), Objective::default(), false).unwrap_err();
         assert!(matches!(err, MilpError::Unservable(ids) if ids.len() == 1));
     }
 
@@ -498,7 +501,7 @@ metrics:
         let facts = facts(&config);
         let w = build_milp_workload(&config, &facts, SCRAPE_MS).unwrap();
         assert_eq!(w.raqes.len(), 2);
-        let err = solve_milp(&w, &facts, &costs(), Objective::default()).unwrap_err();
+        let err = solve_milp(&w, &facts, &costs(), Objective::default(), false).unwrap_err();
         assert!(matches!(err, MilpError::Unservable(ids) if ids == [w.raqes[0].id.clone()]));
     }
 
@@ -510,7 +513,7 @@ metrics:
         let config = config(&groups);
         let facts = facts(&config);
         let w = build_milp_workload(&config, &facts, SCRAPE_MS).unwrap();
-        let solution = solve_milp(&w, &facts, &costs(), Objective::default()).unwrap();
+        let solution = solve_milp(&w, &facts, &costs(), Objective::default(), false).unwrap();
         let chosen: BTreeMap<Capability, &str> = w
             .raqes
             .iter()
@@ -538,7 +541,7 @@ metrics:
         let caps: Vec<_> = w.raqes.iter().map(|r| r.capability).collect();
         assert_eq!(caps.len(), 2);
         assert!(caps.contains(&Capability::Sum) && caps.contains(&Capability::Count));
-        let solution = solve_milp(&w, &facts, &costs(), Objective::default()).unwrap();
+        let solution = solve_milp(&w, &facts, &costs(), Objective::default(), false).unwrap();
         assert_eq!(solution.deployments.len(), 2);
     }
 }