feat(asap-tools): rqe query protocol and frontend switch for e2e runs

asap-tools/experiments/config/config.yaml

@@ -89,6 +89,13 @@ query_engine:
   # dropping them; the run records their size and fails if any write failed.
   # Spill files stay on the node; results syncs skip them.
   spill: false
+  # How the querier sends queries to ASAPQuery: "promql" (PromQL/SQL text) or
+  # "rqe" (each query's RQE expression tree, POSTed to /api/v1/rqe_query).
+  # Baselines always get PromQL/SQL.
+  frontend: "promql"
+  # rqe only: JSON object of query text -> RQE expression tree, e.g. written
+  # by sketch_bench_e2e/sketch_bench_e2e_configs.py.
+  rqe_exprs_file: null
 
 # Controller configuration
 controller:

asap-tools/experiments/constants.py

@@ -23,6 +23,9 @@
 PROMETHEUS_HEALTH_POLLING_INTERVAL = 5
 
 SKETCHDB_EXPERIMENT_NAME = "sketchdb"
+# The querier protocol that POSTs RQE expression trees (prometheus-client's
+# RqeQueryClient).
+RQE_PROTOCOL = "rqe"
 BASELINE_EXPERIMENT_NAME = "baseline"
 INGEST_MONITOR_STOP_FILE = ".ingest_monitor_stop"
 # ClickHouse precompute-strategy arms (issue #491): same raw table/queries as

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

@@ -154,7 +154,7 @@ def pilot_configs(workload, pilot_dir):
     table = pilot_dir / "eval_table.json"
     out_dir = Path(CONFIG_DIR) / workload_set(workload) / workload["id"]
 
