feat(asap-tools): run the e2e planner through the rqe-optimizer MILP (planner.use_optimizer)

asap-planner-rs/docker-compose.yml.j2

@@ -9,15 +9,23 @@ services:
       - TZ=UTC
     volumes:
       - {{ input_config_path }}:/app/input/config.yaml:ro
-      - {{ output_dir }}:/app/output
+      - {{ output_dir }}:/app/output{% if saturation_dir %}
+      - {{ saturation_dir }}:/app/saturation:ro
+      - {{ workload_facts }}:/app/input/workload_facts.yaml:ro{% endif %}
     command: [
       "--input_config", "/app/input/config.yaml",
       "--output_dir", "/app/output",
       "--data-ingestion-interval-ms", "{{ data_ingestion_interval_ms }}",
       "--streaming_engine", "{{ streaming_engine }}",
       "--query-language", "{{ query_language }}"{% if prometheus_url %},
       "--prometheus-url", "{{ prometheus_url }}"{% endif %}{% if punting %},
-      "--enable-punting"{% endif %}
+      "--enable-punting"{% endif %}{% if saturation_dir %},
+      "--planner", "milp",
+      "--workload-facts", "/app/input/workload_facts.yaml",
+      "--atomic-costs", "/app/saturation/optimizer_cost/rqe_atomic_costs.json",
+      "--saturation-dir", "/app/saturation"{% endif %}{% if w_cpu is not none %},
+      "--w-cpu", "{{ w_cpu }}"{% endif %}{% if w_mem is not none %},
+      "--w-mem", "{{ w_mem }}"{% endif %}
     ]
     network_mode: "host"
     restart: no

asap-tools/experiments/config/config.yaml

@@ -91,6 +91,17 @@ controller:
   punting: ???  # REQUIRED: Enable query punting based on performance heuristics (should_be_performant check)
   auto_discover_labels: false  # Query Prometheus for metric labels instead of using config hints
 
+# Planner: sketch-bench's rqe-optimizer MILP (asap-planner --planner milp), or
+# the legacy planner when use_optimizer is false. The MILP needs punting and
+# auto_discover_labels off and no windowing; it picks sketch parameters and
+# cleanup itself.
+planner:
+  use_optimizer: true
+  saturation_dir: null  # Remote path of sketch-bench's saturation-study dir (atomic costs read from optimizer_cost/rqe_atomic_costs.json inside it)
+  workload_facts: null  # Remote path of the workload-facts YAML
+  w_cpu: null  # MILP objective weight on CPU-sec/sec; null = rqe-optimizer's default
+  w_mem: null  # MILP objective weight on memory GiB; null = rqe-optimizer's default
+
 # Optional manual planner windowing override. Omit this block to preserve the
 # default tumbling-window behavior. For tumbling windows, omit
 # slide_interval_ms. For sliding windows, both sizes are explicit.

asap-tools/experiments/experiment_run_clickhouse.py

@@ -339,6 +339,7 @@ def main(cfg: DictConfig) -> None:
         + config.required_cloudlab_params(cfg),
         script_name="experiment_run_clickhouse",
     )
+    config.require_legacy_planner(cfg, "experiment_run_clickhouse")
     config.validate_experiment_config(cfg.experiment_params)
 
     provider = create_provider(cfg)

asap-tools/experiments/experiment_run_e2e.py

@@ -45,6 +45,7 @@ def main(cfg: DictConfig):
     config.validate_config(cfg)
     # Validate experiment configuration
     config.validate_experiment_config(cfg.experiment_params)
+    milp = config.milp_planner(cfg)
 
     # Convert config to args-like object for backward compatibility
     # (also constructs the infrastructure provider, exposed as args.provider)
@@ -173,8 +174,9 @@ def main(cfg: DictConfig):
         config.generate_controller_client_configs(
             cfg.experiment_params,
             local_experiment_root_dir,
-            cfg.aggregate_cleanup,
-            cfg.get("sketch_parameters", None),
+            # The MILP rejects both; it chooses sketch parameters and cleanup itself.
+            None if milp else cfg.aggregate_cleanup,
+            None if milp else cfg.get("sketch_parameters", None),
             cfg.get("windowing", None),
         )
     )
@@ -323,6 +325,7 @@ def main(cfg: DictConfig):
                 punting=args.controller_punting,
                 discovery_backend=discovery_backend,
                 data_ingestion_interval_ms=data_ingestion_interval_ms,
