feat(tools): thread sketch_parameters, rsync init_sql_file, add clickhouse_image_tag (#553)

- Pass optional sketch_parameters through from experiment config into
  the generated SQL planner input YAML.
- Rsync init_sql_file to the remote node before executing it, instead
  of assuming it's already present there.
- Allow overriding the ClickHouse image tag via
  dataset.clickhouse_image_tag (pre-AVX2 CPUs SIGILL on "latest").

Split out of #514.

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

Author: milindsrivastava1997

asap-tools/experiments/experiment_run_clickhouse.py

@@ -93,6 +93,7 @@
 import json
 import os
 import time
+from typing import Optional
 from urllib.parse import urlparse
 
 import hydra
@@ -168,7 +169,7 @@ def _run_query_workload(
     remote_monitor_service: RemoteMonitorService,
     minimum_experiment_running_time: int,
     manual_remote_monitor: bool,
-    query_engine_service: QueryEngineRustService | None,
+    query_engine_service: Optional[QueryEngineRustService],
     profile_query_engine: bool,
     profile_prometheus_time,
 ) -> None:
@@ -306,6 +307,9 @@ def main(cfg: DictConfig) -> None:
         local_experiment_dir=local_experiment_root_dir,
         http_port=clickhouse_http_port,
         database=CLICKHOUSE_DATABASE,
+        # Pre-AVX2 CPUs (e.g. Sandy/Ivy Bridge CloudLab nodes) SIGILL on the
+        # "latest" image; override via dataset.clickhouse_image_tag if needed.
+        image_tag=dataset_cfg.get("clickhouse_image_tag", "latest"),
     )
 
     # --- load data once before the mode loop (DROP + reload) ---
@@ -373,7 +377,7 @@ def main(cfg: DictConfig) -> None:
 
             # Generate and rsync the planner input config to the node
             planner_input_yaml = config.generate_sql_planner_input(
-                ep.query_groups, dataset_cfg
+                ep.query_groups, dataset_cfg, cfg.get("sketch_parameters", None)
             )
             local_planner_input = os.path.join(
                 local_controller_dir, "planner_input.yaml"

asap-tools/experiments/experiment_utils/config.py

@@ -827,13 +827,19 @@ def generate_clickhouse_client_configs(
     return modes
 
 
-def generate_sql_planner_input(query_groups: Any, dataset_cfg: Any) -> str:
+def generate_sql_planner_input(
+    query_groups: Any, dataset_cfg: Any, sketch_parameters: Any = None
+) -> str:
     """Generate the YAML input file for asap-planner in SQL mode.
 
     The planner (``asap-planner --query-language sql``) reads a
     ``SQLControllerConfig`` YAML that contains:
       - ``tables``: schema of the tables being queried
       - ``query_groups``: SQL queries with controller options
+      - ``sketch_parameters``: optional per-sketch-type overrides (e.g.
+        ``DatasketchesKLL.K``), matching ``ControllerConfig``'s PromQL-mode
+        field of the same name (``SketchParameterOverrides`` in
+        asap-planner-rs's ``config/input.rs``).
 
     This function builds that YAML from the experiment config so the runner
     does not need a hand-authored planner input file.
@@ -844,6 +850,10 @@ def generate_sql_planner_input(query_groups: Any, dataset_cfg: Any) -> str:
             ``controller_options`` (``accuracy_sla``, ``latency_sla``).
         dataset_cfg: DictConfig with ``table``/``name``, and ``precompute``
             sub-config (``timestamp_col``, ``value_col``, ``label_cols``).
+        sketch_parameters: Optional DictConfig/dict mirroring ``config.yaml``'s
+            top-level ``sketch_parameters`` section (``CountMinSketch``,
+            ``DatasketchesKLL``, etc.). When ``None``, the planner falls back
+            to its own defaults.
 
     Returns:
         YAML string ready to write to disk and pass to asap-planner.
@@ -899,6 +909,10 @@ def generate_sql_planner_input(query_groups: Any, dataset_cfg: Any) -> str:
         "query_groups": planner_query_groups,
         "aggregate_cleanup": {"policy": "read_based"},
     }
+    if sketch_parameters is not None:
+        if isinstance(sketch_parameters, (DictConfig, ListConfig)):
+            sketch_parameters = OmegaConf.to_container(sketch_parameters, resolve=True)
+        planner_input["sketch_parameters"] = sketch_parameters
     return yaml.dump(planner_input, default_flow_style=False, allow_unicode=True)
 
 

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

@@ -268,8 +268,9 @@ def start(
             table: Target table name. Defaults to the dataset's standard table
                 (``hits`` for clickbench, ``h2o_groupby`` for h2o).
             batch_size: INSERT batch size for H2O loading (default 50 000).
-            init_sql_file: Path to a DDL SQL file *already on the remote node*.
-                When ``None``, the built-in ``*_init.sql`` is rsynced and used.
+            init_sql_file: Path to a local DDL SQL file; rsynced to the remote
+                node and executed. When ``None``, the built-in ``*_init.sql``
+                is rsynced and used instead.
             max_rows: Maximum rows to load (0 = all).
         """
         if remote_data_file is None:
@@ -292,7 +293,12 @@ def start(
 
         if init_sql_file is not None:
             print(f"Running init SQL from {init_sql_file!r}...")
-            self._exec_sql_file(init_sql_file, url)
+            remote_ddl = f"/tmp/{dataset_name}_init_{os.getpid()}.sql"
+            self._rsync_to_remote(init_sql_file, remote_ddl)
+            try:
+                self._exec_sql_file(remote_ddl, url)
+            finally:
+                self._remote_rm(remote_ddl)
         elif dataset_name in self.BUILTIN_DDL_FILES:
             local_ddl = os.path.join(
                 self._ASSETS_DIR, self.BUILTIN_DDL_FILES[dataset_name]