diff --git a/asap-tools/experiments/config/config.yaml b/asap-tools/experiments/config/config.yaml index 97cc3c36..dc056b42 100644 --- a/asap-tools/experiments/config/config.yaml +++ b/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: diff --git a/asap-tools/experiments/constants.py b/asap-tools/experiments/constants.py index ded60da9..03973670 100644 --- a/asap-tools/experiments/constants.py +++ b/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 diff --git a/asap-tools/experiments/datasets/fake_metric/random_workload_artifacts.py b/asap-tools/experiments/datasets/fake_metric/random_workload_artifacts.py index 289236a8..538a9a4f 100644 --- a/asap-tools/experiments/datasets/fake_metric/random_workload_artifacts.py +++ b/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, diff --git a/asap-tools/experiments/experiment_only_ingest_path.py b/asap-tools/experiments/experiment_only_ingest_path.py index 6885167f..d6cc844f 100644 --- a/asap-tools/experiments/experiment_only_ingest_path.py +++ b/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"): diff --git a/asap-tools/experiments/experiment_run_e2e.py b/asap-tools/experiments/experiment_run_e2e.py index 6c139544..d7732ae0 100644 --- a/asap-tools/experiments/experiment_run_e2e.py +++ b/asap-tools/experiments/experiment_run_e2e.py @@ -179,6 +179,7 @@ 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, @@ -186,6 +187,7 @@ def main(cfg: DictConfig): 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, diff --git a/asap-tools/experiments/experiment_run_grafana_demo.py b/asap-tools/experiments/experiment_run_grafana_demo.py index 9af9ca82..c219d29a 100644 --- a/asap-tools/experiments/experiment_run_grafana_demo.py +++ b/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) diff --git a/asap-tools/experiments/experiment_utils/config.py b/asap-tools/experiments/experiment_utils/config.py index 3863f31c..8b1be24d 100644 --- a/asap-tools/experiments/experiment_utils/config.py +++ b/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 = [ diff --git a/asap-tools/experiments/experiment_utils/services/prometheus_client_service.py b/asap-tools/experiments/experiment_utils/services/prometheus_client_service.py index 0b6a3e60..d6ac9749 100644 --- a/asap-tools/experiments/experiment_utils/services/prometheus_client_service.py +++ b/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 diff --git a/asap-tools/experiments/generate_prometheus_client_compose.py b/asap-tools/experiments/generate_prometheus_client_compose.py index 5771b172..337b2e0d 100644 --- a/asap-tools/experiments/generate_prometheus_client_compose.py +++ b/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 diff --git a/asap-tools/experiments/sketch_bench_e2e/sketch_bench_e2e_configs.py b/asap-tools/experiments/sketch_bench_e2e/sketch_bench_e2e_configs.py index 9e6b487d..db659658 100644 --- a/asap-tools/experiments/sketch_bench_e2e/sketch_bench_e2e_configs.py +++ b/asap-tools/experiments/sketch_bench_e2e/sketch_bench_e2e_configs.py @@ -20,12 +20,21 @@ planner configs (e.g. from `asap-planner --planner saved`). - `streaming.forward_unsupported_queries=false`: Prometheus gets no data in sketchdb mode, so a forwarded query would come back empty. +- with `--frontend rqe`, `query_engine.frontend=rqe` and + `query_engine.rqe_exprs_file`: a side file next to the config, + `.rqe_exprs.json`, mapping each query's text in the config to + the table's `expr` tree. The config leaves out the queries without `expr` + (e.g. `rate`), with a note, so every query it runs has a tree; the baseline + mode runs the same queries. It then refuses a + `--prebuilt-output-dir` that is not an `--expr-only` plan (`rqe_configs` + only, no per-query configs): plan those configs with + `asap-planner --planner saved --expr-only`. Usage (from asap-tools/experiments): python sketch_bench_e2e/sketch_bench_e2e_configs.py \\ --workload-table /.../synthetic-templatesclassic-shared1.json \\ --seed 42 --starting-delay-secs 120 --query-duration-secs 3600 \\ - [--prebuilt-output-dir DIR] [--sql-queries FILE] + [--prebuilt-output-dir DIR] [--sql-queries FILE] [--frontend rqe] Then: python experiment_run_e2e.py experiment_type=sketch_bench_e2e/ ... """ @@ -42,6 +51,10 @@ sys.path.insert(0, EXPERIMENTS_DIR) from experiment_utils.config import ( # noqa: E402 + INFERENCE_QUERIES_KEY, + INFERENCE_RQE_CONFIGS_KEY, + QUERY_FRONTEND_PROMQL, + QUERY_FRONTEND_RQE, _load_sql_queries, check_prebuilt_controller_output, ) @@ -60,6 +73,8 @@ ] SKETCHDB_MODE = {"mode": "sketchdb", "server": "sketchdb"} BASELINE_MODE = {"mode": "baseline", "server": "prometheus"} +# The side file next to a generated config: + this. +RQE_EXPRS_SUFFIX = ".rqe_exprs.json" # Seconds the window ending at an evaluation time gets for its data to arrive. DEFAULT_QUERY_TIME_OFFSET_SECS = 10 # Feeder time beyond the query phase: covers flow.steady_state_wait (60 s in @@ -91,8 +106,12 @@ def build_config( records_per_sec: Optional[int] = None, prebuilt_output_dir: Optional[str] = None, sql_queries_file: Optional[str] = None, -) -> Tuple[dict, List[str]]: - """Return (config, notes on what was skipped).""" + frontend: str = QUERY_FRONTEND_PROMQL, +) -> Tuple[dict, List[str], Dict[str, dict]]: + """Return (config, notes on what was skipped, RQE side file: query text in + the config -> its `expr`).""" + if frontend not in (QUERY_FRONTEND_PROMQL, QUERY_FRONTEND_RQE): + raise ValueError(f"unknown frontend {frontend!r}") if sql_queries_file is not None and prebuilt_output_dir is None: # The planner rejects the SQL groups' `time_window_seconds` and can't # plan them; they need configs from `asap-planner --planner saved`. @@ -100,11 +119,10 @@ def build_config( e2e = load_e2e(table_path) notes: List[str] = [] - promql = [ - (q["interval_secs"], q["promql"]) for q in e2e["queries"] if "promql" in q - ] + # (entry, query text in the config) + promql_entries = [(q, q["promql"]) for q in e2e["queries"] if "promql" in q] structured = [q for q in e2e["queries"] if "promql" not in q] - sql: List[Tuple[int, str]] = [] + sql_entries: List[Tuple[dict, str]] = [] if structured and sql_queries_file is None: notes.append( f"skipped {len(structured)} structured (multi-grouping) queries: " @@ -117,11 +135,30 @@ def build_config( f"{sql_queries_file} has {len(statements)} statements; the table " f"has {len(structured)} structured queries" ) - sql = [(q["interval_secs"], s) for q, s in zip(structured, statements)] + sql_entries = list(zip(structured, statements)) + + no_expr: List[str] = [] + if frontend == QUERY_FRONTEND_RQE: + # The rqe querier has no PromQL/SQL fallback. + no_expr = [t for q, t in promql_entries + sql_entries if "expr" not in q] + if no_expr: + notes.append( + f"skipped {len(no_expr)} queries without an RQE expression " + f"(--frontend rqe): {no_expr}" + ) + promql_entries = [(q, t) for q, t in promql_entries if "expr" in q] + sql_entries = [(q, t) for q, t in sql_entries if "expr" in q] + rqe_exprs: Dict[str, dict] = {} + for q, text in promql_entries + sql_entries: + if "expr" not in q: + continue + if rqe_exprs.setdefault(text, q["expr"]) != q["expr"]: + raise ValueError(f"{table_path}: two different exprs for {text!r}") + sql = [(q["interval_secs"], t) for q, t in sql_entries] query_groups = [] for queries_by_interval, extra in ( - (by_interval(promql), {}), + (by_interval([(q["interval_secs"], t) for q, t in promql_entries]), {}), # The SQL templates read only {{ end_time }} / {{ end_datetime }} (the # window is in the SQL); the client requires a window to render them. (by_interval(sql), {"time_window_seconds": 0}), @@ -211,12 +248,56 @@ def build_config( if prebuilt_output_dir is not None: prebuilt_output_dir = os.path.abspath(prebuilt_output_dir) check_prebuilt_controller_output(prebuilt_output_dir) + if frontend == QUERY_FRONTEND_RQE: + check_prebuilt_plan_is_expr_only(prebuilt_output_dir) config["controller"] = {"prebuilt_output_dir": prebuilt_output_dir} - return config, notes + if frontend == QUERY_FRONTEND_RQE: + config["query_engine"] = {"frontend": QUERY_FRONTEND_RQE} + return config, notes, rqe_exprs + + +def check_prebuilt_plan_is_expr_only(prebuilt_output_dir: str) -> None: + """Refuse planner configs that are not an `--expr-only` plan: the rqe + frontend sends only RQE trees, so ASAPQuery would otherwise maintain + sketches for per-query configs it is never asked for, skewing its cost. + An `--expr-only` plan has `rqe_configs` and no per-query configs.""" + path = os.path.join(prebuilt_output_dir, "inference_config.yaml") + with open(path) as f: + inference_config = yaml.safe_load(f) + if not isinstance(inference_config, dict): + inference_config = {} + problems = [] + if not inference_config.get(INFERENCE_RQE_CONFIGS_KEY): + problems.append(f"no {INFERENCE_RQE_CONFIGS_KEY}") + if inference_config.get(INFERENCE_QUERIES_KEY): + problems.append(f"per-query configs under {INFERENCE_QUERIES_KEY!r}") + if problems: + raise ValueError( + f"{path} is not an --expr-only plan ({'; '.join(problems)}), which " + "--frontend rqe needs: plan with " + "`asap-planner --planner saved --expr-only`" + ) + + +def rqe_exprs_path(output_path: str) -> str: + return os.path.splitext(output_path)[0] + RQE_EXPRS_SUFFIX -def write_config(config: dict, table_path: str, output_path: str) -> None: +def write_config( + config: dict, rqe_exprs: Dict[str, dict], table_path: str, output_path: str +) -> None: + """Write the config and, under the rqe frontend, its RQE side file next to + it.""" os.makedirs(os.path.dirname(output_path), exist_ok=True) + query_engine = config.get("query_engine", {}) + if query_engine.get("frontend") == QUERY_FRONTEND_RQE: + side_file = os.path.abspath(rqe_exprs_path(output_path)) + with open(side_file, "w") as f: + json.dump(rqe_exprs, f, indent=1) + config = { + **config, + "query_engine": {**query_engine, "rqe_exprs_file": side_file}, + } with open(output_path, "w") as f: f.write("# @package _global_\n") f.write(f"# Generated by sketch_bench_e2e_configs.py from {table_path}\n") @@ -250,10 +331,16 @@ def main() -> None: parser.add_argument("--records-per-sec", type=int) parser.add_argument("--prebuilt-output-dir") parser.add_argument("--sql-queries") + parser.add_argument( + "--frontend", + choices=[QUERY_FRONTEND_PROMQL, QUERY_FRONTEND_RQE], + default=QUERY_FRONTEND_PROMQL, + help="rqe: set query_engine.frontend=rqe and leave out queries without expr", + ) parser.add_argument("--output", help="default: /
.yaml") args = parser.parse_args() - config, notes = build_config( + config, notes, rqe_exprs = build_config( args.workload_table, seed=args.seed, starting_delay_secs=args.starting_delay_secs, @@ -264,10 +351,11 @@ def main() -> None: records_per_sec=args.records_per_sec, prebuilt_output_dir=args.prebuilt_output_dir, sql_queries_file=args.sql_queries, + frontend=args.frontend, ) stem = os.path.splitext(os.path.basename(args.workload_table))[0] output = args.output or os.path.join(DEFAULT_OUTPUT_DIR, f"{stem}.yaml") - write_config(config, args.workload_table, output) + write_config(config, rqe_exprs, args.workload_table, output) for note in notes: print(f"note: {note}") print(output) diff --git a/asap-tools/experiments/tests/test_prebuilt_controller_output.py b/asap-tools/experiments/tests/test_prebuilt_controller_output.py index c11d242a..b6871d98 100644 --- a/asap-tools/experiments/tests/test_prebuilt_controller_output.py +++ b/asap-tools/experiments/tests/test_prebuilt_controller_output.py @@ -1,5 +1,6 @@ """Tests for running e2e experiments on prebuilt planner configs.""" +import json import os import tempfile import unittest @@ -10,6 +11,18 @@ QUERY = "sum by (job) (data)" SQL_QUERY = "SELECT count() FROM hits" +SUM_ID = "t/sum_by_job/keys/1s" +MAX_ID = "t/max_by_job/keys/1s" +# sum by (job) (data) - max by (job) (data), as an RQE tree. +TREE = { + "op": "-", + "lhs": {"rqe": {"id": SUM_ID, "capability": "sum", "metric": "data"}}, + "rhs": { + "agg": "max", + "by": ["job"], + "of": {"rqe": {"id": MAX_ID, "capability": "max", "metric": "data"}}, + }, +} class CopyPrebuiltControllerOutputTest(unittest.TestCase): @@ -129,5 +142,70 @@ def test_sql_output_refuses_an_invalid_table_schema(self): self.copy_sql() +class CopyPrebuiltRqePlanTest(CopyPrebuiltControllerOutputTest): + """Under query_engine.frontend=rqe, with an --expr-only plan.""" + + def setUp(self): + super().setUp() + self.write( + os.path.join(self.src, "inference_config.yaml"), + {"queries": [], "rqe_configs": {SUM_ID: {}, MAX_ID: {}}}, + ) + self.side_file = os.path.join(self.tmp.name, "rqe_exprs.json") + with open(self.side_file, "w") as f: + json.dump({QUERY: TREE, "2 * " + QUERY: {"scalar": 2.0}}, f) + + def copy(self, src=None, dest=None, scrape_interval_ms=1000): + copy_prebuilt_controller_output( + src or self.src, + dest or self.dest, + self.client, + scrape_interval_ms, + rqe_exprs_file=self.side_file, + ) + + # The promql-only checks don't apply: the querier sends the trees. + test_accepts_arithmetic_over_planned_leaves = None + test_refuses_an_unplanned_query = None + test_sql_output_requires_a_sql_schema = None + test_sql_output_with_a_sql_schema_is_copied = None + test_sql_output_refuses_an_invalid_table_schema = None + + def test_expr_only_plan_without_queries_is_copied(self): + self.write( + os.path.join(self.src, "inference_config.yaml"), + {"rqe_configs": {SUM_ID: {}, MAX_ID: {}}}, + ) + self.copy() + self.assertTrue( + os.path.isfile(os.path.join(self.dest, "inference_config.yaml")) + ) + + def test_refuses_an_unplanned_leaf(self): + self.write( + os.path.join(self.src, "inference_config.yaml"), + {"queries": [{"query": QUERY}], "rqe_configs": {SUM_ID: {}}}, + ) + with self.assertRaisesRegex( + ValueError, r"no rqe_configs entry for these RQE leaves .*\['t/max_by_job" + ): + self.copy() + self.assertFalse(os.path.exists(self.dest)) + + def test_refuses_a_plan_without_rqe_configs(self): + self.write( + os.path.join(self.src, "inference_config.yaml"), + {"queries": [{"query": QUERY}]}, + ) + with self.assertRaisesRegex(ValueError, "t/max_by_job.*t/sum_by_job"): + self.copy() + + def test_refuses_a_malformed_tree(self): + with open(self.side_file, "w") as f: + json.dump({QUERY: {"unknown": 1}}, f) + with self.assertRaisesRegex(ValueError, "not an RQE expression node"): + self.copy() + + if __name__ == "__main__": unittest.main() diff --git a/asap-tools/experiments/tests/test_prometheus_client_compose.py b/asap-tools/experiments/tests/test_prometheus_client_compose.py new file mode 100644 index 00000000..6ef91ed2 --- /dev/null +++ b/asap-tools/experiments/tests/test_prometheus_client_compose.py @@ -0,0 +1,95 @@ +"""Tests for generate_prometheus_client_compose.py's rqe side-file mounts.""" + +import os +import tempfile +import unittest + +import yaml + +import generate_prometheus_client_compose as gen + +TEMPLATE = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(gen.__file__))), + "queriers", + "prometheus-client", + "docker-compose.yml.j2", +) + + +class RqeExprsMountsTest(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.config_file = os.path.join(self.tmp.name, "sketchdb.yaml") + + def tearDown(self): + self.tmp.cleanup() + + def write_config(self, servers): + with open(self.config_file, "w") as f: + yaml.safe_dump({"servers": servers, "query_groups": []}, f) + + def render(self): + output = os.path.join(self.tmp.name, "compose.yml") + gen.generate_compose_file( + TEMPLATE, + output, + container_name="c", + experiment_output_dir="/out", + config_file=self.config_file, + client_output_dir="/out/client", + client_output_file="client.log", + prometheus_host="10.0.0.1", + sketchdb_host="10.0.0.1", + rqe_exprs_mounts=gen.rqe_exprs_mounts(self.config_file), + ) + with open(output) as f: + return yaml.safe_load(f)["services"]["prometheusclient"]["volumes"] + + def test_side_file_keeps_its_place_next_to_the_config(self): + self.write_config( + [ + {"name": "prometheus", "url": "http://p"}, + { + "name": "sketchdb", + "url": "http://s", + "protocol": "rqe", + "rqe_exprs_file": "rqe_exprs.json", + }, + ] + ) + open(os.path.join(self.tmp.name, "rqe_exprs.json"), "w").close() + self.assertIn( + f"{self.tmp.name}/rqe_exprs.json:/app/rqe_exprs.json:ro", self.render() + ) + + def test_missing_side_file_of_an_rqe_server_is_an_error(self): + self.write_config( + [ + { + "name": "sketchdb", + "url": "http://s", + "protocol": "rqe", + "rqe_exprs_file": "rqe_exprs.json", + } + ] + ) + with self.assertRaisesRegex(FileNotFoundError, "rqe server 'sketchdb'"): + gen.rqe_exprs_mounts(self.config_file) + + def test_no_mount_without_servers(self): + with open(self.config_file, "w") as f: + yaml.safe_dump({"query_groups": []}, f) + self.assertEqual(gen.rqe_exprs_mounts(self.config_file), []) + + def test_missing_config_file_is_an_error(self): + with self.assertRaisesRegex(FileNotFoundError, "client config"): + gen.rqe_exprs_mounts(self.config_file) + + def test_no_mount_without_rqe_servers(self): + self.write_config([{"name": "sketchdb", "url": "http://s"}]) + self.assertEqual(gen.rqe_exprs_mounts(self.config_file), []) + self.assertEqual(len(self.render()), 3) + + +if __name__ == "__main__": + unittest.main() diff --git a/asap-tools/experiments/tests/test_prometheus_client_service.py b/asap-tools/experiments/tests/test_prometheus_client_service.py new file mode 100644 index 00000000..9d36de61 --- /dev/null +++ b/asap-tools/experiments/tests/test_prometheus_client_service.py @@ -0,0 +1,49 @@ +"""Tests for the querier container's start command.""" + +import unittest +from unittest import mock + +import constants +from experiment_utils.services.prometheus_client_service import QueryClientService + + +class Provider: + def get_home_dir(self): + return "/home/u" + + def get_node_ip(self, node_offset): + return "10.0.0.1" + + +class StartContainerizedTest(unittest.TestCase): + def start_cmd(self): + service = QueryClientService(Provider(), use_container=True, node_offset=0) + with mock.patch("utils.run_cmd") as run_cmd: + service.start( + constants.SKETCHDB_EXPERIMENT_NAME, + "/out/sketchdb.yaml", + "/out/inference_config.yaml", + "/out", + "client.log", + export_cost_and_latency=False, + profile_query_engine_pid=None, + profile_prometheus_time=None, + parallel=False, + backend_protocol="prometheus", + ) + run_cmd.assert_called_once() + return run_cmd.call_args.args[0] + + def test_failed_compose_generation_stops_the_start(self): + steps = self.start_cmd().split(" && ") + self.assertEqual( + [s.split()[0] for s in steps], ["cd", "mkdir", "python3", "docker"] + ) + self.assertIn("generate_prometheus_client_compose.py", steps[2]) + self.assertTrue(steps[3].endswith("up --no-build -d")) + for step in steps: + self.assertNotIn(";", step) + + +if __name__ == "__main__": + unittest.main() diff --git a/asap-tools/experiments/tests/test_sketch_bench_e2e_configs.py b/asap-tools/experiments/tests/test_sketch_bench_e2e_configs.py index fc38e390..b2b992da 100644 --- a/asap-tools/experiments/tests/test_sketch_bench_e2e_configs.py +++ b/asap-tools/experiments/tests/test_sketch_bench_e2e_configs.py @@ -23,20 +23,56 @@ CONFIG_DIR = os.path.join(sbe.EXPERIMENTS_DIR, "config") PREFIX = "synthetic/templates=classic/shared=1" + +def leaf(capability, rqe_id, lookback_ms, grouping, **extra): + return { + "rqe": { + "capability": capability, + "grouping": grouping, + "id": rqe_id, + "lookback_ms": lookback_ms, + "metric": "data", + **extra, + } + } + + +SUM_BY_JOB = leaf("sum", f"{PREFIX}/sum_by_job/keys/1s", 100, ["job"]) +SUM_OVER_TIME = leaf( + "sum", f"{PREFIX}/sum_over_time/keys/15m", 900000, ["job", "label_0"] +) +QUANTILE_BY_JOB = leaf( + "quantile", f"{PREFIX}/quantile_by_job_p0.5/values/1s", 100, ["job"], phi=0.5 +) +DISTINCT_SRC = { + "rqe": { + "capability": "cardinality", + "grouping": ["dst_subnet"], + "id": "synthetic/templates=all/shared=1/t11_distinct_src/keys/5m", + "lookback_ms": 300000, + "metric": "flows", + "value": "src_ip", + } +} +SQL_STATEMENT = "SELECT dst_subnet, uniq(src_ip) FROM flows GROUP BY dst_subnet" + # A slice of synthetic-templatesclassic-shared1.json's e2e.queries, plus a -# structured entry from the multi-grouping tables. +# structured entry from the multi-grouping tables. `rate` has no `expr`. QUERIES = [ { + "expr": SUM_BY_JOB, "interval_secs": 1, "promql": "sum by (job) (data)", "rqes": [f"{PREFIX}/sum_by_job/keys/1s"], }, { + "expr": SUM_OVER_TIME, "interval_secs": 60, "promql": "sum_over_time(data[15m])", "rqes": [f"{PREFIX}/sum_over_time/keys/15m"], }, { + "expr": QUANTILE_BY_JOB, "interval_secs": 1, "promql": "quantile by (job) (0.5, data)", "rqes": [f"{PREFIX}/quantile_by_job_p0.5/values/1s"], @@ -48,9 +84,11 @@ }, { "capability": "cardinality", + "expr": DISTINCT_SRC, "grouping": ["dst_subnet"], "interval_secs": 60, "rqes": ["synthetic/templates=all/shared=1/t11_distinct_src/keys/5m"], + "sql": "SELECT dst_subnet, COUNT(DISTINCT src_ip) FROM flows GROUP BY dst_subnet", "table": "flows", "value": "src_ip", "window_secs": 300, @@ -87,7 +125,7 @@ def build(self, **kwargs): ) def test_query_groups_by_interval_in_table_order(self): - config, _ = self.build() + config, _, _ = self.build() groups = config["experiment_params"]["query_groups"] self.assertEqual( [(g["id"], g["repetition_delay_ms"], g["queries"]) for g in groups], @@ -110,7 +148,7 @@ def test_query_groups_by_interval_in_table_order(self): ) def test_feeder_sends_the_classic_metrics_without_sql(self): - config, notes = self.build() + config, notes, _ = self.build() params = config["experiment_params"] self.assertEqual( params["metrics"], @@ -144,7 +182,7 @@ def test_feeder_sends_the_classic_metrics_without_sql(self): ) def test_structured_queries_skipped_without_sql(self): - config, notes = self.build() + config, notes, _ = self.build() self.assertEqual(len(config["experiment_params"]["query_groups"]), 2) self.assertIn( "skipped 1 structured (multi-grouping) queries: no --sql-queries file", @@ -155,7 +193,7 @@ def test_structured_queries_from_sql_file(self): sql = os.path.join(self.tmp.name, "multigroup.sql") with open(sql, "w") as f: f.write("SELECT dst_subnet, uniq(src_ip) FROM flows GROUP BY dst_subnet;\n") - config, notes = self.build(sql_queries_file=sql) + config, notes, _ = self.build(sql_queries_file=sql) group = config["experiment_params"]["query_groups"][2] self.assertEqual(group["id"], 3) self.assertEqual(group["repetition_delay_ms"], 60000) @@ -189,10 +227,10 @@ def test_records_rate_only_with_records_metrics(self): sql = os.path.join(self.tmp.name, "multigroup.sql") with open(sql, "w") as f: f.write("SELECT 1;\n") - config, _ = self.build(sql_queries_file=sql, records_per_sec=1000) + config, _, _ = self.build(sql_queries_file=sql, records_per_sec=1000) feeder = config["experiment_params"]["exporters"]["exporter_list"] self.assertEqual(feeder["paced_feeder"]["records_per_sec"], 1000) - config, _ = self.build(records_per_sec=1000) + config, _, _ = self.build(records_per_sec=1000) feeder = config["experiment_params"]["exporters"]["exporter_list"] self.assertNotIn("records_per_sec", feeder["paced_feeder"]) @@ -227,21 +265,77 @@ def test_prebuilt_output_dir(self): self.build(prebuilt_output_dir=prebuilt) for name in experiment_config.PLANNER_OUTPUT_FILES: open(os.path.join(prebuilt, name), "w").close() - config, _ = self.build(prebuilt_output_dir=prebuilt) + config, _, _ = self.build(prebuilt_output_dir=prebuilt) self.assertEqual(config["controller"], {"prebuilt_output_dir": prebuilt}) - def compose_and_validate(self, config): + def write_prebuilt(self, inference_config): + prebuilt = os.path.join(self.tmp.name, "prebuilt") + os.makedirs(prebuilt, exist_ok=True) + open(os.path.join(prebuilt, "streaming_config.yaml"), "w").close() + with open(os.path.join(prebuilt, "inference_config.yaml"), "w") as f: + yaml.safe_dump(inference_config, f) + return prebuilt + + RQE_CONFIGS = { + f"{PREFIX}/sum_by_job/keys/1s": { + "aggregations": [{"aggregation_id": 1, "num_aggregates_to_retain": 1}], + "query_time_aggregations": [], + } + } + + def test_rqe_accepts_an_expr_only_plan(self): + for queries in ({}, {"queries": []}): + prebuilt = self.write_prebuilt({**queries, "rqe_configs": self.RQE_CONFIGS}) + config, _, _ = self.build(prebuilt_output_dir=prebuilt, frontend="rqe") + self.assertEqual(config["controller"], {"prebuilt_output_dir": prebuilt}) + + def test_rqe_refuses_a_plan_with_query_configs(self): + prebuilt = self.write_prebuilt( + { + "queries": [{"query": "rate(data_total[15m])", "aggregations": []}], + "rqe_configs": self.RQE_CONFIGS, + } + ) + with self.assertRaisesRegex( + ValueError, + r"is not an --expr-only plan \(per-query configs under 'queries'\).*" + r"`asap-planner --planner saved --expr-only`", + ): + self.build(prebuilt_output_dir=prebuilt, frontend="rqe") + # PromQL runs the per-query configs, so the same plan is fine there. + self.build(prebuilt_output_dir=prebuilt) + + def test_rqe_refuses_a_plan_without_rqe_configs(self): + for inference_config in ( + {"queries": [{"query": "sum by (job) (data)"}]}, + {"queries": [], "rqe_configs": {}}, + None, + ): + prebuilt = self.write_prebuilt(inference_config) + with self.assertRaisesRegex( + ValueError, + r"is not an --expr-only plan \(no rqe_configs.*" + r"`asap-planner --planner saved --expr-only`", + ): + self.build(prebuilt_output_dir=prebuilt, frontend="rqe") + + def compose_and_validate(self, built, overrides=()): # A dry run of what experiment_run_e2e.py does with the file before it # starts any service; returns the experiment modes. + config, _, rqe_exprs = built config_dir = os.path.join(self.tmp.name, "config") + shutil.rmtree(config_dir, ignore_errors=True) shutil.copytree(CONFIG_DIR, config_dir) sbe.write_config( config, + rqe_exprs, self.table, os.path.join(config_dir, "experiment_type", "sketch_bench_e2e", "t.yaml"), ) with initialize_config_dir(version_base=None, config_dir=config_dir): - cfg = compose("config", overrides=["experiment_type=sketch_bench_e2e/t"]) + cfg = compose( + "config", overrides=["experiment_type=sketch_bench_e2e/t", *overrides] + ) params = cfg.experiment_params experiment_config.validate_experiment_config(params) exporter_config, reason = read_exporter_config(params) @@ -257,19 +351,23 @@ def compose_and_validate(self, config): auto_discover_labels=cfg.controller.auto_discover_labels, ) modes, _ = experiment_config.generate_controller_client_configs( - params, self.tmp.name + params, + self.tmp.name, + rqe_exprs_file=experiment_config.rqe_exprs_file_for_frontend( + cfg.query_engine + ), ) return modes def test_generated_config_composes_and_validates(self): self.assertEqual( - self.compose_and_validate(self.build()[0]), ["sketchdb", "baseline"] + self.compose_and_validate(self.build()), ["sketchdb", "baseline"] ) def test_runner_can_anchor_the_generated_client_config(self): # What experiment_run_e2e.py does to the mode's client config once the # feeder starts: every group must learn where the data starts. - self.compose_and_validate(self.build()[0]) + self.compose_and_validate(self.build()) path = os.path.join(self.tmp.name, "controller_client_configs", "sketchdb.yaml") with open(path) as f: client_config = yaml.safe_load(f) @@ -287,8 +385,225 @@ def test_generated_config_with_records_validates(self): sql = os.path.join(self.tmp.name, "multigroup.sql") with open(sql, "w") as f: f.write("SELECT 1;\n") - config, _ = self.build(sql_queries_file=sql) - self.assertEqual(self.compose_and_validate(config), ["sketchdb"]) + built = self.build(sql_queries_file=sql) + self.assertEqual(self.compose_and_validate(built), ["sketchdb"]) + + def client_config(self, mode): + path = os.path.join(self.tmp.name, "controller_client_configs", f"{mode}.yaml") + with open(path) as f: + return yaml.safe_load(f) + + def write_sql_file(self): + sql = os.path.join(self.tmp.name, "multigroup.sql") + with open(sql, "w") as f: + f.write(SQL_STATEMENT + ";\n") + return sql + + def test_side_file_maps_config_text_to_expr(self): + config, _, rqe_exprs = self.build() + # `rate` has no expr: it stays in the PromQL config, not in the side file. + self.assertEqual( + rqe_exprs, + { + "sum by (job) (data)": SUM_BY_JOB, + "sum_over_time(data[15m])": SUM_OVER_TIME, + "quantile by (job) (0.5, data)": QUANTILE_BY_JOB, + }, + ) + queries = [ + q for g in config["experiment_params"]["query_groups"] for q in g["queries"] + ] + self.assertIn("rate(data_total[15m])", queries) + self.assertNotIn("query_engine", config) + + def test_side_file_keys_sql_by_the_statement_in_the_config(self): + # Not by the table's own `sql`: the config runs the --sql-queries text. + config, _, rqe_exprs = self.build(sql_queries_file=self.write_sql_file()) + self.assertEqual(rqe_exprs[SQL_STATEMENT], DISTINCT_SRC) + self.assertEqual( + config["experiment_params"]["query_groups"][2]["queries"], [SQL_STATEMENT] + ) + self.assertNotIn(QUERIES[4]["sql"], rqe_exprs) + + def test_rqe_frontend_leaves_out_queries_without_expr(self): + config, notes, rqe_exprs = self.build(frontend="rqe") + groups = config["experiment_params"]["query_groups"] + self.assertEqual( + [g["queries"] for g in groups], + [ + ["sum by (job) (data)", "quantile by (job) (0.5, data)"], + ["sum_over_time(data[15m])"], + ], + ) + self.assertIn( + "skipped 1 queries without an RQE expression (--frontend rqe): " + "['rate(data_total[15m])']", + notes, + ) + self.assertEqual(set(rqe_exprs), {q for g in groups for q in g["queries"]}) + self.assertEqual(config["query_engine"], {"frontend": "rqe"}) + + def test_unknown_frontend_rejected(self): + with self.assertRaisesRegex(ValueError, "unknown frontend 'sql'"): + self.build(frontend="sql") + + def test_one_text_with_two_exprs_rejected(self): + with open(self.table, "w") as f: + json.dump( + { + "workloads": [ + { + "e2e": { + "metrics": E2E_METRICS, + "queries": [ + QUERIES[0], + {**QUERIES[0], "expr": QUANTILE_BY_JOB}, + ], + } + } + ] + }, + f, + ) + with self.assertRaisesRegex(ValueError, "two different exprs"): + self.build() + + def test_side_file_written_next_to_the_config(self): + config, _, rqe_exprs = self.build(frontend="rqe") + output = os.path.join(self.tmp.name, "out", "t.yaml") + sbe.write_config(config, rqe_exprs, self.table, output) + side_file = os.path.join(self.tmp.name, "out", "t.rqe_exprs.json") + with open(side_file) as f: + self.assertEqual(json.load(f), rqe_exprs) + with open(output) as f: + written = yaml.safe_load(f) + self.assertEqual( + written["query_engine"], + {"frontend": "rqe", "rqe_exprs_file": side_file}, + ) + + def test_promql_frontend_keeps_the_prometheus_protocol(self): + self.compose_and_validate(self.build()) + for mode in ("sketchdb", "baseline"): + for server in self.client_config(mode)["servers"]: + self.assertNotIn("protocol", server) + self.assertFalse( + os.path.exists( + os.path.join( + self.tmp.name, "controller_client_configs", "rqe_exprs.json" + ) + ) + ) + + def test_rqe_frontend_switches_the_sketchdb_server(self): + _, _, rqe_exprs = built = self.build(frontend="rqe") + self.assertEqual(self.compose_and_validate(built), ["sketchdb", "baseline"]) + self.assertEqual( + self.client_config("sketchdb")["servers"], + [ + { + "name": "sketchdb", + "url": "http://localhost:8088", + "protocol": "rqe", + # Relative to the client config, which it travels with. + "rqe_exprs_file": "rqe_exprs.json", + } + ], + ) + # The baseline keeps PromQL. + self.assertEqual( + self.client_config("baseline")["servers"], + [{"name": "prometheus", "url": "http://localhost:9090"}], + ) + with open( + os.path.join(self.tmp.name, "controller_client_configs", "rqe_exprs.json") + ) as f: + self.assertEqual(json.load(f), rqe_exprs) + + def test_promql_config_has_no_side_file(self): + config, _, rqe_exprs = self.build() + output = os.path.join(self.tmp.name, "out", "t.yaml") + sbe.write_config(config, rqe_exprs, self.table, output) + self.assertFalse(os.path.exists(sbe.rqe_exprs_path(output))) + with open(output) as f: + self.assertNotIn("query_engine", yaml.safe_load(f)) + + def test_rqe_switch_on_a_promql_config_needs_a_side_file(self): + # A promql config has no side file, so it can't switch on the command + # line alone; this fails before any service. + with self.assertRaisesRegex(ValueError, "needs query_engine.rqe_exprs_file"): + self.compose_and_validate( + self.build(), overrides=["query_engine.frontend=rqe"] + ) + + def test_rqe_switch_fails_on_a_query_without_expr(self): + # The rqe querier has no fallback, so this fails before any service. + config, notes, rqe_exprs = self.build() + side_file = os.path.join(self.tmp.name, "exprs.json") + with open(side_file, "w") as f: + json.dump(rqe_exprs, f) + with self.assertRaisesRegex( + ValueError, r"no RQE expression for \['rate\(data_total\[15m\]\)'\]" + ): + self.compose_and_validate( + (config, notes, rqe_exprs), + overrides=[ + "query_engine.frontend=rqe", + f"query_engine.rqe_exprs_file={side_file}", + ], + ) + + +class RqeFrontendConfigTest(unittest.TestCase): + def test_promql_needs_no_side_file(self): + self.assertIsNone( + experiment_config.rqe_exprs_file_for_frontend( + OmegaConf.create({"frontend": "promql", "rqe_exprs_file": None}) + ) + ) + + def test_rqe_needs_a_side_file(self): + with self.assertRaisesRegex(ValueError, "needs query_engine.rqe_exprs_file"): + experiment_config.rqe_exprs_file_for_frontend( + OmegaConf.create({"frontend": "rqe", "rqe_exprs_file": None}) + ) + + def test_unknown_frontend_rejected(self): + with self.assertRaisesRegex(ValueError, "got 'sql'"): + experiment_config.rqe_exprs_file_for_frontend( + OmegaConf.create({"frontend": "sql", "rqe_exprs_file": "x.json"}) + ) + + def test_promql_only_runners_accept_promql(self): + experiment_config.require_promql_frontend( + OmegaConf.create({"frontend": "promql"}), "runner.py" + ) + + def test_promql_only_runners_refuse_rqe(self): + with self.assertRaisesRegex( + ValueError, "runner.py supports only query_engine.frontend='promql'" + ): + experiment_config.require_promql_frontend( + OmegaConf.create({"frontend": "rqe"}), "runner.py" + ) + + def test_rqe_needs_a_sketchdb_mode(self): + with tempfile.TemporaryDirectory() as tmp: + side_file = os.path.join(tmp, "x.json") + with open(side_file, "w") as f: + json.dump({}, f) + params = OmegaConf.create( + { + "experiment": [{"mode": "baseline", "server": "prometheus"}], + "servers": [{"name": "prometheus", "url": "http://x"}], + "query_groups": [], + "metrics": [], + } + ) + with self.assertRaisesRegex(ValueError, "needs a 'sketchdb'"): + experiment_config.generate_controller_client_configs( + params, tmp, rqe_exprs_file=side_file + ) if __name__ == "__main__": diff --git a/asap-tools/queriers/prometheus-client/classes/clickhouse_query_client.py b/asap-tools/queriers/prometheus-client/classes/clickhouse_query_client.py index d42a9bc5..e7ea94f9 100644 --- a/asap-tools/queriers/prometheus-client/classes/clickhouse_query_client.py +++ b/asap-tools/queriers/prometheus-client/classes/clickhouse_query_client.py @@ -1,5 +1,5 @@ import requests -from typing import Any, Dict, Optional, Tuple +from typing import Any, Optional, Tuple from requests.adapters import HTTPAdapter from .query_client import QueryClient, QueryResponse @@ -47,6 +47,7 @@ def execute_query( self, query: str, query_time: Optional[int] = None, + config_query: Optional[str] = None, ) -> QueryResponse: """ Execute SQL query via ClickHouse HTTP interface. @@ -55,6 +56,7 @@ def execute_query( query: SQL query string (may contain template variables already substituted) query_time: Not directly used - time filtering should be done via template substitution before calling this method + config_query: Unused; SQL is sent as rendered Returns: QueryResponse with normalized data @@ -112,16 +114,3 @@ def _has_format_clause(self, query: str) -> bool: # Simple check - look for FORMAT keyword followed by format name upper_query = query.upper() return " FORMAT " in upper_query or "\nFORMAT " in upper_query - - def get_runtime_info(self) -> Optional[Dict[str, Any]]: - """Check ClickHouse availability via ping endpoint.""" - try: - response = self._session.get( - f"{self.server_url}/ping", - timeout=5, - ) - if response.status_code == 200: - return {"status": "ok", "response": response.text.strip()} - except Exception: - pass - return None diff --git a/asap-tools/queriers/prometheus-client/classes/config.py b/asap-tools/queriers/prometheus-client/classes/config.py index 01df6534..88165b24 100644 --- a/asap-tools/queriers/prometheus-client/classes/config.py +++ b/asap-tools/queriers/prometheus-client/classes/config.py @@ -1,5 +1,7 @@ from typing import List, Dict, Any, Optional +from .rqe_query_client import RQE_PROTOCOL + class ServerConfig: def __init__( @@ -11,13 +13,19 @@ def __init__( database: Optional[str], user: Optional[str], password: Optional[str], + # rqe-specific: JSON side file of query text -> RQE expression tree, + # relative to the client config's directory + rqe_exprs_file: Optional[str] = None, ): + if protocol == RQE_PROTOCOL and not rqe_exprs_file: + raise ValueError(f"server {name!r}: the rqe protocol needs rqe_exprs_file") self.name = name self.url = url self.protocol = protocol self.database = database self.user = user self.password = password + self.rqe_exprs_file = rqe_exprs_file @classmethod def from_dict(cls, data: Dict[str, Any]) -> "ServerConfig": @@ -28,6 +36,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "ServerConfig": database=data.get("database"), user=data.get("user"), password=data.get("password"), + rqe_exprs_file=data.get("rqe_exprs_file"), ) diff --git a/asap-tools/queriers/prometheus-client/classes/prometheus_query_client.py b/asap-tools/queriers/prometheus-client/classes/prometheus_query_client.py index e3ad8681..c425246f 100644 --- a/asap-tools/queriers/prometheus-client/classes/prometheus_query_client.py +++ b/asap-tools/queriers/prometheus-client/classes/prometheus_query_client.py @@ -5,6 +5,43 @@ from .query_client import QueryClient, QueryResponse +def normalize_vector(raw_result: list) -> Dict[frozenset, float]: + """ + Convert a Prometheus instant-vector result to normalized format. + + Prometheus response format: + [{"metric": {"label1": "value1", ...}, "value": [timestamp, "value_str"]}, ...] + + Returns: + Dict mapping frozenset of labels to float value + """ + result = {} + for item in raw_result: + metric_labels = frozenset(item.get("metric", {}).items()) + value = item.get("value", [None, None]) + if len(value) >= 2 and value[1] is not None: + try: + result[metric_labels] = float(value[1]) + except (ValueError, TypeError): + # Skip non-numeric values (e.g., NaN represented as string) + pass + return result + + +RUNTIME_INFO_PATH = "/api/v1/status/runtimeinfo" +RUNTIME_INFO_TIMEOUT_SECS = 10 + + +def fetch_runtime_info(server_url: str) -> Dict[str, Any]: + """Query SketchDB/Prometheus runtime info endpoint. Raises on a connection + error, a non-200 status or a body without `data`.""" + response = requests.get( + f"{server_url}{RUNTIME_INFO_PATH}", timeout=RUNTIME_INFO_TIMEOUT_SECS + ) + response.raise_for_status() + return response.json()["data"] + + class PrometheusQueryClient(QueryClient): """Query client for Prometheus HTTP API.""" @@ -42,6 +79,7 @@ def execute_query( self, query: str, query_time: Optional[int] = None, + config_query: Optional[str] = None, ) -> QueryResponse: """ Execute PromQL query via Prometheus HTTP API. @@ -49,6 +87,7 @@ def execute_query( Args: query: PromQL query string query_time: Optional Unix timestamp for point-in-time query + config_query: Unused; PromQL is sent as rendered Returns: QueryResponse with normalized data @@ -62,7 +101,7 @@ def execute_query( raw_result = self._client.custom_query(query=query) # Normalize to Dict[frozenset, float] - normalized = self._normalize_response(raw_result) + normalized = normalize_vector(raw_result) return QueryResponse( success=True, data=normalized, @@ -74,38 +113,3 @@ def execute_query( data=None, error_message=str(e), ) - - def _normalize_response(self, raw_result: list) -> Dict[frozenset, float]: - """ - Convert Prometheus response to normalized format. - - Prometheus response format: - [{"metric": {"label1": "value1", ...}, "value": [timestamp, "value_str"]}, ...] - - Returns: - Dict mapping frozenset of labels to float value - """ - result = {} - for item in raw_result: - metric_labels = frozenset(item.get("metric", {}).items()) - value = item.get("value", [None, None]) - if len(value) >= 2 and value[1] is not None: - try: - result[metric_labels] = float(value[1]) - except (ValueError, TypeError): - # Skip non-numeric values (e.g., NaN represented as string) - pass - return result - - def get_runtime_info(self) -> Optional[Dict[str, Any]]: - """Query SketchDB/Prometheus runtime info endpoint.""" - try: - response = requests.get( - f"{self.server_url}/api/v1/status/runtimeinfo", - timeout=10, - ) - if response.status_code == 200: - return response.json().get("data", {}) - except Exception: - pass - return None diff --git a/asap-tools/queriers/prometheus-client/classes/query_client.py b/asap-tools/queriers/prometheus-client/classes/query_client.py index 3c1ebdc9..8edc92aa 100644 --- a/asap-tools/queriers/prometheus-client/classes/query_client.py +++ b/asap-tools/queriers/prometheus-client/classes/query_client.py @@ -26,30 +26,24 @@ def execute_query( self, query: str, query_time: Optional[int] = None, + config_query: Optional[str] = None, ) -> QueryResponse: """ Execute a query and return normalized response. Args: - query: The query string (PromQL, SQL, etc.) + query: The query string (PromQL, SQL, etc.), templates rendered query_time: Optional Unix timestamp for point-in-time queries + config_query: The query as the config lists it, before template + rendering; None means `query` is that text. Clients that + look queries up by text (rqe) use it, since the rendered + text changes every run Returns: QueryResponse with normalized data """ pass - @abstractmethod - def get_runtime_info(self) -> Optional[Dict[str, Any]]: - """ - Get runtime/status info from the backend. - Used for query alignment with SketchDB. - - Returns: - Dict with backend-specific runtime info, or None if unavailable - """ - pass - @property @abstractmethod def protocol_name(self) -> str: diff --git a/asap-tools/queriers/prometheus-client/classes/query_client_factory.py b/asap-tools/queriers/prometheus-client/classes/query_client_factory.py index b7eefbef..9069685d 100644 --- a/asap-tools/queriers/prometheus-client/classes/query_client_factory.py +++ b/asap-tools/queriers/prometheus-client/classes/query_client_factory.py @@ -3,6 +3,7 @@ from .query_client import QueryClient from .prometheus_query_client import PrometheusQueryClient from .clickhouse_query_client import ClickHouseQueryClient +from .rqe_query_client import RQE_PROTOCOL, RqeQueryClient class QueryClientFactory: @@ -11,6 +12,7 @@ class QueryClientFactory: _registry: Dict[str, Type[QueryClient]] = { "prometheus": PrometheusQueryClient, "clickhouse": ClickHouseQueryClient, + RQE_PROTOCOL: RqeQueryClient, } @classmethod diff --git a/asap-tools/queriers/prometheus-client/classes/rqe_query_client.py b/asap-tools/queriers/prometheus-client/classes/rqe_query_client.py new file mode 100644 index 00000000..2d20687c --- /dev/null +++ b/asap-tools/queriers/prometheus-client/classes/rqe_query_client.py @@ -0,0 +1,137 @@ +import json +import os +from typing import Any, Dict, Optional, Tuple + +import requests + +from .prometheus_query_client import normalize_vector +from .query_client import QueryClient, QueryResponse + +# The `protocol` of a config server answered by RqeQueryClient. +RQE_PROTOCOL = "rqe" +RQE_QUERY_PATH = "/api/v1/rqe_query" +# The only result type the route answers with; results compare as vectors. +VECTOR_RESULT_TYPE = "vector" +# (connect, read): the engine answers from precomputed sketches, but a dead +# server must still fail instead of hanging the querier. +RQE_QUERY_TIMEOUT_SECS: Tuple[float, float] = (10, 300) + + +def load_rqe_exprs(path: str) -> Dict[str, Any]: + """Read a side file: a JSON object mapping config query text to its RQE + expression tree.""" + with open(path) as f: + exprs = json.load(f) + if not isinstance(exprs, dict): + raise ValueError(f"{path}: expected a JSON object of query text -> expr") + return exprs + + +def resolve_rqe_exprs_file(config_file: str, rqe_exprs_file: str) -> str: + """A relative side-file path is relative to the client config's directory, + so the two travel (and are mounted) together.""" + return os.path.join(os.path.dirname(os.path.abspath(config_file)), rqe_exprs_file) + + +class RqeQueryClient(QueryClient): + """Sends each query's RQE expression tree to ASAPQuery's RQE route. + + Queries keep their PromQL/SQL text in the config; the side file maps that + text to the tree, so results are keyed and scored as for PromQL/SQL. The + key is the config's unrendered text (`config_query`), not the query sent: + templated SQL renders differently on every run (its `{{ end_time }}`), so + only the config text is stable across runs. + """ + + def __init__( + self, + server_url: str, + server_name: str, + rqe_exprs_file: str, + timeout: Tuple[float, float] = RQE_QUERY_TIMEOUT_SECS, + **kwargs: Any, + ): + super().__init__(server_url, server_name) + self.rqe_exprs_file = rqe_exprs_file + self._exprs = load_rqe_exprs(rqe_exprs_file) + self.timeout = timeout + self._session = requests.Session() + + @property + def protocol_name(self) -> str: + return RQE_PROTOCOL + + @property + def session(self) -> requests.Session: + """Access to underlying requests Session for mounting debug adapters.""" + return self._session + + def execute_query( + self, + query: str, + query_time: Optional[int] = None, + config_query: Optional[str] = None, + ) -> QueryResponse: + """ + POST the tree for `config_query` (else `query`) evaluated at `query_time`. + + Returns: + QueryResponse with normalized data; raw_response is the + Prometheus-format result list, as from PrometheusQueryClient + """ + key = query if config_query is None else config_query + if key not in self._exprs: + return QueryResponse( + success=False, + data=None, + error_message=f"no RQE expression for query {key!r} in " + f"{self.rqe_exprs_file}", + ) + if query_time is None: + return QueryResponse( + success=False, + data=None, + error_message="RQE queries need an evaluation time", + ) + try: + response = self._session.post( + f"{self.server_url}{RQE_QUERY_PATH}", + json={"expr": self._exprs[key], "time": float(query_time)}, + timeout=self.timeout, + ) + if response.status_code != 200: + return QueryResponse( + success=False, + data=None, + error_message=f"HTTP {response.status_code}: {response.text}", + raw_response=response.text, + ) + body = response.json() + if body.get("status") != "success": + return QueryResponse( + success=False, + data=None, + error_message=f"RQE query failed: {body}", + raw_response=body, + ) + result_type = body["data"].get("resultType") + if result_type != VECTOR_RESULT_TYPE: + return QueryResponse( + success=False, + data=None, + error_message=f"unexpected result type {result_type!r}, " + f"expected {VECTOR_RESULT_TYPE!r}", + raw_response=body, + ) + raw_result = body["data"]["result"] + return QueryResponse( + success=True, + data=normalize_vector(raw_result), + raw_response=raw_result, + ) + except Exception as e: + return QueryResponse( + success=False, + data=None, + error_message=f"{type(e).__name__}: {e}", + ) diff --git a/asap-tools/queriers/prometheus-client/docker-compose.yml.j2 b/asap-tools/queriers/prometheus-client/docker-compose.yml.j2 index 9621b84d..62e2e651 100644 --- a/asap-tools/queriers/prometheus-client/docker-compose.yml.j2 +++ b/asap-tools/queriers/prometheus-client/docker-compose.yml.j2 @@ -15,7 +15,8 @@ services: - "{{ experiment_output_dir }}:/app/outputs" - "{{ client_output_dir }}:/app/prometheus_client_output" - "{{ config_file }}:/app/prometheus_client_config.yaml:ro"{% if query_engine_config_file is defined and query_engine_config_file is not none %} - - "{{ query_engine_config_file}}:/app/query_engine_config_file.yaml:ro"{% endif %} + - "{{ query_engine_config_file}}:/app/query_engine_config_file.yaml:ro"{% endif %}{% for host_path, container_path in rqe_exprs_mounts %} + - "{{ host_path }}:{{ container_path }}:ro"{% endfor %} network_mode: "host" # Allows prometheus client to send requests to localhost, not ideal for production command: [ "--config_file", "/app/prometheus_client_config.yaml", diff --git a/asap-tools/queriers/prometheus-client/main_prometheus_client.py b/asap-tools/queriers/prometheus-client/main_prometheus_client.py index 1382c7c3..efebacc4 100644 --- a/asap-tools/queriers/prometheus-client/main_prometheus_client.py +++ b/asap-tools/queriers/prometheus-client/main_prometheus_client.py @@ -1,7 +1,6 @@ import os import yaml import time -import requests import argparse import datetime import logging @@ -28,8 +27,10 @@ from classes.config import Config from classes.QueryLatencyExporter import QueryLatencyExporter from classes.query_client import QueryClient +from classes.prometheus_query_client import fetch_runtime_info from classes.query_client_factory import QueryClientFactory from classes.query_template import QueryTemplate +from classes.rqe_query_client import RQE_PROTOCOL, resolve_rqe_exprs_file from classes.query_time import ( aligned_eval_time, first_aligned_eval_time, @@ -224,6 +225,7 @@ def execute_single_query( response = server_object.execute_query( query=rendered_query, query_time=query_unix_time, + config_query=query, ) query_end_time = time.time() @@ -489,6 +491,23 @@ def handle_query_group( return result +# inference_config sections: per-query configs, and per-RQE-leaf configs +# (an `--expr-only` plan has only the latter). +INFERENCE_QUERIES_KEY = "queries" +INFERENCE_RQE_CONFIGS_KEY = "rqe_configs" + + +def required_aggregation_ids(query_engine_config: QueryEngineConfig) -> Set[int]: + """Aggregation ids read by any planned query or RQE leaf.""" + configs = list(query_engine_config.get(INFERENCE_QUERIES_KEY) or []) + configs += (query_engine_config.get(INFERENCE_RQE_CONFIGS_KEY) or {}).values() + return { + int(aggregation["aggregation_id"]) + for config in configs + for aggregation in config["aggregations"] + } + + def get_query_start_times( server_url: str, query_engine_config: QueryEngineConfig ) -> QueryStartTimes: @@ -496,29 +515,15 @@ def get_query_start_times( query_aggregation_id_map = {} query_start_time_map = {} - required_aggregation_ids: Set[int] = set() - for query_yaml in query_engine_config["queries"]: - # add all aggregation IDs from the query YAML to the required_aggregation_ids set - required_aggregation_ids.update( - int(aggregation["aggregation_id"]) - for aggregation in query_yaml["aggregations"] - ) - # assert len(query_yaml["aggregations"]) == 1 - # required_aggregation_ids.add( - # int(query_yaml["aggregations"][0]["aggregation_id"]) - # ) - logger.debug("Required aggregation IDs: {}", required_aggregation_ids) + required_ids = required_aggregation_ids(query_engine_config) + logger.debug("Required aggregation IDs: {}", required_ids) - # wait for all required aggregation IDs to be present + # wait for all required aggregation IDs to be present; an unreachable + # server or a bad response raises rather than being retried while True: - server_response = requests.get( - server_url + "/api/v1/status/runtimeinfo", - headers={"Content-Type": "application/json"}, - ) - server_response.raise_for_status() - server_response_json = server_response.json() - logger.debug("Server response: {}", server_response_json) - aggregation_id_start_time_map = server_response_json["data"][ + runtime_info = fetch_runtime_info(server_url) + logger.debug("Server runtime info: {}", runtime_info) + aggregation_id_start_time_map = runtime_info[ "earliest_timestamp_per_aggregation_id" ] @@ -527,12 +532,10 @@ def get_query_start_times( int(k): v for k, v in aggregation_id_start_time_map.items() } - if not set(aggregation_id_start_time_map.keys()).issuperset( - required_aggregation_ids - ): + if not set(aggregation_id_start_time_map.keys()).issuperset(required_ids): logger.debug( "Waiting for aggregation IDs {} to be present", - required_aggregation_ids - set(aggregation_id_start_time_map.keys()), + required_ids - set(aggregation_id_start_time_map.keys()), ) time.sleep(10) else: @@ -542,7 +545,7 @@ def get_query_start_times( # different tumbling windows? How long do we wait here? What happens with multiple query groups? # get query to aggregate ID mapping from query_engine_config - for query_yaml in query_engine_config["queries"]: + for query_yaml in query_engine_config.get(INFERENCE_QUERIES_KEY) or []: # TODO: this assert will fail if there are multiple aggregations in a query YAML, including for DeltaSet, so commenting it out # assert len(query_yaml["aggregations"]) == 1 # for now, just take the first aggregation ID @@ -674,11 +677,16 @@ def main(args: Any) -> None: disable_ssl=True, retry=debug_retry, ) - - # Mount debug adapter for HTTP request logging - debug_adapter = DebugHTTPAdapter(server.name) - client.session.mount("http://", debug_adapter) - client.session.mount("https://", debug_adapter) + elif protocol == RQE_PROTOCOL: + assert server.rqe_exprs_file is not None + client = QueryClientFactory.create( + protocol=protocol, + server_url=server.url, + server_name=server.name, + rqe_exprs_file=resolve_rqe_exprs_file( + args.config_file, server.rqe_exprs_file + ), + ) else: # ClickHouse or other protocols client = QueryClientFactory.create( @@ -690,11 +698,10 @@ def main(args: Any) -> None: password=server.password if server.password else "", ) - # Mount debug adapter for HTTP request logging - debug_adapter = DebugHTTPAdapter(server.name) - client.session.mount("http://", debug_adapter) - client.session.mount("https://", debug_adapter) - + # Mount debug adapter for HTTP request logging + debug_adapter = DebugHTTPAdapter(server.name) + client.session.mount("http://", debug_adapter) + client.session.mount("https://", debug_adapter) servers[server.name] = client logger.debug( "Connected to server {} ({}) with HTTP debug logging enabled", diff --git a/asap-tools/queriers/prometheus-client/tests/test_query_start_times.py b/asap-tools/queriers/prometheus-client/tests/test_query_start_times.py new file mode 100644 index 00000000..6b16e44f --- /dev/null +++ b/asap-tools/queriers/prometheus-client/tests/test_query_start_times.py @@ -0,0 +1,103 @@ +"""Tests for main_prometheus_client.get_query_start_times.""" + +import importlib.util +import unittest +from typing import Any, Dict +from unittest import mock + +import requests + +URL = "http://sketchdb:8088" +QUERY = "sum by (job) (data)" + + +def runtime_info(start_ms_by_id: Dict[int, int]) -> Dict[str, Any]: + return { + "earliest_timestamp_per_aggregation_id": { + str(k): v for k, v in start_ms_by_id.items() + } + } + + +@unittest.skipUnless( + importlib.util.find_spec("promql_utilities") is not None, + "main_prometheus_client needs asap-common's promql_utilities", +) +class GetQueryStartTimesTest(unittest.TestCase): + def setUp(self) -> None: + import main_prometheus_client + + self.main = main_prometheus_client + sleep = mock.patch.object(main_prometheus_client.time, "sleep") + self.sleep = sleep.start() + self.addCleanup(sleep.stop) + + def start_times(self, config: Dict[str, Any], *infos: Any) -> Dict[str, float]: + with mock.patch.object( + self.main, "fetch_runtime_info", side_effect=list(infos) + ) as fetch: + result = self.main.get_query_start_times(URL, config) + fetch.assert_called_with(URL) + self.assertEqual(fetch.call_count, len(infos)) + return result + + def test_waits_for_query_aggregations_then_returns_seconds(self) -> None: + config = { + "queries": [{"query": QUERY, "aggregations": [{"aggregation_id": 3}]}] + } + result = self.start_times( + config, runtime_info({}), runtime_info({3: 1_500, 9: 0}) + ) + self.assertEqual(result, {QUERY: 1.5}) + self.sleep.assert_called_once() + + def test_runtime_info_errors_are_not_retried(self) -> None: + config = { + "queries": [{"query": QUERY, "aggregations": [{"aggregation_id": 3}]}] + } + with self.assertRaises(requests.ConnectionError): + self.start_times(config, requests.ConnectionError("down")) + + def test_expr_only_plan_waits_for_rqe_leaf_aggregations(self) -> None: + # An --expr-only plan: no per-query configs, so no per-query start time. + config = { + "rqe_configs": { + "t/sum_by_job/keys/1s": { + "aggregations": [ + {"aggregation_id": 4, "num_aggregates_to_retain": 1} + ], + "query_time_aggregations": [], + }, + "t/max_by_job/keys/1s": { + "aggregations": [{"aggregation_id": 5, "read_count_threshold": 2}], + "query_time_aggregations": [], + }, + } + } + result = self.start_times( + config, runtime_info({4: 0}), runtime_info({4: 0, 5: 1_000}) + ) + self.assertEqual(result, {}) + self.sleep.assert_called_once() + + def test_waits_for_rqe_leaves_next_to_query_configs(self) -> None: + config = { + "queries": [{"query": QUERY, "aggregations": [{"aggregation_id": 3}]}], + "rqe_configs": { + "t/sum_by_job/keys/1s": {"aggregations": [{"aggregation_id": 4}]} + }, + } + result = self.start_times( + config, runtime_info({3: 1_000}), runtime_info({3: 1_000, 4: 0}) + ) + self.assertEqual(result, {QUERY: 1.0}) + + def test_required_ids_with_empty_sections(self) -> None: + self.assertEqual( + self.main.required_aggregation_ids({"queries": None, "rqe_configs": {}}), + set(), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/asap-tools/queriers/prometheus-client/tests/test_rqe_query_client.py b/asap-tools/queriers/prometheus-client/tests/test_rqe_query_client.py new file mode 100644 index 00000000..2b9a281c --- /dev/null +++ b/asap-tools/queriers/prometheus-client/tests/test_rqe_query_client.py @@ -0,0 +1,258 @@ +"""Tests for the rqe protocol client and its side file.""" + +import importlib.util +import json +import os +import tempfile +import unittest +from typing import Any, Dict +from unittest import mock + +import requests + +from classes.config import ServerConfig +from classes.prometheus_query_client import fetch_runtime_info, normalize_vector +from classes.query_client_factory import QueryClientFactory +from classes.rqe_query_client import ( + RQE_QUERY_PATH, + RqeQueryClient, + load_rqe_exprs, + resolve_rqe_exprs_file, +) + +URL = "http://sketchdb:8088" +PROMQL = "sum by (job) (data)" +SUM_BY_JOB = { + "rqe": { + "id": "t/sum_by_job/keys/1s", + "metric": "data", + "capability": "sum", + "lookback_ms": 100, + "grouping": ["job"], + } +} +SQL_TEMPLATE = "SELECT job, SUM(v) FROM t WHERE ts <= '{{ end_time }}' GROUP BY job" +VECTOR = [ + {"metric": {"job": "a"}, "value": [1791885600, "3"]}, + {"metric": {"job": "b"}, "value": [1791885600, "4.5"]}, +] + + +def response(status_code: int, body: Any) -> mock.Mock: + r = mock.Mock(spec=requests.Response) + r.status_code = status_code + r.json.return_value = body + r.text = json.dumps(body) + return r + + +class RqeQueryClientTest(unittest.TestCase): + def setUp(self) -> None: + self.tmp = tempfile.TemporaryDirectory() + self.side_file = os.path.join(self.tmp.name, "rqe_exprs.json") + self.write_side_file({PROMQL: SUM_BY_JOB, SQL_TEMPLATE: SUM_BY_JOB}) + self.client = RqeQueryClient(URL, "sketchdb", self.side_file) + self.post = mock.Mock( + return_value=response( + 200, + { + "status": "success", + "data": {"resultType": "vector", "result": VECTOR}, + }, + ) + ) + self.client.session.post = self.post # type: ignore[method-assign] + + def tearDown(self) -> None: + self.tmp.cleanup() + + def write_side_file(self, exprs: Any) -> None: + with open(self.side_file, "w") as f: + json.dump(exprs, f) + + def test_posts_the_tree_and_parses_like_prometheus(self) -> None: + result = self.client.execute_query(PROMQL, query_time=1791885600) + self.post.assert_called_once() + args, kwargs = self.post.call_args + self.assertEqual(args, (URL + RQE_QUERY_PATH,)) + self.assertEqual(kwargs["json"], {"expr": SUM_BY_JOB, "time": 1791885600.0}) + self.assertTrue(result.success) + # raw_response is what QueryResult.result stores for Prometheus. + self.assertEqual(result.raw_response, VECTOR) + self.assertEqual(result.data, normalize_vector(VECTOR)) + + def test_looks_up_the_unrendered_config_text(self) -> None: + rendered = SQL_TEMPLATE.replace("{{ end_time }}", "2026-10-13 10:00:00") + result = self.client.execute_query( + rendered, query_time=1791885600, config_query=SQL_TEMPLATE + ) + self.assertTrue(result.success, result.error_message) + + def test_query_missing_from_side_file_fails_without_a_request(self) -> None: + result = self.client.execute_query("rate(data_total[15m])", query_time=1) + self.assertFalse(result.success) + assert result.error_message is not None + self.assertIn( + "no RQE expression for query 'rate(data_total[15m])'", result.error_message + ) + self.assertIn(self.side_file, result.error_message) + self.post.assert_not_called() + + def test_rendered_text_is_not_a_key(self) -> None: + # Only the config's text is a key; a rendered template never matches. + rendered = SQL_TEMPLATE.replace("{{ end_time }}", "2026-10-13 10:00:00") + result = self.client.execute_query(rendered, query_time=1791885600) + self.assertFalse(result.success) + self.post.assert_not_called() + + def test_needs_an_evaluation_time(self) -> None: + result = self.client.execute_query(PROMQL) + self.assertFalse(result.success) + self.post.assert_not_called() + + def test_bad_request_is_a_failure(self) -> None: + self.post.return_value = response( + 400, + {"status": "error", "errorType": "bad_data", "error": "unknown RQE id"}, + ) + result = self.client.execute_query(PROMQL, query_time=1) + self.assertFalse(result.success) + assert result.error_message is not None + self.assertIn("HTTP 400", result.error_message) + self.assertIn("unknown RQE id", result.error_message) + + def test_error_status_in_a_200_is_a_failure(self) -> None: + self.post.return_value = response(200, {"status": "error", "error": "x"}) + self.assertFalse(self.client.execute_query(PROMQL, query_time=1).success) + + def test_connection_error_is_a_failure(self) -> None: + self.post.side_effect = requests.exceptions.ConnectionError("refused") + result = self.client.execute_query(PROMQL, query_time=1) + self.assertFalse(result.success) + assert result.error_message is not None + self.assertIn("ConnectionError", result.error_message) + + def test_non_vector_result_is_a_failure(self) -> None: + for result_type, raw in (("scalar", [1, "3"]), ("matrix", []), (None, [])): + data: Dict[str, Any] = {"result": raw} + if result_type is not None: + data["resultType"] = result_type + self.post.return_value = response(200, {"status": "success", "data": data}) + result = self.client.execute_query(PROMQL, query_time=1) + self.assertFalse(result.success) + self.assertEqual( + result.error_message, + f"unexpected result type {result_type!r}, expected 'vector'", + ) + + def test_empty_result_is_a_success(self) -> None: + self.post.return_value = response( + 200, {"status": "success", "data": {"resultType": "vector", "result": []}} + ) + result = self.client.execute_query(PROMQL, query_time=1) + self.assertTrue(result.success) + self.assertEqual(result.data, {}) + + def test_side_file_must_be_an_object(self) -> None: + self.write_side_file([SUM_BY_JOB]) + with self.assertRaisesRegex(ValueError, "expected a JSON object"): + load_rqe_exprs(self.side_file) + + def test_factory_creates_rqe_clients(self) -> None: + client = QueryClientFactory.create( + "rqe", URL, "sketchdb", rqe_exprs_file=self.side_file + ) + self.assertIsInstance(client, RqeQueryClient) + self.assertEqual(client.protocol_name, "rqe") + + +class RuntimeInfoTest(unittest.TestCase): + def get(self, r: Any) -> Any: + return mock.patch("classes.prometheus_query_client.requests.get", **r) + + def test_returns_the_data(self) -> None: + info: Dict[str, Any] = {"earliest_timestamp_per_aggregation_id": {}} + r = response(200, {"status": "success", "data": info}) + r.raise_for_status.return_value = None + with self.get({"return_value": r}): + self.assertEqual(fetch_runtime_info(URL), info) + + def test_unreachable_server_raises(self) -> None: + with self.get({"side_effect": requests.exceptions.ConnectionError("x")}): + with self.assertRaises(requests.exceptions.ConnectionError): + fetch_runtime_info(URL) + + def test_error_status_raises(self) -> None: + r = response(503, "down") + r.raise_for_status.side_effect = requests.exceptions.HTTPError("503") + with self.get({"return_value": r}): + with self.assertRaises(requests.exceptions.HTTPError): + fetch_runtime_info(URL) + + def test_malformed_body_raises(self) -> None: + r = response(200, {"status": "success"}) + r.raise_for_status.return_value = None + with self.get({"return_value": r}): + with self.assertRaises(KeyError): + fetch_runtime_info(URL) + + +class RqeServerConfigTest(unittest.TestCase): + def test_rqe_needs_a_side_file(self) -> None: + with self.assertRaisesRegex(ValueError, "needs rqe_exprs_file"): + ServerConfig.from_dict({"name": "sketchdb", "url": URL, "protocol": "rqe"}) + + def test_side_file_read_from_server_config(self) -> None: + server = ServerConfig.from_dict( + { + "name": "sketchdb", + "url": URL, + "protocol": "rqe", + "rqe_exprs_file": "rqe_exprs.json", + } + ) + self.assertEqual(server.rqe_exprs_file, "rqe_exprs.json") + + def test_relative_side_file_is_next_to_the_config(self) -> None: + self.assertEqual( + resolve_rqe_exprs_file( + "/app/prometheus_client_config.yaml", "rqe_exprs.json" + ), + "/app/rqe_exprs.json", + ) + self.assertEqual( + resolve_rqe_exprs_file("/app/c.yaml", "/data/rqe_exprs.json"), + "/data/rqe_exprs.json", + ) + + +@unittest.skipUnless( + importlib.util.find_spec("promql_utilities") is not None, + "main_prometheus_client needs asap-common's promql_utilities", +) +class ExecuteSingleQueryTest(unittest.TestCase): + def test_client_gets_the_rendered_and_the_config_text(self) -> None: + from main_prometheus_client import execute_single_query + + client = mock.Mock() + client.execute_query.return_value = mock.Mock(success=True, raw_response=VECTOR) + result = execute_single_query( + "sketchdb", + client, + SQL_TEMPLATE, + query_idx=0, + repetition_idx=0, + query_unix_time=1791885600, + dry_run=False, + query_group_idx=0, + time_window_seconds=0, + ) + kwargs: Dict[str, Any] = client.execute_query.call_args.kwargs + self.assertEqual(kwargs["config_query"], SQL_TEMPLATE) + self.assertNotIn("{{", kwargs["query"]) + # QueryResult keys the vector by label set, as for Prometheus. + self.assertEqual(result.result, normalize_vector(VECTOR)) + + +if __name__ == "__main__": + unittest.main()