+                milp=milp,
             )
             sync.rsync_controller_config_remote_to_local(
                 provider,

asap-tools/experiments/experiment_utils/config.py

@@ -7,12 +7,13 @@
 import os
 import copy
 import yaml
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 from omegaconf import DictConfig, ListConfig, OmegaConf
 
 import constants
 from experiment_utils.providers.factory import create_provider
+from experiment_utils.services.misc import MilpPlanner
 
 # Keys the planner's `ControllerOptions` accepts.
 CONTROLLER_OPTION_KEYS = {"accuracy_sla", "latency_sla_ms"}
@@ -738,6 +739,49 @@ def validate_config(cfg: DictConfig, script_name: str = "experiment_run_e2e"):
         )
 
 
+def milp_planner(cfg: DictConfig) -> Optional[MilpPlanner]:
+    """`planner.*` as the planner's MILP inputs, or None for the legacy planner.
+
+    Fails on settings --planner milp can't honor rather than dropping them.
+    """
+    planner = cfg.get("planner")
+    if planner is None or not planner.get("use_optimizer", False):
+        return None
+    conflicts = [
+        f"planner.{key} is unset"
+        for key in ("saturation_dir", "workload_facts")
+        if not planner.get(key)
+    ]
+    if cfg.controller.punting:
+        conflicts.append("controller.punting=true")
+    if cfg.controller.get("auto_discover_labels", False):
+        conflicts.append("controller.auto_discover_labels=true")
+    if cfg.get("windowing") is not None:
+        conflicts.append("windowing is set")
+    if conflicts:
+        raise ValueError(
+            "planner.use_optimizer=true (asap-planner --planner milp) conflicts with: "
+            + "; ".join(conflicts)
+            + ". Pass planner.use_optimizer=false for the legacy planner."
+        )
+    return MilpPlanner(
+        saturation_dir=planner.saturation_dir,
+        workload_facts=planner.workload_facts,
+        w_cpu=planner.get("w_cpu"),
+        w_mem=planner.get("w_mem"),
+    )
+
+
+def require_legacy_planner(cfg: DictConfig, script_name: str) -> None:
+    """For scripts whose planner runs can't use --planner milp."""
+    planner = cfg.get("planner")
+    if planner is not None and planner.get("use_optimizer", False):
+        raise ValueError(
+            f"{script_name} supports only the legacy planner; "
+            "pass planner.use_optimizer=false."
+        )
+
+
 def _load_sql_queries(sql_file: str) -> List[str]:
     """Read a SQL file and return individual statements, preserving comment lines."""
     with open(sql_file) as f:

asap-tools/experiments/experiment_utils/services/misc.py

@@ -24,6 +24,43 @@ class DiscoveryBackend:
     database: Optional[str]
 
 
+# Atomic cost table under the saturation dir (rqe_optimizer::saturation::COST_TABLE).
+COST_TABLE = "optimizer_cost/rqe_atomic_costs.json"
+
+
+@dataclass
+class MilpPlanner:
+    """Remote inputs for `asap-planner --planner milp` (sketch-bench's rqe-optimizer)."""
+
+    saturation_dir: str
+    workload_facts: str
+    w_cpu: Optional[float] = None
+    w_mem: Optional[float] = None
+
+    def cli_args(self) -> str:
+        args = (
+            f" --planner milp --workload-facts {self.workload_facts}"
+            f" --atomic-costs {os.path.join(self.saturation_dir, COST_TABLE)}"
+            f" --saturation-dir {self.saturation_dir}"
+        )
+        return args + self._weight_args()
+
+    def compose_generator_args(self) -> str:
+        args = (
+            f" --saturation-dir {self.saturation_dir}"
+            f" --workload-facts {self.workload_facts}"
+        )
+        return args + self._weight_args()
+
+    def _weight_args(self) -> str:
+        args = ""
+        if self.w_cpu is not None:
+            args += f" --w-cpu {self.w_cpu}"
+        if self.w_mem is not None:
+            args += f" --w-mem {self.w_mem}"
+        return args
+
+
 class DeathstarService(BaseService):
     """Service for managing DeathStar benchmark."""
 
