diff --git a/benchmarking/README.md b/benchmarking/README.md index 7251e3c383..43f2fe3f8b 100644 --- a/benchmarking/README.md +++ b/benchmarking/README.md @@ -79,6 +79,17 @@ not a local entry point. See [automation/README.md](automation/README.md). python3 runner.py -f tests/.py -t 1m -u 1 --name --dest /tmp/bench ``` +Two flags control the optional post-run measurements described in +[Benchmark output files](#benchmark-output-files): + +* `--cluster-facts` / `--no-cluster-facts`: read node capacity and worker pod + count from the Kubernetes API once the run ends, to derive density frontiers. + On by default. Pass `--no-cluster-facts` on a large cluster, where listing + every node and pod is expensive. +* `--prometheus-url`: the Prometheus to harvest server-side telemetry from. + Defaults to the in-cluster service installed by + [Optional: Prometheus + Grafana](#optional-prometheus--grafana). + Test-specific flags are appended to the same command; see the sections below. ### DurDir Benchmark @@ -114,6 +125,64 @@ You must have enabled otel tracing for your cluster to view traces. You can find trace IDs by viewing the `logs` tab in the Locust UI +## Benchmark output files + +A run writes the following to `--dest`. Each run produces them fresh; none of +them are checked into the repository. + +* `status.json`: `locust_exit_code` and `stats_generated`. Deliberately just + those two keys, because it is what CI orchestration reads to decide whether a + trial ran at all. +* `stats.csv`, `stats_history.csv`, `failures.csv`, `exceptions.csv`: Locust's + own CSV output. +* `logs.txt`, `traces.txt`: the runner log, and the trace IDs seen during the run. +* `stats.jsonl`: one JSON object per line, one per metric. Every row carries + `timestamp`, `tag`, `test_name` and `metric`. +* `server_summary.json`: server-side telemetry harvested from Prometheus, + including the per-sample bin-packing timeseries. + +### Density frontiers + +With cluster discovery enabled, `stats.jsonl` gains a `trial_summary` row +describing how densely actors packed onto the hardware. + +* `raw_configuration`: the measured facts, before any arithmetic: + `machine_type`, `node_count`, `allocatable_cores`, `allocatable_ram_gb` + (GiB), `worker_pod_count`. They are recorded so the ratios below can be + re-derived later, or recomputed against a different denominator. +* `frontiers.actors_per_node`, `frontiers.actors_per_vcpu`, + `frontiers.actors_per_gb_ram`: active users over the matching capacity. +* `frontiers.ap_ratio_p50`, `ap_ratio_p90`, `ap_ratio_p99`: the + actor-to-pod ratio across the steady-state part of the run. Reported as a + distribution rather than one average, because the ratio moves a lot while + users are still ramping up. +* `frontiers.aggregate_failure_ratio`: failures over requests for the run. + +### Server ground truth + +With a reachable Prometheus, `server_summary.json` records what the server +actually did, independent of what the load generator reported. + +* `cluster_packing`: assigned workers over total workers, as a percentile + `summary` plus the per-sample `timeseries` it was computed from. +* `node_psi.cpu_stall_pct`, `mem_stall_pct`, `io_stall_pct`: kernel pressure + stall percentages on the nodes under test. +* `node_psi.cfs_throttled_rate`: CFS quota throttling rate. +* `snapshots.size_p50_mb`, `size_p90_mb`, `size_p95_mb`: actor snapshot sizes. +* `snapshots.size_avg_mb`: mean snapshot size, taken from the histogram's + own sum and count, so it is exact rather than bucket-interpolated. +* `snapshots.checkpoint_p50_s`, `checkpoint_p95_s`, `restore_p50_s`, + `restore_p95_s`: checkpoint and restore latency. +* `snapshots.checkpoints_in_window`, `checkpoints_cumulative`, + `throughput_mb_s`: checkpoint volume over the steady-state window. + +A flattened subset of the same numbers is appended to `stats.jsonl` as a +`server_summary` row, so both metrics can be read from the one file. + +Neither the Kubernetes API nor Prometheus is required. If either is unreachable, +or discovery was skipped, the affected fields are written as `null` and the run +still succeeds. A `null` means the value was not measured. It never means zero. + ## Optional: Prometheus + Grafana Locust provides graphs, statistics, etc. via the UI. However, you @@ -138,3 +207,14 @@ Once installed: code; it manages its own virtual environment under `locust/codegen/venv`. `hack/verify/codegen.sh` fails if the checked-in clients have drifted from the protos. + +### Unit tests + +`locust/unit_tests` covers the runner's helpers and needs no cluster. From the +repository root: + +```bash +python3 -m unittest discover -s benchmarking/locust/unit_tests +``` + +Tests that need the Kubernetes client are skipped when it is not installed. diff --git a/benchmarking/automation/manifests/runner-job.yaml.tmpl b/benchmarking/automation/manifests/runner-job.yaml.tmpl index 5855d1d05a..3eae830b1e 100644 --- a/benchmarking/automation/manifests/runner-job.yaml.tmpl +++ b/benchmarking/automation/manifests/runner-job.yaml.tmpl @@ -41,6 +41,37 @@ roleRef: name: atelet-endpointslices apiGroup: rbac.authorization.k8s.io --- +# Hardware discovery for the runner: runner.py reads nodes and worker pods to +# derive cluster capacity frontiers. Same ClusterRole as the one in +# benchmarking/locust/manifests/locust.yaml, repeated here so this manifest +# stands alone -- the locust Deployment is not part of an automated run. +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + name: locust-hardware-discovery +rules: +- apiGroups: [""] + resources: ["nodes", "pods"] + # list only: the runner reads whole collections, never a single object. + verbs: ["list"] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + # Intentionally not named locust-hardware-discovery: that binding belongs to + # locust.yaml. Two bindings on one ClusterRole add up, but two bindings + # sharing a name overwrite each other, silently revoking whichever subject + # the losing copy listed. + name: benchmark-runner-hardware-discovery +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: ClusterRole + name: locust-hardware-discovery +subjects: +- kind: ServiceAccount + name: benchmark-runner + namespace: benchmarking +--- apiVersion: batch/v1 kind: Job metadata: diff --git a/benchmarking/locust/Dockerfile b/benchmarking/locust/Dockerfile index 6de127ec1d..9def469e36 100644 --- a/benchmarking/locust/Dockerfile +++ b/benchmarking/locust/Dockerfile @@ -44,6 +44,8 @@ COPY benchmarking/locust/common/ /app/common/ COPY benchmarking/locust/shapes/ /app/shapes/ COPY benchmarking/locust/tests/ /app/tests/ COPY benchmarking/locust/runner.py /app/runner.py +COPY benchmarking/locust/cluster_facts.py /app/cluster_facts.py +COPY benchmarking/locust/server_telemetry.py /app/server_telemetry.py ENV PYTHONPATH=/app:/app/deps ENV PYTHONUNBUFFERED=1 diff --git a/benchmarking/locust/cluster_facts.py b/benchmarking/locust/cluster_facts.py new file mode 100644 index 0000000000..b53ae19a73 --- /dev/null +++ b/benchmarking/locust/cluster_facts.py @@ -0,0 +1,272 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +"""Discovers cluster hardware capacity and records per-trial density frontiers. + +Reads allocatable CPU/RAM, node count and worker pod count from the Kubernetes +API, then derives the actor-density frontiers (actors per node / vCPU / GB RAM +and the A/P bin-packing percentiles) for a completed trial. +""" + +from __future__ import annotations + +import argparse +import csv +import json +from pathlib import Path +from typing import Any, TextIO + +from kubernetes import client, config +from kubernetes.client.rest import ApiException +from kubernetes.utils import parse_quantity + +API_TIMEOUT_SECONDS = 5 +WORKER_POOL_NAMESPACE = "benchmark-workloads" +WORKER_POOL_LABEL = "ate.dev/worker-pool" +LIVE_POD_PHASES = ("Running", "Pending") +MACHINE_TYPE_LABEL = "node.kubernetes.io/instance-type" + +# Shape returned when the cluster cannot be read, or when discovery is skipped +# with --no-cluster-facts. Keeping one definition means a trial_summary row has +# the same keys either way, so consumers never have to special-case it. +EMPTY_FACTS: dict[str, Any] = { + "machine_type": None, + "node_count": None, + "allocatable_cores": None, + "allocatable_ram_gb": None, + "worker_pod_count": None, +} + + +def _log(logs: TextIO | None, msg: str) -> None: + """Mirrors runner.tee without importing it, to avoid a circular import.""" + print(msg, flush=True) + if logs is not None: + logs.write(msg + "\n") + logs.flush() + + +def _load_kube_config(logs: TextIO | None = None) -> bool: + """Loads in-cluster credentials, falling back to a local kubeconfig.""" + try: + config.load_incluster_config() + return True + except config.ConfigException: + pass + try: + config.load_kube_config() + return True + except config.ConfigException as e: + _log(logs, f"Notice: no Kubernetes credentials available: {e}") + return False + + +def _count_worker_pods(v1: client.CoreV1Api, logs: TextIO | None = None) -> int | None: + """Counts live pods in the worker pool. + + Prefers the dedicated worker namespace and falls back to a cluster-wide + lookup for clusters that place the pool elsewhere. Listing a namespace that + does not exist returns an empty list rather than an error, so an empty + result is what "the pool is somewhere else" looks like and it has to + trigger the fallback. + + Both listings are filtered server-side by label and served from the watch + cache, but the cluster-wide one still scans every pod, so it is logged + whenever it happens. Use --no-cluster-facts to skip discovery entirely. + """ + namespaced_failed = False + try: + pods = v1.list_namespaced_pod( + namespace=WORKER_POOL_NAMESPACE, + label_selector=WORKER_POOL_LABEL, + resource_version="0", + _request_timeout=API_TIMEOUT_SECONDS, + ).items + except ApiException as e: + _log(logs, f"Notice: could not list pods in {WORKER_POOL_NAMESPACE}: {e.reason}") + pods = [] + namespaced_failed = True + + if not pods: + reason = ( + "the namespaced lookup failed" + if namespaced_failed + else f"no {WORKER_POOL_LABEL} pods in {WORKER_POOL_NAMESPACE}" + ) + _log(logs, f"Notice: {reason}; scanning all namespaces for the worker pool") + try: + pods = v1.list_pod_for_all_namespaces( + label_selector=WORKER_POOL_LABEL, + resource_version="0", + _request_timeout=API_TIMEOUT_SECONDS, + ).items + except ApiException as e: + _log(logs, f"Notice: cluster-wide pod lookup failed: {e.reason}") + return None + + live = [p for p in pods if p.status.phase in LIVE_POD_PHASES] + return len(live) or None + + +def get_cluster_hardware_facts(logs: TextIO | None = None) -> dict[str, Any]: + """Reads allocatable node capacity and worker pod count from the cluster. + + Never raises: a trial must still publish its results when the cluster is + unreadable, so any failure leaves the affected facts as None. + """ + facts: dict[str, Any] = dict(EMPTY_FACTS) + if not _load_kube_config(logs): + return facts + + v1 = client.CoreV1Api() + + # resource_version="0" is served from the apiserver's watch cache rather + # than etcd, which keeps this cheap on large clusters. + try: + nodes = v1.list_node( + resource_version="0", _request_timeout=API_TIMEOUT_SECONDS + ).items + total_cores = 0.0 + total_ram_bytes = 0 + machine_types = set() + for node in nodes: + allocatable = node.status.allocatable or {} + if "cpu" in allocatable: + total_cores += float(parse_quantity(allocatable["cpu"])) + if "memory" in allocatable: + total_ram_bytes += int(parse_quantity(allocatable["memory"])) + labels = (node.metadata.labels or {}) if node.metadata else {} + machine_type = labels.get(MACHINE_TYPE_LABEL) + if machine_type: + machine_types.add(machine_type) + facts["node_count"] = len(nodes) + facts["allocatable_cores"] = round(total_cores, 2) + # GiB, as the apiserver and kubectl quote it. Named _gb for continuity + # with rows already collected; renaming would break consumers. + facts["allocatable_ram_gb"] = round(total_ram_bytes / (1024**3), 2) + # Kept so results stay comparable across hardware changes. A mixed pool + # is a sorted comma-joined list rather than one node picked at random. + facts["machine_type"] = ",".join(sorted(machine_types)) or None + except Exception as e: + _log(logs, f"Notice: could not read node capacity: {e}") + + try: + facts["worker_pod_count"] = _count_worker_pods(v1, logs) + except Exception as e: + _log(logs, f"Notice: could not count worker pods: {e}") + + return facts + + +def append_trial_summary( + jsonl_path: Path, + stats_csv: Path, + stats_history_csv: Path, + args: argparse.Namespace, + data_ts: str, + facts: dict[str, Any], + logs: TextIO | None = None, +) -> None: + active_users = args.users + # Only the facts the frontier math divides by. machine_type is recorded + # but never computed with, so it goes straight into raw_configuration. + node_count = facts.get("node_count") + cores = facts.get("allocatable_cores") + ram_gb = facts.get("allocatable_ram_gb") + pod_count = facts.get("worker_pod_count") + + actors_per_node = round(active_users / node_count, 2) if node_count else None + actors_per_vcpu = round(active_users / cores, 2) if cores else None + actors_per_gb_ram = round(active_users / ram_gb, 2) if ram_gb else None + + # A/P bin-packing percentiles over the steady-state samples. None when + # unmeasurable: a ratio derived from configured user count is not a reading. + ap_p50, ap_p90, ap_p99 = None, None, None + if stats_history_csv.exists() and pod_count and pod_count > 0: + try: + observed: list[float] = [] + with open(stats_history_csv) as f: + for row in csv.DictReader(f): + if row.get("Name", "") not in ("", "Aggregated", "Total"): + continue + try: + u = float(row.get("User Count", "")) + except (TypeError, ValueError): + continue + if u > 0: + observed.append(u) + # Steady state is every sample at or above 90% of the target. A run + # that never got there falls back to every non-zero sample. + steady = [u for u in observed if u >= active_users * 0.9] or observed + if steady: + ratios = sorted(round(u / pod_count, 4) for u in steady) + n = len(ratios) + ap_p50 = round(ratios[int(n * 0.50)], 2) + ap_p90 = round(ratios[min(int(n * 0.90), n - 1)], 2) + ap_p99 = round(ratios[min(int(n * 0.99), n - 1)], 2) + except Exception as e: + _log(logs, f"Notice: Error calculating A/P ratio percentiles: {e}") + + total_requests = 0 + total_failures = 0 + stats_parsed = False + if stats_csv.exists(): + try: + with open(stats_csv) as f: + reader = csv.DictReader(f) + for row in reader: + name = row.get("Name", "") + reqs = int(row.get("Request Count", 0) or 0) + fails = int(row.get("Failure Count", 0) or 0) + if name == "Aggregated": + total_requests = reqs + total_failures = fails + break + total_requests += reqs + total_failures += fails + stats_parsed = True + except Exception as e: + _log(logs, f"Notice: could not parse {stats_csv}: {e}") + else: + _log(logs, f"Notice: {stats_csv} not found; failure ratio unknown") + + # None, not 0.0, when undetermined. A run with zero failures is a real + # result and must not look like one where the stats file was unreadable. + if stats_parsed and total_requests > 0: + failure_ratio = round(total_failures / total_requests, 4) + else: + failure_ratio = None + + summary_entry = { + "timestamp": data_ts, + "tag": args.tag, + "test_name": args.name, + "metric": "trial_summary", + # Keyed off EMPTY_FACTS so the raw block always carries every fact, + # present or not, and the names are declared in one place. + "raw_configuration": {k: facts.get(k) for k in EMPTY_FACTS}, + "frontiers": { + "actors_per_node": actors_per_node, + "actors_per_vcpu": actors_per_vcpu, + "actors_per_gb_ram": actors_per_gb_ram, + "ap_ratio_p50": ap_p50, + "ap_ratio_p90": ap_p90, + "ap_ratio_p99": ap_p99, + "aggregate_failure_ratio": failure_ratio, + }, + } + with open(jsonl_path, "a") as f: + f.write(json.dumps(summary_entry) + "\n") + _log(logs, f"Appended trial_summary to {jsonl_path}") diff --git a/benchmarking/locust/manifests/locust.yaml b/benchmarking/locust/manifests/locust.yaml index 36374e61b4..4ec9404df8 100644 --- a/benchmarking/locust/manifests/locust.yaml +++ b/benchmarking/locust/manifests/locust.yaml @@ -202,3 +202,35 @@ roleRef: kind: Role name: atelet-endpointslices apiGroup: rbac.authorization.k8s.io +--- +# Hardware discovery for benchmarking runner: nodes and worker pods across +# namespaces are read to compute cluster capacity frontiers (actors/vCPU, +# actors/GB RAM, A/P bin-packing ratio). +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + name: locust-hardware-discovery +rules: +- apiGroups: [""] + resources: ["nodes", "pods"] + # list only: the runner reads whole collections, never a single object. + verbs: ["list"] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + name: locust-hardware-discovery +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: ClusterRole + name: locust-hardware-discovery +subjects: +# The locust Deployment runs as `default`; the automated runner Job in +# benchmarking/automation/manifests/runner-job.yaml.tmpl runs as +# `benchmark-runner`. Both execute runner.py, so both need to read capacity. +- kind: ServiceAccount + name: default + namespace: benchmarking +- kind: ServiceAccount + name: benchmark-runner + namespace: benchmarking diff --git a/benchmarking/locust/requirements.txt b/benchmarking/locust/requirements.txt index 033e5a7422..dc7449831c 100644 --- a/benchmarking/locust/requirements.txt +++ b/benchmarking/locust/requirements.txt @@ -24,3 +24,4 @@ opentelemetry-exporter-otlp opentelemetry-instrumentation-grpc opentelemetry-instrumentation-requests google-cloud-storage +kubernetes diff --git a/benchmarking/locust/runner.py b/benchmarking/locust/runner.py index 14bd15629d..e2e057302b 100644 --- a/benchmarking/locust/runner.py +++ b/benchmarking/locust/runner.py @@ -41,8 +41,13 @@ import time from datetime import datetime, timezone from pathlib import Path -from typing import IO, TextIO +from typing import IO, Any, TextIO +from cluster_facts import ( + EMPTY_FACTS, + append_trial_summary, + get_cluster_hardware_facts, +) from common.boomer_config import build_config_json # Path inside the locust image to the boomer-worker binary baked in by @@ -54,6 +59,10 @@ # holds 5557 (master) and 8089 (web UI) in this container. BOOMER_CONFIG_PORT = 5560 +# In-cluster Prometheus that benchmarking/monitoring.yaml deploys. Override +# with --prometheus-url, for example when port-forwarding to a local run. +DEFAULT_PROMETHEUS_URL = "http://prometheus.benchmarking.svc.cluster.local:9090" + # Tab-separated columns written to traces.txt. Order matters — readers split # on \t and index positionally. TRACE_COLUMNS = ("time", "name", "duration_ms", "latency_source", "trace_id", "err") @@ -94,6 +103,26 @@ def parse_args() -> argparse.Namespace: "no request on purpose" ), ) + p.add_argument( + "--cluster-facts", + action=argparse.BooleanOptionalAction, + default=True, + help=( + "Read node capacity and worker pod count from the Kubernetes API " + "after the run to derive density frontiers. Pass " + "--no-cluster-facts to skip those API calls, for example on a " + "large cluster where listing nodes is expensive" + ), + ) + p.add_argument( + "--prometheus-url", + default=DEFAULT_PROMETHEUS_URL, + help=( + "Prometheus to harvest server-side telemetry from after the run. " + "An unreachable Prometheus is not an error: the affected fields " + "are recorded as null" + ), + ) args, extra = p.parse_known_args() args.locust_extra = extra return args @@ -390,6 +419,20 @@ def upload(src: Path, dest: str) -> None: shutil.copy(src, dest_path) +def collect_cluster_facts( + args: argparse.Namespace, logs: TextIO +) -> dict[str, Any]: + """Returns cluster hardware facts, or empty facts when discovery is off. + + --no-cluster-facts short-circuits before any Kubernetes API call, for + clusters where listing nodes and pods is expensive. + """ + if not args.cluster_facts: + tee(logs, "Skipping cluster hardware discovery (--no-cluster-facts)") + return dict(EMPTY_FACTS) + return get_cluster_hardware_facts(logs) + + def main() -> None: args = parse_args() now = datetime.now(timezone.utc) @@ -409,6 +452,7 @@ def main() -> None: logs_path = work_dir / f"{args.name}_logs.txt" traces_path = work_dir / f"{args.name}_traces.txt" status_path = work_dir / f"{args.name}_status.json" + server_summary_json = work_dir / f"{args.name}_server_summary.json" prefix = ( f"{args.dest.rstrip('/')}/runs/{args.name}" @@ -420,6 +464,7 @@ def main() -> None: traces.flush() log_run_config(args, prefix, work_dir, logs) exit_code = run_test(args, csv_prefix, logs, traces) + run_end_ts = int(datetime.now(timezone.utc).timestamp()) stats_generated = False if stats_csv.exists(): @@ -444,6 +489,49 @@ def main() -> None: else: tee(logs, f"Stats CSV {stats_csv} not produced; skipping JSONL") + # Density frontiers and server-side telemetry are additive. They are + # kept out of the block above so that a failure here cannot discard + # the measurements the trial actually came for. + if stats_generated: + stats_history_csv = work_dir / f"{args.name}_stats_history.csv" + # Seeded up front so that a later failure still leaves a usable + # value for the telemetry call below. + facts = dict(EMPTY_FACTS) + try: + facts = collect_cluster_facts(args, logs) + append_trial_summary( + jsonl_path, + stats_csv, + stats_history_csv, + args, + data_ts, + facts, + logs, + ) + except Exception as e: + tee(logs, f"Warning: Failed to record cluster facts: {e}") + + # Harvest server-side ground truth from Prometheus (bin-packing, PSI, snapshots) + try: + from server_telemetry import extract_and_record_server_telemetry + + extract_and_record_server_telemetry( + prom_url=args.prometheus_url, + start_ts=run_ts, + end_ts=run_end_ts, + stats_history_csv=stats_history_csv, + active_users=args.users, + worker_pod_count=facts.get("worker_pod_count"), + output_json_path=server_summary_json, + jsonl_path=jsonl_path, + data_ts=data_ts, + tag=args.tag, + test_name=args.name, + logs=logs, + ) + except Exception as e: + tee(logs, f"Warning: Failed to harvest server telemetry: {e}") + status_path.write_text( json.dumps( {"locust_exit_code": exit_code, "stats_generated": stats_generated} @@ -459,6 +547,7 @@ def main() -> None: (work_dir / f"{args.name}_exceptions.csv", "exceptions.csv"), (work_dir / f"{args.name}_failures.csv", "failures.csv"), (work_dir / f"{args.name}_stats_history.csv", "stats_history.csv"), + (server_summary_json, "server_summary.json"), # TODO: remove after data migration (jsonl_path, f"{args.name}.jsonl"), ] diff --git a/benchmarking/locust/server_telemetry.py b/benchmarking/locust/server_telemetry.py new file mode 100644 index 0000000000..70085606c6 --- /dev/null +++ b/benchmarking/locust/server_telemetry.py @@ -0,0 +1,503 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Harvests server-side Prometheus ground-truth timeseries during benchmark trials. + +Queries Prometheus over [T_start, T_end] and the steady-state window [T_steady, T_end] +to capture dynamic cluster packing, node PSI stalls, and snapshot throughput. +""" + +from __future__ import annotations + +import csv +import json +import math +from pathlib import Path +import sys +from typing import Any, TextIO +import urllib.parse +import urllib.request + + +def query_prometheus_instant( + base_url: str, + query: str, + time_ts: float | int | None = None, + timeout_s: float = 5.0, +) -> list[dict[str, Any]]: + """Executes an instant query against Prometheus /api/v1/query.""" + params = {"query": query} + if time_ts is not None: + params["time"] = str(time_ts) + url = f"{base_url.rstrip('/')}/api/v1/query?{urllib.parse.urlencode(params)}" + try: + req = urllib.request.Request( + url, headers={"User-Agent": "Substrate-Locust-Runner"} + ) + with urllib.request.urlopen(req, timeout=timeout_s) as resp: + data = json.loads(resp.read().decode("utf-8")) + if data.get("status") == "success": + return data.get("data", {}).get("result", []) + except Exception as e: + print(f"Warning: Instant query failed '{query}': {e}", file=sys.stderr) + return [] + + +def query_prometheus_range( + base_url: str, + query: str, + start_ts: int, + end_ts: int, + step: str = "5s", + timeout_s: float = 8.0, +) -> list[dict[str, Any]]: + """Executes a range query against Prometheus /api/v1/query_range.""" + # Guard against Prometheus 400 Bad Request: end must be greater than start + if end_ts <= start_ts: + end_ts = start_ts + 1 + + params = { + "query": query, + "start": str(start_ts), + "end": str(end_ts), + "step": step, + } + url = f"{base_url.rstrip('/')}/api/v1/query_range?{urllib.parse.urlencode(params)}" + try: + req = urllib.request.Request( + url, headers={"User-Agent": "Substrate-Locust-Runner"} + ) + with urllib.request.urlopen(req, timeout=timeout_s) as resp: + data = json.loads(resp.read().decode("utf-8")) + if data.get("status") == "success": + return data.get("data", {}).get("result", []) + except Exception as e: + print(f"Warning: Range query failed '{query}': {e}", file=sys.stderr) + return [] + + +def compute_percentiles(values: list[float]) -> dict[str, float | None]: + """Computes min, p50, p90, p99, max, avg after filtering out NaN and Inf values.""" + clean = sorted([v for v in values if not math.isnan(v) and not math.isinf(v)]) + if not clean: + return { + "min": None, + "p50": None, + "p90": None, + "p99": None, + "max": None, + "avg": None, + } + n = len(clean) + return { + "min": round(clean[0], 4), + "p50": round(clean[int(n * 0.50)], 4), + "p90": round(clean[min(int(n * 0.90), n - 1)], 4), + "p99": round(clean[min(int(n * 0.99), n - 1)], 4), + "max": round(clean[-1], 4), + "avg": round(sum(clean) / n, 4), + } + + +def get_steady_state_window( + stats_history_csv: Path, + active_users: int, + start_ts: int, + end_ts: int, +) -> tuple[int, int]: + """Derives steady-state [T_steady, T_end] where User Count >= 0.9 * active_users.""" + if not stats_history_csv.exists() or active_users <= 0: + return start_ts, end_ts + + steady_ts: int | None = None + try: + with open(stats_history_csv, "r", encoding="utf-8") as f: + reader = csv.DictReader(f) + for row in reader: + name = row.get("Name", "") + if ( + name in ("", "Aggregated", "Total") + and "User Count" in row + and "Timestamp" in row + ): + try: + users = float(row["User Count"]) + ts = int(row["Timestamp"]) + if users >= active_users * 0.9: + steady_ts = ts + break + except (ValueError, TypeError): + continue + except Exception: + pass + + if steady_ts is not None and start_ts <= steady_ts <= end_ts: + return steady_ts, end_ts + return start_ts, end_ts + + +def _parse_instant_float(res: list[dict[str, Any]]) -> float | None: + if res and "value" in res[0]: + try: + v = float(res[0]["value"][1]) + return None if math.isnan(v) or math.isinf(v) else round(v, 4) + except (KeyError, IndexError, TypeError, ValueError): + # Every caller treats None as "no reading". A bad response shape + # must cost this one field, not the whole server_summary.json. + pass + return None + + +def _parse_instant_int(res: list[dict[str, Any]]) -> int | None: + """Returns None, not 0, when the query yielded nothing. + + Prometheus is optional, so an unreachable server is an expected outcome. + Zero is a legitimate reading here, so returning it on failure would leave + a consumer unable to tell "no checkpoints happened" from "we never found + out". Mirrors _parse_instant_float. + """ + val = _parse_instant_float(res) + return int(val) if val is not None else None + + +def _query_quantile_with_fallback( + prom_url: str, + quantile: float, + rate_metric_expr: str, + raw_metric_expr: str, + end_ts: int, + unit_scale: float = 1.0, +) -> float | None: + """Queries histogram quantile using rate(5m), falling back to cumulative buckets if NaN/empty.""" + # _parse_instant_float maps NaN and Inf to None, so None is the only + # "no reading" value either query can produce. + q_rate = ( + f"histogram_quantile({quantile}, sum({rate_metric_expr}) by (le)) / {unit_scale}" + ) + val = _parse_instant_float( + query_prometheus_instant(prom_url, q_rate, time_ts=end_ts) + ) + if val is not None: + return val + + # Fallback to cumulative bucket distribution + q_cum = ( + f"histogram_quantile({quantile}, sum by (le) ({raw_metric_expr})) / {unit_scale}" + ) + return _parse_instant_float( + query_prometheus_instant(prom_url, q_cum, time_ts=end_ts) + ) + + +def harvest_server_telemetry( + prom_url: str, + start_ts: int, + end_ts: int, + steady_start_ts: int, + worker_pod_count: int | None, +) -> dict[str, Any]: + """Harvests all 4 ground truth metric streams from Prometheus.""" + summary: dict[str, Any] = { + "cluster_packing": {}, + "node_psi": {}, + "snapshots": {}, + } + + # 1. Cluster Packing Timeseries (deduping ateapi replicas via max by state) + packing_query = ( + 'max by (ate_worker_state) ' + '(ate_workerpool_workers{ate_workerpool_name="benchmark-ateom"})' + ) + packing_series = query_prometheus_range( + prom_url, packing_query, start_ts, end_ts, step="5s" + ) + + ts_packing_map: dict[int, dict[str, float]] = {} + for series in packing_series: + state = series.get("metric", {}).get("ate_worker_state", "unknown") + for pt in series.get("values", []): + try: + t = int(pt[0]) + val = float(pt[1]) + # Inf as well as NaN: these reach server_summary.json, and + # json.dumps would emit a bare Infinity that strict parsers reject. + if not math.isnan(val) and not math.isinf(val): + ts_packing_map.setdefault(t, {})[state] = val + except (ValueError, IndexError): + continue + + packing_points = [] + steady_packing_ratios = [] + for t in sorted(ts_packing_map.keys()): + states = ts_packing_map[t] + assigned = states.get("assigned", 0.0) + # Real worker pod count is the physical denominator; when unknown, use + # the worker states Prometheus reports rather than assuming a number. + total = ( + float(worker_pod_count) + if worker_pod_count + else (sum(states.values()) or 1.0) + ) + ratio = round(assigned / total, 4) + packing_points.append({ + "timestamp": t, + "assigned_workers": assigned, + "total_workers": total, + "packing_ratio": ratio, + }) + if t >= steady_start_ts: + steady_packing_ratios.append(ratio) + + summary["cluster_packing"] = { + "summary": compute_percentiles(steady_packing_ratios), + "timeseries": packing_points, + } + + # 2. Host Linux Kernel PSI Stalls & CFS Throttling + # + # _waiting_ is PSI "some" (at least one task stalled); cAdvisor also + # exports _stalled_, PSI "full" (all tasks stalled). "some" is the earlier + # warning signal and the CPU full-stall series carries none, so one series + # keeps the three comparable. Both are scraped, so "full" stays available. + def psi_query(resource: str) -> str: + return ( + f'sum by (instance) (rate(container_pressure_{resource}_waiting_seconds_total' + '{container="node"}[1m])) * 100' + ) + + # container!="" drops the per-pod rollups cAdvisor reports alongside each + # container, which would double count. Unlike the memory and CPU series + # these are not dropped at scrape time (see monitoring.yaml). + cfs_throttled_query = ( + 'sum(rate(container_cpu_cfs_throttled_seconds_total' + '{container!=""}[1m]))' + ) + + psi_cpu_res = query_prometheus_range( + prom_url, psi_query("cpu"), start_ts, end_ts, step="5s" + ) + psi_mem_res = query_prometheus_range( + prom_url, psi_query("memory"), start_ts, end_ts, step="5s" + ) + psi_io_res = query_prometheus_range( + prom_url, psi_query("io"), start_ts, end_ts, step="5s" + ) + cfs_res = query_prometheus_range( + prom_url, cfs_throttled_query, start_ts, end_ts, step="5s" + ) + + def extract_steady_values(results: list[dict[str, Any]]) -> list[float]: + vals = [] + for s in results: + for pt in s.get("values", []): + try: + if int(pt[0]) >= steady_start_ts: + v = float(pt[1]) + # Inf needs no filter here: compute_percentiles is the + # only consumer and it drops both NaN and Inf. + if not math.isnan(v): + vals.append(v) + except (ValueError, IndexError): + pass + return vals + + summary["node_psi"] = { + "cpu_stall_pct": compute_percentiles(extract_steady_values(psi_cpu_res)), + "mem_stall_pct": compute_percentiles(extract_steady_values(psi_mem_res)), + "io_stall_pct": compute_percentiles(extract_steady_values(psi_io_res)), + "cfs_throttled_rate": compute_percentiles(extract_steady_values(cfs_res)), + } + + # 3. Snapshot Sizes, Checkpoint Count & Latencies (with histogram fallback) + snap_bucket = "atelet_snapshot_size_bytes_bucket" + snap_rate = f"rate({snap_bucket}[5m])" + snap_p50 = _query_quantile_with_fallback( + prom_url, 0.50, snap_rate, snap_bucket, end_ts, unit_scale=1024 * 1024 + ) + snap_p90 = _query_quantile_with_fallback( + prom_url, 0.90, snap_rate, snap_bucket, end_ts, unit_scale=1024 * 1024 + ) + snap_p95 = _query_quantile_with_fallback( + prom_url, 0.95, snap_rate, snap_bucket, end_ts, unit_scale=1024 * 1024 + ) + + snap_count_query = "sum(atelet_snapshot_size_bytes_count)" + snap_count_start = query_prometheus_instant( + prom_url, snap_count_query, time_ts=steady_start_ts + ) + snap_count_end = query_prometheus_instant( + prom_url, snap_count_query, time_ts=end_ts + ) + c_start = _parse_instant_int(snap_count_start) + c_end = _parse_instant_int(snap_count_end) + + # A missing endpoint, or a counter that went backwards because an atelet + # restarted, makes the delta unknown. 0 would read as "nothing happened". + if c_start is None or c_end is None or c_end < c_start: + window_checkpoints = None + else: + window_checkpoints = c_end - c_start + + # Mean over the same window as the percentiles above, from the histogram's + # own _sum/_count so it is exact rather than bucket interpolated. _sum + # counts up from atelet start, so the end value alone would average the + # whole lifetime, not the test. + snap_sum_query = "sum(atelet_snapshot_size_bytes_sum)" + s_start = _parse_instant_float( + query_prometheus_instant(prom_url, snap_sum_query, time_ts=steady_start_ts) + ) + s_end = _parse_instant_float( + query_prometheus_instant(prom_url, snap_sum_query, time_ts=end_ts) + ) + snap_avg = None + if ( + s_start is not None + and s_end is not None + and s_end >= s_start + and window_checkpoints + ): + snap_avg = round((s_end - s_start) / window_checkpoints / (1024 * 1024), 4) + + def rpc_bucket(method: str) -> str: + return ( + 'rpc_server_call_duration_seconds_bucket' + f'{{rpc_method="atelet.AteomHerder/{method}"}}' + ) + + restore_bucket = rpc_bucket("Restore") + restore_rate = f"rate({restore_bucket}[5m])" + ckpt_bucket = rpc_bucket("Checkpoint") + ckpt_rate = f"rate({ckpt_bucket}[5m])" + + restore_p50 = _query_quantile_with_fallback( + prom_url, 0.50, restore_rate, restore_bucket, end_ts + ) + restore_p95 = _query_quantile_with_fallback( + prom_url, 0.95, restore_rate, restore_bucket, end_ts + ) + ckpt_p50 = _query_quantile_with_fallback( + prom_url, 0.50, ckpt_rate, ckpt_bucket, end_ts + ) + ckpt_p95 = _query_quantile_with_fallback( + prom_url, 0.95, ckpt_rate, ckpt_bucket, end_ts + ) + + steady_duration_s = max(1, end_ts - steady_start_ts) + throughput_mb_s = None + # `is not None`, not truthiness: a genuine 0.0 median or 0 checkpoints is a + # reading, not a missing value, and must not be mistaken for an absent one. + if ( + snap_p50 is not None + and window_checkpoints is not None + and window_checkpoints > 0 + ): + throughput_mb_s = round( + (snap_p50 * window_checkpoints) / steady_duration_s, 2 + ) + + summary["snapshots"] = { + "size_p50_mb": snap_p50, + "size_p90_mb": snap_p90, + "size_p95_mb": snap_p95, + "size_avg_mb": snap_avg, + "checkpoints_in_window": window_checkpoints, + "checkpoints_cumulative": c_end, + "restore_p50_s": restore_p50, + "restore_p95_s": restore_p95, + "checkpoint_p50_s": ckpt_p50, + "checkpoint_p95_s": ckpt_p95, + "throughput_mb_s": throughput_mb_s, + } + + return summary + + +def extract_and_record_server_telemetry( + prom_url: str, + start_ts: int, + end_ts: int, + stats_history_csv: Path, + active_users: int, + worker_pod_count: int | None, + output_json_path: Path, + jsonl_path: Path, + data_ts: str, + tag: str, + test_name: str, + logs: TextIO | None = None, +) -> None: + """Entry point called by runner.py to query Prometheus and persist artifacts.""" + def log(msg: str) -> None: + if logs: + print(f"[ServerTelemetry] {msg}", file=logs, flush=True) + print(f"[ServerTelemetry] {msg}", flush=True) + + log(f"Harvesting Prometheus metrics from {prom_url} over [{start_ts}, {end_ts}]...") + steady_start, steady_end = get_steady_state_window( + stats_history_csv, active_users, start_ts, end_ts + ) + log( + f"Detected steady-state window: [{steady_start}, {steady_end}] " + f"({steady_end - steady_start}s)" + ) + + telemetry = harvest_server_telemetry( + prom_url, start_ts, end_ts, steady_start, worker_pod_count + ) + + full_artifact = { + "metadata": { + "test_name": test_name, + "tag": tag, + "data_timestamp": data_ts, + "prom_url": prom_url, + "start_ts": start_ts, + "end_ts": end_ts, + "steady_start_ts": steady_start, + "steady_end_ts": steady_end, + "worker_pod_count": worker_pod_count, + }, + **telemetry, + } + + output_json_path.write_text(json.dumps(full_artifact, indent=2) + "\n") + log(f"Wrote server summary artifact to {output_json_path}") + + # Append normalized single-row summary into stats.jsonl + packing_s = telemetry.get("cluster_packing", {}).get("summary", {}) + psi = telemetry.get("node_psi", {}) + snaps = telemetry.get("snapshots", {}) + + jsonl_row = { + "timestamp": data_ts, + "tag": tag, + "test_name": test_name, + "metric": "server_summary", + "cluster_packing_p50": packing_s.get("p50"), + "cluster_packing_p90": packing_s.get("p90"), + "cluster_packing_p99": packing_s.get("p99"), + "psi_cpu_stall_p90": psi.get("cpu_stall_pct", {}).get("p90"), + "psi_mem_stall_p90": psi.get("mem_stall_pct", {}).get("p90"), + "psi_io_stall_p90": psi.get("io_stall_pct", {}).get("p90"), + "cfs_throttled_rate_avg": psi.get("cfs_throttled_rate", {}).get("avg"), + "snapshot_size_p50_mb": snaps.get("size_p50_mb"), + "checkpoints_in_window": snaps.get("checkpoints_in_window"), + "restore_p50_s": snaps.get("restore_p50_s"), + "checkpoint_p50_s": snaps.get("checkpoint_p50_s"), + "checkpoint_throughput_mb_s": snaps.get("throughput_mb_s"), + } + + with open(jsonl_path, "a", encoding="utf-8") as f: + f.write(json.dumps(jsonl_row) + "\n") + log(f"Appended server_summary row to {jsonl_path}") diff --git a/benchmarking/locust/unit_tests/test_cluster_facts.py b/benchmarking/locust/unit_tests/test_cluster_facts.py new file mode 100644 index 0000000000..6d9928c1f9 --- /dev/null +++ b/benchmarking/locust/unit_tests/test_cluster_facts.py @@ -0,0 +1,253 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for cluster_facts.py: python3 benchmarking/locust/unit_tests/test_cluster_facts.py + +Never contacts a cluster. Nodes and pods are stand-in objects handed to a +mocked CoreV1Api. Needs the kubernetes client: +pip install -r benchmarking/locust/requirements.txt +""" + +import argparse +import contextlib +import io +import json +from pathlib import Path +import sys +import tempfile +from types import SimpleNamespace +import unittest +from unittest import mock + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +# Probe only the third-party dep, so a bad import in our own modules fails the +# suite instead of skipping it. +try: + import kubernetes # noqa: F401 + + HAS_KUBERNETES = True +except ImportError: # pragma: no cover - depends on the local environment + HAS_KUBERNETES = False + +if HAS_KUBERNETES: + import cluster_facts + import runner + from kubernetes.client.rest import ApiException + +needs_kubernetes = unittest.skipUnless( + HAS_KUBERNETES, + "kubernetes client not installed; pip install -r benchmarking/locust/requirements.txt", +) + +# Real apiserver quantity strings: 3920m -> 3.92 cores, 13591700Ki -> 12.96 GiB. +NODE_CPU, NODE_MEMORY = "3920m", "13591700Ki" + +# One successful discovery. 10 users against 5 worker pods puts A/P at 2.0. +FACTS = {"machine_type": "c3-standard-4", "node_count": 1, + "allocatable_cores": 3.92, "allocatable_ram_gb": 12.96, + "worker_pod_count": 5} + +STATS_HEADER = "Type,Name,Request Count,Failure Count\n" +ARGV = ["runner.py", "-f", "tests/glutton.py", "-t", "1m", "-u", "10", + "--tag", "unit", "--name", "unit-run", "--dest", "/tmp/unit"] + + +def node(machine_type="c3-standard-4"): + labels = {} if machine_type is None else { + cluster_facts.MACHINE_TYPE_LABEL: machine_type} + return SimpleNamespace( + metadata=SimpleNamespace(labels=labels), + status=SimpleNamespace(allocatable={"cpu": NODE_CPU, "memory": NODE_MEMORY})) + + +def pod(phase="Running"): + return SimpleNamespace(status=SimpleNamespace(phase=phase)) + + +def fake_api(nodes=None, ns_pods=None, all_pods=None): + """A stand-in CoreV1Api. None means that call raises 403, which is how the + apiserver answers a ServiceAccount that lacks the ClusterRole.""" + def forbidden(*_args, **_kwargs): + raise ApiException(status=403, reason="Forbidden") + + api = mock.Mock() + for attr, items in (("list_node", nodes), ("list_namespaced_pod", ns_pods), + ("list_pod_for_all_namespaces", all_pods)): + if items is None: + getattr(api, attr).side_effect = forbidden + else: + getattr(api, attr).return_value = SimpleNamespace(items=items) + return api + + +def discover(api): + """Notices print unconditionally, mirroring runner.tee, so stdout is + swallowed to keep the suite quiet.""" + with mock.patch.object(cluster_facts, "_load_kube_config", return_value=True), \ + mock.patch.object(cluster_facts.client, "CoreV1Api", return_value=api), \ + contextlib.redirect_stdout(io.StringIO()): + return cluster_facts.get_cluster_hardware_facts() + + +def summarize(facts, directory, stats=STATS_HEADER + ",Aggregated,100,25\n", + users=10, user_counts=None): + """Writes what append_trial_summary reads, returns the row it emitted.""" + d = Path(directory) + if stats is not None: + (d / "stats.csv").write_text(stats) + if user_counts is None: + user_counts = [users] * 61 + (d / "stats_history.csv").write_text( + "Timestamp,User Count,Type,Name,Requests/s,Failures/s\n" + + "".join(f"{1788914584 + i},{u},,Aggregated,1.0,0.0\n" + for i, u in enumerate(user_counts))) + out = d / "out.jsonl" + with contextlib.redirect_stdout(io.StringIO()): + cluster_facts.append_trial_summary( + out, d / "stats.csv", d / "stats_history.csv", + argparse.Namespace(users=users, tag="unit", name="unit-run"), + "2026-01-01", facts) + return json.loads(out.read_text().splitlines()[0]) + + +def parse(*extra): + with mock.patch.object(sys, "argv", ARGV + list(extra)): + return runner.parse_args() + + +@needs_kubernetes +class ClusterFactsTest(unittest.TestCase): + def test_node_capacity(self): + # An unlabeled node, as a pool can hit mid-upgrade, still counts. + facts = discover(fake_api(nodes=[node(), node(None), node("n2-standard-8")], + ns_pods=[pod()])) + self.assertEqual(facts["node_count"], 3) + self.assertEqual(facts["allocatable_cores"], 11.76) # 3 x 3.92 + # Bytes are summed and rounded once, so this is not 3 x 12.96. + self.assertEqual(facts["allocatable_ram_gb"], 38.89) + # A mixed pool is reported in full rather than attributed to one node. + self.assertEqual(facts["machine_type"], "c3-standard-4,n2-standard-8") + + def test_worker_pod_count(self): + pods = [pod("Running"), pod("Pending"), pod("Succeeded"), pod("Failed")] + self.assertEqual(discover(fake_api(nodes=[node()], ns_pods=pods)) + ["worker_pod_count"], 2) + + # An empty namespace means the pool may live elsewhere, so scan wide. + api = fake_api(nodes=[node()], ns_pods=[], all_pods=[pod(), pod()]) + self.assertEqual(discover(api)["worker_pod_count"], 2) + api.list_pod_for_all_namespaces.assert_called_once() + + api = fake_api(nodes=[node()], ns_pods=[pod()], all_pods=[pod(), pod()]) + self.assertEqual(discover(api)["worker_pod_count"], 1) + api.list_pod_for_all_namespaces.assert_not_called() + + def test_unreadable_facts_are_none(self): + # Nodes denied. Those facts drop out, pods are still counted. + facts = discover(fake_api(nodes=None, ns_pods=[pod(), pod()])) + self.assertIsNone(facts["node_count"]) + self.assertIsNone(facts["allocatable_cores"]) + self.assertEqual(facts["worker_pod_count"], 2) + + # No pod carries the pool label. A guess here would skew every A/P ratio. + facts = discover(fake_api(nodes=[node(), node()], ns_pods=[], all_pods=[])) + self.assertIsNone(facts["worker_pod_count"]) + self.assertEqual(facts["node_count"], 2) + + # Everything denied, and no credentials at all. + self.assertEqual(discover(fake_api()), cluster_facts.EMPTY_FACTS) + with mock.patch.object(cluster_facts, "_load_kube_config", return_value=False): + self.assertEqual(cluster_facts.get_cluster_hardware_facts(), + cluster_facts.EMPTY_FACTS) + + def test_flags(self): + self.assertEqual(parse().prometheus_url, runner.DEFAULT_PROMETHEUS_URL) + self.assertEqual(parse("--prometheus-url", "http://x:9090").prometheus_url, + "http://x:9090") + # Neither flag is ours to hand on to locust. + extra = parse("--no-cluster-facts", "--prometheus-url", "http://x:9090") + self.assertNotIn("--no-cluster-facts", extra.locust_extra) + self.assertNotIn("--prometheus-url", extra.locust_extra) + + def test_no_cluster_facts_skips_the_api(self): + def tripwire(*_args, **_kwargs): + raise AssertionError("Kubernetes was contacted with --no-cluster-facts") + + with mock.patch.object(cluster_facts.client, "CoreV1Api", tripwire), \ + mock.patch.object(cluster_facts.config, "load_incluster_config", tripwire), \ + mock.patch.object(cluster_facts.config, "load_kube_config", tripwire), \ + contextlib.redirect_stdout(io.StringIO()): + facts = runner.collect_cluster_facts(parse("--no-cluster-facts"), + io.StringIO()) + self.assertEqual(facts, cluster_facts.EMPTY_FACTS) + + # Guards the assertion above: broken discovery also makes no API call. + with mock.patch.object(runner, "get_cluster_hardware_facts", + return_value={"node_count": 1}) as discovery, \ + contextlib.redirect_stdout(io.StringIO()): + runner.collect_cluster_facts(parse(), io.StringIO()) + discovery.assert_called_once() + + def test_trial_summary(self): + with tempfile.TemporaryDirectory() as td: + row = summarize(FACTS, td) + self.assertEqual(row["metric"], "trial_summary") + # Raw readings sit beside the derived numbers, so ratios can be re-derived. + self.assertEqual(row["raw_configuration"], FACTS) + f = row["frontiers"] + self.assertEqual(f["actors_per_node"], 10.0) # 10 users / 1 node + self.assertEqual(f["actors_per_vcpu"], 2.55) # 10 / 3.92 + self.assertEqual(f["actors_per_gb_ram"], 0.77) # 10 / 12.96 + + # Skipped or unreadable: same keys, so consumers need no special case. + with tempfile.TemporaryDirectory() as td: + row = summarize(dict(cluster_facts.EMPTY_FACTS), td) + self.assertEqual(set(row["raw_configuration"]), set(cluster_facts.EMPTY_FACTS)) + for key in ("actors_per_node", "actors_per_vcpu", "actors_per_gb_ram", + "ap_ratio_p50", "ap_ratio_p90", "ap_ratio_p99"): + self.assertIsNone(row["frontiers"][key]) + + def test_ap_ratio_percentiles(self): + # 200 steady samples behind three ramp-up ones the filter must drop, + # each a different ratio so an off-by-one shows up. + with tempfile.TemporaryDirectory() as td: + row = summarize(FACTS, td, users=100, + user_counts=[1, 50, 89] + list(range(100, 300))) + f = row["frontiers"] + self.assertEqual([f["ap_ratio_p50"], f["ap_ratio_p90"], f["ap_ratio_p99"]], + [40.0, 56.0, 59.6]) # users 200, 280, 298 over 5 pods + # p99 is not max: the top sample, 299 users, is 59.8. + + # No usable sample: a computed ratio here would be one nobody measured. + with tempfile.TemporaryDirectory() as td: + row = summarize(FACTS, td, users=10, user_counts=[]) + for key in ("ap_ratio_p50", "ap_ratio_p90", "ap_ratio_p99"): + self.assertIsNone(row["frontiers"][key]) + + def test_failure_ratio(self): + # 0.0 is the most optimistic value here, so only a real count may give it. + def ratio(stats): + with tempfile.TemporaryDirectory() as td: + return summarize(FACTS, td, stats)["frontiers"]["aggregate_failure_ratio"] + + self.assertEqual(ratio(STATS_HEADER + ",Aggregated,100,25\n"), 0.25) + self.assertEqual(ratio(STATS_HEADER + ",Aggregated,1708,0\n"), 0.0) + self.assertIsNone(ratio(STATS_HEADER + "not,a,valid\n")) # truncated + self.assertIsNone(ratio(STATS_HEADER + ",Aggregated,0,0\n")) # 0/0 + self.assertIsNone(ratio(None)) # file absent + + +if __name__ == "__main__": + unittest.main() diff --git a/benchmarking/locust/unit_tests/test_server_telemetry.py b/benchmarking/locust/unit_tests/test_server_telemetry.py new file mode 100644 index 0000000000..9bc19d10ef --- /dev/null +++ b/benchmarking/locust/unit_tests/test_server_telemetry.py @@ -0,0 +1,281 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for server_telemetry.py: python3 benchmarking/locust/unit_tests/test_server_telemetry.py""" + +import csv +import json +from pathlib import Path +import sys +import tempfile +import unittest +from unittest import mock + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +import server_telemetry + +NO_PERCENTILES = {"min": None, "p50": None, "p90": None, + "p99": None, "max": None, "avg": None} + +# Every harvest test shares this five-second window; only the pod count varies. +WINDOW = {"prom_url": "http://localhost:9090", "start_ts": 100, + "end_ts": 105, "steady_start_ts": 100} + +# packing, CPU PSI, memory PSI, IO PSI, CFS throttling. +EMPTY_RANGES = [[], [], [], [], []] + + +def harvest(worker_pod_count=5): + return server_telemetry.harvest_server_telemetry( + worker_pod_count=worker_pod_count, **WINDOW + ) + + +def snapshot_instants(size_p50, size_p90, c_start="100", c_end="150", + size_p95="0", size_sum_start=None, size_sum_end=None): + """The eleven instant queries the snapshot block issues, in order. + + Order matters: these are consumed as a mock side_effect. + """ + return [ + [{"value": [105, size_p50]}], # snapshot size p50 + [{"value": [105, size_p90]}], # snapshot size p90 + [{"value": [105, size_p95]}], # snapshot size p95 + [{"value": [100, c_start]}], # checkpoint count at window start + [{"value": [105, c_end]}], # checkpoint count at window end + [{"value": [100, size_sum_start]}] if size_sum_start is not None else [], + [{"value": [105, size_sum_end]}] if size_sum_end is not None else [], + [{"value": [105, "0.08"]}], # restore p50 + [{"value": [105, "0.15"]}], # restore p95 + [{"value": [105, "0.12"]}], # checkpoint p50 + [{"value": [105, "0.22"]}], # checkpoint p95 + ] + + +def history_csv(rows): + """A stats_history.csv built from (timestamp, user count) pairs.""" + f = tempfile.NamedTemporaryFile("w", delete=False, suffix=".csv") + writer = csv.DictWriter(f, fieldnames=["Timestamp", "Name", "User Count"]) + writer.writeheader() + for ts, users in rows: + writer.writerow({"Timestamp": ts, "Name": "Aggregated", "User Count": users}) + f.close() + return Path(f.name) + + +class ServerTelemetryTest(unittest.TestCase): + def test_compute_percentiles(self): + # 200 distinct samples, so an off-by-one in the indexing shows up. + res = server_telemetry.compute_percentiles([float(i) for i in range(1, 201)]) + self.assertEqual( + (res["min"], res["p50"], res["p90"], res["p99"], res["max"], res["avg"]), + (1.0, 101.0, 181.0, 199.0, 200.0, 100.5)) + + # With a single sample, every percentile is that sample. + res = server_telemetry.compute_percentiles([7.5]) + self.assertEqual((res["p50"], res["p90"], res["p99"]), (7.5, 7.5, 7.5)) + + res = server_telemetry.compute_percentiles( + [1.0, float("nan"), 2.0, float("inf"), float("-inf"), 3.0] + ) + self.assertEqual((res["min"], res["p50"], res["max"], res["avg"]), + (1.0, 2.0, 3.0, 2.0)) + + # Nothing left to measure is None, not zero. + self.assertEqual(server_telemetry.compute_percentiles([]), NO_PERCENTILES) + self.assertEqual(server_telemetry.compute_percentiles([float("nan")]), + NO_PERCENTILES) + + def test_steady_state_window(self): + # Steady starts at the first sample at or above 90% of target (6.3). + path = history_csv([("100", "2"), ("110", "4"), ("120", "7"), ("130", "7")]) + try: + self.assertEqual( + server_telemetry.get_steady_state_window( + path, active_users=7, start_ts=100, end_ts=150), + (120, 150), + ) + finally: + path.unlink() + + # Target never reached, so the whole run is used rather than nothing. + path = history_csv([("100", "2")]) + try: + self.assertEqual( + server_telemetry.get_steady_state_window( + path, active_users=10, start_ts=100, end_ts=150), + (100, 150), + ) + finally: + path.unlink() + + @mock.patch("urllib.request.urlopen") + def test_range_query_window_guard(self, mock_urlopen): + # Prometheus answers 400 when end is not after start, so a zero-length + # window is widened by a second. Assert on the URL actually sent. + resp = mock.MagicMock() + resp.read.return_value = json.dumps({ + "status": "success", + "data": {"result": [{"metric": {}, "values": [[100, "1.0"]]}]}, + }).encode("utf-8") + mock_urlopen.return_value.__enter__.return_value = resp + + res = server_telemetry.query_prometheus_range( + "http://localhost:9090", "up", 100, 100) + url = mock_urlopen.call_args[0][0].full_url + self.assertIn("start=100", url) + self.assertIn("end=101", url) + self.assertEqual(len(res), 1) + + server_telemetry.query_prometheus_range( + "http://localhost:9090", "up", 100, 160) + url = mock_urlopen.call_args[0][0].full_url + self.assertIn("start=100", url) + self.assertIn("end=160", url) + + @mock.patch("server_telemetry.query_prometheus_instant") + def test_quantile_fallback(self, mock_instant): + # A rate query over a quiet window yields NaN: fall back to cumulative. + mock_instant.side_effect = [ + [{"value": [100, "NaN"]}], + [{"value": [100, "11.625"]}], + ] + val = server_telemetry._query_quantile_with_fallback( + "http://localhost:9090", 0.50, "rate(bucket[5m])", "bucket", end_ts=100) + self.assertEqual(val, 11.625) + + # A malformed response costs one field, it must not raise. + for bad in ([{"value": None}], [{"value": []}], [{"value": [100, None]}]): + self.assertIsNone(server_telemetry._parse_instant_float(bad)) + + @mock.patch("server_telemetry.query_prometheus_range") + @mock.patch("server_telemetry.query_prometheus_instant") + def test_packing_and_checkpoint_math(self, mock_instant, mock_range): + # The Inf sample must be dropped: json.dumps would write a bare + # Infinity that strict parsers reject. + assigned = {"metric": {"ate_worker_state": "assigned"}, + "values": [[100, "4.0"], [105, "4.0"], [110, "Inf"]]} + quiet = [{"values": [[100, "0.0"], [105, "0.0"]]}] + mock_range.side_effect = [[assigned], quiet, quiet, quiet, quiet] + mock_instant.side_effect = snapshot_instants("11.5", "12.0") + + summary = harvest(worker_pod_count=5) + + packing = summary["cluster_packing"] + self.assertEqual(packing["summary"]["p50"], 0.8) # 4 assigned / 5 pods + self.assertEqual(packing["timeseries"][0]["total_workers"], 5.0) + self.assertEqual(len(packing["timeseries"]), 2) # the Inf point is gone + self.assertNotIn("Infinity", json.dumps(summary)) + + snapshots = summary["snapshots"] + self.assertEqual(snapshots["checkpoints_in_window"], 50) # 150 - 100 + self.assertEqual(snapshots["checkpoints_cumulative"], 150) + + @mock.patch("server_telemetry.query_prometheus_range") + @mock.patch("server_telemetry.query_prometheus_instant") + def test_unknown_pod_count_uses_observed_workers(self, mock_instant, mock_range): + # cluster_facts returns None rather than guessing, so the denominator + # is what Prometheus reports. + mock_range.side_effect = [ + [ + {"metric": {"ate_worker_state": "assigned"}, + "values": [[100, "4.0"]]}, + {"metric": {"ate_worker_state": "idle"}, + "values": [[100, "16.0"]]}, + ], + [{"values": [[100, "0.0"]]}], [{"values": [[100, "0.0"]]}], + [{"values": [[100, "0.0"]]}], [{"values": [[100, "0.0"]]}], + ] + mock_instant.return_value = [] + + point = harvest(worker_pod_count=None)["cluster_packing"]["timeseries"][0] + self.assertEqual(point["total_workers"], 20.0) # 4 + 16 observed + self.assertEqual(point["packing_ratio"], 0.2) + + @mock.patch("server_telemetry.query_prometheus_range") + @mock.patch("server_telemetry.query_prometheus_instant") + def test_snapshot_fields_are_null_not_zero(self, mock_instant, mock_range): + # "no checkpoints happened" and "we could not find out" must differ. + mock_range.side_effect = EMPTY_RANGES + mock_instant.return_value = [] + snaps = harvest()["snapshots"] + self.assertIsNone(snaps["checkpoints_in_window"]) + self.assertIsNone(snaps["checkpoints_cumulative"]) + self.assertIsNone(snaps["throughput_mb_s"]) + self.assertIsNone(snaps["size_p95_mb"]) + self.assertIsNone(snaps["size_avg_mb"]) + + # A 0.0 median is falsy but real, and must still produce throughput. + mock_range.side_effect = EMPTY_RANGES + mock_instant.side_effect = snapshot_instants("0.0", "0.0") + snaps = harvest()["snapshots"] + self.assertEqual(snaps["size_p50_mb"], 0.0) + self.assertEqual(snaps["checkpoints_in_window"], 50) + self.assertEqual(snaps["throughput_mb_s"], 0.0) + + # Counter went backwards, so an atelet restarted and the delta is + # unknowable. 0 would read as "nothing was checkpointed". + mock_range.side_effect = EMPTY_RANGES + mock_instant.side_effect = snapshot_instants("1.0", "1.0", + c_start="900", c_end="150") + snaps = harvest()["snapshots"] + self.assertIsNone(snaps["checkpoints_in_window"]) + self.assertIsNone(snaps["throughput_mb_s"]) + self.assertEqual(snaps["checkpoints_cumulative"], 150) + + # Windowed delta over windowed count: 100 MiB across 50 checkpoints is + # 2.0 MB. A lifetime mean would read 400/150 = 2.667 instead. + mock_range.side_effect = EMPTY_RANGES + mock_instant.side_effect = snapshot_instants( + "1.0", "1.5", size_p95="1.75", + size_sum_start=str(300 * 1024 * 1024), + size_sum_end=str(400 * 1024 * 1024), + ) + snaps = harvest()["snapshots"] + self.assertEqual(snaps["size_p95_mb"], 1.75) + self.assertEqual(snaps["size_avg_mb"], 2.0) + + # The mock replies by position, so only the query text proves p95 was + # asked for. restore_p95 also uses 0.95, hence the metric name too. + queries = [c.args[1] for c in mock_instant.call_args_list] + self.assertTrue(any("histogram_quantile(0.95" in q + and "atelet_snapshot_size_bytes" in q + for q in queries)) + + # Likewise time_ts, or both sum endpoints could read the same instant. + sum_times = {c.kwargs.get("time_ts") for c in mock_instant.call_args_list + if c.args[1] == "sum(atelet_snapshot_size_bytes_sum)"} + self.assertEqual(len(sum_times), 2) + + # Sum went backwards even though the counts rose, so a restart again. + mock_range.side_effect = EMPTY_RANGES + mock_instant.side_effect = snapshot_instants( + "1.0", "1.0", + size_sum_start=str(400 * 1024 * 1024), + size_sum_end=str(300 * 1024 * 1024), + ) + self.assertIsNone(harvest()["snapshots"]["size_avg_mb"]) + + # A zero count means nothing was snapshotted: unknown, not 0 MB. + mock_range.side_effect = EMPTY_RANGES + mock_instant.side_effect = snapshot_instants( + "1.0", "1.0", c_start="0", c_end="0", + size_sum_start="0", size_sum_end="0", + ) + self.assertIsNone(harvest()["snapshots"]["size_avg_mb"]) + + +if __name__ == "__main__": + unittest.main()