@@ -5,8 +5,7 @@ use promql_utilities::query_logics::enums::Statistic;
55use rqe_optimizer:: candidates:: { build_all_candidates, eligible_deployments_for} ;
66use rqe_optimizer:: milp:: { minimize, MilpSolution , Objective } ;
77use rqe_optimizer:: {
8- validate_facts, AccuracyDirection , AtomicCostEntry , Capability , Deployment , LabelSet , Raqe ,
9- WorkloadFacts ,
8+ validate_facts, AccuracyDirection , AtomicCostEntry , Capability , LabelSet , Raqe , WorkloadFacts ,
109} ;
1110use thiserror:: Error ;
1211
@@ -37,20 +36,22 @@ pub enum MilpError {
3736 query : String ,
3837 statistics : Vec < Statistic > ,
3938 } ,
39+ #[ error( "query {query:?}: topk must rank by value (sum_over_time) or count (count_over_time)" ) ]
40+ TopkWeightingUnknown { query : String } ,
4041 #[ error( "no eligible deployment for raqes: {0:?}" ) ]
4142 Unservable ( Vec < String > ) ,
4243 #[ error( "MILP solve failed: {0}" ) ]
4344 Solver ( String ) ,
4445}
4546
46- /// Candidate deployments and the MILP's choice among them . Only families in
47- /// sketch-bench's `DEPLOYABLE_FAMILIES` are candidates.
47+ /// The cheapest plan . Only families in sketch-bench's `DEPLOYABLE_FAMILIES`
48+ /// are candidates.
4849pub fn solve_milp (
4950 workload : & MilpWorkload ,
5051 facts : & WorkloadFacts ,
5152 costs : & [ AtomicCostEntry ] ,
5253 objective : Objective ,
53- ) -> Result < ( Vec < Deployment > , MilpSolution ) , MilpError > {
54+ ) -> Result < MilpSolution , MilpError > {
5455 let deployments = build_all_candidates ( & workload. raqes , costs, facts, false ) ;
5556 tracing:: debug!(
5657 candidates = deployments. len( ) ,
@@ -59,7 +60,7 @@ pub fn solve_milp(
5960 ) ;
6061 let mut missing = Vec :: new ( ) ;
6162 for raqe in & workload. raqes {
62- let eligible = eligible_deployments_for ( raqe, & deployments) . len ( ) ;
63+ let eligible = eligible_deployments_for ( raqe, & deployments, facts ) . len ( ) ;
6364 tracing:: debug!( raqe = %raqe. id, eligible, "milp: eligible deployments" ) ;
6465 if eligible == 0 {
6566 missing. push ( raqe. id . clone ( ) ) ;
@@ -74,10 +75,10 @@ pub fn solve_milp(
7475 objective = objective. value( & solution. plan_cost) ,
7576 cpu_secs_per_sec = solution. plan_cost. cpu_secs_per_sec( ) ,
7677 memory_bytes = solution. plan_cost. memory_bytes( ) ,
77- mapping = ? solution. mapping ,
78+ deployments = solution. deployments . len ( ) ,
7879 "milp: solved"
7980 ) ;
80- Ok ( ( deployments , solution) )
81+ Ok ( solution)
8182}
8283
8384/// The MILP's view of a workload.
@@ -166,7 +167,9 @@ fn item_to_raqe(item: &OptimizerItem) -> Result<Raqe, MilpError> {
166167 statistics : req. statistics . clone ( ) ,
167168 } ) ;
168169 } ;
169- let capability = capability ( * statistic) ;
170+ let Some ( capability) = capability ( * statistic, req. topk_count_events ) else {
171+ return Err ( MilpError :: TopkWeightingUnknown { query } ) ;
172+ } ;
170173 if !( 0.0 ..=1.0 ) . contains ( & item. accuracy_sla ) {
171174 return Err ( MilpError :: AccuracySlaOutOfRange {
172175 query,
@@ -179,7 +182,7 @@ fn item_to_raqe(item: &OptimizerItem) -> Result<Raqe, MilpError> {
179182 // TopK keeps one heap per `topk by` bucket; its `grouping_labels` is the
180183 // output label set (every label), which would cost one heap per series.
181184 let grouping_labels: LabelSet = match capability {
182- Capability :: TopK => req
185+ Capability :: TopKByValue | Capability :: TopKByCount => req
183186 . topk_by_labels
184187 . as_ref ( )
185188 . map ( |labels| labels. labels . iter ( ) . cloned ( ) . collect ( ) )
@@ -202,16 +205,23 @@ fn item_to_raqe(item: &OptimizerItem) -> Result<Raqe, MilpError> {
202205 } )
203206}
204207
205- fn capability ( statistic : Statistic ) -> Capability {
206- match statistic {
207- Statistic :: Sum | Statistic :: Count => Capability :: SumOrCount ,
208+ /// `None` for a topk whose weighting is unknown: value-ranked
209+ /// (`sum_over_time`) and count-ranked (`count_over_time`) top-k need separate
210+ /// heaps.
211+ fn capability ( statistic : Statistic , topk_count_events : Option < bool > ) -> Option < Capability > {
212+ Some ( match statistic {
213+ Statistic :: Sum => Capability :: Sum ,
214+ Statistic :: Count => Capability :: Count ,
208215 Statistic :: Rate | Statistic :: Increase => Capability :: RateOrIncrease ,
209216 Statistic :: Min => Capability :: Min ,
210217 Statistic :: Max => Capability :: Max ,
211218 Statistic :: Quantile => Capability :: Quantile ,
212219 Statistic :: Cardinality => Capability :: Cardinality ,
213- Statistic :: Topk => Capability :: TopK ,
214- }
220+ Statistic :: Topk => match topk_count_events? {
221+ true => Capability :: TopKByCount ,
222+ false => Capability :: TopKByValue ,
223+ } ,
224+ } )
215225}
216226
217227/// `accuracy_sla` is required accuracy: error metrics must stay within
@@ -222,13 +232,14 @@ fn accuracy_target(
222232) -> ( & ' static str , f64 , AccuracyDirection ) {
223233 let max_error = 1.0 - accuracy_sla + SLA_EPSILON ;
224234 match capability {
225- Capability :: SumOrCount
235+ Capability :: Sum
236+ | Capability :: Count
226237 | Capability :: Min
227238 | Capability :: Max
228239 | Capability :: RateOrIncrease
229240 | Capability :: Cardinality => ( RELATIVE_ERROR , max_error, AccuracyDirection :: LowerIsBetter ) ,
230241 Capability :: Quantile => ( MAX_RANK_ERROR , max_error, AccuracyDirection :: LowerIsBetter ) ,
231- Capability :: TopK => (
242+ Capability :: TopKByValue | Capability :: TopKByCount => (
232243 PRECISION_AT_K ,
233244 accuracy_sla - SLA_EPSILON ,
234245 AccuracyDirection :: HigherIsBetter ,
@@ -352,10 +363,10 @@ metrics:
352363 }
353364
354365 #[ test]
355- fn sum_maps_to_sum_or_count_with_relative_error_ceiling ( ) {
366+ fn sum_maps_to_sum_with_relative_error_ceiling ( ) {
356367 let w = workload ( & group ( "sum by (job) (http_requests_total)" , 0.99 ) ) ;
357368 let r = & w. raqes [ 0 ] ;
358- assert_eq ! ( r. capability, Capability :: SumOrCount ) ;
369+ assert_eq ! ( r. capability, Capability :: Sum ) ;
359370 assert_eq ! ( r. grouping_labels, labels( & [ "job" ] ) ) ;
360371 assert_eq ! ( r. accuracy_metric, RELATIVE_ERROR ) ;
361372 assert_eq ! ( r. accuracy_direction, AccuracyDirection :: LowerIsBetter ) ;
@@ -364,6 +375,34 @@ metrics:
364375 assert_eq ! ( r. latency_sla_ms, None ) ;
365376 }
366377
378+ #[ test]
379+ fn count_and_topk_weighting_map_to_their_own_capabilities ( ) {
380+ let cap = |query| workload ( & group ( query, 0.99 ) ) . raqes [ 0 ] . capability ;
381+ assert_eq ! (
382+ cap( "sum by (job) (count_over_time(http_requests_total[1m]))" ) ,
383+ Capability :: Count
384+ ) ;
385+ assert_eq ! (
386+ cap( "topk(5, sum_over_time(http_requests_total[1m]))" ) ,
387+ Capability :: TopKByValue
388+ ) ;
389+ assert_eq ! (
390+ cap( "topk(5, count_over_time(http_requests_total[1m]))" ) ,
391+ Capability :: TopKByCount
392+ ) ;
393+ }
394+
395+ #[ test]
396+ fn topk_with_unknown_weighting_is_an_error ( ) {
397+ let w = workload ( & group ( "topk(5, http_requests_total)" , 0.99 ) ) ;
398+ let mut item = w. items [ 0 ] . clone ( ) ;
399+ item. requirements . topk_count_events = None ;
400+ assert ! ( matches!(
401+ item_to_raqe( & item) ,
402+ Err ( MilpError :: TopkWeightingUnknown { .. } )
403+ ) ) ;
404+ }
405+
367406 #[ test]
368407 fn quantile_uses_max_rank_error ( ) {
369408 let w = workload ( & group (
@@ -380,7 +419,7 @@ metrics:
380419 fn topk_requires_precision_and_groups_by_its_buckets ( ) {
381420 let w = workload ( & group ( "topk by (job) (5, http_requests_total)" , 0.9 ) ) ;
382421 let r = & w. raqes [ 0 ] ;
383- assert_eq ! ( r. capability, Capability :: TopK ) ;
422+ assert_eq ! ( r. capability, Capability :: TopKByValue ) ;
384423 assert_eq ! ( r. accuracy_metric, PRECISION_AT_K ) ;
385424 assert_eq ! ( r. accuracy_direction, AccuracyDirection :: HigherIsBetter ) ;
386425 assert ! ( ( r. accuracy_sla - 0.9 ) . abs( ) < 1e-6 ) ;
@@ -427,10 +466,13 @@ metrics:
427466 let config = config ( & ( group ( query, 0.99 ) + & group ( query, 0.99 ) ) ) ;
428467 let facts = facts ( & config) ;
429468 let w = build_milp_workload ( & config, & facts, SCRAPE_MS ) . unwrap ( ) ;
430- let ( deployments, solution) =
431- solve_milp ( & w, & facts, & costs ( ) , Objective :: default ( ) ) . unwrap ( ) ;
432- assert_eq ! ( solution. mapping[ 0 ] , solution. mapping[ 1 ] ) ;
433- assert_eq ! ( deployments[ solution. mapping[ 0 ] ] . config. sketch, "exact-sum" ) ;
469+ let solution = solve_milp ( & w, & facts, & costs ( ) , Objective :: default ( ) ) . unwrap ( ) ;
470+ assert_eq ! ( solution. deployments. len( ) , 1 ) ;
471+ assert_eq ! ( solution. raqes[ 0 ] . deployment, solution. raqes[ 1 ] . deployment) ;
472+ assert_eq ! (
473+ solution. deployments[ 0 ] . deployment. config. sketch,
474+ "exact-sum"
475+ ) ;
434476 }
435477
436478 #[ test]
@@ -454,16 +496,35 @@ metrics:
454496 let config = config ( & groups) ;
455497 let facts = facts ( & config) ;
456498 let w = build_milp_workload ( & config, & facts, SCRAPE_MS ) . unwrap ( ) ;
457- let ( deployments, solution) =
458- solve_milp ( & w, & facts, & costs ( ) , Objective :: default ( ) ) . unwrap ( ) ;
499+ let solution = solve_milp ( & w, & facts, & costs ( ) , Objective :: default ( ) ) . unwrap ( ) ;
459500 let chosen: BTreeMap < Capability , & str > = w
460501 . raqes
461502 . iter ( )
462- . zip ( & solution. mapping )
463- . map ( |( r, & d) | ( r. capability , deployments[ d] . config . sketch . as_str ( ) ) )
503+ . zip ( & solution. raqes )
504+ . map ( |( r, p) | {
505+ let sketch = & solution. deployments [ p. deployment ] . deployment . config . sketch ;
506+ ( r. capability , sketch. as_str ( ) )
507+ } )
464508 . collect ( ) ;
465- assert_eq ! ( chosen[ & Capability :: SumOrCount ] , "exact-sum" ) ;
509+ assert_eq ! ( chosen[ & Capability :: Sum ] , "exact-sum" ) ;
466510 assert_eq ! ( chosen[ & Capability :: Quantile ] , "kll-percall" ) ;
467- assert_eq ! ( chosen[ & Capability :: TopK ] , "cms-heap-topk-fastpath-vector2d" ) ;
511+ assert_eq ! (
512+ chosen[ & Capability :: TopKByValue ] ,
513+ "cms-heap-topk-fastpath-vector2d"
514+ ) ;
515+ }
516+
517+ #[ test]
518+ fn avg_sum_and_count_get_separate_deployments ( ) {
519+ // One exact-sum accumulator can't answer both halves of the avg
520+ // rewrite; sharing would undercount ingest and memory.
521+ let config = config ( & group ( "avg by (job) (http_requests_total)" , 0.99 ) ) ;
522+ let facts = facts ( & config) ;
523+ let w = build_milp_workload ( & config, & facts, SCRAPE_MS ) . unwrap ( ) ;
524+ let caps: Vec < _ > = w. raqes . iter ( ) . map ( |r| r. capability ) . collect ( ) ;
525+ assert_eq ! ( caps. len( ) , 2 ) ;
526+ assert ! ( caps. contains( & Capability :: Sum ) && caps. contains( & Capability :: Count ) ) ;
527+ let solution = solve_milp ( & w, & facts, & costs ( ) , Objective :: default ( ) ) . unwrap ( ) ;
528+ assert_eq ! ( solution. deployments. len( ) , 2 ) ;
468529 }
469530}
0 commit comments