From 685ea4a439cab9bef363321c1623d4dee4cad9a1 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Mon, 12 Oct 2026 14:35:02 -0400 Subject: [PATCH 1/4] refactor(planner): thread a language-tagged MilpConfig through the MILP path Co-Authored-By: Claude Opus 5.5 --- asap-planner-rs/src/bin/optimizer_cli.rs | 5 +- asap-planner-rs/src/main.rs | 10 +- asap-planner-rs/src/optimizer/milp.rs | 103 +++++++++++++------ asap-planner-rs/src/optimizer/milp_output.rs | 60 ++++++----- asap-planner-rs/src/optimizer/mod.rs | 4 +- 5 files changed, 117 insertions(+), 65 deletions(-) diff --git a/asap-planner-rs/src/bin/optimizer_cli.rs b/asap-planner-rs/src/bin/optimizer_cli.rs index e4ba6e78..44b5595e 100644 --- a/asap-planner-rs/src/bin/optimizer_cli.rs +++ b/asap-planner-rs/src/bin/optimizer_cli.rs @@ -4,8 +4,8 @@ use std::path::PathBuf; use asap_planner::optimizer::{ - parse_weight, plan_milp, plan_to_planner_output, reject_unwritable_queries, MilpInputs, - MilpPlan, + parse_weight, plan_milp, plan_to_planner_output, reject_unwritable_queries, MilpConfig, + MilpInputs, MilpPlan, }; use asap_planner::ControllerConfig; use clap::Parser; @@ -82,6 +82,7 @@ fn main() -> anyhow::Result<()> { } fn run_milp(args: &Args, config: &ControllerConfig) -> anyhow::Result<()> { + let config = MilpConfig::Promql(config); // Fail before solving when the plan would be written but can't be. if args.output_dir.is_some() { reject_unwritable_queries(config)?; diff --git a/asap-planner-rs/src/main.rs b/asap-planner-rs/src/main.rs index c4b2e20b..69f5a95e 100644 --- a/asap-planner-rs/src/main.rs +++ b/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, MilpConfig, + MilpInputs, }; use asap_planner::{ Controller, ControllerConfig, ElasticController, ElasticRuntimeOptions, RuntimeOptions, @@ -247,10 +248,11 @@ fn run_milp(args: &Args) -> anyhow::Result<()> { "--data-ingestion-interval-ms must be positive" ); let config: ControllerConfig = serde_yaml::from_str(&std::fs::read_to_string(config_path)?)?; + let config = MilpConfig::Promql(&config); // Fail before solving: the plan is always written. - reject_unwritable_queries(&config)?; + reject_unwritable_queries(config)?; let plan = plan_milp( - &config, + config, &MilpInputs { workload_facts: args .workload_facts @@ -270,7 +272,7 @@ fn run_milp(args: &Args) -> anyhow::Result<()> { allow_undeployable_families: false, }, )?; - plan_to_planner_output(&config, &plan.workload, &plan.solution)? + plan_to_planner_output(config, &plan.workload, &plan.solution)? .write_to_dir(&args.output_dir)?; Ok(()) } diff --git a/asap-planner-rs/src/optimizer/milp.rs b/asap-planner-rs/src/optimizer/milp.rs index 71d8d83c..60ec899c 100644 --- a/asap-planner-rs/src/optimizer/milp.rs +++ b/asap-planner-rs/src/optimizer/milp.rs @@ -11,7 +11,7 @@ use rqe_optimizer::{ }; use thiserror::Error; -use crate::config::input::ControllerConfig; +use crate::config::input::{ControllerConfig, MetricDefinition}; use std::path::Path; @@ -119,6 +119,49 @@ pub fn solve_milp( Ok(solution) } +/// The workload config the MILP plans, by query language. +#[derive(Debug, Clone, Copy)] +pub enum MilpConfig<'a> { + Promql(&'a ControllerConfig), +} + +impl MilpConfig<'_> { + /// Each workload metric's labels. + fn hints(self) -> Result, MilpError> { + match self { + MilpConfig::Promql(config) => { + config.metrics.clone().ok_or(MilpError::MissingMetricHints) + } + } + } + + fn items(self, scrape_interval_ms: u64) -> Result, MilpError> { + match self { + MilpConfig::Promql(config) => { + reject_fields([ + ("windowing", config.windowing.is_some()), + ("sketch_parameters", config.sketch_parameters.is_some()), + ("aggregate_cleanup", config.aggregate_cleanup.is_some()), + ])?; + extract_hinted_items(config, scrape_interval_ms) + } + } + } +} + +/// The MILP chooses windows and sketch parameters itself. +fn reject_fields(fields: [(&'static str, bool); 3]) -> Result<(), MilpError> { + let unsupported: Vec<&'static str> = fields + .into_iter() + .filter_map(|(field, set)| set.then_some(field)) + .collect(); + if unsupported.is_empty() { + Ok(()) + } else { + Err(MilpError::UnsupportedFields(unsupported)) + } +} + /// What a MILP planning run reads besides the workload config. pub struct MilpInputs<'a> { pub workload_facts: &'a Path, @@ -141,12 +184,11 @@ pub struct MilpPlan { /// Loads the facts, costs and saturation curves and solves for the cheapest /// plan. -pub fn plan_milp(config: &ControllerConfig, inputs: &MilpInputs) -> anyhow::Result { - 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)?; +pub fn plan_milp(config: MilpConfig, inputs: &MilpInputs) -> anyhow::Result { + let MilpConfig::Promql(promql) = config; + promql.warn_default_slas(); + let hints = config.hints()?; + let facts = load_workload_facts(inputs.workload_facts, &hints, inputs.scrape_interval_ms)?; let curves = SaturationCurves::load(inputs.saturation_dir).with_context(|| { format!( "loading saturation curves from {}", @@ -240,23 +282,11 @@ pub struct MilpWorkload { /// dashboards running the same query at the same cadence become two /// identical Raqes, which the MILP serves from one shared deployment. pub fn build_milp_workload( - config: &ControllerConfig, + config: MilpConfig, facts: &WorkloadFacts, scrape_interval_ms: u64, ) -> Result { - // The MILP chooses windows and sketch parameters itself. - let unsupported: Vec<&'static str> = [ - ("windowing", config.windowing.is_some()), - ("sketch_parameters", config.sketch_parameters.is_some()), - ("aggregate_cleanup", config.aggregate_cleanup.is_some()), - ] - .into_iter() - .filter_map(|(field, set)| set.then_some(field)) - .collect(); - if !unsupported.is_empty() { - return Err(MilpError::UnsupportedFields(unsupported)); - } - let mut items = extract_hinted_items(config, scrape_interval_ms)?; + let mut items = config.items(scrape_interval_ms)?; // Stable Raqe order and ids across runs. items.sort_by(|a, b| { (&a.query_strings, a.t_repeat_ms) @@ -500,7 +530,7 @@ metrics: fn workload(groups: &str) -> MilpWorkload { let config = config(groups); - build_milp_workload(&config, &facts(&config), SCRAPE_MS).unwrap() + build_milp_workload(MilpConfig::Promql(&config), &facts(&config), SCRAPE_MS).unwrap() } fn labels(names: &[&str]) -> LabelSet { @@ -673,7 +703,8 @@ metrics: #[test] fn accuracy_sla_above_one_is_rejected() { let config = config(&group("sum by (job) (http_requests_total)", 1.5)); - let err = build_milp_workload(&config, &facts(&config), SCRAPE_MS).unwrap_err(); + let err = build_milp_workload(MilpConfig::Promql(&config), &facts(&config), SCRAPE_MS) + .unwrap_err(); assert!(matches!(err, MilpError::AccuracySlaOutOfRange { .. })); } @@ -683,7 +714,8 @@ metrics: "sum by (job) (http_requests_total{job='api'})", 0.99, )); - let err = build_milp_workload(&config, &facts(&config), SCRAPE_MS).unwrap_err(); + let err = build_milp_workload(MilpConfig::Promql(&config), &facts(&config), SCRAPE_MS) + .unwrap_err(); let MilpError::InvalidInputs(problems) = err else { panic!("expected InvalidInputs, got {err:?}"); }; @@ -693,7 +725,8 @@ metrics: #[test] fn missing_cardinality_is_rejected() { let config = config(&group("sum by (instance) (http_requests_total)", 0.99)); - let err = build_milp_workload(&config, &facts(&config), SCRAPE_MS).unwrap_err(); + let err = build_milp_workload(MilpConfig::Promql(&config), &facts(&config), SCRAPE_MS) + .unwrap_err(); assert!(matches!(err, MilpError::InvalidInputs(_))); } @@ -705,7 +738,8 @@ metrics: config.windowing = Some(serde_yaml::from_str("{type: tumbling, window_size_ms: 60000}").unwrap()); config.sketch_parameters = Some(Default::default()); - let err = build_milp_workload(&config, &facts(&config), SCRAPE_MS).unwrap_err(); + let err = build_milp_workload(MilpConfig::Promql(&config), &facts(&config), SCRAPE_MS) + .unwrap_err(); assert!( matches!(&err, MilpError::UnsupportedFields(f) if f == &["windowing", "sketch_parameters"]), "{err}" @@ -713,7 +747,8 @@ metrics: config.windowing = None; config.sketch_parameters = None; config.aggregate_cleanup = Some(serde_yaml::from_str("{}").unwrap()); - let err = build_milp_workload(&config, &facts(&config), SCRAPE_MS).unwrap_err(); + let err = build_milp_workload(MilpConfig::Promql(&config), &facts(&config), SCRAPE_MS) + .unwrap_err(); assert!( matches!(&err, MilpError::UnsupportedFields(f) if f == &["aggregate_cleanup"]), "{err}" @@ -725,7 +760,7 @@ metrics: let query = "sum by (job) (http_requests_total)"; 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 w = build_milp_workload(MilpConfig::Promql(&config), &facts, SCRAPE_MS).unwrap(); let solution = solve_milp( &w, &facts, @@ -756,7 +791,7 @@ metrics: .unwrap() .data_shape .clear(); - let w = build_milp_workload(&both, &facts, SCRAPE_MS).unwrap(); + let w = build_milp_workload(MilpConfig::Promql(&both), &facts, SCRAPE_MS).unwrap(); let err = solve_milp( &w, &facts, @@ -787,7 +822,7 @@ metrics: .unwrap_err(); assert!(matches!(err, MilpError::Unservable(_)), "{err:?}"); let exact_only = config(&group(sum, 0.99)); - let w = build_milp_workload(&exact_only, &facts, SCRAPE_MS).unwrap(); + let w = build_milp_workload(MilpConfig::Promql(&exact_only), &facts, SCRAPE_MS).unwrap(); assert!(solve_milp( &w, &facts, @@ -807,7 +842,7 @@ metrics: 0.999, )); let facts = facts(&config); - let w = build_milp_workload(&config, &facts, SCRAPE_MS).unwrap(); + let w = build_milp_workload(MilpConfig::Promql(&config), &facts, SCRAPE_MS).unwrap(); let err = solve_milp( &w, &facts, @@ -825,7 +860,7 @@ metrics: let quantile = group("quantile_over_time(0.99, http_requests_total[5m])", 0.999); let config = config(&(quantile.clone() + &quantile)); let facts = facts(&config); - let w = build_milp_workload(&config, &facts, SCRAPE_MS).unwrap(); + let w = build_milp_workload(MilpConfig::Promql(&config), &facts, SCRAPE_MS).unwrap(); assert_eq!(w.raqes.len(), 2); let err = solve_milp( &w, @@ -846,7 +881,7 @@ metrics: + &group("topk(5, http_requests_total)", 0.99); let config = config(&groups); let facts = facts(&config); - let w = build_milp_workload(&config, &facts, SCRAPE_MS).unwrap(); + let w = build_milp_workload(MilpConfig::Promql(&config), &facts, SCRAPE_MS).unwrap(); let solution = solve_milp( &w, &facts, @@ -879,7 +914,7 @@ metrics: // rewrite; sharing would undercount ingest and memory. let config = config(&group("avg by (job) (http_requests_total)", 0.99)); let facts = facts(&config); - let w = build_milp_workload(&config, &facts, SCRAPE_MS).unwrap(); + let w = build_milp_workload(MilpConfig::Promql(&config), &facts, SCRAPE_MS).unwrap(); 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)); diff --git a/asap-planner-rs/src/optimizer/milp_output.rs b/asap-planner-rs/src/optimizer/milp_output.rs index 8f934931..6298b727 100644 --- a/asap-planner-rs/src/optimizer/milp_output.rs +++ b/asap-planner-rs/src/optimizer/milp_output.rs @@ -13,7 +13,6 @@ use rqe_optimizer::{Capability, Deployment}; use serde_json::Value; use thiserror::Error; -use crate::config::input::ControllerConfig; use crate::error::ControllerError; use crate::generator::{GeneratorOutput, QueryPlanEntry}; use crate::planner::agg_config::IntermediateAggConfig; @@ -22,7 +21,7 @@ use crate::planner_output::PlannerOutput; use crate::promql::generator::{build_inference_yaml, build_streaming_yaml}; use super::aqe_extractor::contains_avg; -use super::milp::MilpWorkload; +use super::milp::{MilpConfig, MilpWorkload}; use super::solution::OptimizerItem; #[derive(Debug, Error)] @@ -63,7 +62,8 @@ pub enum MilpOutputError { /// Errors on queries whose plan can be costed but not written as configs: /// avg queries, and range-query overrides, which only size retention. -pub fn reject_unwritable_queries(config: &ControllerConfig) -> Result<(), MilpOutputError> { +pub fn reject_unwritable_queries(config: MilpConfig) -> Result<(), MilpOutputError> { + let MilpConfig::Promql(config) = config; if let Some(query) = config .query_groups .iter() @@ -88,7 +88,7 @@ pub fn reject_unwritable_queries(config: &ControllerConfig) -> Result<(), MilpOu /// plan's deployment order, starting at 1. No cleanup policy yet, so the /// retained instance counts are not emitted. pub fn plan_to_planner_output( - config: &ControllerConfig, + config: MilpConfig, workload: &MilpWorkload, solution: &MilpSolution, ) -> Result { @@ -152,11 +152,19 @@ pub fn plan_to_planner_output( .enumerate() .map(|(index, key)| (key.clone(), index as u32 + 1)) .collect(); - let schema = config.schema_from_hints(); + let (streaming_yaml, inference_yaml) = match config { + MilpConfig::Promql(config) => { + let schema = config.schema_from_hints(); + ( + build_streaming_yaml(&aggregations, &id_map, &schema)?, + build_inference_yaml(CleanupPolicy::NoCleanup, &queries, &id_map, &schema)?, + ) + } + }; Ok(PlannerOutput::from_output(GeneratorOutput { punted_queries: Vec::new(), - streaming_yaml: build_streaming_yaml(&aggregations, &id_map, &schema)?, - inference_yaml: build_inference_yaml(CleanupPolicy::NoCleanup, &queries, &id_map, &schema)?, + streaming_yaml, + inference_yaml, aggregation_count: aggregations.len(), query_count: queries.len(), })) @@ -315,13 +323,14 @@ mod tests { use serde_json::json; use super::*; + use crate::config::input::ControllerConfig; use crate::optimizer::milp::tests::{config, costs, facts, group, table_accuracy, SCRAPE_MS}; use crate::optimizer::milp::{build_milp_workload, solve_milp}; fn plan(groups: &str) -> (ControllerConfig, MilpWorkload, MilpSolution) { let config = config(groups); let facts = facts(&config); - let workload = build_milp_workload(&config, &facts, SCRAPE_MS).unwrap(); + let workload = build_milp_workload(MilpConfig::Promql(&config), &facts, SCRAPE_MS).unwrap(); let mut costs = costs(); for cost in &mut costs { cost.sketch_config["params"] = match cost.sketch.as_str() { @@ -390,7 +399,8 @@ mod tests { .map(|q| group(q, 0.99)) .collect(); let (config, workload, solution) = plan(&groups); - let output = plan_to_planner_output(&config, &workload, &solution).unwrap(); + let output = + plan_to_planner_output(MilpConfig::Promql(&config), &workload, &solution).unwrap(); assert_eq!(output.streaming_aggregation_count(), 5); let a = aggregation_for(&output, sum); @@ -481,7 +491,8 @@ mod tests { } else { assert_eq!(heap, needed, "k = {k}: m · k"); } - let output = plan_to_planner_output(&config, &workload, &solution).unwrap(); + let output = + plan_to_planner_output(MilpConfig::Promql(&config), &workload, &solution).unwrap(); assert_eq!( aggregation_for(&output, &query).parameters["heapsize"], json!(heap), @@ -500,7 +511,8 @@ mod tests { assert_eq!(solution.deployments.len(), 1); let heap = planned_heap(&workload, &solution, large); assert!(heap >= heap_needed(&workload, &solution, large)); - let output = plan_to_planner_output(&config, &workload, &solution).unwrap(); + let output = + plan_to_planner_output(MilpConfig::Promql(&config), &workload, &solution).unwrap(); for query in [small, large] { assert_eq!( aggregation_for(&output, query).parameters["heapsize"], @@ -526,7 +538,7 @@ mod tests { let query = "quantile_over_time(0.99, http_requests_total[5m])"; let config = config(&group(query, 0.99)); let facts = facts(&config); - let workload = build_milp_workload(&config, &facts, SCRAPE_MS).unwrap(); + let workload = build_milp_workload(MilpConfig::Promql(&config), &facts, SCRAPE_MS).unwrap(); let costs = vec![dd_cost(json!({"alpha": 0.02}))]; let solution = solve_milp( &workload, @@ -538,7 +550,8 @@ mod tests { ) .unwrap(); - let output = plan_to_planner_output(&config, &workload, &solution).unwrap(); + let output = + plan_to_planner_output(MilpConfig::Promql(&config), &workload, &solution).unwrap(); let aggregation = aggregation_for(&output, query); assert_eq!(aggregation.aggregation_type, AggregationType::DDSketch); assert_eq!(aggregation.parameters["alpha"], json!(0.02)); @@ -549,7 +562,7 @@ mod tests { let query = "quantile_over_time(0.99, http_requests_total[5m])"; let config = config(&group(query, 0.99)); let facts = facts(&config); - let workload = build_milp_workload(&config, &facts, SCRAPE_MS).unwrap(); + let workload = build_milp_workload(MilpConfig::Promql(&config), &facts, SCRAPE_MS).unwrap(); let costs = vec![dd_cost(json!({}))]; let solution = solve_milp( &workload, @@ -562,7 +575,7 @@ mod tests { .unwrap(); assert!(matches!( - plan_to_planner_output(&config, &workload, &solution), + plan_to_planner_output(MilpConfig::Promql(&config), &workload, &solution), Err(MilpOutputError::MissingParam { variant, param: "alpha", @@ -575,7 +588,7 @@ mod tests { let query = "quantile_over_time(0.99, http_requests_total[5m])"; let config = config(&group(query, 0.99)); let facts = facts(&config); - let workload = build_milp_workload(&config, &facts, SCRAPE_MS).unwrap(); + let workload = build_milp_workload(MilpConfig::Promql(&config), &facts, SCRAPE_MS).unwrap(); for alpha in [0.0, -0.01, 1.0] { let costs = vec![dd_cost(json!({"alpha": alpha}))]; @@ -589,7 +602,7 @@ mod tests { ) .unwrap(); assert!(matches!( - plan_to_planner_output(&config, &workload, &solution), + plan_to_planner_output(MilpConfig::Promql(&config), &workload, &solution), Err(MilpOutputError::InvalidDdsAlpha { .. }) )); } @@ -602,7 +615,8 @@ mod tests { let deployment = &mut solution.deployments[0].deployment; deployment.window_ms = 300_000; deployment.slide_ms = 60_000; - let output = plan_to_planner_output(&config, &workload, &solution).unwrap(); + let output = + plan_to_planner_output(MilpConfig::Promql(&config), &workload, &solution).unwrap(); let a = aggregation_for(&output, query); assert_eq!(a.window_type, WindowType::Sliding); assert_eq!((a.window_size_ms, a.slide_interval_ms), (300_000, 60_000)); @@ -612,7 +626,7 @@ mod tests { fn avg_query_is_an_error() { let query = "avg by (job) (http_requests_total)"; let (config, workload, solution) = plan(&group(query, 0.99)); - let err = plan_to_planner_output(&config, &workload, &solution) + let err = plan_to_planner_output(MilpConfig::Promql(&config), &workload, &solution) .err() .unwrap(); assert!(matches!(err, MilpOutputError::AvgQuery(q) if q == query)); @@ -626,7 +640,7 @@ mod tests { for field in ["step_ms: 60000", "range_duration_ms: 3600000"] { let (config, workload, solution) = plan(&format!("{} {field}\n", group(query, 0.99))); - let err = plan_to_planner_output(&config, &workload, &solution) + let err = plan_to_planner_output(MilpConfig::Promql(&config), &workload, &solution) .err() .unwrap(); assert!( @@ -644,7 +658,7 @@ mod tests { 0.99, )); solution.deployments[0].deployment.config.sketch = "hydra-kll".into(); - let err = plan_to_planner_output(&config, &workload, &solution) + let err = plan_to_planner_output(MilpConfig::Promql(&config), &workload, &solution) .err() .unwrap(); assert!( @@ -662,7 +676,7 @@ mod tests { .deployment .grouping_labels .insert("instance".into()); - let err = plan_to_planner_output(&config, &workload, &solution) + let err = plan_to_planner_output(MilpConfig::Promql(&config), &workload, &solution) .err() .unwrap(); assert!( @@ -684,7 +698,7 @@ mod tests { other.deployment.slide_ms *= 2; solution.deployments.push(other); solution.raqes[1].deployment = 1; - let err = plan_to_planner_output(&config, &workload, &solution) + let err = plan_to_planner_output(MilpConfig::Promql(&config), &workload, &solution) .err() .unwrap(); assert!(matches!(err, MilpOutputError::QueryOnTwoDeployments(q) if q == query)); diff --git a/asap-planner-rs/src/optimizer/mod.rs b/asap-planner-rs/src/optimizer/mod.rs index 7731c14c..b3321be1 100644 --- a/asap-planner-rs/src/optimizer/mod.rs +++ b/asap-planner-rs/src/optimizer/mod.rs @@ -10,8 +10,8 @@ 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, MilpConfig, MilpError, MilpInputs, + MilpPlan, MilpWorkload, }; pub use milp_output::{plan_to_planner_output, reject_unwritable_queries, MilpOutputError}; pub use solution::OptimizerItem; From 351d5f1d71cb486405f0c30d62ad686f27bc7ae4 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Mon, 12 Oct 2026 16:06:55 -0400 Subject: [PATCH 2/4] feat(planner): plan SQL workloads with the MILP Co-Authored-By: Claude Opus 5.5 --- .../src/optimizer/aqe_extractor.rs | 237 +++++++++++++++--- asap-planner-rs/src/optimizer/milp.rs | 76 ++++-- asap-planner-rs/src/optimizer/milp_output.rs | 200 ++++++++++++--- asap-planner-rs/src/optimizer/solution.rs | 4 + asap-planner-rs/src/planner/sql.rs | 126 +++++----- asap-planner-rs/src/sql/generator.rs | 4 +- 6 files changed, 493 insertions(+), 154 deletions(-) diff --git a/asap-planner-rs/src/optimizer/aqe_extractor.rs b/asap-planner-rs/src/optimizer/aqe_extractor.rs index 1a7b3262..a993904c 100644 --- a/asap-planner-rs/src/optimizer/aqe_extractor.rs +++ b/asap-planner-rs/src/optimizer/aqe_extractor.rs @@ -1,14 +1,18 @@ use std::collections::HashMap; +use std::time::{SystemTime, UNIX_EPOCH}; use asap_types::query_requirements::{build_query_requirements_promql, QueryRequirements}; +use asap_types::utils::normalize_spatial_filter; use asap_types::PromQLSchema; use promql_parser::parser::token::{self, TokenType}; use promql_parser::parser::{AggregateExpr, Call, Expr, Function}; use promql_utilities::data_model::KeyByLabelNames; use promql_utilities::query_logics::enums::Statistic; +use crate::config::input::TableDefinition; use crate::planner::patterns::build_patterns; use crate::planner::promql::{parse_binary_arms, BinaryArm}; +use crate::planner::sql::{parse_sql_query, sql_data_range_ms, sql_statistics}; use super::error::OptimizerError; use super::solution::OptimizerItem; @@ -67,54 +71,154 @@ pub fn extract_aqes( metric_schema: &PromQLSchema, scrape_interval_ms: u64, ) -> Result, OptimizerError> { - let mut acc: HashMap, usize)> = - HashMap::new(); - + let mut leaves = Vec::new(); for rqe in rqes { - if rqe.t_repeat_ms == 0 { - return Err(OptimizerError::InvalidRepeatInterval { - query: rqe.query_string.clone(), - }); + check_repeat_interval(rqe)?; + for leaf in decompose_to_leaves(&rqe.query_string) { + let req = extract_requirements(&leaf, metric_schema, scrape_interval_ms).map_err( + |reason| OptimizerError::UnsupportedLeaf { + query: rqe.query_string.clone(), + leaf: leaf.clone(), + reason, + }, + )?; + let k = topk_k(&leaf); + leaves.push((req, leaf, k, rqe)); } + } + Ok(merge_leaves(leaves)) +} - let leaves = decompose_to_leaves(&rqe.query_string); - - for leaf in leaves { - match extract_requirements(&leaf, metric_schema, scrape_interval_ms) { - Ok(req) => { - let key = OptimizerItemKey::from_rqe(&req, rqe); - let entry = acc.entry(key).or_insert_with(|| (req, Vec::new(), 0)); - if !entry.1.contains(&leaf) { - entry.1.push(leaf); - } - entry.2 += 1; - } - Err(reason) => { - return Err(OptimizerError::UnsupportedLeaf { - query: rqe.query_string.clone(), - leaf, - reason, - }) +/// `extract_aqes` for SQL: each query is one leaf, and its metric is +/// `table.value_column` (see `sql_metric`). +pub fn extract_sql_aqes( + rqes: &[RQE], + tables: &[TableDefinition], + scrape_interval_ms: u64, +) -> Result, OptimizerError> { + // Only the queries' durations are used, which don't depend on it. + let eval_time = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("clock after 1970") + .as_secs_f64(); + let mut leaves = Vec::new(); + for rqe in rqes { + check_repeat_interval(rqe)?; + let query = &rqe.query_string; + let (req, k) = + sql_requirements(query, tables, scrape_interval_ms, eval_time).map_err(|reason| { + OptimizerError::UnsupportedLeaf { + query: query.clone(), + leaf: query.clone(), + reason, } - } + })?; + leaves.push((req, query.clone(), k, rqe)); + } + Ok(merge_leaves(leaves)) +} + +fn check_repeat_interval(rqe: &RQE) -> Result<(), OptimizerError> { + if rqe.t_repeat_ms == 0 { + return Err(OptimizerError::InvalidRepeatInterval { + query: rqe.query_string.clone(), + }); + } + Ok(()) +} + +/// Merges leaf occurrences, each `(requirements, leaf query, top-k k, RQE)`, +/// into optimizer items. +fn merge_leaves(leaves: Vec<(QueryRequirements, String, Option, &RQE)>) -> Vec { + struct Acc { + requirements: QueryRequirements, + query_strings: Vec, + occurrences: usize, + topk_k: Option, + } + let mut acc: HashMap = HashMap::new(); + for (req, leaf, k, rqe) in leaves { + let key = OptimizerItemKey::from_rqe(&req, rqe); + let entry = acc.entry(key).or_insert_with(|| Acc { + requirements: req, + query_strings: Vec::new(), + occurrences: 0, + topk_k: k, + }); + if !entry.query_strings.contains(&leaf) { + entry.query_strings.push(leaf); } + entry.occurrences += 1; + entry.topk_k = entry.topk_k.zip(k).map(|(a, b)| a.max(b)); } - Ok(acc - .into_iter() - .map( - |(key, (requirements, query_strings, occurrences))| OptimizerItem { - requirements, - query_strings, - // Hz (queries per real second): 1000.0 / ms, not 1.0 / ms. - query_frequency_hz: occurrences as f64 * 1000.0 / key.t_repeat_ms as f64, - occurrences, - t_repeat_ms: key.t_repeat_ms, - accuracy_sla: f64::from_bits(key.accuracy_sla_bits), - latency_sla_ms: key.latency_sla_ms_bits.map(f64::from_bits), - }, - ) - .collect()) + acc.into_iter() + .map(|(key, entry)| OptimizerItem { + requirements: entry.requirements, + query_strings: entry.query_strings, + // Hz (queries per real second): 1000.0 / ms, not 1.0 / ms. + query_frequency_hz: entry.occurrences as f64 * 1000.0 / key.t_repeat_ms as f64, + occurrences: entry.occurrences, + t_repeat_ms: key.t_repeat_ms, + accuracy_sla: f64::from_bits(key.accuracy_sla_bits), + latency_sla_ms: key.latency_sla_ms_bits.map(f64::from_bits), + topk_k: entry.topk_k, + }) + .collect() +} + +/// The MILP's name for a SQL value column. +pub(crate) fn sql_metric(table: &str, value_column: &str) -> String { + format!("{table}.{value_column}") +} + +/// A SQL query's requirements and top-k k. A top-k must GROUP BY every +/// metadata column: it then ranks whole series, which is what sketch-bench +/// models. +fn sql_requirements( + query: &str, + tables: &[TableDefinition], + scrape_interval_ms: u64, + eval_time: f64, +) -> Result<(QueryRequirements, Option), String> { + let (data, topk) = parse_sql_query(query, tables, scrape_interval_ms, eval_time) + .map_err(|error| error.to_string())?; + let statistics = sql_statistics(&data, topk).map_err(|error| error.to_string())?; + let table = tables + .iter() + .find(|table| table.name == data.metric) + .ok_or_else(|| format!("unknown table {:?}", data.metric))?; + let grouping_labels = KeyByLabelNames::new(data.labels.iter().cloned().collect()); + if topk.is_some() && grouping_labels != KeyByLabelNames::new(table.metadata_columns.clone()) { + return Err(format!( + "top-k must GROUP BY every metadata column of {} ({:?})", + table.name, table.metadata_columns + )); + } + let requirements = QueryRequirements { + metric: sql_metric(&table.name, data.aggregation_info.get_value_column_name()), + statistics, + data_range_ms: sql_data_range_ms(&data), + grouping_labels, + spatial_filter_normalized: normalize_spatial_filter(""), + topk_count_events: topk.map(|t| t.count_events()), + topk_by_labels: None, + }; + Ok((requirements, topk.map(|t| t.k))) +} + +/// A `topk` query's literal `k`; `None` for anything else. +fn topk_k(query: &str) -> Option { + let Ok(Expr::Aggregate(aggregate)) = promql_parser::parser::parse(query) else { + return None; + }; + if aggregate.op.id() != token::T_TOPK { + return None; + } + match aggregate.param.as_deref()? { + Expr::NumberLiteral(number) if number.val >= 1.0 => Some(number.val as u64), + _ => None, + } } fn normalized_f64_bits(value: f64) -> u64 { @@ -228,6 +332,57 @@ fn extract_requirements( mod tests { use super::*; + #[test] + fn topk_k_reads_the_literal() { + assert_eq!(topk_k("topk(7, http_requests_total)"), Some(7)); + assert_eq!(topk_k("topk by (job) (3, http_requests_total)"), Some(3)); + assert_eq!(topk_k("sum(http_requests_total)"), None); + assert_eq!(topk_k("topk(scalar(up), http_requests_total)"), None); + } + + /// Sketch-bench ranks whole series, so a SQL top-k must GROUP BY every + /// metadata column. + #[test] + fn sql_topk_must_group_by_every_metadata_column() { + let tables: Vec = serde_yaml::from_str( + "[{name: t, time_column: time, value_columns: [v], metadata_columns: [host, region]}]", + ) + .unwrap(); + let topk = |by: &str| { + format!( + "SELECT {by}, SUM(v) AS s FROM t WHERE time BETWEEN DATEADD(s, -60, NOW()) \ + AND NOW() GROUP BY {by} ORDER BY s DESC LIMIT 5" + ) + }; + let aqes = + extract_sql_aqes(&[rqe(&topk("host, region"), 60_000)], &tables, 15_000).unwrap(); + assert_eq!(aqes[0].requirements.metric, "t.v"); + assert_eq!(aqes[0].topk_k, Some(5)); + assert!(matches!( + extract_sql_aqes(&[rqe(&topk("host"), 60_000)], &tables, 15_000), + Err(OptimizerError::UnsupportedLeaf { reason, .. }) if reason.contains("top-k") + )); + } + + #[test] + fn merged_topk_takes_the_largest_k_and_none_wins() { + let rqes = [ + rqe("topk(3, metric)", 60_000), + rqe("topk(7, metric)", 60_000), + ]; + let aqes = extract_aqes(&rqes, &empty_schema(), 15_000).unwrap(); + assert_eq!(aqes.len(), 1); + assert_eq!(aqes[0].topk_k, Some(7)); + + let rqes = [ + rqe("topk(3, metric)", 60_000), + rqe("topk(scalar(up), metric)", 60_000), + ]; + let aqes = extract_aqes(&rqes, &empty_schema(), 15_000).unwrap(); + assert_eq!(aqes.len(), 1); + assert_eq!(aqes[0].topk_k, None); + } + fn empty_schema() -> PromQLSchema { PromQLSchema::new() } diff --git a/asap-planner-rs/src/optimizer/milp.rs b/asap-planner-rs/src/optimizer/milp.rs index 60ec899c..1cdb3df6 100644 --- a/asap-planner-rs/src/optimizer/milp.rs +++ b/asap-planner-rs/src/optimizer/milp.rs @@ -11,14 +11,14 @@ use rqe_optimizer::{ }; use thiserror::Error; -use crate::config::input::{ControllerConfig, MetricDefinition}; +use crate::config::input::{ControllerConfig, MetricDefinition, SQLControllerConfig}; use std::path::Path; use anyhow::Context; use rqe_optimizer::saturation::SaturationCurves; -use super::aqe_extractor::{extract_aqes, RQE}; +use super::aqe_extractor::{extract_aqes, extract_sql_aqes, sql_metric, RQE}; use super::atomic_costs::load_flat_atomic_cost_table; use super::error::OptimizerError; use super::solution::OptimizerItem; @@ -38,6 +38,10 @@ pub enum MilpError { MissingMetricHints, #[error("workload metrics missing from `metrics:` hints: {0:?}")] MetricsWithoutHints(Vec), + #[error("table {0:?} lists no metadata_columns; the MILP planner doesn't discover them")] + TableWithoutMetadataColumns(String), + #[error("two SQL value columns both name metric {0:?}")] + DuplicateSqlMetric(String), #[error(transparent)] Extraction(#[from] OptimizerError), #[error("query {query:?}: accuracy_sla {accuracy_sla} must be between 0 and 1")] @@ -123,15 +127,36 @@ pub fn solve_milp( #[derive(Debug, Clone, Copy)] pub enum MilpConfig<'a> { Promql(&'a ControllerConfig), + Sql(&'a SQLControllerConfig), } impl MilpConfig<'_> { - /// Each workload metric's labels. - fn hints(self) -> Result, MilpError> { + /// Each workload metric's labels. A SQL table's value columns are its + /// metrics (named by `sql_metric`), labelled by its metadata columns. + pub(crate) fn hints(self) -> Result, MilpError> { match self { MilpConfig::Promql(config) => { config.metrics.clone().ok_or(MilpError::MissingMetricHints) } + MilpConfig::Sql(config) => { + let mut hints: Vec = Vec::new(); + for table in &config.tables { + if table.metadata_columns.is_empty() { + return Err(MilpError::TableWithoutMetadataColumns(table.name.clone())); + } + for column in &table.value_columns { + let metric = sql_metric(&table.name, column); + if hints.iter().any(|hint| hint.metric == metric) { + return Err(MilpError::DuplicateSqlMetric(metric)); + } + hints.push(MetricDefinition { + metric, + labels: table.metadata_columns.clone(), + }); + } + } + Ok(hints) + } } } @@ -145,6 +170,26 @@ impl MilpConfig<'_> { ])?; extract_hinted_items(config, scrape_interval_ms) } + MilpConfig::Sql(config) => { + reject_fields([ + ("windowing", config.windowing.is_some()), + ("sketch_parameters", config.sketch_parameters.is_some()), + ("aggregate_cleanup", config.aggregate_cleanup.is_some()), + ])?; + let rqes: Vec = config + .query_groups + .iter() + .flat_map(|qg| { + qg.queries.iter().map(|q| RQE { + query_string: q.clone(), + t_repeat_ms: qg.repetition_delay_ms, + accuracy_sla: qg.controller_options.accuracy_sla, + latency_sla_ms: qg.controller_options.latency_sla_ms, + }) + }) + .collect(); + Ok(extract_sql_aqes(&rqes, &config.tables, scrape_interval_ms)?) + } } } } @@ -185,8 +230,9 @@ pub struct MilpPlan { /// Loads the facts, costs and saturation curves and solves for the cheapest /// plan. pub fn plan_milp(config: MilpConfig, inputs: &MilpInputs) -> anyhow::Result { - let MilpConfig::Promql(promql) = config; - promql.warn_default_slas(); + if let MilpConfig::Promql(promql) = config { + promql.warn_default_slas(); + } let hints = config.hints()?; let facts = load_workload_facts(inputs.workload_facts, &hints, inputs.scrape_interval_ms)?; let curves = SaturationCurves::load(inputs.saturation_dir).with_context(|| { @@ -412,15 +458,11 @@ fn item_to_raqe(item: &OptimizerItem) -> Result { // A top-k Raqe asks for its query's literal k (the largest, should an // item carry several), which sizes and prices the heap. let topk_k = match capability { - Capability::TopKByValue | Capability::TopKByCount => Some( - item.query_strings - .iter() - .map(|q| super::milp_output::topk_k(q)) - .try_fold(0, |max, k| Some(max.max(k?))) - .ok_or_else(|| MilpError::TopkWithoutK { - query: query.clone(), - })?, - ), + Capability::TopKByValue | Capability::TopKByCount => { + Some(item.topk_k.ok_or_else(|| MilpError::TopkWithoutK { + query: query.clone(), + })?) + } _ => None, }; @@ -495,7 +537,7 @@ pub(super) mod tests { pub(crate) const SCRAPE_MS: u64 = 15_000; - const FACTS: &str = r#" + pub(crate) const FACTS: &str = r#" metrics: - metric: http_requests_total value_range: [1.0, 1000.0] @@ -652,7 +694,7 @@ metrics: )); assert_eq!(w.raqes[0].topk_k, Some(7)); let mut item = w.items[0].clone(); - item.query_strings = vec!["topk(scalar(up), http_requests_total)".into()]; + item.topk_k = None; assert!(matches!( item_to_raqe(&item), Err(MilpError::TopkWithoutK { .. }) diff --git a/asap-planner-rs/src/optimizer/milp_output.rs b/asap-planner-rs/src/optimizer/milp_output.rs index 6298b727..3529266d 100644 --- a/asap-planner-rs/src/optimizer/milp_output.rs +++ b/asap-planner-rs/src/optimizer/milp_output.rs @@ -5,7 +5,6 @@ use std::collections::{HashMap, HashSet}; use asap_types::enums::{CleanupPolicy, WindowType}; use indexmap::map::Entry; use indexmap::IndexMap; -use promql_parser::parser::{token, Expr}; use promql_utilities::data_model::KeyByLabelNames; use promql_utilities::query_logics::enums::AggregationType; use rqe_optimizer::milp::MilpSolution; @@ -19,8 +18,9 @@ use crate::planner::agg_config::IntermediateAggConfig; use crate::planner::labels::set_subpopulation_labels; use crate::planner_output::PlannerOutput; use crate::promql::generator::{build_inference_yaml, build_streaming_yaml}; +use crate::sql::generator::{build_sql_inference_yaml, build_sql_streaming_yaml}; -use super::aqe_extractor::contains_avg; +use super::aqe_extractor::{contains_avg, sql_metric}; use super::milp::{MilpConfig, MilpWorkload}; use super::solution::OptimizerItem; @@ -61,9 +61,11 @@ pub enum MilpOutputError { } /// Errors on queries whose plan can be costed but not written as configs: -/// avg queries, and range-query overrides, which only size retention. +/// PromQL avg queries, and range-query overrides, which only size retention. pub fn reject_unwritable_queries(config: MilpConfig) -> Result<(), MilpOutputError> { - let MilpConfig::Promql(config) = config; + let MilpConfig::Promql(config) = config else { + return Ok(()); + }; if let Some(query) = config .query_groups .iter() @@ -119,10 +121,33 @@ pub fn plan_to_planner_output( } } + // A SQL aggregation names its table and value column, and its metric is + // the table. + let sql_columns: HashMap = match config { + MilpConfig::Promql(_) => HashMap::new(), + MilpConfig::Sql(config) => config + .tables + .iter() + .flat_map(|table| { + table.value_columns.iter().map(move |column| { + ( + sql_metric(&table.name, column), + (table.name.as_str(), column.as_str()), + ) + }) + }) + .collect(), + }; + let mut aggregations: IndexMap = 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)?; + let mut aggregation = aggregation_config(&planned.deployment, items)?; + if let Some(&(table, column)) = sql_columns.get(&aggregation.metric) { + aggregation.metric = table.to_string(); + aggregation.table_name = Some(table.to_string()); + aggregation.value_column = Some(column.to_string()); + } let key = aggregation.identifying_key(); aggregations.entry(key.clone()).or_insert(aggregation); deployment_keys.push(key); @@ -160,6 +185,10 @@ pub fn plan_to_planner_output( build_inference_yaml(CleanupPolicy::NoCleanup, &queries, &id_map, &schema)?, ) } + MilpConfig::Sql(config) => ( + build_sql_streaming_yaml(config, &aggregations, &id_map)?, + build_sql_inference_yaml(config, CleanupPolicy::NoCleanup, &queries, &id_map)?, + ), }; Ok(PlannerOutput::from_output(GeneratorOutput { punted_queries: Vec::new(), @@ -298,20 +327,6 @@ fn heap_size(deployment: &Deployment) -> Result { .ok_or_else(|| missing_param(&deployment.config.sketch, "heap")) } -/// A `topk` query's literal `k`; `None` for anything else. -pub(crate) fn topk_k(query: &str) -> Option { - let Ok(Expr::Aggregate(aggregate)) = promql_parser::parser::parse(query) else { - return None; - }; - if aggregate.op.id() != token::T_TOPK { - return None; - } - match aggregate.param.as_deref()? { - Expr::NumberLiteral(number) if number.val >= 1.0 => Some(number.val as u64), - _ => None, - } -} - #[cfg(test)] mod tests { use std::collections::BTreeMap; @@ -319,18 +334,25 @@ mod tests { use asap_types::aggregation_config::AggregationConfig; use asap_types::enums::QueryLanguage; use rqe_optimizer::milp::Objective; - use rqe_optimizer::AtomicCostEntry; + use rqe_optimizer::{AtomicCostEntry, WorkloadFacts}; use serde_json::json; use super::*; - use crate::config::input::ControllerConfig; - use crate::optimizer::milp::tests::{config, costs, facts, group, table_accuracy, SCRAPE_MS}; + use crate::config::input::{ControllerConfig, SQLControllerConfig}; + use crate::optimizer::milp::tests::{ + config, costs, facts, group, table_accuracy, FACTS, SCRAPE_MS, + }; use crate::optimizer::milp::{build_milp_workload, solve_milp}; + use crate::optimizer::workload_facts::parse_workload_facts; fn plan(groups: &str) -> (ControllerConfig, MilpWorkload, MilpSolution) { let config = config(groups); - let facts = facts(&config); - let workload = build_milp_workload(MilpConfig::Promql(&config), &facts, SCRAPE_MS).unwrap(); + let (workload, solution) = solve(MilpConfig::Promql(&config), &facts(&config)); + (config, workload, solution) + } + + fn solve(config: MilpConfig, facts: &WorkloadFacts) -> (MilpWorkload, MilpSolution) { + let workload = build_milp_workload(config, facts, SCRAPE_MS).unwrap(); let mut costs = costs(); for cost in &mut costs { cost.sketch_config["params"] = match cost.sketch.as_str() { @@ -345,20 +367,28 @@ mod tests { } let solution = solve_milp( &workload, - &facts, + facts, &costs, Objective::default(), false, &table_accuracy, ) .unwrap(); - (config, workload, solution) + (workload, solution) } /// The aggregation each query names, parsed back the way the engine reads it. fn aggregation_for(output: &PlannerOutput, query: &str) -> AggregationConfig { - let inference = output.to_inference_config(QueryLanguage::promql).unwrap(); - let streaming = output.to_streaming_config(QueryLanguage::promql).unwrap(); + aggregation_in(output, query, QueryLanguage::promql) + } + + fn aggregation_in( + output: &PlannerOutput, + query: &str, + language: QueryLanguage, + ) -> AggregationConfig { + let inference = output.to_inference_config(language).unwrap(); + let streaming = output.to_streaming_config(language).unwrap(); let query_config = inference .query_configs .iter() @@ -446,6 +476,113 @@ mod tests { } } + /// What the engine builds from an aggregation, without its id, metric + /// and SQL columns. + fn layout(a: &AggregationConfig) -> String { + let parameters: BTreeMap<_, _> = a.parameters.iter().collect(); + format!( + "{:?}", + ( + a.aggregation_type, + &a.aggregation_sub_type, + parameters, + &a.grouping_labels, + &a.aggregated_labels, + &a.rollup_labels, + a.window_size_ms, + a.slide_interval_ms, + a.window_type, + &a.spatial_filter_normalized, + a.num_aggregates_to_retain, + ) + ) + } + + /// A workload written in SQL, on a table whose value column has + /// `http_requests_total`'s labels and facts, plans exactly as in PromQL. + #[test] + fn sql_twin_plans_like_promql() { + let last = |secs: u32| format!("time BETWEEN DATEADD(s, -{secs}, NOW()) AND NOW()"); + let twins = [ + ( + "sum by (job) (http_requests_total)", + format!( + "SELECT job, SUM(requests) FROM http WHERE {} GROUP BY job", + last(15) + ), + ), + ( + "sum by (job) (count_over_time(http_requests_total[1m]))", + format!( + "SELECT job, COUNT(requests) FROM http WHERE {} GROUP BY job", + last(60) + ), + ), + ( + "quantile_over_time(0.99, http_requests_total[5m])", + format!( + "SELECT instance, job, quantile(0.99)(requests) FROM http WHERE {} \ + GROUP BY instance, job", + last(300) + ), + ), + ( + "topk(5, sum_over_time(http_requests_total[1m]))", + format!( + "SELECT instance, job, SUM(requests) AS s FROM http WHERE {} \ + GROUP BY instance, job ORDER BY s DESC LIMIT 5", + last(60) + ), + ), + ( + "topk(5, count_over_time(http_requests_total[1m]))", + format!( + "SELECT instance, job, COUNT(requests) AS c FROM http WHERE {} \ + GROUP BY instance, job ORDER BY c DESC LIMIT 5", + last(60) + ), + ), + ]; + let promql_groups: String = twins.iter().map(|(q, _)| group(q, 0.99)).collect(); + let (promql, promql_workload, promql_solution) = plan(&promql_groups); + + let sql_groups: String = twins.iter().map(|(_, q)| group(q, 0.99)).collect(); + let sql: SQLControllerConfig = serde_yaml::from_str(&format!( + "tables:\n - name: http\n time_column: time\n value_columns: [requests]\n \ + metadata_columns: [instance, job]\nquery_groups:\n{sql_groups}" + )) + .unwrap(); + let facts = parse_workload_facts( + &FACTS.replace("http_requests_total", "http.requests"), + &MilpConfig::Sql(&sql).hints().unwrap(), + SCRAPE_MS, + ) + .unwrap(); + let (sql_workload, sql_solution) = solve(MilpConfig::Sql(&sql), &facts); + + let (promql_cost, sql_cost) = (&promql_solution.plan_cost, &sql_solution.plan_cost); + assert_eq!(sql_cost.cpu_secs_per_sec(), promql_cost.cpu_secs_per_sec()); + assert_eq!(sql_cost.memory_bytes(), promql_cost.memory_bytes()); + + let promql_output = plan_to_planner_output( + MilpConfig::Promql(&promql), + &promql_workload, + &promql_solution, + ) + .unwrap(); + let sql_output = + plan_to_planner_output(MilpConfig::Sql(&sql), &sql_workload, &sql_solution).unwrap(); + assert_eq!(sql_output.streaming_aggregation_count(), 5); + for (promql_query, sql_query) in &twins { + let expected = aggregation_for(&promql_output, promql_query); + let a = aggregation_in(&sql_output, sql_query, QueryLanguage::sql); + assert_eq!(layout(&a), layout(&expected), "{sql_query}"); + assert_eq!(a.metric, "http.requests"); + assert_eq!(a.table_name.as_deref(), Some("http")); + assert_eq!(a.value_column.as_deref(), Some("requests")); + } + } + /// The heap `query`'s deployment needs for it: `m · k` for the windows /// its lookback merges and its own `k` (`Deployment::heap_needed`). fn heap_needed(workload: &MilpWorkload, solution: &MilpSolution, query: &str) -> u64 { @@ -711,11 +848,4 @@ mod tests { )); assert!(!contains_avg("sum by (job) (http_requests_total) * 2")); } - - #[test] - fn topk_k_reads_the_literal() { - assert_eq!(topk_k("topk(7, http_requests_total)"), Some(7)); - assert_eq!(topk_k("topk by (job) (3, http_requests_total)"), Some(3)); - assert_eq!(topk_k("sum(http_requests_total)"), None); - } } diff --git a/asap-planner-rs/src/optimizer/solution.rs b/asap-planner-rs/src/optimizer/solution.rs index 44d79537..cc083f71 100644 --- a/asap-planner-rs/src/optimizer/solution.rs +++ b/asap-planner-rs/src/optimizer/solution.rs @@ -27,4 +27,8 @@ pub struct OptimizerItem { /// Maximum query latency (ms) for every RQE contributing to this item; /// `None` means no limit. pub latency_sla_ms: Option, + + /// The largest literal `k` of the item's top-k queries; `None` when one + /// has no literal `k`, or for anything but top-k. + pub topk_k: Option, } diff --git a/asap-planner-rs/src/planner/sql.rs b/asap-planner-rs/src/planner/sql.rs index ad120503..bbf2e67d 100644 --- a/asap-planner-rs/src/planner/sql.rs +++ b/asap-planner-rs/src/planner/sql.rs @@ -3,7 +3,9 @@ use std::collections::HashSet; use asap_types::enums::{CleanupPolicy, WindowType}; use promql_utilities::data_model::KeyByLabelNames; use promql_utilities::query_logics::enums::{AggregationType, QueryTreatmentType, Statistic}; -use sql_utilities::ast_matching::sqlhelper::{detect_sql_topk, Table, TimeInfo}; +use sql_utilities::ast_matching::sqlhelper::{ + detect_sql_topk, SQLQueryData, SqlTopk, Table, TimeInfo, +}; use sql_utilities::ast_matching::sqlpattern_matcher::SQLPatternMatcher; use sql_utilities::ast_matching::sqlpattern_parser::SQLPatternParser; use sql_utilities::ast_matching::SQLSchema; @@ -58,58 +60,27 @@ impl SQLSingleQueryProcessor { &self, query_evaluation_time: f64, ) -> Result<(Vec, Option), ControllerError> { - let schema = build_sql_schema(&self.table_definitions); - - // Parse SQL - let stmts = SqlParser::parse_sql(&ClickHouseDialect {}, &self.query_string) - .map_err(|e| ControllerError::SqlParse(e.to_string()))?; - - // Parse query into SQLQueryData - let qdata = SQLPatternParser::new(&schema, query_evaluation_time) - .parse_query(&stmts) - .ok_or_else(|| { - ControllerError::SqlParse(format!( - "Failed to parse SQL query: {}", - self.query_string - )) - })?; - - // Match query to pattern - // SQLPatternMatcher.scrape_interval is in seconds (SQL timestamps are seconds-based). - let sql_query = SQLPatternMatcher::new( - schema.clone(), - self.data_ingestion_interval_ms as f64 / 1000.0, - ) - .query_info_to_pattern(&qdata); - - if !sql_query.is_valid() { - return Err(ControllerError::SqlParse(sql_query.msg.unwrap_or_default())); - } - - let n = sql_query.query_data.len(); - - if n != 1 { - return Err(ControllerError::SqlParse(format!( - "Nested SQL queries (n={}) are not supported", - n - ))); - } + let (query_data, sql_topk) = parse_sql_query( + &self.query_string, + &self.table_definitions, + self.data_ingestion_interval_ms, + query_evaluation_time, + )?; // Determine fields from query vecs - let agg_info = &sql_query.query_data[0].aggregation_info; - let labels = &sql_query.query_data[0].labels; - let table_name = &sql_query.query_data[0].metric; + let agg_info = &query_data.aggregation_info; + let labels = &query_data.labels; + let table_name = &query_data.metric; let value_column = agg_info.get_value_column_name().to_string(); // Compute window let mut window_cfg = compute_sql_window( - &sql_query.query_data[0].time_info, + &query_data.time_info, self.data_ingestion_interval_ms, self.t_repeat_ms, )?; - let data_range_ms = - (sql_query.query_data[0].time_info.get_duration() * 1000.0).round() as u64; + let data_range_ms = sql_data_range_ms(&query_data); crate::planner::window::apply_windowing_override( &mut window_cfg, data_range_ms, @@ -122,15 +93,8 @@ impl SQLSingleQueryProcessor { // Label routing let spatial_output = KeyByLabelNames::new(labels.iter().cloned().collect::>()); - // Top-k needs ORDER BY / LIMIT from the parser; SQLPatternMatcher drops them - // when building `sql_query.query_data[0]`, so use `qdata` not query_data[0]. - let sql_topk = detect_sql_topk(&qdata); let treatment_type = get_sql_treatment_type(agg_info.get_name()); - let statistics = if sql_topk.is_some() { - vec![Statistic::Topk] - } else { - get_sql_statistics(agg_info.get_name())? - }; + let statistics = sql_statistics(&query_data, sql_topk)?; let rollup = if statistics.contains(&Statistic::Cardinality) { // Distinct target is value_column, not a rollup label dimension. KeyByLabelNames::empty() @@ -182,18 +146,11 @@ impl SQLSingleQueryProcessor { } } - // SQLPatternParser always produces second-based durations; convert to ms. - // For a single-scrape-interval query this equals data_ingestion_interval_ms - // by construction (the matcher's classification boundary), so this is a - // plain unconditional formula, not a special case. - let t_lookback_ms = - (sql_query.query_data[0].time_info.get_duration() * 1000.0).round() as u64; - let cleanup_param = if self.cleanup_policy == CleanupPolicy::NoCleanup { None } else { Some( - get_sql_cleanup_param(self.cleanup_policy, t_lookback_ms, self.t_repeat_ms) + get_sql_cleanup_param(self.cleanup_policy, data_range_ms, self.t_repeat_ms) .map_err(ControllerError::PlannerError)?, ) }; @@ -202,6 +159,57 @@ impl SQLSingleQueryProcessor { } } +/// Parses `query` and matches it against `tables`: a single-level query's +/// data, and its top-k clause if it is one. +pub(crate) fn parse_sql_query( + query: &str, + tables: &[TableDefinition], + data_ingestion_interval_ms: u64, + query_evaluation_time: f64, +) -> Result<(SQLQueryData, Option), ControllerError> { + let schema = build_sql_schema(tables); + let stmts = SqlParser::parse_sql(&ClickHouseDialect {}, query) + .map_err(|e| ControllerError::SqlParse(e.to_string()))?; + let qdata = SQLPatternParser::new(&schema, query_evaluation_time) + .parse_query(&stmts) + .ok_or_else(|| ControllerError::SqlParse(format!("Failed to parse SQL query: {query}")))?; + + // SQLPatternMatcher.scrape_interval is in seconds (SQL timestamps are seconds-based). + let sql_query = SQLPatternMatcher::new(schema, data_ingestion_interval_ms as f64 / 1000.0) + .query_info_to_pattern(&qdata); + if !sql_query.is_valid() { + return Err(ControllerError::SqlParse(sql_query.msg.unwrap_or_default())); + } + let n = sql_query.query_data.len(); + let Ok([query_data]) = <[SQLQueryData; 1]>::try_from(sql_query.query_data) else { + return Err(ControllerError::SqlParse(format!( + "Nested SQL queries (n={n}) are not supported" + ))); + }; + // Top-k needs ORDER BY / LIMIT from the parser; SQLPatternMatcher drops + // them from its query data. + Ok((query_data, detect_sql_topk(&qdata))) +} + +/// The statistics `query_data` needs: top-k, or its aggregation's. +pub(crate) fn sql_statistics( + query_data: &SQLQueryData, + topk: Option, +) -> Result, ControllerError> { + match topk { + Some(_) => Ok(vec![Statistic::Topk]), + None => get_sql_statistics(query_data.aggregation_info.get_name()), + } +} + +/// SQLPatternParser always produces second-based durations; convert to ms. +/// For a single-scrape-interval query this equals data_ingestion_interval_ms +/// by construction (the matcher's classification boundary), so this is a +/// plain unconditional formula, not a special case. +pub(crate) fn sql_data_range_ms(query_data: &SQLQueryData) -> u64 { + (query_data.time_info.get_duration() * 1000.0).round() as u64 +} + fn build_sql_schema(tables: &[TableDefinition]) -> SQLSchema { let table_vec: Vec = tables .iter() diff --git a/asap-planner-rs/src/sql/generator.rs b/asap-planner-rs/src/sql/generator.rs index 020d49ab..aa6df271 100644 --- a/asap-planner-rs/src/sql/generator.rs +++ b/asap-planner-rs/src/sql/generator.rs @@ -182,7 +182,7 @@ fn build_tables_yaml(config: &SQLControllerConfig) -> Vec { .collect() } -fn build_sql_streaming_yaml( +pub(crate) fn build_sql_streaming_yaml( config: &SQLControllerConfig, dedup_map: &IndexMap, id_map: &HashMap, @@ -205,7 +205,7 @@ fn build_sql_streaming_yaml( Ok(YamlValue::Mapping(root)) } -fn build_sql_inference_yaml( +pub(crate) fn build_sql_inference_yaml( config: &SQLControllerConfig, cleanup_policy: CleanupPolicy, query_plan_map: &IndexMap, From 1483156fe5d0145cac51f56fb3e4d2073f870cf9 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Mon, 12 Oct 2026 17:48:38 -0400 Subject: [PATCH 3/4] feat(planner): accept --query-language sql with --planner milp Co-Authored-By: Claude Opus 5.5 --- asap-planner-rs/src/main.rs | 29 ++++++++++++++++++++--------- 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/asap-planner-rs/src/main.rs b/asap-planner-rs/src/main.rs index 69f5a95e..62d71e96 100644 --- a/asap-planner-rs/src/main.rs +++ b/asap-planner-rs/src/main.rs @@ -4,7 +4,7 @@ use asap_planner::optimizer::{ }; use asap_planner::{ Controller, ControllerConfig, ElasticController, ElasticRuntimeOptions, RuntimeOptions, - SQLController, SQLRuntimeOptions, StreamingEngine, + SQLController, SQLControllerConfig, SQLRuntimeOptions, StreamingEngine, }; use asap_types::enums::QueryLanguage; use clap::Parser; @@ -57,8 +57,9 @@ struct Args { #[arg(long = "clickhouse-database", required = false)] clickhouse_database: Option, - /// `milp` plans with sketch-bench's rqe-optimizer: PromQL with - /// --input_config only, labels from its `metrics:` hints. + /// `milp` plans with sketch-bench's rqe-optimizer: PromQL or SQL with + /// --input_config only, labels from its `metrics:` hints or tables' + /// `metadata_columns`. #[arg(long, value_enum, default_value = "legacy")] planner: PlannerArg, @@ -216,16 +217,19 @@ fn main() -> anyhow::Result<()> { fn run_milp(args: &Args) -> anyhow::Result<()> { anyhow::ensure!( - matches!(args.query_language, QueryLanguage::promql), - "--planner milp supports only --query-language promql" + matches!( + args.query_language, + QueryLanguage::promql | QueryLanguage::sql + ), + "--planner milp supports only --query-language promql or sql" ); anyhow::ensure!( args.query_log.is_none(), - "--planner milp needs --input_config: query logs have no `metrics:` hints" + "--planner milp needs --input_config: query logs have no label hints" ); anyhow::ensure!( args.prometheus_url.is_none(), - "--planner milp takes labels from the config's `metrics:` hints, not --prometheus-url" + "--planner milp takes labels from the config, not --prometheus-url" ); anyhow::ensure!( !args.enable_punting && args.range_duration_ms == 0 && args.step_ms == 0, @@ -247,8 +251,15 @@ fn run_milp(args: &Args) -> anyhow::Result<()> { 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)?)?; - let config = MilpConfig::Promql(&config); + let yaml = std::fs::read_to_string(config_path)?; + let (promql, sql): (ControllerConfig, SQLControllerConfig); + let config = if args.query_language == QueryLanguage::sql { + sql = serde_yaml::from_str(&yaml)?; + MilpConfig::Sql(&sql) + } else { + promql = serde_yaml::from_str(&yaml)?; + MilpConfig::Promql(&promql) + }; // Fail before solving: the plan is always written. reject_unwritable_queries(config)?; let plan = plan_milp( From 553b49711342009354965f9730e6257f12421cbd Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Mon, 12 Oct 2026 21:04:23 -0400 Subject: [PATCH 4/4] fix(planner): address review on --planner milp for SQL Co-Authored-By: Claude Opus 5.5 --- .../src/optimizer/aqe_extractor.rs | 62 ++++++-- asap-planner-rs/src/optimizer/milp.rs | 140 ++++++++++-------- asap-planner-rs/src/optimizer/milp_output.rs | 22 +-- 3 files changed, 129 insertions(+), 95 deletions(-) diff --git a/asap-planner-rs/src/optimizer/aqe_extractor.rs b/asap-planner-rs/src/optimizer/aqe_extractor.rs index a993904c..0ef3c933 100644 --- a/asap-planner-rs/src/optimizer/aqe_extractor.rs +++ b/asap-planner-rs/src/optimizer/aqe_extractor.rs @@ -1,5 +1,4 @@ use std::collections::HashMap; -use std::time::{SystemTime, UNIX_EPOCH}; use asap_types::query_requirements::{build_query_requirements_promql, QueryRequirements}; use asap_types::utils::normalize_spatial_filter; @@ -82,7 +81,9 @@ pub fn extract_aqes( reason, }, )?; - let k = topk_k(&leaf); + let k = (req.statistics == [Statistic::Topk]) + .then(|| topk_k(&leaf)) + .flatten(); leaves.push((req, leaf, k, rqe)); } } @@ -96,23 +97,17 @@ pub fn extract_sql_aqes( tables: &[TableDefinition], scrape_interval_ms: u64, ) -> Result, OptimizerError> { - // Only the queries' durations are used, which don't depend on it. - let eval_time = SystemTime::now() - .duration_since(UNIX_EPOCH) - .expect("clock after 1970") - .as_secs_f64(); let mut leaves = Vec::new(); for rqe in rqes { check_repeat_interval(rqe)?; let query = &rqe.query_string; - let (req, k) = - sql_requirements(query, tables, scrape_interval_ms, eval_time).map_err(|reason| { - OptimizerError::UnsupportedLeaf { - query: query.clone(), - leaf: query.clone(), - reason, - } - })?; + let (req, k) = sql_requirements(query, tables, scrape_interval_ms, SQL_EVAL_TIME).map_err( + |reason| OptimizerError::UnsupportedLeaf { + query: query.clone(), + leaf: query.clone(), + reason, + }, + )?; leaves.push((req, query.clone(), k, rqe)); } Ok(merge_leaves(leaves)) @@ -167,6 +162,10 @@ fn merge_leaves(leaves: Vec<(QueryRequirements, String, Option, &RQE)>) -> .collect() } +/// The time `NOW()` means when SQL queries are planned: fixed, so plans +/// don't depend on the clock. +const SQL_EVAL_TIME: f64 = 1_700_000_000.0; + /// The MILP's name for a SQL value column. pub(crate) fn sql_metric(table: &str, value_column: &str) -> String { format!("{table}.{value_column}") @@ -184,10 +183,21 @@ fn sql_requirements( let (data, topk) = parse_sql_query(query, tables, scrape_interval_ms, eval_time) .map_err(|error| error.to_string())?; let statistics = sql_statistics(&data, topk).map_err(|error| error.to_string())?; + let aggregation = data.aggregation_info.get_name(); + if statistics.len() > 1 { + return Err(format!("{aggregation} is not supported by --planner milp")); + } let table = tables .iter() .find(|table| table.name == data.metric) .ok_or_else(|| format!("unknown table {:?}", data.metric))?; + let column = data.aggregation_info.get_value_column_name(); + if !table.value_columns.iter().any(|c| c == column) { + return Err(format!( + "{aggregation} reads {column:?}, which is not a value column of {}", + table.name + )); + } let grouping_labels = KeyByLabelNames::new(data.labels.iter().cloned().collect()); if topk.is_some() && grouping_labels != KeyByLabelNames::new(table.metadata_columns.clone()) { return Err(format!( @@ -196,7 +206,7 @@ fn sql_requirements( )); } let requirements = QueryRequirements { - metric: sql_metric(&table.name, data.aggregation_info.get_value_column_name()), + metric: sql_metric(&table.name, column), statistics, data_range_ms: sql_data_range_ms(&data), grouping_labels, @@ -364,6 +374,26 @@ mod tests { )); } + #[test] + fn sql_rejects_non_value_columns_and_avg() { + let tables: Vec = serde_yaml::from_str( + "[{name: t, time_column: time, value_columns: [v], metadata_columns: [host, region]}]", + ) + .unwrap(); + let reason = |select: &str| { + let query = format!( + "SELECT host, {select} FROM t WHERE time BETWEEN DATEADD(s, -60, NOW()) \ + AND NOW() GROUP BY host" + ); + match extract_sql_aqes(&[rqe(&query, 60_000)], &tables, 15_000) { + Err(OptimizerError::UnsupportedLeaf { reason, .. }) => reason, + other => panic!("{select}: {other:?}"), + } + }; + assert!(reason("COUNT(DISTINCT region)").contains("not a value column of t")); + assert!(reason("AVG(v)").contains("not supported by --planner milp")); + } + #[test] fn merged_topk_takes_the_largest_k_and_none_wins() { let rqes = [ diff --git a/asap-planner-rs/src/optimizer/milp.rs b/asap-planner-rs/src/optimizer/milp.rs index 1cdb3df6..7480c543 100644 --- a/asap-planner-rs/src/optimizer/milp.rs +++ b/asap-planner-rs/src/optimizer/milp.rs @@ -11,7 +11,9 @@ use rqe_optimizer::{ }; use thiserror::Error; -use crate::config::input::{ControllerConfig, MetricDefinition, SQLControllerConfig}; +use crate::config::input::{ + ControllerConfig, ControllerOptions, MetricDefinition, SQLControllerConfig, TableDefinition, +}; use std::path::Path; @@ -130,64 +132,65 @@ pub enum MilpConfig<'a> { Sql(&'a SQLControllerConfig), } -impl MilpConfig<'_> { +impl<'a> MilpConfig<'a> { /// Each workload metric's labels. A SQL table's value columns are its - /// metrics (named by `sql_metric`), labelled by its metadata columns. + /// metrics, labelled by its metadata columns. pub(crate) fn hints(self) -> Result, MilpError> { - match self { - MilpConfig::Promql(config) => { - config.metrics.clone().ok_or(MilpError::MissingMetricHints) - } - MilpConfig::Sql(config) => { - let mut hints: Vec = Vec::new(); - for table in &config.tables { - if table.metadata_columns.is_empty() { - return Err(MilpError::TableWithoutMetadataColumns(table.name.clone())); - } - for column in &table.value_columns { - let metric = sql_metric(&table.name, column); - if hints.iter().any(|hint| hint.metric == metric) { - return Err(MilpError::DuplicateSqlMetric(metric)); - } - hints.push(MetricDefinition { - metric, - labels: table.metadata_columns.clone(), - }); - } + let MilpConfig::Promql(config) = self else { + let mut hints: Vec = Vec::new(); + for (metric, table, _) in self.sql_columns() { + if table.metadata_columns.is_empty() { + return Err(MilpError::TableWithoutMetadataColumns(table.name.clone())); + } + if hints.iter().any(|hint| hint.metric == metric) { + return Err(MilpError::DuplicateSqlMetric(metric)); } - Ok(hints) + hints.push(MetricDefinition { + metric, + labels: table.metadata_columns.clone(), + }); } - } + return Ok(hints); + }; + config.metrics.clone().ok_or(MilpError::MissingMetricHints) + } + + /// Every SQL value column as `(metric, table, column)`; empty for PromQL. + pub(crate) fn sql_columns(self) -> Vec<(String, &'a TableDefinition, &'a str)> { + let MilpConfig::Sql(config) = self else { + return Vec::new(); + }; + config + .tables + .iter() + .flat_map(|table| { + table + .value_columns + .iter() + .map(move |column| (sql_metric(&table.name, column), table, column.as_str())) + }) + .collect() } fn items(self, scrape_interval_ms: u64) -> Result, MilpError> { + let (windowing, sketch_parameters, aggregate_cleanup) = match self { + MilpConfig::Promql(c) => (&c.windowing, &c.sketch_parameters, &c.aggregate_cleanup), + MilpConfig::Sql(c) => (&c.windowing, &c.sketch_parameters, &c.aggregate_cleanup), + }; + reject_fields([ + ("windowing", windowing.is_some()), + ("sketch_parameters", sketch_parameters.is_some()), + ("aggregate_cleanup", aggregate_cleanup.is_some()), + ])?; match self { - MilpConfig::Promql(config) => { - reject_fields([ - ("windowing", config.windowing.is_some()), - ("sketch_parameters", config.sketch_parameters.is_some()), - ("aggregate_cleanup", config.aggregate_cleanup.is_some()), - ])?; - extract_hinted_items(config, scrape_interval_ms) - } + MilpConfig::Promql(config) => extract_hinted_items(config, scrape_interval_ms), MilpConfig::Sql(config) => { - reject_fields([ - ("windowing", config.windowing.is_some()), - ("sketch_parameters", config.sketch_parameters.is_some()), - ("aggregate_cleanup", config.aggregate_cleanup.is_some()), - ])?; - let rqes: Vec = config - .query_groups - .iter() - .flat_map(|qg| { - qg.queries.iter().map(|q| RQE { - query_string: q.clone(), - t_repeat_ms: qg.repetition_delay_ms, - accuracy_sla: qg.controller_options.accuracy_sla, - latency_sla_ms: qg.controller_options.latency_sla_ms, - }) - }) - .collect(); + let rqes = config_to_rqes( + config + .query_groups + .iter() + .map(|qg| (&qg.queries, qg.repetition_delay_ms, &qg.controller_options)), + ); Ok(extract_sql_aqes(&rqes, &config.tables, scrape_interval_ms)?) } } @@ -395,7 +398,12 @@ fn extract_hinted_items( return Err(MilpError::MissingMetricHints); } let schema = config.schema_from_hints(); - let rqes = config_to_rqes(config); + let rqes = config_to_rqes( + config + .query_groups + .iter() + .map(|qg| (&qg.queries, qg.repetition_delay_ms, &qg.controller_options)), + ); let aqes = extract_aqes(&rqes, &schema, scrape_interval_ms)?; // Requirement extraction treats an unknown metric as having no labels, @@ -413,18 +421,19 @@ fn extract_hinted_items( Ok(aqes) } -/// Convert a `ControllerConfig`'s query groups into a flat list of RQEs. -/// Each (query, repetition_delay_ms) pair becomes one RQE. -fn config_to_rqes(config: &ControllerConfig) -> Vec { - config - .query_groups - .iter() - .flat_map(|qg| { - qg.queries.iter().map(|q| RQE { +/// Flattens query groups, each `(queries, repetition_delay_ms, options)`, +/// into one RQE per query. +fn config_to_rqes<'g>( + groups: impl IntoIterator, u64, &'g ControllerOptions)>, +) -> Vec { + groups + .into_iter() + .flat_map(|(queries, t_repeat_ms, options)| { + queries.iter().map(move |q| RQE { query_string: q.clone(), - t_repeat_ms: qg.repetition_delay_ms, - accuracy_sla: qg.controller_options.accuracy_sla, - latency_sla_ms: qg.controller_options.latency_sla_ms, + t_repeat_ms, + accuracy_sla: options.accuracy_sla, + latency_sla_ms: options.latency_sla_ms, }) }) .collect() @@ -1025,7 +1034,12 @@ metrics: group_at("sum_over_time(a[5m])", 60_000), group_at("sum_over_time(b[5m])", 30_000) )); - let rqes = config_to_rqes(&config); + let rqes = config_to_rqes( + config + .query_groups + .iter() + .map(|qg| (&qg.queries, qg.repetition_delay_ms, &qg.controller_options)), + ); assert_eq!(rqes.len(), 2); assert_eq!(rqes[0].t_repeat_ms, 60_000); assert_eq!(rqes[1].t_repeat_ms, 30_000); diff --git a/asap-planner-rs/src/optimizer/milp_output.rs b/asap-planner-rs/src/optimizer/milp_output.rs index 3529266d..dbe77f08 100644 --- a/asap-planner-rs/src/optimizer/milp_output.rs +++ b/asap-planner-rs/src/optimizer/milp_output.rs @@ -20,7 +20,7 @@ use crate::planner_output::PlannerOutput; use crate::promql::generator::{build_inference_yaml, build_streaming_yaml}; use crate::sql::generator::{build_sql_inference_yaml, build_sql_streaming_yaml}; -use super::aqe_extractor::{contains_avg, sql_metric}; +use super::aqe_extractor::contains_avg; use super::milp::{MilpConfig, MilpWorkload}; use super::solution::OptimizerItem; @@ -123,21 +123,11 @@ pub fn plan_to_planner_output( // A SQL aggregation names its table and value column, and its metric is // the table. - let sql_columns: HashMap = match config { - MilpConfig::Promql(_) => HashMap::new(), - MilpConfig::Sql(config) => config - .tables - .iter() - .flat_map(|table| { - table.value_columns.iter().map(move |column| { - ( - sql_metric(&table.name, column), - (table.name.as_str(), column.as_str()), - ) - }) - }) - .collect(), - }; + let sql_columns: HashMap = config + .sql_columns() + .into_iter() + .map(|(metric, table, column)| (metric, (table.name.as_str(), column))) + .collect(); let mut aggregations: IndexMap = IndexMap::new(); let mut deployment_keys = Vec::with_capacity(served.len());