fix(planner): reject avg before solving and write configs atomically

- asap-optimizer-cli --output-dir rejects avg queries before the MILP
  solve; print-only runs still plan them.
- Both YAML strings are serialized before either file is written.
- Each item is visited once per deployment, not once per occurrence.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: milindsrivastava1997

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

@@ -8,8 +8,8 @@ use std::path::PathBuf;
 
 use asap_planner::optimizer::{
     build_milp_workload, load_flat_atomic_cost_table, load_optional_selected_atomic_cost_table,
-    load_workload_facts, plan_to_planner_output, run_greedy_pipeline, solve_milp, AtomicCostTable,
-    LabelSetFacts, LabelSetFactsError,
+    load_workload_facts, plan_to_planner_output, reject_avg_queries, run_greedy_pipeline,
+    solve_milp, AtomicCostTable, LabelSetFacts, LabelSetFactsError,
 };
 use asap_planner::ControllerConfig;
 use clap::Parser;
@@ -200,6 +200,10 @@ fn run_milp(args: &Args, config: &ControllerConfig) -> anyhow::Result<()> {
     let objective = Objective::AUCCost { w_cpu, w_mem };
     tracing::debug!(?objective, cost_rows = costs.len(), "milp: inputs loaded");
 
+    // Fail before solving when the plan would be written but can't be.
+    if args.output_dir.is_some() {
+        reject_avg_queries(config)?;
+    }
     let workload = build_milp_workload(config, &facts, args.data_ingestion_interval_ms)?;
     let solution = solve_milp(
         &workload,
@@ -259,15 +263,13 @@ fn run_milp(args: &Args, config: &ControllerConfig) -> anyhow::Result<()> {
 
     if let Some(dir) = &args.output_dir {
         let output = plan_to_planner_output(config, &workload, &solution)?;
+        // Serialize both before writing either, so a failure can't leave a
+        // new streaming config next to a stale inference config.
+        let streaming = output.to_streaming_yaml_string()?;
+        let inference = output.to_inference_yaml_string()?;
         std::fs::create_dir_all(dir)?;
-        std::fs::write(
-            dir.join("streaming_config.yaml"),
-            output.to_streaming_yaml_string()?,
-        )?;
-        std::fs::write(
-            dir.join("inference_config.yaml"),
-            output.to_inference_yaml_string()?,
-        )?;
+        std::fs::write(dir.join("streaming_config.yaml"), streaming)?;
+        std::fs::write(dir.join("inference_config.yaml"), inference)?;
         println!("\nwrote configs to {}", dir.display());
     }
     Ok(())

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

@@ -1,6 +1,6 @@
 //! Turns a solved MILP plan into the planner's streaming and inference YAML.
 
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 use asap_types::enums::{CleanupPolicy, WindowType};
 use indexmap::map::Entry;
@@ -47,6 +47,19 @@ pub enum MilpOutputError {
     Generator(#[from] ControllerError),
 }
 
+/// Errors on the first avg query: its plan can be costed but not deployed.
+pub fn reject_avg_queries(config: &ControllerConfig) -> Result<(), MilpOutputError> {
+    match config
+        .query_groups
+        .iter()
+        .flat_map(|group| &group.queries)
+        .find(|query| contains_avg(query))
+    {
+        Some(query) => Err(MilpOutputError::AvgQuery(query.clone())),
+        None => Ok(()),
+    }
+}
+
 /// Streaming and inference YAML for `solution`. Aggregation ids follow the
 /// plan's deployment order, starting at 1. No cleanup policy yet, so the
 /// retained instance counts are not emitted.
@@ -55,19 +68,16 @@ pub fn plan_to_planner_output(
     workload: &MilpWorkload,
     solution: &MilpSolution,
 ) -> Result<PlannerOutput, MilpOutputError> {
-    if let Some(query) = config
-        .query_groups
-        .iter()
-        .flat_map(|group| &group.queries)
-        .find(|query| contains_avg(query))
-    {
-        return Err(MilpOutputError::AvgQuery(query.clone()));
-    }
+    reject_avg_queries(config)?;
 
     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 seen = HashSet::new();
     for (raqe, planned) in solution.raqes.iter().enumerate() {
-        served[planned.deployment].push(item_of(raqe));
+        if seen.insert((planned.deployment, workload.raqe_items[raqe])) {
+            served[planned.deployment].push(item_of(raqe));
+        }
     }
 
     let mut aggregations: IndexMap<String, IntermediateAggConfig> = IndexMap::new();

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

@@ -26,7 +26,7 @@ pub use error::{OptimizerError, UnservableItem};
 pub use greedy::greedy_assign;
 pub use label_set_facts::{ItemFacts, LabelSetFacts, LabelSetFactsError, LabelSetKey, SeriesKey};
 pub use milp::{build_milp_workload, solve_milp, MilpError, MilpWorkload};
-pub use milp_output::{plan_to_planner_output, MilpOutputError};
+pub use milp_output::{plan_to_planner_output, reject_avg_queries, MilpOutputError};
 pub use pipeline::{run_greedy_pipeline, OptimizerPipelineError};
 pub use sketch_properties::{sketch_properties, SketchProperties};
 pub use solution::{AQEAssignment, OptimizerItem, OptimizerSolution, QueryMethod};