feat(planner): add asap-planner --planner legacy|milp

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

@@ -4,12 +4,10 @@
 use std::path::PathBuf;
 
 use asap_planner::optimizer::{
-    build_milp_workload, load_flat_atomic_cost_table, load_workload_facts, plan_to_planner_output,
-    reject_avg_queries, solve_milp, MilpError,
+    parse_weight, plan_milp, plan_to_planner_output, reject_avg_queries, MilpInputs, MilpPlan,
 };
 use asap_planner::ControllerConfig;
 use clap::Parser;
-use rqe_optimizer::milp::Objective;
 
 #[derive(Parser, Debug)]
 #[command(
@@ -74,44 +72,25 @@ fn main() -> anyhow::Result<()> {
     run_milp(&args, &config)
 }
 
-/// Objective weights must be finite and non-negative: a negative weight
-/// rewards cost, and NaN poisons every coefficient.
-fn parse_weight(s: &str) -> Result<f64, String> {
-    match s.parse::<f64>() {
-        Ok(w) if w.is_finite() && w >= 0.0 => Ok(w),
-        Ok(w) => Err(format!("must be finite and >= 0, got {w}")),
-        Err(e) => Err(e.to_string()),
-    }
-}
-
 fn run_milp(args: &Args, config: &ControllerConfig) -> anyhow::Result<()> {
-    config.warn_default_slas();
-    let Some(hints) = config.metrics.as_deref() else {
-        return Err(MilpError::MissingMetricHints.into());
-    };
-    let facts = load_workload_facts(&args.workload_facts, hints, args.data_ingestion_interval_ms)?;
-    let costs = load_flat_atomic_cost_table(&args.atomic_costs)?;
-    let Objective::AUCCost { w_cpu, w_mem } = Objective::default();
-    let (w_cpu, w_mem) = (args.w_cpu.unwrap_or(w_cpu), args.w_mem.unwrap_or(w_mem));
-    // All-zero weights make every plan cost 0, so the solver's pick is arbitrary.
-    anyhow::ensure!(
-        w_cpu > 0.0 || w_mem > 0.0,
-        "--w-cpu and --w-mem are both 0; at least one must be positive"
-    );
-    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,
+    let MilpPlan {
+        workload,
+        solution,
         objective,
-        args.allow_undeployable_families,
+    } = plan_milp(
+        config,
+        &MilpInputs {
+            workload_facts: &args.workload_facts,
+            atomic_costs: &args.atomic_costs,
+            scrape_interval_ms: args.data_ingestion_interval_ms,
+            w_cpu: args.w_cpu,
+            w_mem: args.w_mem,
+            allow_undeployable_families: args.allow_undeployable_families,
+        },
     )?;
 
     println!("=== Deployments: {} ===", solution.deployments.len());
@@ -163,29 +142,8 @@ fn run_milp(args: &Args, config: &ControllerConfig) -> anyhow::Result<()> {
     }
 
     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)?;
+        plan_to_planner_output(config, &workload, &solution)?.write_to_dir(dir)?;
         println!("\nwrote configs to {}", dir.display());
     }
     Ok(())
 }
-
-#[cfg(test)]
-mod tests {
-    use super::parse_weight;
-
-    #[test]
-    fn weights_must_be_finite_and_non_negative() {
-        assert_eq!(parse_weight("0.5"), Ok(0.5));
-        assert_eq!(parse_weight("0"), Ok(0.0));
-        for bad in ["-1", "NaN", "inf", "x"] {
-            assert!(parse_weight(bad).is_err(), "{bad}");
-        }
-    }
-}

asap-planner-rs/src/main.rs

@@ -1,6 +1,9 @@
+use asap_planner::optimizer::{
+    parse_weight, plan_milp, plan_to_planner_output, reject_avg_queries, MilpInputs,
+};
 use asap_planner::{
-    Controller, ElasticController, ElasticRuntimeOptions, RuntimeOptions, SQLController,
-    SQLRuntimeOptions, StreamingEngine,
+    Controller, ControllerConfig, ElasticController, ElasticRuntimeOptions, RuntimeOptions,
+    SQLController, SQLRuntimeOptions, StreamingEngine,
 };
 use asap_types::enums::QueryLanguage;
 use clap::Parser;
