feat(planner): asap-planner --planner saved writes a sketch-bench eval plan without re-solving

asap-planner-rs/Cargo.toml

@@ -37,6 +37,7 @@ indexmap.workspace = true
 chrono.workspace = true
 promql-parser = "0.5.0"
 reqwest = { version = "0.11", features = ["blocking", "json"] }
+flate2 = "1.1"
 rqe-optimizer = { git = "https://github.com/ProjectASAP/sketch-bench", branch = "main" }
 
 [dev-dependencies]

asap-planner-rs/src/main.rs

@@ -1,5 +1,6 @@
 use asap_planner::optimizer::{
-    parse_weight, plan_milp, plan_to_planner_output, reject_unwritable_queries, MilpInputs,
+    parse_weight, plan_milp, plan_to_planner_output, reject_unwritable_queries, saved_plan_output,
+    MilpInputs, PlanSelector,
 };
 use asap_planner::{
     Controller, ControllerConfig, ElasticController, ElasticRuntimeOptions, RuntimeOptions,
@@ -57,10 +58,33 @@ struct Args {
     clickhouse_database: Option<String>,
 
     /// `milp` plans with sketch-bench's rqe-optimizer: PromQL with
-    /// --input_config only, labels from its `metrics:` hints.
+    /// --input_config only, labels from its `metrics:` hints. `saved` writes
+    /// a plan sketch-bench's `autosketch_vs_asap` already solved.
     #[arg(long, value_enum, default_value = "legacy")]
     planner: PlannerArg,
 
+    /// Saved only. sketch-bench `autosketch_vs_asap` plan file
+    /// (`<table>.plans.json.gz`).
+    #[arg(long, required_if_eq("planner", "saved"))]
+    plan: Option<PathBuf>,
+
+    /// Saved only. The eval workload table the plan was solved for; its
+    /// `e2e` block gives the PromQL queries and metric labels.
+    #[arg(long = "workload-table", required_if_eq("planner", "saved"))]
+    workload_table: Option<PathBuf>,
+
+    /// Saved only. The plan's method, e.g. `asap`.
+    #[arg(long, required_if_eq("planner", "saved"))]
+    method: Option<String>,
+
+    /// Saved only. The plan's objective weights, e.g. `cpu` or `fargate`.
+    #[arg(long, required_if_eq("planner", "saved"))]
+    weights: Option<String>,
+
+    /// Saved only. The plan at this latency SLA; without it, the cheapest.
+    #[arg(long = "sla-ms")]
+    sla_ms: Option<f64>,
+
     /// MILP only. YAML workload facts: per metric, positive `value_range` and
     /// `cardinality` per label set, including all labels (the series count),
     /// plus the `shape` of each grouping sketches may serve.
@@ -99,6 +123,7 @@ enum EngineArg {
 enum PlannerArg {
     Legacy,
     Milp,
+    Saved,
 }
 
 fn main() -> anyhow::Result<()> {
@@ -116,20 +141,36 @@ fn main() -> anyhow::Result<()> {
         EngineArg::Precompute => StreamingEngine::Precompute,
     };
 
-    if args.planner == PlannerArg::Milp {
-        run_milp(&args)?;
+    if args.planner != PlannerArg::Saved {
+        anyhow::ensure!(
+            args.plan.is_none()
+                && args.workload_table.is_none()
+                && args.method.is_none()
+                && args.weights.is_none()
+                && args.sla_ms.is_none(),
+            "--plan, --workload-table, --method, --weights and --sla-ms require --planner saved"
+        );
+    }
+    if args.planner != PlannerArg::Milp {
+        anyhow::ensure!(
+            args.workload_facts.is_none()
+                && args.atomic_costs.is_none()
+                && args.saturation_dir.is_none()
+                && args.w_cpu.is_none()
+                && args.w_mem.is_none(),
+            "--workload-facts, --atomic-costs, --saturation-dir, --w-cpu and --w-mem require \
+             --planner milp"
+        );
+    }
+    if args.planner != PlannerArg::Legacy {
+        match args.planner {
+            PlannerArg::Milp => run_milp(&args)?,
+            PlannerArg::Saved => run_saved(&args)?,
+            PlannerArg::Legacy => unreachable!(),
+        }
         println!("Generated configs in {}", args.output_dir.display());
         return Ok(());
     }
-    anyhow::ensure!(
-        args.workload_facts.is_none()
-            && args.atomic_costs.is_none()
-            && args.saturation_dir.is_none()
-            && args.w_cpu.is_none()
-            && args.w_mem.is_none(),
-        "--workload-facts, --atomic-costs, --saturation-dir, --w-cpu and --w-mem require \
-         --planner milp"
-    );
 
     match args.query_language {
         QueryLanguage::promql => {
@@ -213,39 +254,43 @@ fn main() -> anyhow::Result<()> {
     Ok(())
 }
 
-fn run_milp(args: &Args) -> anyhow::Result<()> {
+/// Checks the flags `--planner milp` and `--planner saved` share: PromQL
+/// with labels from `metrics:` hints, no legacy-only flags. Returns the
+/// positive scrape interval, which per-series sample rates divide by.
+fn check_optimizer_flags(args: &Args, planner: &str) -> anyhow::Result<u64> {
     anyhow::ensure!(
         matches!(args.query_language, QueryLanguage::promql),
-        "--planner milp supports only --query-language promql"
+        "--planner {planner} supports only --query-language promql"
     );
     anyhow::ensure!(
         args.query_log.is_none(),
-        "--planner milp needs --input_config: query logs have no `metrics:` hints"
+        "--planner {planner} doesn't take --query-log: 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"
+        "--planner {planner} takes labels from `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"
+        "--enable-punting, --range-duration-ms and --step-ms don't apply to --planner {planner}"
     );
     anyhow::ensure!(
         args.clickhouse_url.is_none() && args.clickhouse_database.is_none(),
-        "--clickhouse-url and --clickhouse-database don't apply to --planner milp"
+        "--clickhouse-url and --clickhouse-database don't apply to --planner {planner}"
     );
+    args.data_ingestion_interval_ms
+        .filter(|&ms| ms > 0)
+        .ok_or_else(|| {
+            anyhow::anyhow!("--planner {planner} requires a positive --data-ingestion-interval-ms")
+        })
+}
+
+fn run_milp(args: &Args) -> anyhow::Result<()> {
+    let scrape_interval_ms = check_optimizer_flags(args, "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!("--planner milp requires --data-ingestion-interval-ms"))?;
-    // Per-series sample rates divide by it.
-    anyhow::ensure!(
-        scrape_interval_ms > 0,
-        "--data-ingestion-interval-ms must be positive"
-    );
     let config: ControllerConfig = serde_yaml::from_str(&std::fs::read_to_string(config_path)?)?;
     // Fail before solving: the plan is always written.
     reject_unwritable_queries(&config)?;
@@ -274,3 +319,27 @@ fn run_milp(args: &Args) -> anyhow::Result<()> {
         .write_to_dir(&args.output_dir)?;
     Ok(())
 }
+
+fn run_saved(args: &Args) -> anyhow::Result<()> {
+    let scrape_interval_ms = check_optimizer_flags(args, "saved")?;
+    anyhow::ensure!(
+        args.input_config.is_none(),
+        "--planner saved reads its queries from --workload-table, not --input_config"
+    );
+    saved_plan_output(
+        args.workload_table
+            .as_deref()
+            .expect("clap requires --workload-table with --planner saved"),
+        args.plan
+            .as_deref()
+            .expect("clap requires --plan with --planner saved"),
+        &PlanSelector {
+            method: args.method.as_deref().expect("clap requires --method"),
+            weights: args.weights.as_deref().expect("clap requires --weights"),
+            sla_ms: args.sla_ms,
+        },
+        scrape_interval_ms,
+    )?
+    .write_to_dir(&args.output_dir)?;
+    Ok(())
+}

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

@@ -132,7 +132,7 @@ fn normalized_f64_bits(value: f64) -> u64 {
 /// dropped — they contribute no AQE. Only arithmetic operators are split;
 /// comparison and set operators are left as-is (treated as opaque leaves).
 /// An `avg` leaf becomes its sum and count leaves (see `rewrite_avg`).
-fn decompose_to_leaves(query: &str) -> Vec<String> {
+pub(crate) fn decompose_to_leaves(query: &str) -> Vec<String> {
     let (lhs, rhs) = match parse_binary_arms(query) {
         Some(arms) => arms,
         None => return rewrite_avg(query),

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

@@ -243,6 +243,17 @@ pub fn build_milp_workload(
     config: &ControllerConfig,
     facts: &WorkloadFacts,
     scrape_interval_ms: u64,
+) -> Result<MilpWorkload, MilpError> {
+    let workload = workload_from_config(config, scrape_interval_ms)?;
+    validate_facts(&workload.raqes, facts).map_err(MilpError::InvalidInputs)?;
+    Ok(workload)
+}
+
+/// [`build_milp_workload`] without the facts check, for a plan solved
+/// elsewhere.
+pub fn workload_from_config(
+    config: &ControllerConfig,
+    scrape_interval_ms: u64,
 ) -> Result<MilpWorkload, MilpError> {
     // The MILP chooses windows and sketch parameters itself.
     let unsupported: Vec<&'static str> = [
@@ -298,10 +309,8 @@ pub fn build_milp_workload(
     tracing::debug!(
         items = items.len(),
         raqes = raqes.len(),
-        metrics_with_facts = facts.len(),
         "milp inputs: built raqes"
     );
-    validate_facts(&raqes, facts).map_err(MilpError::InvalidInputs)?;
     Ok(MilpWorkload {
         items,
         raqes,

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

@@ -92,14 +92,36 @@ pub fn plan_to_planner_output(
     workload: &MilpWorkload,
     solution: &MilpSolution,
 ) -> Result<PlannerOutput, MilpOutputError> {
+    let deployments: Vec<Deployment> = solution
+        .deployments
+        .iter()
+        .map(|planned| planned.deployment.clone())
+        .collect();
+    let served_by: Vec<usize> = solution.raqes.iter().map(|r| r.deployment).collect();
+    deployments_to_planner_output(config, workload, &deployments, &served_by)
+}
+
+/// [`plan_to_planner_output`] for a plan given as its deployments and, per
+/// Raqe of `workload`, the index of the deployment serving it.
+pub fn deployments_to_planner_output(
+    config: &ControllerConfig,
+    workload: &MilpWorkload,
+    deployments: &[Deployment],
+    served_by: &[usize],
+) -> Result<PlannerOutput, MilpOutputError> {
+    assert_eq!(
+        served_by.len(),
+        workload.raqes.len(),
+        "one deployment per Raqe"
+    );
     reject_unwritable_queries(config)?;
     // sketch-bench lets a deployment serve a RAQE grouped by a subset of its
     // grouping (a roll-up, sketch-bench#190). The engine reads a deployment
     // only at its own grouping, and `aggregation_config` takes the labels
     // from the first RAQE served, so refuse a roll-up rather than write
     // configs at the wrong grouping.
-    for (raqe, planned) in workload.raqes.iter().zip(&solution.raqes) {
-        let deployment = &solution.deployments[planned.deployment].deployment;
+    for (raqe, &index) in workload.raqes.iter().zip(served_by) {
+        let deployment = &deployments[index];
         if raqe.grouping_labels != deployment.grouping_labels {
             return Err(MilpOutputError::RollupUnsupported {
                 query: raqe.id.clone(),
@@ -111,26 +133,26 @@ pub fn plan_to_planner_output(
 
     let item_of = |raqe: usize| &workload.items[workload.raqe_items[raqe]];
     // Each item once per deployment, however many occurrences it has.
-    let mut served: Vec<Vec<&OptimizerItem>> = vec![Vec::new(); solution.deployments.len()];
+    let mut served: Vec<Vec<&OptimizerItem>> = vec![Vec::new(); deployments.len()];
     let mut seen = HashSet::new();
-    for (raqe, planned) in solution.raqes.iter().enumerate() {
-        if seen.insert((planned.deployment, workload.raqe_items[raqe])) {
-            served[planned.deployment].push(item_of(raqe));
+    for (raqe, &index) in served_by.iter().enumerate() {
+        if seen.insert((index, workload.raqe_items[raqe])) {
+            served[index].push(item_of(raqe));
         }
     }
 
     let mut aggregations: IndexMap<String, IntermediateAggConfig> = IndexMap::new();
     let mut deployment_keys = Vec::with_capacity(served.len());
-    for (planned, items) in solution.deployments.iter().zip(&served) {
-        let aggregation = aggregation_config(&planned.deployment, items)?;
+    for (deployment, items) in deployments.iter().zip(&served) {
+        let aggregation = aggregation_config(deployment, items)?;
         let key = aggregation.identifying_key();
         aggregations.entry(key.clone()).or_insert(aggregation);
         deployment_keys.push(key);
     }
 
     let mut queries: IndexMap<String, QueryPlanEntry> = IndexMap::new();
-    for (raqe, planned) in solution.raqes.iter().enumerate() {
-        let key = &deployment_keys[planned.deployment];
+    for (raqe, &index) in served_by.iter().enumerate() {
+        let key = &deployment_keys[index];
         for query in &item_of(raqe).query_strings {
             match queries.entry(query.clone()) {
                 Entry::Vacant(entry) => {

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

@@ -3,16 +3,21 @@ pub mod atomic_costs;
 pub mod error;
 pub mod milp;
 pub mod milp_output;
+pub mod saved_plan;
 pub mod solution;
 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, parse_weight, plan_milp, solve_milp, MilpError, MilpInputs, MilpPlan,
-    MilpWorkload,
+    build_milp_workload, parse_weight, plan_milp, solve_milp, workload_from_config, MilpError,
+    MilpInputs, MilpPlan, MilpWorkload,
 };
-pub use milp_output::{plan_to_planner_output, reject_unwritable_queries, MilpOutputError};
+pub use milp_output::{
+    deployments_to_planner_output, plan_to_planner_output, reject_unwritable_queries,
+    MilpOutputError,
+};
+pub use saved_plan::{saved_plan_output, PlanSelector};
 pub use solution::OptimizerItem;
 pub use workload_facts::{load_workload_facts, parse_workload_facts, WorkloadFactsError};