@@ -206,6 +243,7 @@ def start(
         discovery_backend: Optional[DiscoveryBackend],
         data_ingestion_interval_ms: int,
         query_language: str = "promql",
+        milp: Optional[MilpPlanner] = None,
         **kwargs,
     ) -> None:
         """
@@ -222,6 +260,7 @@ def start(
                 SQL mode:    DiscoveryBackend(type="clickhouse", url=<url>, database=<db>)
             data_ingestion_interval_ms: Data ingestion interval in milliseconds (required for all modes)
             query_language: 'promql' (default) or 'sql'
+            milp: Plan with --planner milp instead of the legacy planner.
             **kwargs: Additional configuration
         """
         if self.use_container:
@@ -233,6 +272,7 @@ def start(
                 discovery_backend,
                 data_ingestion_interval_ms,
                 query_language,
+                milp,
             )
         else:
             return self._start_bare_metal(
@@ -243,6 +283,7 @@ def start(
                 discovery_backend,
                 data_ingestion_interval_ms,
                 query_language,
+                milp,
             )
 
     @staticmethod
@@ -268,6 +309,7 @@ def _start_bare_metal(
         discovery_backend: Optional[DiscoveryBackend],
         data_ingestion_interval_ms: int,
         query_language: str,
+        milp: Optional[MilpPlanner] = None,
     ) -> None:
         controller_log = os.path.join(controller_remote_output_dir, "controller.log")
         # Force UTC so naive (no Z/offset) datetime-string time literals in SQL
@@ -286,6 +328,8 @@ def _start_bare_metal(
         cmd += self._discovery_args(discovery_backend)
         if punting:
             cmd += " --enable-punting"
+        if milp is not None:
+            cmd += milp.cli_args()
         cmd += " -v"
         cmd += f" > {controller_log} 2>&1"
         cmd_dir = os.path.join(self.provider.get_home_dir(), "code", "asap-planner-rs")
@@ -307,6 +351,7 @@ def _start_containerized(
         discovery_backend: Optional[DiscoveryBackend],
         data_ingestion_interval_ms: int,
         query_language: str,
+        milp: Optional[MilpPlanner] = None,
     ):
         controller_dir = os.path.join(
             self.provider.get_home_dir(), "code", "asap-planner-rs"
@@ -338,6 +383,8 @@ def _start_containerized(
         generate_cmd += self._discovery_args(discovery_backend)
         if punting:
             generate_cmd += " --punting"
+        if milp is not None:
+            generate_cmd += milp.compose_generator_args()
 
         controller_log = os.path.join(controller_remote_output_dir, "controller.log")
         cmd = (

asap-tools/experiments/generate_controller_compose.py

@@ -21,6 +21,10 @@ def generate_compose_file(
     punting: bool,
     prometheus_url: Optional[str],
     query_language: str,
+    saturation_dir: Optional[str] = None,
+    workload_facts: Optional[str] = None,
+    w_cpu: Optional[float] = None,
+    w_mem: Optional[float] = None,
 ):
     """Generate docker-compose.yml from template with provided variables."""
 
@@ -46,6 +50,10 @@ def generate_compose_file(
         "punting": punting,
         "prometheus_url": prometheus_url,
         "query_language": query_language,
+        "saturation_dir": saturation_dir,
+        "workload_facts": workload_facts,
+        "w_cpu": w_cpu,
+        "w_mem": w_mem,
     }
 
     # Render the template
@@ -134,7 +142,17 @@ def main():
         help="Data ingestion interval in milliseconds",
     )
 
+    # Set together: the planner runs --planner milp (sketch-bench's rqe-optimizer).
+    parser.add_argument(
+        "--saturation-dir", help="Host path of sketch-bench's saturation-study dir"
+    )
+    parser.add_argument("--workload-facts", help="Host path of the workload-facts YAML")
+    parser.add_argument("--w-cpu", type=float, help="MILP objective weight on CPU")
+    parser.add_argument("--w-mem", type=float, help="MILP objective weight on memory")
+
     args = parser.parse_args()
+    if (args.saturation_dir is None) != (args.workload_facts is None):
+        parser.error("--saturation-dir and --workload-facts go together")
 
     generate_compose_file(
         template_path=args.template_path,
@@ -148,6 +166,10 @@ def main():
         punting=args.punting,
         prometheus_url=args.prometheus_url,
         query_language=args.query_language,
+        saturation_dir=args.saturation_dir,
+        workload_facts=args.workload_facts,
+        w_cpu=args.w_cpu,
+        w_mem=args.w_mem,
     )