@@ -53,6 +56,29 @@ struct Args {
     #[arg(long = "clickhouse-database", required = false)]
     clickhouse_database: Option<String>,
 
+    /// `milp` plans with sketch-bench's rqe-optimizer: PromQL with
+    /// --input_config only, labels from its `metrics:` hints.
+    #[arg(long, value_enum, default_value = "legacy")]
+    planner: PlannerArg,
+
+    /// MILP only. YAML workload facts: per metric, positive `value_range` and
+    /// `cardinality` per label set, including all labels (the series count).
+    #[arg(long = "workload-facts", required_if_eq("planner", "milp"))]
+    workload_facts: Option<PathBuf>,
+
+    /// MILP only. The flat cost table `export_rqe_optimizer_costs.sh` writes
+    /// (`rqe_atomic_costs.json`).
+    #[arg(long = "atomic-costs", required_if_eq("planner", "milp"))]
+    atomic_costs: Option<PathBuf>,
+
+    /// MILP only. Objective weight on CPU-sec/sec. Default: rqe-optimizer's.
+    #[arg(long = "w-cpu", value_parser = parse_weight)]
+    w_cpu: Option<f64>,
+
+    /// MILP only. Objective weight on memory GiB. Default: rqe-optimizer's.
+    #[arg(long = "w-mem", value_parser = parse_weight)]
+    w_mem: Option<f64>,
+
     #[arg(short, long, action = clap::ArgAction::Count)]
     verbose: u8,
 }
@@ -62,6 +88,12 @@ enum EngineArg {
     Precompute,
 }
 