-    e2e, _ = build_config(
+    e2e, _, _ = build_config(
         str(table),
         seed,
         starting_delay,

asap-tools/experiments/experiment_only_ingest_path.py

@@ -36,6 +36,7 @@ def main(cfg: DictConfig):
     config.validate_config(cfg)
     # Validate experiment configuration (queries not required for ingest path experiments)
     config.validate_experiment_config(cfg.experiment_params, require_queries=False)
+    config.require_promql_frontend(cfg.query_engine, "experiment_only_ingest_path.py")
 
     # Check that experiment_duration is specified
     if not hasattr(cfg.experiment_params, "experiment_duration"):

asap-tools/experiments/experiment_run_e2e.py

@@ -179,13 +179,15 @@ def main(cfg: DictConfig):
         )
 
     sync.copy_experiment_config(cfg.experiment_params, local_experiment_root_dir)
+    rqe_exprs_file = config.rqe_exprs_file_for_frontend(cfg.query_engine)
     experiment_modes, metrics_to_remote_write = (
         config.generate_controller_client_configs(
             cfg.experiment_params,
             local_experiment_root_dir,
             cfg.aggregate_cleanup,
             cfg.get("sketch_parameters", None),
             cfg.get("windowing", None),
+            rqe_exprs_file=rqe_exprs_file,
         )
     )
     sync.rsync_controller_client_configs(
@@ -377,6 +379,7 @@ def main(cfg: DictConfig):
                         f"{experiment_mode}_controller_input.yaml",
                     ),
                     data_ingestion_interval_ms,
+                    rqe_exprs_file=rqe_exprs_file,
                 )
                 sync.rsync_streaming_configs(
                     provider,

asap-tools/experiments/experiment_run_grafana_demo.py

@@ -45,6 +45,7 @@ def main(cfg: DictConfig):
     config.validate_config(cfg)
     # Validate experiment configuration
     config.validate_experiment_config(cfg.experiment_params)
+    config.require_promql_frontend(cfg.query_engine, "experiment_run_grafana_demo.py")
     # Convert config to args-like object for backward compatibility
     # (also constructs the infrastructure provider, exposed as args.provider)
     args = config.Args(cfg)

asap-tools/experiments/experiment_utils/config.py

@@ -8,8 +8,9 @@
 import re
 import shutil
 import copy
+import json
 import yaml
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import hydra.utils
 from omegaconf import DictConfig, ListConfig, OmegaConf
@@ -20,6 +21,70 @@
 # Keys the planner's `ControllerOptions` accepts.
 CONTROLLER_OPTION_KEYS = {"accuracy_sla", "latency_sla_ms"}
 
+# `query_engine.frontend`: how the querier sends queries to ASAPQuery.
+QUERY_FRONTEND_PROMQL = "promql"
+QUERY_FRONTEND_RQE = "rqe"
+# The side file's name next to the client configs, which the querier resolves
+# relative to its config.
+RQE_EXPRS_FILENAME = "rqe_exprs.json"
+# inference_config.yaml sections: per-query configs and per-RQE-leaf configs
+# (an `--expr-only` plan has only the latter, keyed by leaf id).
+INFERENCE_QUERIES_KEY = "queries"
+INFERENCE_RQE_CONFIGS_KEY = "rqe_configs"
+
+
+def rqe_exprs_file_for_frontend(query_engine_cfg: Any) -> Optional[str]:
+    """The side file (query text -> RQE expression tree) the sketchdb server's
+    querier reads under `query_engine.frontend=rqe`; None under promql."""
+    frontend = query_engine_cfg.frontend
+    if frontend == QUERY_FRONTEND_PROMQL:
+        return None
+    if frontend != QUERY_FRONTEND_RQE:
+        raise ValueError(
+            f"query_engine.frontend must be {QUERY_FRONTEND_PROMQL!r} or "
+            f"{QUERY_FRONTEND_RQE!r}, got {frontend!r}"
+        )
+    path = query_engine_cfg.get("rqe_exprs_file")
+    if not path:
+        raise ValueError(
+            "query_engine.frontend=rqe needs query_engine.rqe_exprs_file "
+            "(written by sketch_bench_e2e_configs.py)"
+        )
+    return str(path)
+
+
+def require_promql_frontend(query_engine_cfg: Any, script_name: str) -> None:
+    """Fail unless `query_engine.frontend` is promql, for runners that only
+    send PromQL: an rqe run would otherwise measure the PromQL frontend."""
+    frontend = query_engine_cfg.frontend
+    if frontend != QUERY_FRONTEND_PROMQL:
+        raise ValueError(
+            f"{script_name} supports only query_engine.frontend="
+            f"{QUERY_FRONTEND_PROMQL!r}, got {frontend!r}; run "
+            "experiment_run_e2e.py for the rqe frontend"
+        )
+
+
+def _copy_rqe_exprs(
+    rqe_exprs_file: str, query_groups: List[Dict[str, Any]], output_dir: str
+) -> None:
+    """Copy the side file next to the client configs after checking it has a
+    tree for every query: the rqe querier has no PromQL/SQL fallback."""
+    with open(rqe_exprs_file) as f:
+        exprs = json.load(f)
+    if not isinstance(exprs, dict):
+        raise ValueError(
+            f"{rqe_exprs_file}: expected a JSON object of query text -> expr"
+        )
+    missing = [q for g in query_groups for q in g["queries"] if q not in exprs]
+    if missing:
+        raise ValueError(
+            f"{rqe_exprs_file} has no RQE expression for {missing}; drop these "
+            "queries for query_engine.frontend=rqe (sketch_bench_e2e_configs.py "
+            "--frontend rqe does)"
+        )
+    shutil.copyfile(rqe_exprs_file, os.path.join(output_dir, RQE_EXPRS_FILENAME))
+
 
 def validate_basic_config(
     cfg: DictConfig,
@@ -370,8 +435,13 @@ def generate_controller_client_configs(
     aggregate_cleanup: DictConfig = None,
     sketch_parameters: DictConfig = None,
     windowing: DictConfig = None,
+    rqe_exprs_file: Optional[str] = None,
 ) -> Tuple[List[str], List[str]]:
-    """Generate controller client configurations from experiment parameters."""
+    """Generate controller client configurations from experiment parameters.
+
+    With `rqe_exprs_file` (see `rqe_exprs_file_for_frontend`), the sketchdb
+    mode's server uses the rqe protocol and a copy of that side file.
+    """
     # experiment_params is already loaded by Hydra
     experiment_config = OmegaConf.to_container(experiment_params, resolve=True)
     assert experiment_config is not None and isinstance(experiment_config, dict)
@@ -413,6 +483,16 @@ def generate_controller_client_configs(
         "existing_streaming_config",
     }
 
+    if rqe_exprs_file is not None:
+        if not any(
+            m["mode"] == constants.SKETCHDB_EXPERIMENT_NAME for m in experiment_modes
+        ):
+            raise ValueError(
+                "query_engine.frontend=rqe needs a "
+                f"{constants.SKETCHDB_EXPERIMENT_NAME!r} experiment mode"
+            )
+        _copy_rqe_exprs(rqe_exprs_file, experiment_config["query_groups"], output_dir)
+
     for experiment_mode in experiment_modes:
         full_config = copy.deepcopy(experiment_config)
         del full_config["experiment"]
@@ -429,6 +509,23 @@ def generate_controller_client_configs(
         ):
             full_config["servers"] = servers_config
 
+        if (
+            rqe_exprs_file is not None
+            and experiment_mode["mode"] == constants.SKETCHDB_EXPERIMENT_NAME
+        ):
+            full_config["servers"] = [
+                (
+                    {
+                        **server,
+                        "protocol": constants.RQE_PROTOCOL,
+                        "rqe_exprs_file": RQE_EXPRS_FILENAME,
+                    }
+                    if server["name"] == experiment_mode["server"]
+                    else server
+                )
+                for server in full_config["servers"]
+            ]
+
         # Full config — used by prometheus_client (needs "servers", etc.)
         with open(
             os.path.join(output_dir, "{}.yaml".format(experiment_mode["mode"])), "w"
@@ -627,15 +724,58 @@ def _validate_sql_inference_config(inference_config):
             )
 
 
+def rqe_leaf_ids(expr):
+    """Ids of the `rqe` leaves of an RQE expression tree, e.g.
+    {"agg": "sum", "by": [], "of": {"rqe": {"id": "a", ...}}} -> ["a"]."""
+    if not isinstance(expr, dict):
+        raise ValueError(f"not an RQE expression node: {expr!r}")
+    if "rqe" in expr:
+        return [expr["rqe"]["id"]]
+    if "of" in expr:
+        return rqe_leaf_ids(expr["of"])
+    if "lhs" in expr and "rhs" in expr:
+        return rqe_leaf_ids(expr["lhs"]) + rqe_leaf_ids(expr["rhs"])
+    if "scalar" in expr:
+        return []
+    raise ValueError(f"not an RQE expression node: {expr!r}")
+
+
+def _check_rqe_leaves_planned(inference_config, rqe_exprs_file):
+    """Raise unless every leaf of every tree in the side file has an
+    `rqe_configs` entry: ASAPQuery answers a leaf only from its plan."""
+    with open(rqe_exprs_file) as f:
+        exprs = json.load(f)
+    rqe_configs = (
+        inference_config.get(INFERENCE_RQE_CONFIGS_KEY)
+        if isinstance(inference_config, dict)
+        else None
+    ) or {}
+    missing = sorted(
+        {i for e in exprs.values() for i in rqe_leaf_ids(e) if i not in rqe_configs}
+    )
+    if missing:
+        raise ValueError(
+            "controller.prebuilt_output_dir inference_config.yaml has no "
+            f"{INFERENCE_RQE_CONFIGS_KEY} entry for these RQE leaves of "
+            f"{rqe_exprs_file}: {missing}; plan with "
+            "`asap-planner --planner saved --expr-only`"
+        )
+
+
 def copy_prebuilt_controller_output(
     src,
     dest,
     controller_input_config,
     scrape_interval_ms,
     query_language=QUERY_LANGUAGE_PROMQL,
+    rqe_exprs_file=None,
 ):
     """Replaces `dest` with the planner configs in `src`, after checking they
-    plan every query of `controller_input_config` at this scrape interval."""
+    plan every query of `controller_input_config` at this scrape interval.
+
+    With `rqe_exprs_file` (query_engine.frontend=rqe) the querier sends those
+    trees instead of the queries' text, so every tree leaf must be planned
+    under `rqe_configs` instead."""
     check_prebuilt_controller_output(src)
     src_real, dest_real = os.path.realpath(src), os.path.realpath(dest)
     if os.path.commonpath([src_real, dest_real]) in (src_real, dest_real):
@@ -647,20 +787,27 @@ def copy_prebuilt_controller_output(
         inference_config = yaml.safe_load(f)
     if query_language == QUERY_LANGUAGE_SQL:
         _validate_sql_inference_config(inference_config)
-    if not isinstance(inference_config, dict) or not isinstance(
-        inference_config.get("queries"), list
-    ):
-        raise ValueError(
-            "controller.prebuilt_output_dir inference_config.yaml has invalid queries"
-        )
-    planned = {q["query"] for q in inference_config["queries"]}
-    with open(controller_input_config) as f:
-        groups = yaml.safe_load(f)["query_groups"]
-    unplanned = [q for g in groups for q in g["queries"] if not _is_planned(q, planned)]
-    if unplanned:
-        raise ValueError(
-            f"controller.prebuilt_output_dir doesn't plan these queries: {unplanned}"
-        )
+    if rqe_exprs_file is not None:
+        _check_rqe_leaves_planned(inference_config, rqe_exprs_file)
+    else:
+        if not isinstance(inference_config, dict) or not isinstance(
+            inference_config.get("queries"), list
+        ):
+            raise ValueError(
+                "controller.prebuilt_output_dir inference_config.yaml has invalid "
+                "queries"
+            )
+        planned = {q["query"] for q in inference_config["queries"]}
+        with open(controller_input_config) as f:
+            groups = yaml.safe_load(f)["query_groups"]
+        unplanned = [
+            q for g in groups for q in g["queries"] if not _is_planned(q, planned)
+        ]
+        if unplanned:
+            raise ValueError(
+                "controller.prebuilt_output_dir doesn't plan these queries: "
+                f"{unplanned}"
+            )
     with open(os.path.join(src, "streaming_config.yaml")) as f:
         aggregations = yaml.safe_load(f)["aggregations"]
     bad = [

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

@@ -138,9 +138,10 @@ def _start_containerized(
         if profile_prometheus_time is not None:
             gen_compose_cmd += f" --profile-prometheus-time {profile_prometheus_time}"
 
-        cmd = f"mkdir -p {output_dir}; {gen_compose_cmd}; docker compose -f {remote_compose_file} up --no-build -d"
+        # `&&`: a failed step (e.g. compose generation) stops the start.
+        cmd = f"mkdir -p {output_dir} && {gen_compose_cmd} && docker compose -f {remote_compose_file} up --no-build -d"
         try:
-            utils.run_cmd(f"cd {prometheus_client_dir}; {cmd}", popen=False)
+            utils.run_cmd(f"cd {prometheus_client_dir} && {cmd}", popen=False)
         except Exception as e:
             print(f"Failed to start PrometheusClient container: {e}")
             raise

asap-tools/experiments/generate_prometheus_client_compose.py

@@ -7,8 +7,49 @@
 import argparse
 import os
 import sys
+from typing import List, Tuple
+
+import yaml
 from jinja2 import Template
 
+import constants
+
+# Where docker-compose.yml.j2 mounts the client config.
+CONTAINER_CONFIG_DIR = "/app"
+
+
+def rqe_exprs_mounts(config_file: str) -> List[Tuple[str, str]]:
+    """(host path, container path) of each rqe server's side file.
+
+    The client resolves a relative side file against its config's directory,
+    so the file keeps that relative place inside the container. A config
+    that lists no servers has no rqe server to mount for; a missing config
+    is an error, since its rqe servers' side files would go unmounted.
+    """
+    if not os.path.isfile(config_file):
+        raise FileNotFoundError(
+            f"client config {config_file} does not exist; cannot tell which "
+            "rqe side files to mount"
+        )
+    with open(config_file) as f:
+        servers = (yaml.safe_load(f) or {}).get("servers") or []
+    config_dir = os.path.dirname(os.path.abspath(config_file))
+    mounts = []
+    for server in servers:
+        if server.get("protocol") != constants.RQE_PROTOCOL:
+            continue
+        host_path = os.path.join(config_dir, server["rqe_exprs_file"])
+        if not os.path.isfile(host_path):
+            # Docker would mount an empty directory in its place.
+            raise FileNotFoundError(
+                f"rqe server {server['name']!r} in {config_file}: side file "
+                f"{host_path} does not exist"
+            )
+        mounts.append(
+            (host_path, os.path.join(CONTAINER_CONFIG_DIR, server["rqe_exprs_file"]))
+        )
+    return mounts
+
 
 def generate_compose_file(template_path: str, output_path: str, **template_vars):
     """Generate docker-compose.yml from template with provided variables."""
@@ -154,6 +195,7 @@ def main():
         "dry_run": args.dry_run,
         "compare_results": args.compare_results,
         "parallel": args.parallel,
+        "rqe_exprs_mounts": rqe_exprs_mounts(args.config_file),
     }
 
     # Only add optional args if they have values