diff --git a/benchmarking/README.md b/benchmarking/README.md index 7251e3c383..d8d0beddd1 100644 --- a/benchmarking/README.md +++ b/benchmarking/README.md @@ -79,6 +79,16 @@ 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` to skip Kubernetes API discovery. +* `--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 +124,83 @@ 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 + the same five keys: `timestamp`, `tag`, `test_name`, `metric`, and a flat + `measurements` map holding that metric's numbers. +* `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. Its `measurements` +map holds the raw facts and the derived numbers side by side. + +* `machine_type`, `node_count`, `allocatable_cores`, `allocatable_ram_gb` + (GiB), `worker_pod_count`: the measured facts, before any arithmetic. + Capacity covers the nodes the worker pods are running on rather than the + whole cluster, so a separate infrastructure pool is not counted. They are + recorded so the ratios below can be re-derived later, or recomputed against + a different denominator. +* `actors_per_node`, `actors_per_vcpu`, `actors_per_gb_ram`: the most users + Locust reported running, over the matching capacity. The `-u` flag only + stands in when no sample was read. +* `actors_per_pod_p50`, `actors_per_pod_p90`, `actors_per_pod_p99`: users per + worker pod across the run. Reported as a distribution rather than one + average, and it spans ramp-up too, because a custom load shape has no + single user count to call steady. +* `aggregate_failure_ratio`: failures over requests for the run. +* `_failure_ratio`: the same ratio for every operation Locust + reported, so each test carries its own names through. The operation name is + lowercased with underscores, so `DurDirWrite` becomes + `dur_dir_write_failure_ratio`. A key is absent when the test has no such + row, and null when the row ran no requests. + +### 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. +* `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`: checkpoint + volume over the steady-state window. +* `snapshots.checkpoint_mb_s`: bytes written over the seconds spent writing + them, taken from the histogram sums, so it reads as how fast a checkpoint + writes rather than how many bytes the cluster moved per second of wall + clock. + +A flattened subset of the same numbers goes into the `measurements` map of a +`server_summary` row in `stats.jsonl`, so both metrics can be read from the +one file. + +The `metadata.start_ts` the file records is when the runner started, not when +load did, so the window it covers includes setup. That is deliberate, it gives +the percentiles an idle stretch to sit against, but it does mean the window is +a little longer than the test. The steady-state window is reported separately. + +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 +225,12 @@ 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 +``` diff --git a/benchmarking/automation/manifests/runner-job.yaml.tmpl b/benchmarking/automation/manifests/runner-job.yaml.tmpl index 5855d1d05a..c43d678261 100644 --- a/benchmarking/automation/manifests/runner-job.yaml.tmpl +++ b/benchmarking/automation/manifests/runner-job.yaml.tmpl @@ -41,6 +41,53 @@ roleRef: name: atelet-endpointslices apiGroup: rbac.authorization.k8s.io --- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + name: benchmark-runner-hardware-discovery +rules: +- apiGroups: [""] + resources: ["nodes"] + # list only: the runner reads whole collections, never a single object. + verbs: ["list"] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + name: benchmark-runner-hardware-discovery +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: ClusterRole + name: benchmark-runner-hardware-discovery +subjects: +- kind: ServiceAccount + name: benchmark-runner + namespace: benchmarking +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: Role +metadata: + name: benchmark-runner-worker-pods + namespace: benchmark-workloads +rules: +- apiGroups: [""] + resources: ["pods"] + verbs: ["list"] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: RoleBinding +metadata: + name: benchmark-runner-worker-pods + namespace: benchmark-workloads +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: Role + name: benchmark-runner-worker-pods +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..731232c3e8 --- /dev/null +++ b/benchmarking/locust/cluster_facts.py @@ -0,0 +1,266 @@ +# 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 actors-per-pod percentiles) for a completed trial. +""" + +import argparse +import csv +import json +import re +from pathlib import Path +from typing import Any, TextIO + +from kubernetes import client, config +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 hardware 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 _list_worker_pods( + v1: client.CoreV1Api, logs: TextIO | None = None +) -> list[Any] | None: + """Lists the live pods of the worker pool. + + The pool lives in one namespace by convention and the `WorkerPool` CRD is + namespaced, so this is a single scoped read. The listing is filtered + server-side by label and served from the watch cache. Use + --no-cluster-facts to skip discovery entirely. + + Returns None only when the read failed. An empty list is a reading: the + namespace holds no live worker pods. + """ + try: + pods = v1.list_namespaced_pod( + namespace=WORKER_POOL_NAMESPACE, + label_selector=WORKER_POOL_LABEL, + resource_version="0", + _request_timeout=API_TIMEOUT_SECONDS, + ).items + return [p for p in pods if p.status.phase in LIVE_POD_PHASES] + except Exception as e: + # An ApiException prints its whole HTTP response, so log the reason on + # its own. Anything without one logs itself. + reason = getattr(e, "reason", e) + _log(logs, + f"Notice: could not list pods in {WORKER_POOL_NAMESPACE}: {reason}") + return None + + +def get_cluster_hardware_facts(logs: TextIO | None = None) -> dict[str, Any]: + """Reads the worker pool size and the capacity of the nodes it runs on. + + Capacity is scoped to the nodes carrying worker pods, so a cluster that + keeps its infrastructure on a separate pool does not count that pool's + cores and memory against the density frontiers. + + 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() + + pods = _list_worker_pods(v1, logs) + + if pods is None: + # Without a pod set there is no node set, so capacity stays unmeasured + # rather than falling back to every node in the cluster. + return facts + + facts["worker_pod_count"] = len(pods) + + try: + # A Pending pod may not be scheduled yet, so it counts toward the pool + # size without contributing a node. + worker_nodes = {p.spec.node_name for p in pods if p.spec.node_name} + # resource_version="0" is served from the apiserver's watch cache + # rather than etcd, avoiding a quorum read on large clusters. + nodes = v1.list_node( + resource_version="0", _request_timeout=API_TIMEOUT_SECONDS + ).items + node_count = 0 + total_cores = 0.0 + total_ram_bytes = 0 + machine_types = set() + for node in nodes: + metadata = node.metadata + if metadata is None or metadata.name not in worker_nodes: + continue + node_count += 1 + allocatable = node.status.allocatable or {} + total_cores += float(parse_quantity(allocatable["cpu"])) + total_ram_bytes += int(parse_quantity(allocatable["memory"])) + machine_type = (metadata.labels or {}).get(MACHINE_TYPE_LABEL) + if machine_type: + machine_types.add(machine_type) + facts["node_count"] = node_count + facts["allocatable_cores"] = round(total_cores, 2) + # GiB, as the apiserver and kubectl quote it. + 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: + reason = getattr(e, "reason", e) + _log(logs, f"Notice: could not read node capacity: {reason}") + + 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: + # Locust's own User Count samples. The -u flag is a request; under a custom + # load shape what actually ran is whatever the shape asked for. + observed: list[float] = [] + if stats_history_csv.exists(): + try: + with open(stats_history_csv, encoding="utf-8") 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) + except Exception as e: + _log(logs, f"Notice: could not read user counts: {e}") + # A read that threw partway leaves a truncated sample behind, and + # a truncated sample understates the peak without looking wrong. + observed = [] + + # The flag stands in only when no sample was read at all. + peak_users = max(observed) if observed else args.users + + 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(peak_users / node_count, 2) if node_count else None + actors_per_vcpu = round(peak_users / cores, 2) if cores else None + actors_per_gb_ram = round(peak_users / ram_gb, 2) if ram_gb else None + + # Actors per pod across every sample, ramp-up included. Under a load shape + # there is no one target to measure steadiness against, so the + # distribution covers the whole run. + actors_per_pod_p50, actors_per_pod_p90, actors_per_pod_p99 = None, None, None + if observed and pod_count: + ratios = sorted(round(u / pod_count, 4) for u in observed) + n = len(ratios) + actors_per_pod_p50 = round(ratios[int(n * 0.50)], 2) + actors_per_pod_p90 = round(ratios[min(int(n * 0.90), n - 1)], 2) + actors_per_pod_p99 = round(ratios[min(int(n * 0.99), n - 1)], 2) + + # One ratio per row Locust reported, so a test's own operation names carry + # through. Absent when the test has no such row, null when it ran nothing. + failure_ratios: dict[str, float | None] = {"aggregate_failure_ratio": None} + if stats_csv.exists(): + try: + with open(stats_csv, encoding="utf-8") as f: + for row in csv.DictReader(f): + name = row.get("Name", "") + reqs = row.get("Request Count") + fails = row.get("Failure Count") + # Both columns required so a missing one is not read as zero. + if not name or reqs is None or fails is None: + continue + requests, failures = int(reqs), int(fails) + ratio = round(failures / requests, 4) if requests else None + if name == "Aggregated": + failure_ratios["aggregate_failure_ratio"] = ratio + continue + key = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", name) + key = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1_\2", key) + key = re.sub(r"[^a-z0-9]+", "_", key.lower()).strip("_") + failure_ratios[f"{key}_failure_ratio"] = ratio + except Exception as e: + _log(logs, f"Notice: could not parse {stats_csv}: {e}") + failure_ratios = {"aggregate_failure_ratio": None} + else: + _log(logs, f"Notice: {stats_csv} not found; failure ratios unknown") + + summary_entry = { + "timestamp": data_ts, + "tag": args.tag, + "test_name": args.name, + "metric": "trial_summary", + "measurements": { + **{k: facts.get(k) for k in EMPTY_FACTS}, + "actors_per_node": actors_per_node, + "actors_per_vcpu": actors_per_vcpu, + "actors_per_gb_ram": actors_per_gb_ram, + "actors_per_pod_p50": actors_per_pod_p50, + "actors_per_pod_p90": actors_per_pod_p90, + "actors_per_pod_p99": actors_per_pod_p99, + **failure_ratios, + }, + } + with open(jsonl_path, "a", encoding="utf-8") 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..d2cf6b3c51 100644 --- a/benchmarking/locust/manifests/locust.yaml +++ b/benchmarking/locust/manifests/locust.yaml @@ -202,3 +202,50 @@ roleRef: kind: Role name: atelet-endpointslices apiGroup: rbac.authorization.k8s.io +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + name: locust-hardware-discovery +rules: +- apiGroups: [""] + resources: ["nodes"] + # 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: +- kind: ServiceAccount + name: default + namespace: benchmarking +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: Role +metadata: + name: locust-worker-pods + namespace: benchmark-workloads +rules: +- apiGroups: [""] + resources: ["pods"] + verbs: ["list"] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: RoleBinding +metadata: + name: locust-worker-pods + namespace: benchmark-workloads +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: Role + name: locust-worker-pods +subjects: +- kind: ServiceAccount + name: default + 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 1996e91d0d..32b55a2a56 100644 --- a/benchmarking/locust/runner.py +++ b/benchmarking/locust/runner.py @@ -41,9 +41,15 @@ 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 +from server_telemetry import extract_and_record_server_telemetry # Path inside the locust image to the boomer-worker binary baked in by # benchmarking/locust/Dockerfile. @@ -54,6 +60,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") @@ -108,6 +118,25 @@ def parse_args() -> argparse.Namespace: "default of 1." ), ) + 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 Kubernetes API discovery" + ), + ) + 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 @@ -406,6 +435,16 @@ 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.""" + 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) @@ -425,6 +464,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}" @@ -436,6 +476,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(): @@ -460,6 +501,53 @@ 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. + 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) + except Exception as e: + tee(logs, f"Warning: Failed to read cluster facts: {e}") + + # The frontiers divide by user counts, so they need the CSV. + if stats_generated: + try: + 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}") + + # Server telemetry is a Prometheus time-window query, so it runs + # either way. A run too loaded to write a CSV is the one its + # bin-packing, PSI and snapshot numbers matter most for. + try: + 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, + 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} @@ -475,6 +563,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..9a825be4e3 --- /dev/null +++ b/benchmarking/locust/server_telemetry.py @@ -0,0 +1,547 @@ +# 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. +""" + +import csv +import json +import math +import sys +import urllib.parse +import urllib.request +from pathlib import Path +from typing import Any, TextIO + + +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, + start_ts: int, + end_ts: int, +) -> tuple[int, int]: + """Derives the steady-state window [T_steady, T_end] from Locust's samples. + + T_steady is the first sample reaching 90% of the highest user count the run + reached. The target is read from the samples rather than the `-u` flag, + because a custom load shape ignores the flag. Unlike the frontier + percentiles this has to stay a single threshold: the result is one + contiguous span, which is what a Prometheus range query takes. + """ + if not stats_history_csv.exists(): + return start_ts, end_ts + + samples: list[tuple[int, float]] = [] + try: + with open(stats_history_csv, 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: + samples.append((int(row["Timestamp"]), + float(row["User Count"]))) + except (ValueError, TypeError): + continue + except Exception: + # A read that threw partway leaves only the rows before the fault, + # which drags the peak down and opens the window during ramp-up. + samples = [] + + peak = max((users for _, users in samples), default=0.0) + if peak <= 0: + return start_ts, end_ts + + steady_ts: int | None = None + for ts, users in samples: + if users >= peak * 0.9: + steady_ts = ts + break + + 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 (IndexError, TypeError, ValueError): + # Malformed response yields None for this field without failing harvest. + 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 round(val) if val is not None else None + + +def _query_rate_quantile( + prom_url: str, + quantile: float, + rate_metric_expr: str, + end_ts: int, + unit_scale: float = 1.0, +) -> float | None: + """Reads a histogram quantile over the trailing rate window. + + Scoped to the window the rate covers, so a quiet window reports nothing + rather than the distribution accumulated since the atelet started. + """ + query = ( + f"histogram_quantile({quantile}, sum({rate_metric_expr}) by (le))" + f" / {unit_scale}" + ) + return _parse_instant_float( + query_prometheus_instant(prom_url, query, time_ts=end_ts) + ) + + +def _query_checkpoint_throughput( + prom_url: str, steady_start_ts: int, end_ts: int +) -> float | None: + """Checkpoint write throughput in MiB/s across the steady-state window. + + Bytes written over the seconds spent writing them, both read from the + histogram sums, so this is how fast a checkpoint writes rather than how + many bytes the cluster moved per second of wall clock. The duration is + pinned to the `total` phase because the sub-phases overlap and an + unpinned sum counts the same checkpoint more than once, and failed saves + are excluded because they spend time without writing anything. + """ + window_s = end_ts - steady_start_ts + if window_s <= 0: + return None + written = _parse_instant_float( + query_prometheus_instant( + prom_url, + f"sum(increase(atelet_snapshot_size_bytes_sum[{window_s}s]))", + time_ts=end_ts, + ) + ) + spent = _parse_instant_float( + query_prometheus_instant( + prom_url, + "sum(increase(ate_actor_checkpoint_duration_seconds_sum" + '{ate_snapshot_phase="total", ' + 'ate_failure_reason!="FAILED_SAVE_SNAPSHOT"}' + f"[{window_s}s]))", + time_ts=end_ts, + ) + ) + # No seconds recorded against the checkpoints leaves the rate unknowable, + # where 0.0 would read as a stalled disk. + if written is None or spent is None or spent <= 0: + return None + return round((written / (1024 * 1024)) / spent, 2) + + +def _psi_query(resource: str) -> str: + """The stall percentage query for one PSI resource, summed by instance. + + `_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. + """ + return ( + f'sum by (instance) (rate(container_pressure_{resource}_waiting_seconds_total' + '{container="node"}[1m])) * 100' + ) + + +def _steady_values( + results: list[dict[str, Any]], steady_start_ts: int +) -> list[float]: + """Every range series value at or after `steady_start_ts`.""" + vals = [] + for s in results: + for pt in s.get("values", []): + try: + if int(pt[0]) >= steady_start_ts: + # NaN and Inf are filtered downstream by compute_percentiles. + vals.append(float(pt[1])) + except (ValueError, IndexError): + pass + return vals + + +def _rpc_rate(method: str) -> str: + """The bucket rate expression for one AteomHerder RPC.""" + return ( + 'rate(rpc_server_call_duration_seconds_bucket' + f'{{rpc_method="atelet.AteomHerder/{method}"}}[5m])' + ) + + +def _harvest_cluster_packing( + prom_url: str, + start_ts: int, + end_ts: int, + steady_start_ts: int, + worker_pod_count: int | None, +) -> dict[str, Any]: + """Assigned workers over total workers, per sample and as percentiles.""" + # Dedupes 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 is not None + else sum(states.values()) + ) + # The total is recorded either way, but a ratio over zero workers is + # not a reading, so it stays null rather than taking a stand-in. + ratio = round(assigned / total, 4) if total > 0 else None + packing_points.append({ + "timestamp": t, + "assigned_workers": assigned, + "total_workers": total, + "packing_ratio": ratio, + }) + if t >= steady_start_ts and ratio is not None: + steady_packing_ratios.append(ratio) + + return { + "summary": compute_percentiles(steady_packing_ratios), + "timeseries": packing_points, + } + + +def _harvest_node_psi( + prom_url: str, start_ts: int, end_ts: int, steady_start_ts: int +) -> dict[str, Any]: + """Host kernel pressure stall percentages over the steady-state window.""" + 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" + ) + + return { + "cpu_stall_pct": compute_percentiles( + _steady_values(psi_cpu_res, steady_start_ts) + ), + "mem_stall_pct": compute_percentiles( + _steady_values(psi_mem_res, steady_start_ts) + ), + "io_stall_pct": compute_percentiles( + _steady_values(psi_io_res, steady_start_ts) + ), + } + + +def _harvest_snapshots( + prom_url: str, steady_start_ts: int, end_ts: int +) -> dict[str, Any]: + """Snapshot sizes, checkpoint volume, latencies and write throughput.""" + snap_rate = "rate(atelet_snapshot_size_bytes_bucket[5m])" + snap_p50 = _query_rate_quantile( + prom_url, 0.50, snap_rate, end_ts, unit_scale=1024 * 1024 + ) + snap_p90 = _query_rate_quantile( + prom_url, 0.90, snap_rate, end_ts, unit_scale=1024 * 1024 + ) + snap_p95 = _query_rate_quantile( + prom_url, 0.95, snap_rate, 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) + + restore_rate = _rpc_rate("Restore") + ckpt_rate = _rpc_rate("Checkpoint") + + restore_p50 = _query_rate_quantile(prom_url, 0.50, restore_rate, end_ts) + restore_p95 = _query_rate_quantile(prom_url, 0.95, restore_rate, end_ts) + ckpt_p50 = _query_rate_quantile(prom_url, 0.50, ckpt_rate, end_ts) + ckpt_p95 = _query_rate_quantile(prom_url, 0.95, ckpt_rate, end_ts) + + checkpoint_mb_s = _query_checkpoint_throughput(prom_url, steady_start_ts, end_ts) + + return { + "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, + "checkpoint_mb_s": checkpoint_mb_s, + } + + +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 three ground truth metric streams from Prometheus.""" + return { + "cluster_packing": _harvest_cluster_packing( + prom_url, start_ts, end_ts, steady_start_ts, worker_pod_count + ), + "node_psi": _harvest_node_psi(prom_url, start_ts, end_ts, steady_start_ts), + "snapshots": _harvest_snapshots(prom_url, steady_start_ts, end_ts), + } + + +def extract_and_record_server_telemetry( + prom_url: str, + start_ts: int, + end_ts: int, + stats_history_csv: Path, + 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, 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", encoding="utf-8" + ) + 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", + # Flat here, nested in server_summary.json: the jsonl is the graph + # feed and every row in it carries the same five keys. + "measurements": { + "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"), + "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_mb_s": snaps.get("checkpoint_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..1e40354e33 --- /dev/null +++ b/benchmarking/locust/unit_tests/test_cluster_facts.py @@ -0,0 +1,280 @@ +# 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. + +Run via: 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 +import sys +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest import mock + +from kubernetes.client.rest import ApiException + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +import cluster_facts +import runner + +# 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 actors per pod 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", name="node-a"): + labels = {} if machine_type is None else { + cluster_facts.MACHINE_TYPE_LABEL: machine_type} + return SimpleNamespace( + metadata=SimpleNamespace(labels=labels, name=name), + status=SimpleNamespace(allocatable={"cpu": NODE_CPU, "memory": NODE_MEMORY})) + + +def pod(phase="Running", node_name="node-a"): + return SimpleNamespace(spec=SimpleNamespace(node_name=node_name), + status=SimpleNamespace(phase=phase)) + + +def fake_api(nodes=None, ns_pods=None): + """Stand-in CoreV1Api; None raises 403 Forbidden.""" + 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)): + if items is None: + getattr(api, attr).side_effect = forbidden + else: + getattr(api, attr).return_value = SimpleNamespace(items=items) + return api + + +def discover(api): + 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 CSV inputs and returns the emitted trial_summary row.""" + 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() + + +class ClusterFactsTest(unittest.TestCase): + def test_node_capacity(self): + facts = discover(fake_api( + nodes=[node(name="node-a"), node(None, name="node-b"), + node("n2-standard-8", name="node-c")], + ns_pods=[pod(node_name="node-a"), pod(node_name="node-b"), + pod(node_name="node-c")])) + self.assertEqual(facts["node_count"], 3) + self.assertEqual(facts["allocatable_cores"], 11.76) # 3 x 3.92 + self.assertEqual(facts["allocatable_ram_gb"], 38.89) + 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")] + api = fake_api(nodes=[node()], ns_pods=pods) + self.assertEqual(discover(api)["worker_pod_count"], 2) + api.list_namespaced_pod.assert_called_once_with( + namespace="benchmark-workloads", + label_selector="ate.dev/worker-pool", + resource_version="0", + _request_timeout=5, + ) + api.list_node.assert_called_once_with( + resource_version="0", + _request_timeout=5, + ) + + def test_capacity_is_scoped_to_nodes_running_worker_pods(self): + # Excludes nodes without worker pods (e.g. infra-a). + facts = discover(fake_api( + nodes=[node(name="node-a"), node(name="node-b"), + node(name="infra-a")], + ns_pods=[pod(node_name="node-a"), pod(node_name="node-b"), + pod(node_name="node-b")])) + self.assertEqual(facts["worker_pod_count"], 3) + self.assertEqual(facts["node_count"], 2) + self.assertEqual(facts["allocatable_cores"], 7.84) # 2 x 3.92 + + def test_zero_worker_pods_is_a_reading(self): + # Empty pool -> 0; denied read -> None. + facts = discover(fake_api(nodes=[node()], ns_pods=[])) + self.assertEqual(facts["worker_pod_count"], 0) + self.assertEqual(facts["node_count"], 0) + + denied = discover(fake_api(nodes=[node()], ns_pods=None)) + self.assertIsNone(denied["worker_pod_count"]) + self.assertIsNone(denied["node_count"]) + + def test_unreadable_facts_are_none(self): + # Nodes denied -> capacity None, pod count still recorded. + 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) + + # Everything denied or no credentials -> EMPTY_FACTS. + 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): + extra = parse("--no-cluster-facts", "--max-wait-time", "1.0") + self.assertNotIn("--no-cluster-facts", extra.locust_extra) + self.assertEqual(extra.locust_extra, ["--max-wait-time", "1.0"]) + + 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) + + 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") + self.assertEqual(set(row), {"timestamp", "tag", "test_name", "metric", + "measurements"}) + m = row["measurements"] + self.assertEqual({k: m[k] for k in FACTS}, FACTS) + self.assertEqual(m["actors_per_node"], 10.0) # 10 users / 1 node + self.assertEqual(m["actors_per_vcpu"], 2.55) # 10 / 3.92 + self.assertEqual(m["actors_per_gb_ram"], 0.77) # 10 / 12.96 + + # Unmeasured facts keep the same keys with None values. + with tempfile.TemporaryDirectory() as td: + row = summarize(dict(cluster_facts.EMPTY_FACTS), td) + self.assertLessEqual(set(cluster_facts.EMPTY_FACTS), + set(row["measurements"])) + for key in ("actors_per_node", "actors_per_vcpu", "actors_per_gb_ram", + "actors_per_pod_p50", "actors_per_pod_p90", + "actors_per_pod_p99"): + self.assertIsNone(row["measurements"][key]) + + def test_actors_per_pod_percentiles(self): + with tempfile.TemporaryDirectory() as td: + row = summarize(FACTS, td, users=100, + user_counts=[1, 50, 89] + list(range(100, 300))) + m = row["measurements"] + self.assertEqual([m["actors_per_pod_p50"], m["actors_per_pod_p90"], + m["actors_per_pod_p99"]], + [39.6, 55.8, 59.4]) # users 198, 279, 297 over 5 pods + + # Uses observed peak (60) rather than requested -u (10). + with tempfile.TemporaryDirectory() as td: + row = summarize(FACTS, td, users=10, user_counts=[20, 40, 60]) + self.assertEqual(row["measurements"]["actors_per_node"], 60.0) + self.assertEqual(row["measurements"]["actors_per_pod_p50"], 8.0) + + # Empty history falls back to -u for frontiers and None for percentiles. + with tempfile.TemporaryDirectory() as td: + row = summarize(FACTS, td, users=10, user_counts=[]) + for key in ("actors_per_pod_p50", "actors_per_pod_p90", + "actors_per_pod_p99"): + self.assertIsNone(row["measurements"][key]) + self.assertEqual(row["measurements"]["actors_per_node"], 10.0) + + def test_failure_ratio(self): + def ratio(stats): + with tempfile.TemporaryDirectory() as td: + return summarize(FACTS, td, stats)["measurements"][ + "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 + ",Aggregated,100\n")) # truncated + self.assertIsNone(ratio(STATS_HEADER + ",Aggregated,bad,5\n")) # corrupt int + self.assertIsNone(ratio(STATS_HEADER + ",Aggregated,0,0\n")) # 0/0 + self.assertIsNone(ratio(None)) # file absent + + def test_per_rpc_failure_ratios(self): + # First resume is its own Locust row, so it gets its own key. + stats = (STATS_HEADER + + "grpc,ResumeActor,100,2\n" + + "grpc,ResumeActorFirstResume,100,99\n" + + "grpc,SuspendActor,200,0\n" + + ",Aggregated,400,101\n") + with tempfile.TemporaryDirectory() as td: + f = summarize(FACTS, td, stats)["measurements"] + self.assertEqual(f["resume_actor_failure_ratio"], 0.02) # 2 / 100 + self.assertEqual(f["resume_actor_first_resume_failure_ratio"], 0.99) # 99 / 100 + self.assertEqual(f["suspend_actor_failure_ratio"], 0.0) + self.assertEqual(f["aggregate_failure_ratio"], 0.2525) # 101 / 400 + + # An RPC the test never ran has no key at all. + with tempfile.TemporaryDirectory() as td: + row = summarize(FACTS, td, STATS_HEADER + ",Aggregated,10,0\n") + self.assertNotIn("resume_actor_failure_ratio", row["measurements"]) + self.assertNotIn("suspend_actor_failure_ratio", row["measurements"]) + + +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..9ca1489214 --- /dev/null +++ b/benchmarking/locust/unit_tests/test_server_telemetry.py @@ -0,0 +1,369 @@ +# 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. + +Run via: python3 benchmarking/locust/unit_tests/test_server_telemetry.py +""" + +import contextlib +import csv +import io +import json +import sys +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +import runner +import server_telemetry + +NO_PERCENTILES = {"min": None, "p50": None, "p90": None, + "p99": None, "max": None, "avg": None} + +WINDOW = {"prom_url": "http://localhost:9090", "start_ts": 100, + "end_ts": 105, "steady_start_ts": 100} + +# packing, CPU PSI, memory PSI, IO PSI. +EMPTY_RANGES = [[], [], [], []] + +ARGV = ["runner.py", "-f", "tests/glutton.py", "-t", "1m", "-u", "10", + "--tag", "unit", "--name", "unit-run", "--dest", "/tmp"] + + +def harvest(worker_pod_count=5): + return server_telemetry.harvest_server_telemetry( + worker_pod_count=worker_pod_count, **WINDOW + ) + + +def snapshots(): + """The snapshot section alone, which issues no range queries.""" + return server_telemetry._harvest_snapshots( + WINDOW["prom_url"], WINDOW["steady_start_ts"], WINDOW["end_ts"] + ) + + +def parse(*extra): + with mock.patch.object(sys, "argv", ARGV + list(extra)): + return runner.parse_args() + + +def snapshot_instants(size_p50, size_p90, c_start="100", c_end="150", + size_p95="0", size_sum_start=None, size_sum_end=None, + ckpt_bytes=None, ckpt_seconds=None): + """Mock responses for the 13 instant queries in _harvest_snapshots.""" + 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 + [{"value": [105, ckpt_bytes]}] if ckpt_bytes is not None else [], + [{"value": [105, ckpt_seconds]}] if ckpt_seconds is not None else [], + ] + + +def history_csv(rows): + """A stats_history.csv built from (timestamp, user count) pairs.""" + f = tempfile.NamedTemporaryFile( + "w", delete=False, suffix=".csv", encoding="utf-8" + ) + 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): + 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)) + + 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)) + + self.assertEqual(server_telemetry.compute_percentiles([]), NO_PERCENTILES) + self.assertEqual(server_telemetry.compute_percentiles([float("nan")]), + NO_PERCENTILES) + + def test_steady_state_window(self): + # 90% of peak (7) is 6.3. + path = history_csv([ + ("100", "2"), ("110", "4"), ("120", "6.3"), ("130", "7") + ]) + try: + self.assertEqual( + server_telemetry.get_steady_state_window( + path, start_ts=100, end_ts=150), + (120, 150), + ) + finally: + path.unlink() + + # Uses observed peak (15) rather than requested -u (5). + path = history_csv([("100", "5"), ("110", "10"), ("120", "15"), + ("130", "15")]) + try: + self.assertEqual( + server_telemetry.get_steady_state_window( + path, start_ts=100, end_ts=150), + (120, 150), + ) + finally: + path.unlink() + + # Zero users falls back to [start_ts, end_ts]. + path = history_csv([("120", "0")]) + try: + self.assertEqual( + server_telemetry.get_steady_state_window( + path, 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): + # Zero-length window [100, 100] is widened to [100, 101]. + 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) + + def test_malformed_instant_response(self): + 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): + # Inf samples are filtered out before JSON serialization. + 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] + mock_instant.side_effect = snapshot_instants( + "11.5", "12.0", ckpt_bytes=str(100 * 1024 * 1024), ckpt_seconds="50" + ) + + with tempfile.TemporaryDirectory() as td: + out_json = Path(td) / "server_summary.json" + out_jsonl = Path(td) / "stats.jsonl" + with contextlib.redirect_stdout(io.StringIO()): + server_telemetry.extract_and_record_server_telemetry( + prom_url="http://localhost:9090", + start_ts=100, + end_ts=105, + stats_history_csv=Path(td) / "missing.csv", + worker_pod_count=5, + output_json_path=out_json, + jsonl_path=out_jsonl, + data_ts="2026-01-01", + tag="unit", + test_name="unit-run", + ) + summary = json.loads(out_json.read_text()) + row = json.loads(out_jsonl.read_text().splitlines()[0]) + + self.assertEqual( + set(summary), {"metadata", "cluster_packing", "node_psi", "snapshots"} + ) + 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) + self.assertNotIn("Infinity", json.dumps(summary)) + + snapshots = summary["snapshots"] + self.assertEqual(snapshots["checkpoints_in_window"], 50) # 150 - 100 + self.assertEqual(snapshots["checkpoints_cumulative"], 150) + + self.assertEqual(row["metric"], "server_summary") + self.assertEqual(row["measurements"]["checkpoint_mb_s"], 2.0) + + @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): + # Falls back to observed worker sum when worker_pod_count is None. + 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"]]}], + ] + 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_missing_denominator_is_null(self, mock_instant, mock_range): + # Zero total workers yields packing_ratio=None. + mock_range.side_effect = [ + [{"metric": {"ate_worker_state": "idle"}, "values": [[100, "0.0"]]}], + *EMPTY_RANGES[1:], + ] + mock_instant.return_value = [] + point = harvest(worker_pod_count=None)["cluster_packing"]["timeseries"][0] + self.assertEqual(point["total_workers"], 0.0) + self.assertIsNone(point["packing_ratio"]) + + # Zero-duration window yields checkpoint_mb_s=None. + mock_range.side_effect = EMPTY_RANGES + mock_instant.side_effect = snapshot_instants("1.0", "1.0") + snaps = server_telemetry.harvest_server_telemetry( + prom_url="http://localhost:9090", start_ts=100, end_ts=100, + steady_start_ts=100, worker_pod_count=5)["snapshots"] + self.assertEqual(snaps["checkpoints_in_window"], 50) + self.assertIsNone(snaps["checkpoint_mb_s"]) + + @mock.patch("server_telemetry.query_prometheus_instant") + def test_snapshot_fields_are_null_not_zero(self, mock_instant): + # Empty response -> all fields None. + mock_instant.return_value = [] + snaps = snapshots() + self.assertIsNone(snaps["checkpoints_in_window"]) + self.assertIsNone(snaps["checkpoints_cumulative"]) + self.assertIsNone(snaps["checkpoint_mb_s"]) + self.assertIsNone(snaps["size_p95_mb"]) + self.assertIsNone(snaps["size_avg_mb"]) + + # 0 bytes written over 4s -> measured 0.0. + mock_instant.side_effect = snapshot_instants( + "0.0", "0.0", ckpt_bytes="0", ckpt_seconds="4") + snaps = snapshots() + self.assertEqual(snaps["size_p50_mb"], 0.0) + self.assertEqual(snaps["checkpoints_in_window"], 50) + self.assertEqual(snaps["checkpoint_mb_s"], 0.0) + + # Counter reset (c_end < c_start) -> window delta None. + mock_instant.side_effect = snapshot_instants("1.0", "1.0", + c_start="900", c_end="150") + snaps = snapshots() + self.assertIsNone(snaps["checkpoints_in_window"]) + self.assertEqual(snaps["checkpoints_cumulative"], 150) + + # Windowed average size: 100 MiB across 50 checkpoints = 2.0 MB. + 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 = snapshots() + self.assertEqual(snaps["size_p95_mb"], 1.75) + self.assertEqual(snaps["size_avg_mb"], 2.0) + + # Assert outbound query strings and timestamps. + 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)) + 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 reset (s_end < s_start) or zero count -> size_avg_mb None. + 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(snapshots()["size_avg_mb"]) + + 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(snapshots()["size_avg_mb"]) + + @mock.patch("server_telemetry.query_prometheus_instant") + def test_checkpoint_throughput_from_counter_sums(self, mock_instant): + # 100 MiB / 50s = 2.0 MB/s. + mock_instant.side_effect = snapshot_instants( + "9.0", "9.0", ckpt_bytes=str(100 * 1024 * 1024), ckpt_seconds="50") + self.assertEqual(snapshots()["checkpoint_mb_s"], 2.0) + + # Zero duration -> None. + mock_instant.side_effect = snapshot_instants( + "9.0", "9.0", ckpt_bytes=str(8 * 1024 * 1024), ckpt_seconds="0") + self.assertIsNone(snapshots()["checkpoint_mb_s"]) + + # Duration query must filter by ate_snapshot_phase="total". + spent = [c.args[1] for c in mock_instant.call_args_list + if "ate_actor_checkpoint_duration_seconds_sum" in c.args[1]] + self.assertIn('ate_snapshot_phase="total"', spent[0]) + + def test_prometheus_url_flag(self): + self.assertEqual(parse().prometheus_url, runner.DEFAULT_PROMETHEUS_URL) + self.assertEqual(parse("--prometheus-url", "http://x:9090").prometheus_url, + "http://x:9090") + extra = parse("--prometheus-url", "http://x:9090", "--max-wait-time", "1.0") + self.assertNotIn("--prometheus-url", extra.locust_extra) + self.assertEqual(extra.locust_extra, ["--max-wait-time", "1.0"]) + + @mock.patch.object(runner, "extract_and_record_server_telemetry") + @mock.patch.object(runner, "upload") + @mock.patch.object(runner, "run_test", return_value=1) + def test_telemetry_survives_a_missing_stats_csv(self, _run, _up, telemetry): + with mock.patch.object(sys, "argv", ARGV + ["--no-cluster-facts", + "--allow-empty-stats"]), \ + contextlib.redirect_stdout(io.StringIO()): + runner.main() + self.assertTrue(telemetry.called) + + +if __name__ == "__main__": + unittest.main()