+#[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
+enum PlannerArg {
+    Legacy,
+    Milp,
+}
+
 fn main() -> anyhow::Result<()> {
     let args = Args::parse();
 
@@ -77,6 +109,19 @@ fn main() -> anyhow::Result<()> {
         EngineArg::Precompute => StreamingEngine::Precompute,
     };
 
+    if args.planner == PlannerArg::Milp {
+        run_milp(&args)?;
+        println!("Generated configs in {}", args.output_dir.display());
+        return Ok(());
+    }
+    anyhow::ensure!(
+        args.workload_facts.is_none()
+            && args.atomic_costs.is_none()
+            && args.w_cpu.is_none()
+            && args.w_mem.is_none(),
+        "--workload-facts, --atomic-costs, --w-cpu and --w-mem require --planner milp"
+    );
+
     match args.query_language {
         QueryLanguage::promql => {
             let scrape_interval_ms = args.data_ingestion_interval_ms.ok_or_else(|| {
@@ -158,3 +203,50 @@ fn main() -> anyhow::Result<()> {
     println!("Generated configs in {}", args.output_dir.display());
     Ok(())
 }
+
+fn run_milp(args: &Args) -> anyhow::Result<()> {
+    anyhow::ensure!(
+        matches!(args.query_language, QueryLanguage::promql),
+        "--planner milp supports only --query-language promql"
+    );
+    anyhow::ensure!(
+        args.query_log.is_none(),
+        "--planner milp needs --input_config: query logs have no `metrics:` hints"
+    );
+    anyhow::ensure!(
+        args.prometheus_url.is_none(),
+        "--planner milp takes labels from the config's `metrics:` hints, not --prometheus-url"
+    );
+    anyhow::ensure!(
+        !args.enable_punting && args.range_duration_ms == 0 && args.step_ms == 0,
+        "--enable-punting, --range-duration-ms and --step-ms don't apply to --planner milp"
+    );
+    let config_path = args
+        .input_config
+        .as_deref()
+        .ok_or_else(|| anyhow::anyhow!("--planner milp requires --input_config"))?;
+    let scrape_interval_ms = args.data_ingestion_interval_ms.ok_or_else(|| {
+        anyhow::anyhow!("--data-ingestion-interval-ms is required for PromQL mode")
+    })?;
+    let config: ControllerConfig = serde_yaml::from_str(&std::fs::read_to_string(config_path)?)?;
+    // Fail before solving: an avg query can't be written as a config.
+    reject_avg_queries(&config)?;
+    let plan = plan_milp(
+        &config,
+        &MilpInputs {
+            workload_facts: args
+                .workload_facts
+                .as_deref()
+                .expect("clap requires --workload-facts with --planner milp"),
+            atomic_costs: args
+                .atomic_costs
+                .as_deref()
+                .expect("clap requires --atomic-costs with --planner milp"),
+            scrape_interval_ms,
+            w_cpu: args.w_cpu,
+            w_mem: args.w_mem,
+            allow_undeployable_families: false,
+        },
+    )?;
+    plan_to_planner_output(&config, &plan.workload, &plan.solution)?.write_to_dir(&args.output_dir)
+}

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

@@ -12,16 +12,25 @@ use thiserror::Error;
 
 use crate::config::input::ControllerConfig;
 
+use std::path::Path;
+
 use super::aqe_extractor::{extract_aqes, RQE};
+use super::atomic_costs::load_flat_atomic_cost_table;
 use super::error::OptimizerError;
 use super::solution::OptimizerItem;
+use super::workload_facts::load_workload_facts;
 
 /// Slack on accuracy tolerances so `1 - sla` rounding (`1 - 0.9 =
 /// 0.0999...98`) doesn't reject a row measured exactly at the boundary.
 const SLA_EPSILON: f64 = 1e-9;
 
 #[derive(Debug, Error)]
 pub enum MilpError {
+    #[error(
+        "query groups set step_ms or range_duration_ms, which the MILP planner doesn't \
+         support yet (#800): {0:?}"
+    )]
+    RangeQueryOverrides(Vec<String>),
     #[error(
         "workload config has no `metrics:` hints; they are required to resolve grouping labels"
     )]
@@ -100,6 +109,73 @@ pub fn solve_milp(
     Ok(solution)
 }
 
+/// What a MILP planning run reads besides the workload config.
+pub struct MilpInputs<'a> {
+    pub workload_facts: &'a Path,
+    pub atomic_costs: &'a Path,
+    pub scrape_interval_ms: u64,
+    /// Objective weights; `None` takes rqe-optimizer's default.
+    pub w_cpu: Option<f64>,
+    pub w_mem: Option<f64>,
+    pub allow_undeployable_families: bool,
+}
+
+pub struct MilpPlan {
+    pub workload: MilpWorkload,
+    pub solution: MilpSolution,
+    pub objective: Objective,
+}
+
+/// Loads the facts and costs and solves for the cheapest plan.
+pub fn plan_milp(config: &ControllerConfig, inputs: &MilpInputs) -> anyhow::Result<MilpPlan> {
+    config.warn_default_slas();
+    let Some(hints) = config.metrics.as_deref() else {
+        return Err(MilpError::MissingMetricHints.into());
+    };
+    let facts = load_workload_facts(inputs.workload_facts, hints, inputs.scrape_interval_ms)?;
+    let costs = load_flat_atomic_cost_table(inputs.atomic_costs)?;
+    let objective = objective(inputs.w_cpu, inputs.w_mem)?;
+    tracing::debug!(?objective, cost_rows = costs.len(), "milp: inputs loaded");
+
+    let workload = build_milp_workload(config, &facts, inputs.scrape_interval_ms)?;
+    let solution = solve_milp(
+        &workload,
+        &facts,
+        &costs,
+        objective,
+        inputs.allow_undeployable_families,
+    )?;
+    Ok(MilpPlan {
+        workload,
+        solution,
+        objective,
+    })
+}
+
+fn objective(w_cpu: Option<f64>, w_mem: Option<f64>) -> anyhow::Result<Objective> {
+    let Objective::AUCCost {
+        w_cpu: default_cpu,
+        w_mem: default_mem,
+    } = Objective::default();
+    let (w_cpu, w_mem) = (w_cpu.unwrap_or(default_cpu), w_mem.unwrap_or(default_mem));
+    // All-zero weights make every plan cost 0, so the solver's pick is arbitrary.
+    anyhow::ensure!(
+        w_cpu > 0.0 || w_mem > 0.0,
+        "--w-cpu and --w-mem are both 0; at least one must be positive"
+    );
+    Ok(Objective::AUCCost { w_cpu, w_mem })
+}
+
+/// Parses an objective weight flag. Weights must be finite and non-negative:
+/// a negative weight rewards cost, and NaN poisons every coefficient.
+pub fn parse_weight(s: &str) -> Result<f64, String> {
+    match s.parse::<f64>() {
+        Ok(w) if w.is_finite() && w >= 0.0 => Ok(w),
+        Ok(w) => Err(format!("must be finite and >= 0, got {w}")),
+        Err(e) => Err(e.to_string()),
+    }
+}
+
 /// The MILP's view of a workload.
 #[derive(Debug)]
 pub struct MilpWorkload {
@@ -117,6 +193,16 @@ pub fn build_milp_workload(
     facts: &WorkloadFacts,
     scrape_interval_ms: u64,
 ) -> Result<MilpWorkload, MilpError> {
+    // They only size retention, which MILP configs don't set yet.
+    let range_queries: Vec<String> = config
+        .query_groups
+        .iter()
+        .filter(|qg| qg.step_ms.is_some() || qg.range_duration_ms.is_some())
+        .flat_map(|qg| qg.queries.iter().cloned())
+        .collect();
+    if !range_queries.is_empty() {
+        return Err(MilpError::RangeQueryOverrides(range_queries));
+    }
     let mut items = extract_hinted_items(config, scrape_interval_ms)?;
     // Stable Raqe order and ids across runs.
     items.sort_by(|a, b| {
@@ -511,6 +597,22 @@ metrics:
         assert!(matches!(err, MilpError::InvalidInputs(_)));
     }
 
+    /// Their retention isn't modeled, so planning them would ignore them silently.
+    #[test]
+    fn range_query_overrides_are_rejected() {
+        for field in ["step_ms: 60000", "range_duration_ms: 3600000"] {
+            let config = config(&format!(
+                "{}    {field}\n",
+                group("sum(http_requests_total)", 0.99)
+            ));
+            let err = build_milp_workload(&config, &facts(&config), SCRAPE_MS).unwrap_err();
+            assert!(
+                matches!(&err, MilpError::RangeQueryOverrides(q) if q == &["sum(http_requests_total)"]),
+                "{field}: {err}"
+            );
+        }
+    }
+
     #[test]
     fn solve_shares_one_deployment_across_repeated_queries() {
         let query = "sum by (job) (http_requests_total)";
@@ -590,6 +692,15 @@ metrics:
         assert_eq!(solution.deployments.len(), 2);
     }
 
+    #[test]
+    fn weights_must_be_finite_and_non_negative() {
+        assert_eq!(parse_weight("0.5"), Ok(0.5));
+        assert_eq!(parse_weight("0"), Ok(0.0));
+        for bad in ["-1", "NaN", "inf", "x"] {
+            assert!(parse_weight(bad).is_err(), "{bad}");
+        }
+    }
+
     #[test]
     fn workload_requires_metric_hints() {
         let config: ControllerConfig = serde_yaml::from_str(&format!(

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

@@ -9,7 +9,10 @@ pub mod workload_facts;
 pub use aqe_extractor::{extract_aqes, RQE};
 pub use atomic_costs::{load_flat_atomic_cost_table, AtomicCostEntry, AtomicCostTable};
 pub use error::OptimizerError;
-pub use milp::{build_milp_workload, solve_milp, MilpError, MilpWorkload};
+pub use milp::{
+    build_milp_workload, parse_weight, plan_milp, solve_milp, MilpError, MilpInputs, MilpPlan,
+    MilpWorkload,
+};
 pub use milp_output::{plan_to_planner_output, reject_avg_queries, MilpOutputError};
 pub use solution::OptimizerItem;
 pub use workload_facts::{load_workload_facts, parse_workload_facts, WorkloadFactsError};