feat(asap-tools): sketchdb arm for sketch-bench ClickHouse workloads

asap-tools/experiments/datasets/fake_metric/sketch_bench_workloads.py

@@ -12,15 +12,21 @@
 - <stem>/init.sql: the metric tables (baseline);
 - <stem>/init_mv.sql: the metric tables and their materialized views
   (baseline_mv);
-- <stem>/interval_<T>s_<mode>.sql: one statement per query that repeats every
-  T s, for the baseline (raw table) and baseline_mv (views) arms;
-- <stem>/asapquery.sql: each structured entry's `sql` (ASAPQuery's SQL, as
-  sketch-bench's exporter writes it) with NOW() replaced by the floored eval
-  time;
+- <stem>/interval_<T>s[_structured]_<mode>.sql: one statement per query that
+  repeats every T s, for the baseline (raw table) and baseline_mv (views)
+  arms, and for structured entries also sketchdb's: the entry's `sql`
+  (ASAPQuery's SQL, as sketch-bench's exporter writes it) with NOW() replaced
+  by the floored eval time;
 - <stem>_clickhouse.yaml, <stem>_clickhouse_mv.yaml: one
   experiment_run_clickhouse.py config per arm, each creating the tables
   through its own init SQL and feeding every metric with the paced feeder,
-  with one query group per T, repeating every T seconds.
+  with one query group per T and kind (structured groups first), repeating
+  every T seconds;
+- <stem>_clickhouse_sketchdb.yaml, with --prebuilt-output-dir: baseline and
+  sketchdb in one experiment, so baseline's results score sketchdb's.
+  sketchdb runs the structured groups only (PromQL entries have no ASAPQuery
+  SQL) on those planner configs; since they come first, its query indices
+  are baseline's.
 
 baseline_mv keeps exact mergeable states, no sketches: one view per (metric,
 value, statistic, T) holds per T-second bucket and group the state of sums,
@@ -84,12 +90,16 @@
     os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 )
 import constants  # noqa: E402
+from experiment_utils.config import check_prebuilt_controller_output  # noqa: E402
 
 BASELINE = constants.BASELINE_EXPERIMENT_NAME
 BASELINE_MV = constants.BASELINE_MV_EXPERIMENT_NAME
+SKETCHDB = constants.SKETCHDB_EXPERIMENT_NAME
 
 OUT_CONFIG_DIR = "config/experiment_type/sketch_bench_clickhouse"
 CLICKHOUSE_URL = "http://localhost:8123"
+# ASAPQuery's ClickHouse adapter serves SQL here, not at ClickHouse's `/`.
+ASAP_URL = "http://localhost:8088/clickhouse/query"
 # Seconds after an eval time the client sends it, so data up to the eval
 # time has arrived.
 QUERY_TIME_OFFSET_SECS = 10
