diff --git a/asap-planner-rs/docker-compose.yml.j2 b/asap-planner-rs/docker-compose.yml.j2 index 7212c5e4..7c79682d 100644 --- a/asap-planner-rs/docker-compose.yml.j2 +++ b/asap-planner-rs/docker-compose.yml.j2 @@ -9,7 +9,9 @@ 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", @@ -17,7 +19,13 @@ services: "--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 diff --git a/asap-tools/experiments/config/config.yaml b/asap-tools/experiments/config/config.yaml index dfd4abd9..04deecb3 100644 --- a/asap-tools/experiments/config/config.yaml +++ b/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. diff --git a/asap-tools/experiments/experiment_run_clickhouse.py b/asap-tools/experiments/experiment_run_clickhouse.py index faf14952..427e1285 100644 --- a/asap-tools/experiments/experiment_run_clickhouse.py +++ b/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) diff --git a/asap-tools/experiments/experiment_run_e2e.py b/asap-tools/experiments/experiment_run_e2e.py index 9b288ad8..3e8b3d3e 100644 --- a/asap-tools/experiments/experiment_run_e2e.py +++ b/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, diff --git a/asap-tools/experiments/experiment_utils/config.py b/asap-tools/experiments/experiment_utils/config.py index a731c005..9a041603 100644 --- a/asap-tools/experiments/experiment_utils/config.py +++ b/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: diff --git a/asap-tools/experiments/experiment_utils/services/misc.py b/asap-tools/experiments/experiment_utils/services/misc.py index 10285ef9..915a27d0 100644 --- a/asap-tools/experiments/experiment_utils/services/misc.py +++ b/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=, database=) 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 = ( diff --git a/asap-tools/experiments/generate_controller_compose.py b/asap-tools/experiments/generate_controller_compose.py index 526cbe0b..d84595e8 100644 --- a/asap-tools/experiments/generate_controller_compose.py +++ b/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, ) diff --git a/asap-tools/experiments/tests/test_controller_service.py b/asap-tools/experiments/tests/test_controller_service.py index 176bf2bc..61a8fc8e 100644 --- a/asap-tools/experiments/tests/test_controller_service.py +++ b/asap-tools/experiments/tests/test_controller_service.py @@ -4,7 +4,13 @@ import unittest from pathlib import Path -from experiment_utils.services.misc import ControllerService, DiscoveryBackend +import yaml + +from experiment_utils.services.misc import ( + ControllerService, + DiscoveryBackend, + MilpPlanner, +) from generate_controller_compose import generate_compose_file @@ -80,5 +86,93 @@ def test_compose_generator_omits_discovery_url_when_not_provided(self): self.assertNotIn("--prometheus-url", compose) +MILP = MilpPlanner(saturation_dir="/sat", workload_facts="/facts.yaml", w_cpu=0.5) + + +class MilpPlannerTest(unittest.TestCase): + def start(self, use_container, milp): + provider = RecordingProvider() + ControllerService(provider, use_container, node_offset=0).start( + controller_input_file="/tmp/input.yaml", + streaming_engine="precompute", + controller_remote_output_dir="/tmp/output", + punting=False, + discovery_backend=None, + data_ingestion_interval_ms=1000, + milp=milp, + ) + return provider.commands[0] + + def test_bare_metal_passes_milp_flags(self): + cmd = self.start(use_container=False, milp=MILP) + self.assertIn( + " --planner milp --workload-facts /facts.yaml" + " --atomic-costs /sat/optimizer_cost/rqe_atomic_costs.json" + " --saturation-dir /sat --w-cpu 0.5", + cmd, + ) + self.assertNotIn("--w-mem", cmd) + + def test_container_passes_milp_inputs_to_compose_generator(self): + cmd = self.start(use_container=True, milp=MILP) + self.assertIn( + " --saturation-dir /sat --workload-facts /facts.yaml --w-cpu 0.5", cmd + ) + + def test_legacy_planner_has_no_milp_flags(self): + for use_container in (False, True): + cmd = self.start(use_container=use_container, milp=None) + self.assertNotIn("--planner", cmd) + self.assertNotIn("--saturation-dir", cmd) + + def render(self, **milp_kwargs): + template_path = ( + Path(__file__).resolve().parents[3] + / "asap-planner-rs/docker-compose.yml.j2" + ) + with tempfile.TemporaryDirectory() as temp_dir: + output_path = Path(temp_dir) / "controller-compose.yml" + generate_compose_file( + template_path=str(template_path), + output_path=str(output_path), + controller_dir="/tmp/planner", + container_name="controller", + input_config_path="/tmp/input.yaml", + output_dir="/tmp/output", + data_ingestion_interval_ms=1000, + streaming_engine="precompute", + punting=False, + prometheus_url=None, + query_language="promql", + **milp_kwargs, + ) + return yaml.safe_load(output_path.read_text())["services"]["controller"] + + def test_compose_mounts_milp_inputs_and_points_flags_at_them(self): + controller = self.render( + saturation_dir="/sat", workload_facts="/facts.yaml", w_mem=3.0 + ) + self.assertIn("/sat:/app/saturation:ro", controller["volumes"]) + self.assertIn( + "/facts.yaml:/app/input/workload_facts.yaml:ro", controller["volumes"] + ) + command = controller["command"] + flags = dict(zip(command[::2], command[1::2])) + self.assertEqual(flags["--planner"], "milp") + self.assertEqual(flags["--workload-facts"], "/app/input/workload_facts.yaml") + self.assertEqual( + flags["--atomic-costs"], + "/app/saturation/optimizer_cost/rqe_atomic_costs.json", + ) + self.assertEqual(flags["--saturation-dir"], "/app/saturation") + self.assertEqual(flags["--w-mem"], "3.0") + self.assertNotIn("--w-cpu", flags) + + def test_compose_without_milp_is_unchanged(self): + controller = self.render() + self.assertEqual(len(controller["volumes"]), 2) + self.assertNotIn("--planner", controller["command"]) + + if __name__ == "__main__": unittest.main() diff --git a/asap-tools/experiments/tests/test_planner_config.py b/asap-tools/experiments/tests/test_planner_config.py new file mode 100644 index 00000000..7416019f --- /dev/null +++ b/asap-tools/experiments/tests/test_planner_config.py @@ -0,0 +1,61 @@ +"""Tests for planner.* -> --planner milp inputs.""" + +import unittest + +from omegaconf import OmegaConf + +from experiment_utils.config import milp_planner, require_legacy_planner +from experiment_utils.services.misc import MilpPlanner + + +def make_cfg(**overrides): + cfg = OmegaConf.create( + { + "controller": {"punting": False, "auto_discover_labels": False}, + "planner": { + "use_optimizer": True, + "saturation_dir": "/sat", + "workload_facts": "/facts.yaml", + "w_cpu": None, + "w_mem": 2.0, + }, + } + ) + return OmegaConf.merge(cfg, OmegaConf.create(overrides)) + + +class MilpPlannerConfigTest(unittest.TestCase): + def test_builds_inputs_when_enabled(self): + self.assertEqual( + milp_planner(make_cfg()), + MilpPlanner(saturation_dir="/sat", workload_facts="/facts.yaml", w_mem=2.0), + ) + + def test_none_when_disabled(self): + cfg = make_cfg(planner={"use_optimizer": False, "saturation_dir": None}) + self.assertIsNone(milp_planner(cfg)) + + def test_rejects_each_conflict(self): + cases = { + "planner.saturation_dir is unset": {"planner": {"saturation_dir": None}}, + "planner.workload_facts is unset": {"planner": {"workload_facts": None}}, + "controller.punting=true": {"controller": {"punting": True}}, + "controller.auto_discover_labels=true": { + "controller": {"auto_discover_labels": True} + }, + "windowing is set": {"windowing": {"type": "sliding"}}, + } + for message, overrides in cases.items(): + with self.subTest(message), self.assertRaisesRegex(ValueError, message): + milp_planner(make_cfg(**overrides)) + + def test_legacy_only_scripts_reject_optimizer(self): + with self.assertRaisesRegex(ValueError, "experiment_run_clickhouse"): + require_legacy_planner(make_cfg(), "experiment_run_clickhouse") + require_legacy_planner( + make_cfg(planner={"use_optimizer": False}), "experiment_run_clickhouse" + ) + + +if __name__ == "__main__": + unittest.main()