From b7d3eef87e050c2a7ef76fa6c78114de733a61d7 Mon Sep 17 00:00:00 2001 From: Arunkumar Eli Date: Thu, 12 Nov 2026 00:11:01 +0000 Subject: [PATCH] Support benchmark results on S3 and Azure with shared storage Use obstore for runner uploads and result retrieval while preserving the existing run layout. Keep provider credentials in the environment and propagate upload failures. Signed-off-by: Arunkumar Eli --- benchmarking/automation/README.md | 11 ++- benchmarking/automation/ateapi-cold-resume.md | 13 ++- .../automation/manifests/cronjob.yaml.tmpl | 2 +- benchmarking/locust/Dockerfile | 7 +- benchmarking/locust/codegen/requirements.txt | 2 +- benchmarking/locust/requirements.txt | 2 +- benchmarking/locust/runner.py | 28 ++---- .../unit_tests/test_server_telemetry.py | 10 +++ benchmarking/nighthawk-ingress/Dockerfile | 9 +- .../nighthawk-ingress/requirements.txt | 2 +- benchmarking/nighthawk-ingress/run-dev.sh | 15 ++-- benchmarking/nighthawk-ingress/runner.py | 28 ++---- .../nighthawk-ingress/tests/test_runner.py | 55 ++++++++++++ benchmarking/requirements.txt | 15 ++++ benchmarking/result_store.py | 90 +++++++++++++++++++ benchmarking/test_result_store.py | 76 ++++++++++++++++ docs/guides/benchmarking.md | 27 +++++- hack/verify/benchmarking-python.sh | 1 + 18 files changed, 328 insertions(+), 65 deletions(-) create mode 100644 benchmarking/nighthawk-ingress/tests/test_runner.py create mode 100644 benchmarking/requirements.txt create mode 100644 benchmarking/result_store.py create mode 100644 benchmarking/test_result_store.py diff --git a/benchmarking/automation/README.md b/benchmarking/automation/README.md index 0a5952d1be..181a4bddaa 100644 --- a/benchmarking/automation/README.md +++ b/benchmarking/automation/README.md @@ -3,13 +3,19 @@ Scheduled, repeatable benchmark runs of a substrate branch. A CronJob on an **orchestration cluster** drives the full build/deploy/run/teardown cycle against a separate **test cluster**, once per entry in [tests.yaml](tests.yaml). -Each run uploads results to GCS via `benchmarking/locust/runner.py` +Each run uploads results to object storage via `benchmarking/locust/runner.py` (`type: locust`) or `benchmarking/nighthawk-ingress/runner.py` (`type: nighthawk-ingress`, the router capacity benchmark — see [benchmarking/nighthawk-ingress/README.md](../nighthawk-ingress/README.md)). To run the cold resume tests from your workstation instead, see [ateapi-cold-resume.md](ateapi-cold-resume.md). +Pass `--dest s3://benchmark-results/substrate` (or a `gs://` or +`az://container/prefix` URL). See the +[storage configuration](../../docs/guides/benchmarking.md#headless-automation-only) +for credentials and local-path limitations. The cluster provisioning steps +below still describe GKE; result storage does not need to be GCS. + ## How it works 1. CronJob fires on the orchestration cluster; pod starts (orchestrator + @@ -131,7 +137,8 @@ gcloud container clusters create \ ## IAM prerequisites -This setup assumes both clusters and the destination GCS bucket already exist. +This GKE setup assumes both clusters already exist. Create a result bucket or +container in your chosen provider; the GCS grants below apply only to GCS destinations. Two Workload Identity bindings are needed. Both ServiceAccounts are created by the manifests (`cronjob.yaml` for the diff --git a/benchmarking/automation/ateapi-cold-resume.md b/benchmarking/automation/ateapi-cold-resume.md index 6264d520e7..25cd521db5 100644 --- a/benchmarking/automation/ateapi-cold-resume.md +++ b/benchmarking/automation/ateapi-cold-resume.md @@ -117,7 +117,7 @@ PY $S/benchmarking/automation/venv/bin/python -u $D/run.py \ --repo file://$S --branch $(git branch --show-current) \ - --dest gs:/// \ + --dest s3:/// \ --tests $S/benchmarking/automation/tests.yaml \ --test-name ateapi_cold_resume_pause_fake_smoke \ --target-cluster-dir $D/target-clusters \ @@ -125,6 +125,10 @@ $S/benchmarking/automation/venv/bin/python -u $D/run.py \ --junit-output $D/junit.xml 2>&1 | tee $D/orchestrator.log ``` +Configure storage credentials on the runner Job and on your workstation for +retrieval; see [result storage](../../docs/guides/benchmarking.md#headless-automation-only). +The example uses S3; GCS and Azure destinations work with the same commands. + Keep your workstation awake for the whole run. If it suspends, the orchestrator pauses too, and the teardown and the next test wait for it. For long ladders, run the command under `setsid nohup` or in a VM. @@ -137,8 +141,11 @@ explains the per-step rows and where the run path is in the log. lists the files. To copy them all down: ```bash -st=$(grep -o 'gs://[^ ]*/status.json' $D/orchestrator.log | head -1) -mkdir -p $D/results && gcloud storage cp "${st%/status.json}/*" $D/results/ +python3 -m venv "$D/results-venv" +"$D/results-venv/bin/pip" install -r "$S/benchmarking/requirements.txt" +st=$(grep -Eo '(gs|s3|az)://[^ ]*/status.json' "$D/orchestrator.log" | head -1) +"$D/results-venv/bin/python" "$S/benchmarking/result_store.py" download \ + "${st%/status.json}" "$D/results" ``` The capacity is the highest step whose `resume_rps` keeps up with its diff --git a/benchmarking/automation/manifests/cronjob.yaml.tmpl b/benchmarking/automation/manifests/cronjob.yaml.tmpl index 0b1b3f9b68..3a13dcb4b7 100644 --- a/benchmarking/automation/manifests/cronjob.yaml.tmpl +++ b/benchmarking/automation/manifests/cronjob.yaml.tmpl @@ -68,7 +68,7 @@ spec: - "--branch" - "" - "--dest" - - "" + - "" env: - name: DOCKER_HOST value: tcp://localhost:2375 diff --git a/benchmarking/locust/Dockerfile b/benchmarking/locust/Dockerfile index 8fb8f61a6a..fb13f51fce 100644 --- a/benchmarking/locust/Dockerfile +++ b/benchmarking/locust/Dockerfile @@ -18,8 +18,9 @@ # Stage 1: Python dependency builder FROM python:3.11-slim AS pydeps WORKDIR /app -COPY benchmarking/locust/requirements.txt . -RUN pip install --no-cache-dir --target=/app/deps -r requirements.txt +COPY benchmarking/requirements.txt /requirements/requirements.txt +COPY benchmarking/locust/requirements.txt /requirements/locust/requirements.txt +RUN pip install --no-cache-dir --target=/app/deps -r /requirements/locust/requirements.txt # Stage 2: Python gRPC clients, generated from the protos rather than checked # in. See codegen/generate.sh. @@ -64,5 +65,7 @@ COPY benchmarking/locust/ladder_steps.py /app/ladder_steps.py COPY benchmarking/locust/turnover_window.py /app/turnover_window.py COPY benchmarking/analysis/phase_report.py /app/phase_report.py +COPY benchmarking/result_store.py /app/result_store.py + ENV PYTHONPATH=/app:/app/deps ENV PYTHONUNBUFFERED=1 diff --git a/benchmarking/locust/codegen/requirements.txt b/benchmarking/locust/codegen/requirements.txt index e5617d172d..5ba7fb579e 100644 --- a/benchmarking/locust/codegen/requirements.txt +++ b/benchmarking/locust/codegen/requirements.txt @@ -15,7 +15,7 @@ # Build-time dependencies for generate.sh and for the nighthawk-ingress # image's descriptor build. Kept apart from the load test's runtime # requirements.txt so that generating the protos does not pull in locust, the -# opentelemetry exporters and google-cloud-storage. +# opentelemetry exporters. # # grpcio-tools bundles protoc along with the Python and gRPC codegen plugins. # It is pinned because the generated files stamp the version that produced diff --git a/benchmarking/locust/requirements.txt b/benchmarking/locust/requirements.txt index a76e292a41..ff138da96b 100644 --- a/benchmarking/locust/requirements.txt +++ b/benchmarking/locust/requirements.txt @@ -26,5 +26,5 @@ grpcio>=1.84.1 opentelemetry-exporter-otlp opentelemetry-instrumentation-grpc opentelemetry-instrumentation-requests -google-cloud-storage +-r ../requirements.txt kubernetes diff --git a/benchmarking/locust/runner.py b/benchmarking/locust/runner.py index 58b215b56a..f5f44b6a52 100644 --- a/benchmarking/locust/runner.py +++ b/benchmarking/locust/runner.py @@ -15,7 +15,7 @@ """Headless locust runner that publishes results as JSONL + CSVs. Runs locust with the given flags, converts the resulting stats CSV -to JSONL, and uploads everything to either GCS or local disk under +to JSONL, and uploads everything to object storage or local disk under /runs///. When the test target is glutton.py, also spawns the boomer-worker Go @@ -33,7 +33,6 @@ import json import os import re -import shutil import signal import subprocess import sys @@ -43,6 +42,10 @@ from pathlib import Path from typing import IO, Any, TextIO +# Also find the shared module when running directly from the source tree. +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from result_store import upload + from cluster_facts import ( EMPTY_FACTS, append_trial_summary, @@ -110,7 +113,7 @@ def parse_args() -> argparse.Namespace: p.add_argument( "--dest", required=True, - help="Root destination (gs://bucket/path or local path)", + help="Root destination (gs://bucket/path, s3://bucket/path, az://container/path, or local path)", ) p.add_argument( "--allow-empty-stats", @@ -494,25 +497,6 @@ def stats_to_jsonl(stats_csv: Path, jsonl_path: Path, timestamp: str, tag: str, return rows_written -def upload_to_gcs(local_path: Path, gcs_uri: str) -> None: - # Imported here so non-GCS use doesn't require google-cloud-storage. - from google.cloud import storage - - bucket_name, _, blob_path = gcs_uri[len("gs://"):].partition("/") - storage.Client().bucket(bucket_name).blob(blob_path).upload_from_filename( - str(local_path) - ) - - -def upload(src: Path, dest: str) -> None: - if dest.startswith("gs://"): - upload_to_gcs(src, dest) - else: - dest_path = Path(dest) - dest_path.parent.mkdir(parents=True, exist_ok=True) - shutil.copy(src, dest_path) - - def collect_cluster_facts( args: argparse.Namespace, logs: TextIO ) -> dict[str, Any]: diff --git a/benchmarking/locust/unit_tests/test_server_telemetry.py b/benchmarking/locust/unit_tests/test_server_telemetry.py index 044fad79c9..10dfb10d82 100644 --- a/benchmarking/locust/unit_tests/test_server_telemetry.py +++ b/benchmarking/locust/unit_tests/test_server_telemetry.py @@ -785,6 +785,16 @@ def test_telemetry_survives_a_missing_stats_csv(self, _run, _up, telemetry): self.assertEqual(telemetry.call_args.kwargs["lag_s"], 70) self.assertNotIn("worker_pod_count", telemetry.call_args.kwargs) + @mock.patch.object(runner, "extract_and_record_server_telemetry") + @mock.patch.object(runner, "upload", side_effect=OSError("storage unavailable")) + @mock.patch.object(runner, "run_test", return_value=0) + def test_upload_failure_is_not_success(self, _run, _upload, _telemetry): + with mock.patch.object(sys, "argv", ARGV + ["--no-cluster-facts", "--allow-empty-stats"]), \ + contextlib.redirect_stdout(io.StringIO()): + # An uncaught upload error makes the runner process exit non-zero. + with self.assertRaisesRegex(OSError, "storage unavailable"): + runner.main() + # Ladder rows follow the telemetry harvest, whose wait covers their # last samples, and a failure there still uploads the run. @mock.patch.object(runner, "upload") diff --git a/benchmarking/nighthawk-ingress/Dockerfile b/benchmarking/nighthawk-ingress/Dockerfile index 9cfbb836a8..bdbabc39a1 100644 --- a/benchmarking/nighthawk-ingress/Dockerfile +++ b/benchmarking/nighthawk-ingress/Dockerfile @@ -74,14 +74,15 @@ FROM envoyproxy/nighthawk-dev@${NIGHTHAWK_IMAGE_DIGEST} # Deps must be installed with this image's own interpreter: grpcio/protobuf # native extensions are python-version-specific. -COPY benchmarking/nighthawk-ingress/requirements.txt /tmp/requirements.txt +COPY benchmarking/requirements.txt /requirements/requirements.txt +COPY benchmarking/nighthawk-ingress/requirements.txt /requirements/nighthawk-ingress/requirements.txt RUN apt-get update \ && apt-get install -y --no-install-recommends python3 python3-pip ca-certificates \ && python3 -m pip install --no-cache-dir --break-system-packages \ - --target=/app/deps -r /tmp/requirements.txt \ + --target=/app/deps -r /requirements/nighthawk-ingress/requirements.txt \ && apt-get purge -y python3-pip \ && apt-get autoremove -y \ - && rm -rf /var/lib/apt/lists/* /tmp/requirements.txt + && rm -rf /var/lib/apt/lists/* /requirements WORKDIR /app COPY --from=protogen /out/nighthawk.desc /app/nighthawk.desc @@ -99,6 +100,8 @@ COPY benchmarking/nighthawk-ingress/runner.py \ /app/ COPY benchmarking/nighthawk-ingress/tests /app/tests +COPY benchmarking/result_store.py /app/result_store.py + ENV PYTHONPATH=/app:/app/deps ENV PYTHONUNBUFFERED=1 ENV NIGHTHAWK_DESC=/app/nighthawk.desc diff --git a/benchmarking/nighthawk-ingress/requirements.txt b/benchmarking/nighthawk-ingress/requirements.txt index 244e0c78d0..29dd679e5f 100644 --- a/benchmarking/nighthawk-ingress/requirements.txt +++ b/benchmarking/nighthawk-ingress/requirements.txt @@ -18,4 +18,4 @@ protobuf>=7.35.1 grpcio>=1.84.1 requests -google-cloud-storage +-r ../requirements.txt diff --git a/benchmarking/nighthawk-ingress/run-dev.sh b/benchmarking/nighthawk-ingress/run-dev.sh index ade748fd2d..0450e3f8ed 100755 --- a/benchmarking/nighthawk-ingress/run-dev.sh +++ b/benchmarking/nighthawk-ingress/run-dev.sh @@ -39,7 +39,7 @@ Usage: $0 [options] --actors N actor fleet size; needs at least one worker Running (default: ${ACTORS}) --tail-latency-slo-ms N SLO bound; 0 disables (default: ${TAIL_LATENCY_SLO_MS}) --atespace NAME actor namespace (default: ${ATESPACE}) - --dest gs://... results root (default: gs://\$BUCKET_NAME/nighthawk-ingress-results) + --dest URL results root (default: gs://\$BUCKET_NAME/nighthawk-ingress-results) EOF exit "${1:-1}" } @@ -68,10 +68,13 @@ source hack/util/venv.sh } # shellcheck disable=SC1091 source .ate-dev-env.sh -for var in PROJECT_ID CLUSTER_NAME CLUSTER_LOCATION KO_DOCKER_REPO BUCKET_NAME; do +for var in PROJECT_ID CLUSTER_NAME CLUSTER_LOCATION KO_DOCKER_REPO; do [[ -n "${!var:-}" ]] || { echo "ERROR: ${var} not set by .ate-dev-env.sh" >&2; exit 1; } done -DEST="${DEST:-gs://${BUCKET_NAME}/nighthawk-ingress-results}" +if [[ -z "${DEST}" ]]; then + : "${BUCKET_NAME:?Set BUCKET_NAME or pass --dest}" + DEST="gs://${BUCKET_NAME}/nighthawk-ingress-results" +fi docker info >/dev/null 2>&1 || { echo "ERROR: docker daemon not running (the runner image is built locally)" >&2 @@ -110,6 +113,7 @@ if ! "${VENV}/bin/python" -c 'import yaml' 2>/dev/null; then "${VENV}/bin/pip" install --quiet pyyaml || "${VENV}/bin/pip" install --quiet --index-url https://pypi.org/simple/ pyyaml fi +venv_sync_requirements "${VENV}" "benchmarking/requirements.txt" PY="${VENV}/bin/python" # --- pin the router if its current cpu differs -------------------------------- @@ -194,10 +198,11 @@ kill "${LOGS_PID}" 2>/dev/null || true # --- verdict -------------------------------------------------------------------- if [[ "${STATUS}" == "complete" ]]; then - CAPACITY="$(gcloud storage ls "${DEST}/runs/${NAME}/**/run_tag=${TAG}/capacity.json" 2>/dev/null | tail -1)" + CAPACITY="$("${PY}" benchmarking/result_store.py list "${DEST}/runs/${NAME}" \ + --match "*/run_tag=${TAG}/capacity.json" | tail -1)" if [[ -n "${CAPACITY}" ]]; then echo ">>> ${CAPACITY}" - gcloud storage cat "${CAPACITY}" + "${PY}" benchmarking/result_store.py cat "${CAPACITY}" fi kubectl delete job "${JOB}" -n "${NAMESPACE}" >/dev/null echo ">>> ${NAME}: complete (tag ${TAG})" diff --git a/benchmarking/nighthawk-ingress/runner.py b/benchmarking/nighthawk-ingress/runner.py index c746f6933b..4ca67e947d 100644 --- a/benchmarking/nighthawk-ingress/runner.py +++ b/benchmarking/nighthawk-ingress/runner.py @@ -30,7 +30,6 @@ import json import os import re -import shutil import socket import subprocess import sys @@ -40,6 +39,10 @@ from pathlib import Path from typing import IO, TextIO +# Also find the shared module when running directly from the source tree. +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from result_store import upload + import actors import output as output_mod import spec as spec_mod @@ -62,12 +65,12 @@ def parse_duration_seconds(s: str) -> int: def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser(description=__doc__) - p.add_argument("--name", required=True, help="Test name (GCS path component)") + p.add_argument("--name", required=True, help="Test name (result path component)") p.add_argument("--tag", required=True, help="Tag for this run (commit sha)") p.add_argument( "--dest", required=True, - help="Root destination (gs://bucket/path or local path)", + help="Root destination (gs://bucket/path, s3://bucket/path, az://container/path, or local path)", ) p.add_argument("--envoy-cpu", required=True, type=int, dest="envoy_cpu") p.add_argument("--actors", type=int, default=100) @@ -179,25 +182,6 @@ def run_adaptive_client( return proc.wait() -def upload_to_gcs(local_path: Path, gcs_uri: str) -> None: - # Imported here so non-GCS use doesn't require google-cloud-storage. - from google.cloud import storage - - bucket_name, _, blob_path = gcs_uri[len("gs://"):].partition("/") - storage.Client().bucket(bucket_name).blob(blob_path).upload_from_filename( - str(local_path) - ) - - -def upload(src: Path, dest: str) -> None: - if dest.startswith("gs://"): - upload_to_gcs(src, dest) - else: - dest_path = Path(dest) - dest_path.parent.mkdir(parents=True, exist_ok=True) - shutil.copy(src, dest_path) - - def log_run_config(args: argparse.Namespace, prefix: str, logs: TextIO) -> None: lines = [ "==== Run config ====", diff --git a/benchmarking/nighthawk-ingress/tests/test_runner.py b/benchmarking/nighthawk-ingress/tests/test_runner.py new file mode 100644 index 0000000000..dac9484463 --- /dev/null +++ b/benchmarking/nighthawk-ingress/tests/test_runner.py @@ -0,0 +1,55 @@ +# Copyright 2026 The Agent Substrate Authors +# +# 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. + +import importlib.util +import sys +from pathlib import Path +from unittest import mock + +import pytest + + +@pytest.mark.parametrize("upload_fails", [False, True]) +def test_converged_run_requires_successful_uploads(tmp_path, monkeypatch, upload_fails): + # Actor RPCs are outside this test; their generated clients only exist in + # the runner image. Keep all workload and network operations mocked. + actors = mock.Mock() + spec = importlib.util.spec_from_file_location( + "nighthawk_runner", Path(__file__).resolve().parents[1] / "runner.py" + ) + runner = importlib.util.module_from_spec(spec) + with mock.patch.dict(sys.modules, {"actors": actors}): + spec.loader.exec_module(runner) + monkeypatch.setattr(sys, "argv", [ + "runner.py", "--name", "test", "--tag", "test", "--dest", str(tmp_path / "results"), + "--envoy-cpu", "1", + ]) + # Place the otherwise timestamp-named work directory under this test. + monkeypatch.setattr(runner, "Path", lambda _: tmp_path / "work") + monkeypatch.setattr(runner.spec_mod, "build_spec_dict", lambda **_: {}) + monkeypatch.setattr(runner.spec_mod, "spec_dict_to_textproto", lambda *_: "spec") + monkeypatch.setattr(runner, "start_service", lambda *_: mock.Mock()) + def adaptive(_spec, output, _logs): + output.write_text("output") + return 0 + monkeypatch.setattr(runner, "run_adaptive_client", adaptive) + monkeypatch.setattr(runner.output_mod, "parse_output_textproto", lambda *_: {}) + monkeypatch.setattr(runner.output_mod, "stage_rows", lambda *_: [{"stage": "testing"}]) + monkeypatch.setattr(runner.output_mod, "stats_records", lambda *_: []) + monkeypatch.setattr(runner.output_mod, "capacity_summary", lambda *_, **__: {}) + if upload_fails: + monkeypatch.setattr(runner, "upload", mock.Mock(side_effect=OSError("storage unavailable"))) + with pytest.raises(SystemExit) as exc: + runner.main() + assert exc.value.code == (1 if upload_fails else 0) diff --git a/benchmarking/requirements.txt b/benchmarking/requirements.txt new file mode 100644 index 0000000000..bbae55626b --- /dev/null +++ b/benchmarking/requirements.txt @@ -0,0 +1,15 @@ +# Copyright 2026 The Agent Substrate Authors +# +# 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. + +obstore==0.12.1 diff --git a/benchmarking/result_store.py b/benchmarking/result_store.py new file mode 100644 index 0000000000..c5b435ff03 --- /dev/null +++ b/benchmarking/result_store.py @@ -0,0 +1,90 @@ +# Copyright 2026 The Agent Substrate Authors +# +# 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. + +"""Read and publish benchmark artifacts using local paths or object-store URLs.""" + +import argparse +import fnmatch +import sys +from pathlib import Path +from urllib.parse import urlsplit + +import obstore +from obstore.store import LocalStore, from_url + + +class ResultStore: + """An artifact root. Paths passed to its methods are relative to this root.""" + + def __init__(self, root: str): + url = urlsplit(root) + if not url.scheme: + self.store = LocalStore(Path(root).resolve(), mkdir=True) + elif url.scheme in {"gs", "s3", "az", "memory"}: + if url.query or url.fragment: + raise ValueError("Result-store URLs cannot contain a query or fragment") + self.store = from_url(root) + else: + raise ValueError(f"Unsupported result-store scheme: {url.scheme}") + + def write(self, path: str, source: Path) -> None: + obstore.put(self.store, path, source) + + def read(self, path: str) -> bytes: + return bytes(obstore.get(self.store, path).bytes()) + + def list(self, prefix: str = ""): + """Yield artifact paths relative to the root, recursively.""" + for batch in obstore.list(self.store, prefix=prefix): + for item in batch: + yield item["path"] + + +def upload(source: Path, destination: str) -> None: + """Upload one artifact; propagate failures to the runner's exit status.""" + root, _, name = destination.rpartition("/") + ResultStore(root or ("/" if destination.startswith("/") else ".")).write(name, source) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + commands = parser.add_subparsers(dest="command", required=True) + listing = commands.add_parser("list", help="List artifact URLs under a root") + listing.add_argument("root") + listing.add_argument("--match", default="*", help="Glob matched against relative paths") + cat = commands.add_parser("cat", help="Write one artifact to stdout") + cat.add_argument("url") + download = commands.add_parser("download", help="Download all artifacts under a root") + download.add_argument("root") + download.add_argument("directory", type=Path) + args = parser.parse_args() + if args.command == "cat": + root, _, name = args.url.rpartition("/") + sys.stdout.buffer.write(ResultStore(root or ("/" if args.url.startswith("/") else ".")).read(name)) + return + store = ResultStore(args.root) + for path in store.list(): + if args.command == "list": + if fnmatch.fnmatchcase(path, args.match): + print(f"{args.root.rstrip('/')}/{path}") + else: + target = (args.directory / path).resolve() + if not target.is_relative_to(args.directory.resolve()): + raise ValueError(f"Artifact path escapes the download directory: {path}") + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(store.read(path)) + + +if __name__ == "__main__": + main() diff --git a/benchmarking/test_result_store.py b/benchmarking/test_result_store.py new file mode 100644 index 0000000000..99f846252e --- /dev/null +++ b/benchmarking/test_result_store.py @@ -0,0 +1,76 @@ +# Copyright 2026 The Agent Substrate Authors +# +# 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. + +import subprocess +import sys +from pathlib import Path + +import pytest + +from result_store import ResultStore, upload + + +@pytest.mark.parametrize("backend", ["local", "memory"]) +def test_round_trip(tmp_path, backend): + store = ResultStore(str(tmp_path / "results") if backend == "local" else "memory:///") + source = tmp_path / "source" + source.write_bytes(b"first\x00\xff") + path = "runs/test/run_date=2026-11-12/run_ts=123/run_tag=abc/stats.jsonl" + store.write(path, source) + store.write("unrelated/status.json", source) + assert list(store.list("runs/")) == [path] + assert store.read(path) == source.read_bytes() + source.write_bytes(b"replacement") + store.write(path, source) + assert store.read(path) == b"replacement" + with pytest.raises(FileNotFoundError): + store.read("missing") + + +def test_upload_local_and_failure(tmp_path): + source = tmp_path / "source" + source.write_bytes(b"stats") + destination = tmp_path / "nested" / "stats.jsonl" + upload(source, str(destination)) + assert destination.read_bytes() == b"stats" + with pytest.raises(FileNotFoundError): + upload(tmp_path / "missing", str(destination)) + with pytest.raises(ValueError, match="Unsupported"): + upload(source, "https://example.com/stats.jsonl") + + +def test_cli(tmp_path): + source = tmp_path / "source" + source.write_bytes(b"stats") + root = tmp_path / "results" + upload(source, f"{root}/runs/test/stats.jsonl") + script = Path(__file__).with_name("result_store.py") + def run(*args): + return subprocess.run([sys.executable, str(script), *map(str, args)], capture_output=True) + listing = run("list", root, "--match", "runs/*/stats.jsonl") + assert listing.returncode == 0, listing.stderr + assert listing.stdout.decode().strip() == f"{root}/runs/test/stats.jsonl" + assert run("cat", f"{root}/runs/test/stats.jsonl").stdout == b"stats" + destination = tmp_path / "download" + result = run("download", root, destination) + assert result.returncode == 0, result.stderr + assert (destination / "runs/test/stats.jsonl").read_bytes() == b"stats" + assert run("cat", f"{root}/missing").returncode != 0 + + +def test_azure_url_selects_container_and_prefix(monkeypatch): + monkeypatch.setenv("AZURE_STORAGE_ACCOUNT_NAME", "testaccount") + store = ResultStore("az://results/substrate") + assert store.store.config["container_name"] == "results" + assert store.store.prefix == "substrate" diff --git a/docs/guides/benchmarking.md b/docs/guides/benchmarking.md index 35c4caff6e..957a9e4eb7 100644 --- a/docs/guides/benchmarking.md +++ b/docs/guides/benchmarking.md @@ -122,6 +122,29 @@ not a local entry point. See [automation/README.md](../../benchmarking/automatio python3 runner.py -f tests/.py -t 1m -u 1 --name --dest /tmp/bench ``` +Results can use `gs://bucket/prefix`, `s3://bucket/prefix`, +`az://container/prefix`, or a local path. For example, pass +`--dest s3://benchmark-results/substrate` to the runner or orchestrator. +The `runs//run_date=/run_ts=/run_tag=` layout is the same for every backend. +Local paths inside a Job disappear with the pod unless backed by a volume. + +All runners use [result_store.py](../../benchmarking/result_store.py), backed by +[obstore](https://developmentseed.org/obstore/latest/). Install +`benchmarking/requirements.txt` to use its `list`, `cat`, and `download` commands +from a workstation. New runners should import its `upload` function; consumers +can use `ResultStore` to list and read paths relative to a result root. + +Configure credentials on the runner pod through workload identity or obstore's +[provider configuration](https://developmentseed.org/obstore/latest/api/store/). +For an S3-compatible service, set `AWS_ENDPOINT`, `AWS_REGION`, +`AWS_ACCESS_KEY_ID`, and `AWS_SECRET_ACCESS_KEY`; use `AWS_ALLOW_HTTP=true` only +for a development endpoint using HTTP. These settings must reach the Job, not +just the orchestrator. Workstation retrieval uses the workstation's credentials. +For Azure, `az://` names the container; set `AZURE_STORAGE_ACCOUNT_NAME` for +the account and configure credentials using obstore's Azure environment settings. +The generic Job templates take only the destination; supply provider settings +through your cluster's pod configuration. + Four flags control the optional post-run measurements described in [Benchmark output files](#benchmark-output-files): @@ -673,11 +696,11 @@ To pull the step rows out of a run's results: ```bash RUN=/runs//run_date=/run_ts=/run_tag= -gcloud storage cat "$RUN/stats.jsonl" | grep '"metric": "ladder_step_' +python3 benchmarking/result_store.py cat "$RUN/stats.jsonl" | grep '"metric": "ladder_step_' ``` `` is the orchestrator's `--dest`, and its log has the full `RUN` path -on each `Uploaded ... -> gs://...` line. +on each `Uploaded ... -> ...` line. ### Isolate Mode diff --git a/hack/verify/benchmarking-python.sh b/hack/verify/benchmarking-python.sh index 000cb29a5c..fe4f38aa73 100755 --- a/hack/verify/benchmarking-python.sh +++ b/hack/verify/benchmarking-python.sh @@ -29,6 +29,7 @@ source hack/util/venv.sh # . A suite without a # requirements.txt gets pytest alone. SUITES=( + "benchmarking test_result_store.py" "benchmarking/analysis ." "benchmarking/automation ." "benchmarking/locust unit_tests"