@@ -609,6 +619,7 @@ def convert(
     duration_secs: int,
     samples_per_sec: Optional[int] = None,
     records_per_sec: Optional[int] = None,
+    prebuilt_output_dir: Optional[str] = None,
 ) -> List[str]:
     """Writes the SQL files and one config per arm; returns the config paths.
     Everything is rendered before anything is written, so a table that fails
@@ -653,19 +664,34 @@ def convert(
             metrics, list(views.values()), mode == BASELINE_MV
         )
 
-    by_interval: Dict[int, Dict[str, List[str]]] = {}
-    asapquery: List[str] = []
+    # Structured groups come first: sketchdb runs only those, and results pair
+    # across modes by query index, so its queries must be baseline's first.
+    groups: Dict[Tuple[bool, int], Dict[str, List[str]]] = {}
+    sketchdb_arm = prebuilt_output_dir is not None
     for entry, q in entries:
-        files = by_interval.setdefault(q.interval_secs, {m: [] for m in modes})
-        files[BASELINE].append(f"-- {describe(entry)}\n{raw_sql(q, labels)}")
+        structured = "promql" not in entry
+        files = groups.setdefault(
+            (not structured, q.interval_secs),
+            {
+                m: []
+                for m in modes + ([SKETCHDB] if structured and sketchdb_arm else [])
+            },
+        )
+        comment = f"-- {describe(entry)}\n"
+        files[BASELINE].append(comment + raw_sql(q, labels))
         if BASELINE_MV in files:
-            sql = mv_sql(q, views, labels)
-            files[BASELINE_MV].append(f"-- {describe(entry)}\n{sql}")
-        if "promql" not in entry:
-            asapquery.append(f"-- {describe(entry)}\n{asapquery_sql(entry)}")
-    if asapquery:
-        # For the ASAPQuery arm; the runner does not read it.
-        outputs[os.path.join(sql_dir, "asapquery.sql")] = statements_text(asapquery)
+            files[BASELINE_MV].append(comment + mv_sql(q, views, labels))
+        if structured:
+            # Rendered either way, so a table without ASAPQuery SQL fails alike.
+            sql = asapquery_sql(entry)
+            if SKETCHDB in files:
+                files[SKETCHDB].append(comment + sql)
+    has_structured = any(not promql for promql, _ in groups)
+    if sketchdb_arm and not has_structured:
+        raise ValueError(
+            f"{table_path}: --prebuilt-output-dir given, but no structured entry "
+            "has ASAPQuery SQL for a sketchdb arm"
+        )
     max_window = max((q.window_secs for _, q in entries), default=0)
     if starting_delay_secs < QUERY_TIME_OFFSET_SECS + max_window:
         # Same data for every arm, so arms stay comparable; absolutes are biased.
@@ -676,8 +702,8 @@ def convert(
         )
 
     query_groups = []
-    for interval, files in sorted(by_interval.items()):
-        name = f"interval_{interval}s"
+    for (promql, interval), files in sorted(groups.items()):
+        name = f"interval_{interval}s" + ("" if promql else "_structured")
         sql_files = {}
         for mode, statements in files.items():
             sql_files[mode] = os.path.join(sql_dir, f"{name}_{mode}.sql")
@@ -709,34 +735,56 @@ def convert(
     if records_per_sec is not None:
         feeder["records_per_sec"] = records_per_sec
 
-    out_paths = []
-    for mode in modes:
-        suffix = "_mv" if mode == BASELINE_MV else ""
-        out = {
-            "experiment": [{"mode": mode, "server": "clickhouse"}],
-            "servers": [{"name": "clickhouse", "url": CLICKHOUSE_URL}],
+    def experiment_params(arm_modes: List[str], init_mode: str) -> dict:
+        return {
+            "experiment": [
+                {"mode": m, "server": "asap" if m == SKETCHDB else "clickhouse"}
+                for m in arm_modes
+            ],
+            "servers": [{"name": "clickhouse", "url": CLICKHOUSE_URL}]
+            + ([{"name": "asap", "url": ASAP_URL}] if SKETCHDB in arm_modes else []),
             "dataset": {
                 "name": "custom",
-                "init_sql_file": init_paths[mode],
+                "init_sql_file": init_paths[init_mode],
                 "paced_feeder": feeder,
             },
-            # Both arms' SQL, so the saved config labels queries by baseline's text.
+            # Every arm's SQL, so the saved config labels queries by baseline's text.
             "query_groups": query_groups,
             "data_ingestion_interval_ms": SCRAPE_INTERVAL_SECONDS * 1000,
         }
+
+    out_paths = []
+    header = f"# Generated from {table_path}\n"
+    for mode in modes:
+        suffix = "_mv" if mode == BASELINE_MV else ""
         out_path = os.path.join(OUT_CONFIG_DIR, f"{stem}_clickhouse{suffix}.yaml")
         outputs[out_path] = (
-            f"# @package experiment_params\n# Generated from {table_path}\n"
-            + yaml.safe_dump(out, sort_keys=False)
+            "# @package experiment_params\n"
+            + header
+            + yaml.safe_dump(experiment_params([mode], mode), sort_keys=False)
+        )
+        out_paths.append(out_path)
+    if sketchdb_arm:
+        prebuilt_output_dir = os.path.abspath(prebuilt_output_dir)
+        check_prebuilt_controller_output(prebuilt_output_dir)
+        # With baseline in the same experiment, so its results score sketchdb's.
+        out = {
+            "experiment_params": experiment_params([BASELINE, SKETCHDB], BASELINE),
+            "controller": {"prebuilt_output_dir": prebuilt_output_dir},
+        }
+        out_path = os.path.join(OUT_CONFIG_DIR, f"{stem}_clickhouse_sketchdb.yaml")
+        outputs[out_path] = (
+            "# @package _global_\n" + header + yaml.safe_dump(out, sort_keys=False)
         )
         out_paths.append(out_path)
 
-    # Clear files an earlier run wrote, e.g. a baseline_mv arm or an interval
-    # this table no longer has.
+    # Clear files an earlier run wrote, e.g. a baseline_mv or sketchdb arm or
+    # an interval this table no longer has.
     shutil.rmtree(sql_dir, ignore_errors=True)
-    mv_config = os.path.join(OUT_CONFIG_DIR, f"{stem}_clickhouse_mv.yaml")
-    if os.path.exists(mv_config):
-        os.remove(mv_config)
+    for suffix in ("_mv", "_sketchdb"):
+        stale = os.path.join(OUT_CONFIG_DIR, f"{stem}_clickhouse{suffix}.yaml")
+        if os.path.exists(stale):
+            os.remove(stale)
     os.makedirs(sql_dir)
     for path, content in outputs.items():
         with open(path, "w") as f:
@@ -754,6 +802,11 @@ def main() -> None:
     )
     parser.add_argument("--samples-per-sec", type=int, help="gauges and counters")
     parser.add_argument("--records-per-sec", type=int, help="records metrics")
+    parser.add_argument(
+        "--prebuilt-output-dir",
+        help="planner configs for the sketchdb arm (e.g. from asap-planner "
+        "--planner saved); without it there is no sketchdb arm",
+    )
     args = parser.parse_args()
     for path in convert(
         args.table,
@@ -762,6 +815,7 @@ def main() -> None:
         args.duration_secs,
         args.samples_per_sec,
         args.records_per_sec,
+        args.prebuilt_output_dir,
     ):
         print(path)
 

asap-tools/experiments/tests/test_sketch_bench_workloads.py

@@ -30,7 +30,11 @@
 )
 
 import sketch_bench_workloads as sbw  # noqa: E402
-from experiment_utils.config import validate_experiment_config  # noqa: E402
+from experiment_utils.config import (  # noqa: E402
+    copy_prebuilt_controller_output,
+    generate_clickhouse_client_configs,
+    validate_experiment_config,
+)
 from omegaconf import OmegaConf  # noqa: E402
 
 SERIES_LABELS = [{"name": "job"}, {"name": "label_0", "child_of": "job"}]
@@ -327,7 +331,11 @@ def test_writes_config_grouped_by_interval(self):
                 validate_experiment_config(OmegaConf.create(cfg))
                 _, init = cfgs["baseline"]
                 groups = {g["name"]: g for g in cfg["query_groups"]}
-                self.assertEqual(sorted(groups), ["interval_1s", "interval_60s"])
+                # Structured first, so sketchdb's queries are baseline's first.
+                self.assertEqual(
+                    list(groups),
+                    ["interval_60s_structured", "interval_1s", "interval_60s"],
+                )
                 g60 = groups["interval_60s"]
                 self.assertEqual(g60["repetition_delay_ms"], 60000)
                 self.assertEqual(
@@ -346,7 +354,7 @@ def test_writes_config_grouped_by_interval(self):
                 for path in g60["sql_file"].values():
                     with open(path) as f:
                         statements = [s for s in f.read().split(";") if s.strip()]
-                    self.assertEqual(len(statements), 9)
+                    self.assertEqual(len(statements), 8)
                 self.assertIn(
                     "endpoint String, status String, latency Float64, "
                     "user_id Float64)",
@@ -364,16 +372,13 @@ def test_writes_config_grouped_by_interval(self):
                         "sketch_bench_views.data_total_value_counter_60s"
                     ),
                 )
-                sql_dir = os.path.dirname(g60["sql_file"]["baseline"])
-                # The table's sql as written, with only NOW() substituted.
-                end = "'" + sbw.floored_end_datetime(60) + "'"
-                with open(os.path.join(sql_dir, "asapquery.sql")) as f:
-                    self.assertEqual(
-                        f.read(),
-                        "-- quantile(0.99, latency) by (service) over 300s\n"
-                        + e2e["queries"][-1]["sql"].replace("NOW()", end)
-                        + ";\n",
-                    )
+                # No --prebuilt-output-dir, so no sketchdb arm or SQL.
+                structured = groups["interval_60s_structured"]
+                self.assertEqual(
+                    sorted(structured["sql_file"]), ["baseline", "baseline_mv"]
+                )
+                sql_dir = os.path.dirname(structured["sql_file"]["baseline"])
+                self.assertFalse([n for n in os.listdir(sql_dir) if "sketchdb" in n])
             finally:
                 os.chdir(cwd)
 
@@ -438,6 +443,131 @@ def test_unservable_query_skips_mv_arm(self):
             finally:
                 os.chdir(cwd)
 
+    def test_sketchdb_arm_runs_the_structured_queries_on_the_saved_plan(self):
+        structured = [
+            mg_entry("quantile", ["service"], 300),
+            mg_entry("cardinality", [], 60),
+        ]
+        promql = {"promql": "sum_over_time(data[1m])", "interval_secs": 60, "rqes": []}
+        with tempfile.TemporaryDirectory() as tmp:
+            table = os.path.join(tmp, "mixed.json")
+            with open(table, "w") as f:
+                e2e = {"metrics": METRICS, "queries": [promql] + structured}
+                json.dump({"workloads": [{"e2e": e2e}]}, f)
+            # As `asap-planner --planner saved` writes it: the table's sql, on NOW().
+            plan = os.path.join(tmp, "plan")
+            os.makedirs(plan)
+            with open(os.path.join(plan, "inference_config.yaml"), "w") as f:
+                yaml.safe_dump(
+                    {
+                        "tables": [
+                            {
+                                "name": "http",
+                                "time_column": "ts",
+                                "value_columns": ["latency", "user_id"],
+                                "metadata_columns": ["region", "service"],
+                            }
+                        ],
+                        "queries": [{"query": e["sql"]} for e in structured],
+                    },
+                    f,
+                )
+            with open(os.path.join(plan, "streaming_config.yaml"), "w") as f:
+                yaml.safe_dump(
+                    {"aggregations": [{"aggregationId": 1, "slideIntervalMs": 1000}]},
+                    f,
+                )
+            cwd = os.getcwd()
+            os.chdir(tmp)
+            try:
+                paths = sbw.convert(table, 42, 310, 600, prebuilt_output_dir=plan)
+                sketchdb_path = os.path.join(
+                    sbw.OUT_CONFIG_DIR, "mixed_clickhouse_sketchdb.yaml"
+                )
+                self.assertIn(sketchdb_path, paths)
+                with open(sketchdb_path) as f:
+                    # Sets controller.prebuilt_output_dir, outside experiment_params.
+                    self.assertTrue(f.readline().startswith("# @package _global_"))
+                    cfg = yaml.safe_load(f)
+                self.assertEqual(
+                    cfg["controller"], {"prebuilt_output_dir": os.path.abspath(plan)}
+                )
+                ep = cfg["experiment_params"]
+                # Baseline in the same experiment scores sketchdb's results.
+                self.assertEqual(
+                    ep["experiment"],
+                    [
+                        {"mode": "baseline", "server": "clickhouse"},
+                        {"mode": "sketchdb", "server": "asap"},
+                    ],
+                )
+                # The engine serves SQL at /clickhouse/query and 404s at `/`.
+                self.assertIn(
+                    {"name": "asap", "url": "http://localhost:8088/clickhouse/query"},
+                    ep["servers"],
+                )
+                validate_experiment_config(OmegaConf.create(ep))
+
+                out = os.path.join(tmp, "out")
+                servers = {s["name"]: s["url"] for s in ep["servers"]}
+                generate_clickhouse_client_configs(
+                    ep["query_groups"],
+                    out,
+                    {m["mode"]: servers[m["server"]] for m in ep["experiment"]},
+                )
+                client = {}
+                for mode in ("baseline", "sketchdb"):
+                    path = os.path.join(
+                        out, "controller_client_configs", f"{mode}.yaml"
+                    )
+                    with open(path) as f:
+                        client[mode] = yaml.safe_load(f)["query_groups"]
+                # sketchdb runs the structured groups, which come first, so its
+                # query indices are baseline's.
+                self.assertEqual([g["id"] for g in client["baseline"]], [0, 1])
+                self.assertEqual([g["id"] for g in client["sketchdb"]], [0])
+                self.assertEqual(len(client["baseline"][0]["queries"]), 2)
+                self.assertEqual(len(client["sketchdb"][0]["queries"]), 2)
+                # The table's sql as written, with only NOW() substituted.
+                end = "'" + sbw.floored_end_datetime(60) + "'"
+                self.assertEqual(
+                    client["sketchdb"][0]["queries"][0],
+                    "-- quantile(0.99, latency) by (service) over 300s\n"
+                    + structured[0]["sql"].replace("NOW()", end),
+                )
+                # The runner's check that the saved plan covers sketchdb's queries.
+                copy_prebuilt_controller_output(
+                    plan,
+                    os.path.join(tmp, "deployed"),
+                    os.path.join(out, "controller_client_configs", "sketchdb.yaml"),
+                    1000,
+                    query_language="sql",
+                )
+
+                # A re-run without the plan leaves no stale sketchdb arm.
+                self.assertNotIn(sketchdb_path, sbw.convert(table, 42, 310, 600))
+                self.assertFalse(os.path.exists(sketchdb_path))
+            finally:
+                os.chdir(cwd)
+
+    def test_sketchdb_arm_without_structured_entries_raises(self):
+        promql = {"promql": "sum_over_time(data[1m])", "interval_secs": 60, "rqes": []}
+        with tempfile.TemporaryDirectory() as tmp:
+            table = os.path.join(tmp, "t.json")
+            with open(table, "w") as f:
+                json.dump(
+                    {"workloads": [{"e2e": {"metrics": METRICS, "queries": [promql]}}]},
+                    f,
+                )
+            cwd = os.getcwd()
+            os.chdir(tmp)
+            try:
+                with self.assertRaisesRegex(ValueError, "no structured entry"):
+                    sbw.convert(table, 42, 310, 600, prebuilt_output_dir=tmp)
+                self.assertFalse(os.path.exists(sbw.OUT_CONFIG_DIR))
+            finally:
+                os.chdir(cwd)
+
     def test_failed_reconvert_keeps_earlier_files(self):
         # Cleared before validating, a failing re-run left the earlier config
         # pointing at deleted SQL files.