From 9538806adb653657f0d6f4d1b11617155325a417 Mon Sep 17 00:00:00 2001 From: Jeff Newman Date: Sun, 13 Sep 2026 14:55:51 -0500 Subject: [PATCH 1/8] Add generalized ActivitySim benchmark runner and CI --- .github/workflows/ci.yml | 69 ++++ .gitignore | 3 + .pre-commit-config.yaml | 13 + LICENSE | 1 + MANIFEST.in | 3 + README.md | 209 +++++++++++ pyproject.toml | 35 ++ src/abench/Dockerfile | 10 + src/abench/__init__.py | 3 + src/abench/__main__.py | 3 + src/abench/cli.py | 436 ++++++++++++++++++++++ src/abench/common.py | 11 + src/abench/profiles.py | 139 +++++++ src/abench/profiles/mtc.yaml | 46 +++ src/abench/profiles/sandag.yaml | 65 ++++ src/abench/report.py | 436 ++++++++++++++++++++++ src/abench/runtime/__init__.py | 1 + src/abench/runtime/build_sources.py | 109 ++++++ src/abench/runtime/instrumentation.py | 61 +++ src/abench/runtime/worker.py | 255 +++++++++++++ src/abench/sources.py | 93 +++++ tests/fixtures/tiny/benchmark.yaml | 20 + tests/fixtures/tiny/configs/settings.yaml | 15 + tests/fixtures/tiny/data/households.csv | 5 + tests/fixtures/tiny/tiny_extension.py | 54 +++ tests/requirements.txt | 6 + tests/test_builder.py | 82 ++++ tests/test_docker.py | 61 +++ tests/test_measurement.py | 364 ++++++++++++++++++ tests/test_profiles_sources.py | 165 ++++++++ 30 files changed, 2773 insertions(+) create mode 100644 .github/workflows/ci.yml create mode 100644 .pre-commit-config.yaml create mode 100644 MANIFEST.in create mode 100644 README.md create mode 100644 pyproject.toml create mode 100644 src/abench/Dockerfile create mode 100644 src/abench/__init__.py create mode 100644 src/abench/__main__.py create mode 100644 src/abench/cli.py create mode 100644 src/abench/common.py create mode 100644 src/abench/profiles.py create mode 100644 src/abench/profiles/mtc.yaml create mode 100644 src/abench/profiles/sandag.yaml create mode 100644 src/abench/report.py create mode 100644 src/abench/runtime/__init__.py create mode 100644 src/abench/runtime/build_sources.py create mode 100644 src/abench/runtime/instrumentation.py create mode 100644 src/abench/runtime/worker.py create mode 100644 src/abench/sources.py create mode 100644 tests/fixtures/tiny/benchmark.yaml create mode 100644 tests/fixtures/tiny/configs/settings.yaml create mode 100644 tests/fixtures/tiny/data/households.csv create mode 100644 tests/fixtures/tiny/tiny_extension.py create mode 100644 tests/requirements.txt create mode 100644 tests/test_builder.py create mode 100644 tests/test_docker.py create mode 100644 tests/test_measurement.py create mode 100644 tests/test_profiles_sources.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..df4d8d4 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,69 @@ +name: Tests +on: + push: + pull_request: + workflow_dispatch: +permissions: + contents: read +concurrency: + group: tests-${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true +jobs: + unit: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python: ['3.10', '3.11', '3.12'] + timeout-minutes: 20 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python }} + cache: pip + - uses: actions/setup-node@v4 + with: + node-version: '22' + - run: python -m pip install '.[test]' -r tests/requirements.txt + - run: ruff check . && ruff format --check . + - run: python -m pytest -m 'not docker' -q + - run: python -m build + - name: Verify installed wheel and packaged resources + run: | + python -m venv /tmp/abench-wheel + /tmp/abench-wheel/bin/pip install dist/*.whl + cd /tmp + /tmp/abench-wheel/bin/abench --help + /tmp/abench-wheel/bin/python -c "from pathlib import Path; import abench; p=Path(abench.__file__).parent; assert (p/'Dockerfile').is_file(); assert (p/'profiles/mtc.yaml').is_file(); assert (p/'runtime/build_sources.py').is_file()" + docker: + runs-on: ubuntu-latest + timeout-minutes: 30 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.11' + cache: pip + - run: python -m pip install . pytest + - name: Build exact sources and exercise serial and multiprocess runs + env: + ABENCH_DOCKER_TESTS: '1' + run: | + mkdir -p experiments + python -m pytest tests/test_docker.py -v --basetemp=experiments/ci + - name: Retain reports and diagnostics + if: always() + uses: actions/upload-artifact@v4 + with: + name: docker-experiments + path: | + experiments/ci/**/experiment/*.html + experiments/ci/**/experiment/*.json + experiments/ci/**/experiment/*.txt + experiments/ci/**/experiment/*.log + experiments/ci/**/experiment/measured/*.json* + experiments/ci/**/experiment/measured/*.csv + experiments/ci/**/experiment/measured/*.log + experiments/ci/**/experiment/warmup/*.log + if-no-files-found: ignore diff --git a/.gitignore b/.gitignore index 83972fa..f7bfe9e 100644 --- a/.gitignore +++ b/.gitignore @@ -216,3 +216,6 @@ __marimo__/ # Streamlit .streamlit/secrets.toml + +# Local benchmark results and caches +experiments/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..3baa17b --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,13 @@ +repos: + - repo: local + hooks: + - id: ruff-check + name: Ruff check + entry: ruff check + language: system + types: [python] + - id: ruff-format + name: Ruff format + entry: ruff format --check + language: system + types: [python] diff --git a/LICENSE b/LICENSE index c0d18ef..9e8a46f 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,7 @@ BSD 3-Clause License Copyright (c) 2026, ActivitySim +Copyright (c) 2024, ActivitySim Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..eb5bcdb --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,3 @@ +include LICENSE README.md .pre-commit-config.yaml +recursive-include tests *.py *.yaml *.csv *.txt +recursive-include .github *.yml diff --git a/README.md b/README.md new file mode 100644 index 0000000..17f1df6 --- /dev/null +++ b/README.md @@ -0,0 +1,209 @@ +# abench + +Run reproducible ActivitySim runtime and memory experiments in Linux Docker, +from macOS or Linux. One runner supports MTC, SANDAG ABM3, and other models through +small YAML profiles. ActivitySim itself does not need to be installed on the host. + +```bash +python -m pip install /path/to/abench +abench run --model-dir /path/to/sandag-abm3-example --profile sandag \ + --source activitysim=ActivitySim/activitysim@ \ + --source sharrow=ActivitySim/sharrow@ \ + --multiprocess --processes 4 --sharrow --households 28365 \ + --memory 32g --shm-size 8g --output-dir /path/to/experiments/sandag +``` + +Use `--profile mtc` for MTC; both profiles ship with the package. SANDAG defaults +to its small `benchmarking-data`, **not full-scale skims**. MTC defaults to +`data_full`. `--data-dir` overrides either. Model directories need not be Git +repositories; Git revision/status are recorded where available and model files +are always snapshotted. Existing example scripts and normal configs are untouched. + +The host needs Python 3.10+, PyYAML (installed with abench), and a Linux Docker +engine with cgroup v2 and `memory.peak`. Docker Desktop must have enough VM RAM +for the chosen memory limit plus VM overhead. The default container is Debian +Bookworm/Python 3.11. Current instrumentation requires ActivitySim's +`workflow.State` API (1.4-era or newer); arbitrary historical revisions are not +promised to work. Build/runtime failures retain diagnostics and a failure report. + +## Run controls + +- `--single-process` (default), or `--multiprocess --processes N`. The count applies + to every sliced stage; coordinators are additional processes. +- `--sharrow` (default) or `--no-sharrow`. Sharrow enabled requires its source pin. +- `--households N` (default 1,000); zero uses the original full input population. + abench never replicates households. Positive samples must match realized output. +- `--config-overlay configs_explicit_chunk` adds config directories in listed + priority order. Relative overlay paths are relative to the model directory. +- `--memory 16g`, `--shm-size 8g`, `--interval 0.5`, and optional `--platform`. +- `--output-dir` must be new. `--label` names an experiment, and `--compare` accepts + earlier experiment directories. `--cache-from` seeds compatible generated flows; + a complete warmup still runs. + +`abench validate` accepts the same experiment arguments without `--output-dir`. +It checks the profile, required inputs, CSV population size, source pin syntax, +and Docker capabilities without building an image or running the model. It does +not prove Git commit availability, package compatibility, or skim consistency; +those are checked by the build and model run. + +## Any dependency from GitHub source + +Repeat `--source` for any Python distribution, including add-on extensions: + +```bash +--source 'my-addon[fast]=ExampleOrg/model-addon@#subdirectory=python/addon' +``` + +The left side is the **distribution name** (which can differ from its import +module). Each source uses an exact full SHA and an `organization/repository` name. +Extras and a repository subdirectory are optional. Profiles can declare the same +entries as strings or mappings: + +```yaml +sources: + - name: my-addon + repository: ExampleOrg/model-addon + commit: '0123456789abcdef0123456789abcdef01234567' + extras: [fast] + subdirectory: python/addon +``` + +CLI sources override profile sources by normalized distribution name. Duplicate +CLI entries are errors. `--activitysim-commit` and `--sharrow-commit` remain aliases +for the official repositories; conflicting alias/source declarations are errors. + +Inside Docker, abench verifies each checkout's Git object, builds a wheel, checks +its distribution name, then installs all source wheels together with other +requirements. It runs `pip check` and verifies installed wheel identities. Exact +source commits, resolved versions, and the full dependency environment are saved. +All source dependencies must agree: conflicting requirements fail the build. + +Use profile `requirements` for additional registry requirements and `constraints` +for resolver bounds or exact transitive pins. Profiles may select `python_image` +(a compatible Debian-based image, optionally pinned by digest). The default image +includes a compiler and HDF5 headers. Packages requiring other system libraries +can use a prebuilt compatible base image. Private GitHub authentication and custom +OS provisioning are outside the initial interface. + +Source pins do not freeze unpinned transitive/build dependencies or base images. +Retain images and dependency manifests for strict reproduction. Build isolation +may fetch build requirements; constraints currently govern the final environment, +not those isolated build environments. + +## Add another model + +Create `benchmark.yaml` in its model directory, then use `--model-dir`: + +```yaml +schema_version: 1 +name: My regional model +configs: [configs] +mp_configs: [configs_mp] +snapshot: [configs, configs_mp, extensions] +extensions: [extensions] +data_dir: data +required_inputs: + - [households.csv, households.parquet] + - persons.csv + - land_use.csv + - skims.omx +settings: + use_shadow_pricing: false + rng_base_seed: 0 +input_tables: + households: {} + persons: {} + land_use: + totals: [TOTPOP, TOTHH, TOTEMP] + zone_columns: [TAZ] +output_tables: + households: {} + persons: {} + tours: {} + trips: + categories: [trip_mode, primary_purpose] +``` + +Paths in a profile are relative to the model directory (except `data_dir`, which +may be absolute). `configs` and `mp_configs` are ordered, highest priority first. +`snapshot` must cover config files, local extension modules, and adapter modules; +entries must not overlap or contain directory symlinks. Data is mounted read-only +and should not change during a run. Profiles are trusted model code/configuration. + +Settings precedence, lowest to highest: **normal model configs → profile settings +→ user overlays → required CLI controls** (sample, SP/MP, worker counts, Sharrow, +and fail-fast). Generated inheriting profile configs preserve this ordering when +ActivitySim reconstructs worker settings. Component overlays remain independent. + +Optional profile fields: + +| Field | Purpose | +|---|---| +| `models_from`, `exclude_models` | Take a YAML `models` list and explicitly omit diagnostic steps. Otherwise use normal settings. | +| `mp_settings` | Read `multiprocess_steps` from a separate YAML file. | +| `extensions` | Import modules through ActivitySim's registration mechanism, including spawned workers. Installing a package alone does not register its components. | +| `adapter: module:function` | Optional `function(state, spec, phase)` initialization hook, called once in the model parent before execution. Use `state.import_extensions` for worker setup; parent-only mutations are not automatically worker initialization. | +| `input_tables`, `output_tables` | Logical table names mapped to summary options: `file` (stem), `totals`, `categories`, and `zone_columns` for land use. CSV and Parquet are supported. Use logical `households` for sample validation. | +| `output_prefix` | Default `final_`; applied to output file stems. | +| `household_table` | Default `households.csv`, used for early CSV sample validation. Parquet samples are checked after execution. | +| `zone_label` | Display label for land-use rows, such as zones or MAZs. | + +For specialized data formats, an adapter can arrange compatible CSV/Parquet +summary outputs. The initial generic reader does not interpret arbitrary binary +model outputs. + +## Measurement and reports + +Sharrow runs first complete a full matching warmup in a separate container. The +measured run uses the same sample, seed, layout, and cache path. Numba compilation +of generated flow overloads is rejected during measurement; ordinary non-flow +compilation and disk-cache loading remain included. This guard uses private +Numba internals and is covered by real disk-cache hit/miss tests. + +Memory is the whole-container cgroup v2 charge, counting shared pages once. +Blue is `memory.current` (including file cache, shared memory, and kernel costs). +Green dashed is `anon + shmem`, a subset excluding ordinary file cache and kernel +costs. Never add the lines. Swap is recorded separately and disabled by equal +memory/memory+swap limits. `/dev/shm` capacity is within that limit. + +Component runtimes show worker mean, population SD, count, and maximum. Worker SD +is imbalance, not confidence across repeated trials. The component dropdown +highlights each worker's actual window; overlaps darken and gaps remain clear. +Overall elapsed includes startup, coordination, and checkpoint writes between +components. Kernel peak includes startup; warmup and post-run summaries are +excluded. Shared VM page-cache ownership can influence container charges: this +is not a cold-input I/O benchmark. + +```bash +abench report --compare /path/to/run-a /path/to/run-b \ + --output-dir /path/to/comparison.html +``` + +Reports are offline HTML/SVG/JavaScript plus normalized JSON. They retain fastest +successful component highlighting and use common memory axes. Existing MTC and +SANDAG schema-1 experiments remain readable, including approximate legacy timing +windows where only completion logs exist. New experiments use schema 2 and include +source manifests, resolved profile/settings, model and harness snapshots, file +hashes, Docker details, and source installation provenance. Input file size/mtime +records are provenance hints, not content hashes of large skims. + +## Development and CI + +```bash +python -m pip install -e '.[test]' -r tests/requirements.txt +python -m pytest -m 'not docker' +pre-commit run --all-files +python -m build +ABENCH_DOCKER_TESTS=1 python -m pytest tests/test_docker.py -v +``` + +Install Node.js to exercise the offline chart selector test. GitHub Actions runs +unit tests on Python 3.10–3.12, lint/format checks, wheel packaging checks, and Linux +Docker integration. The Docker tests build pinned ActivitySim/Sharrow sources and +run a four-household extension workflow in serial and multiprocess modes, checking +warmup, generated-code cache reuse, per-worker timings, merged outputs, memory, +and reports. They do not require either example repository or large datasets. + +Adapted from the MTC and SANDAG benchmark harnesses developed in this workspace. +The original measurement approach was informed by WSP's Lighthouse production +benchmark. See LICENSE for the retained BSD license. diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..4f18484 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,35 @@ +[build-system] +requires = ["setuptools>=77"] +build-backend = "setuptools.build_meta" + +[project] +name = "abench" +version = "0.1.0" +description = "Reproducible ActivitySim runtime and memory benchmarks in Linux containers" +readme = "README.md" +requires-python = ">=3.10" +license = "BSD-3-Clause" +license-files = ["LICENSE"] +dependencies = ["PyYAML>=6"] + +[project.optional-dependencies] +test = ["pytest>=8", "pandas<3", "pyarrow", "numba", "ruff", "build"] + +[project.scripts] +abench = "abench.cli:entrypoint" + +[tool.setuptools.package-data] +abench = ["Dockerfile", "profiles/*.yaml"] + +[tool.ruff] +line-length = 88 + +[tool.pytest.ini_options] +testpaths = ["tests"] +markers = ["docker: Linux Docker integration tests (opt in with ABENCH_DOCKER_TESTS=1)"] + +[tool.ruff.lint] +select = ["E4", "E7", "E9", "F", "I"] + +[tool.ruff.lint.isort] +known-first-party = ["abench"] diff --git a/src/abench/Dockerfile b/src/abench/Dockerfile new file mode 100644 index 0000000..0705752 --- /dev/null +++ b/src/abench/Dockerfile @@ -0,0 +1,10 @@ +ARG PYTHON_IMAGE=python:3.11-slim-bookworm +FROM ${PYTHON_IMAGE} +RUN apt-get update && apt-get install -y --no-install-recommends git build-essential libhdf5-dev \ + && rm -rf /var/lib/apt/lists/* +COPY build_sources.py dependencies.json /opt/build/ +RUN python /opt/build/build_sources.py /opt/build/dependencies.json +ENV PYTHONUNBUFFERED=1 OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 \ + NUMEXPR_NUM_THREADS=1 NUMBA_NUM_THREADS=1 PYTHONHASHSEED=0 DASK_SCHEDULER=synchronous +WORKDIR /model +ENTRYPOINT ["python", "/benchmark/worker.py"] diff --git a/src/abench/__init__.py b/src/abench/__init__.py new file mode 100644 index 0000000..334e474 --- /dev/null +++ b/src/abench/__init__.py @@ -0,0 +1,3 @@ +"""Reproducible ActivitySim experiments in Linux containers.""" + +__version__ = "0.1.0" diff --git a/src/abench/__main__.py b/src/abench/__main__.py new file mode 100644 index 0000000..920b4b9 --- /dev/null +++ b/src/abench/__main__.py @@ -0,0 +1,3 @@ +from .cli import entrypoint + +entrypoint() diff --git a/src/abench/cli.py b/src/abench/cli.py new file mode 100644 index 0000000..de67962 --- /dev/null +++ b/src/abench/cli.py @@ -0,0 +1,436 @@ +"""Run reproducible ActivitySim benchmarks with model profiles.""" + +import argparse +import hashlib +import json +import math +import re +import shutil +import subprocess +import sys +import tempfile +import uuid +from datetime import datetime, timezone +from pathlib import Path + +from . import __version__ +from .common import read_json, write_json +from .profiles import load_profile, validate_model +from .report import load_run, report +from .sources import resolve_sources + +PACKAGE = Path(__file__).resolve().parent + + +def command(args, log=None): + """Keep build/run output on disk and propagate failures to the caller.""" + if log: + with log.open("w") as stream: + subprocess.run(args, stdout=stream, stderr=subprocess.STDOUT, check=True) + else: + return subprocess.check_output(args, text=True).strip() + + +def git_info(root, args): + """Non-Git model directories are supported; snapshots still capture their files.""" + result = subprocess.run( + ["git", "-C", str(root), *args], capture_output=True, text=True + ) + return result.stdout.strip() if result.returncode == 0 else None + + +def commit(value): + if not re.fullmatch(r"[0-9a-fA-F]{40}", value): + raise argparse.ArgumentTypeError("provide the full 40-character Git commit SHA") + return value.lower() + + +def positive(value): + value = float(value) + if not math.isfinite(value) or value <= 0: + raise argparse.ArgumentTypeError("must be finite and positive") + return value + + +def parser(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--model-dir", type=Path, default=Path.cwd()) + p.add_argument("--profile", default="benchmark.yaml") + p.add_argument( + "--source", + action="append", + default=[], + help="distribution=organization/repository@40-character-SHA", + ) + p.add_argument("--activitysim-commit", type=commit) + p.add_argument("--sharrow-commit", type=commit) + mode = p.add_mutually_exclusive_group() + mode.add_argument("--single-process", dest="multiprocess", action="store_false") + mode.add_argument("--multiprocess", action="store_true") + p.set_defaults(multiprocess=False) + p.add_argument("--processes", type=int, help="required for --multiprocess") + p.add_argument("--sharrow", action=argparse.BooleanOptionalAction, default=True) + p.add_argument( + "--households", type=int, default=1000, help="0 means full population" + ) + p.add_argument("--data-dir", type=Path, default=None) + p.add_argument( + "--config-overlay", + type=Path, + nargs="+", + default=[], + help="extra config directories, highest priority first", + ) + p.add_argument( + "--cache-from", + type=Path, + help="seed flow cache from an experiment with the same revisions and dependencies", + ) + p.add_argument( + "--output-dir", + type=Path, + required=False, + help="new experiment directory, or report HTML with --report-only", + ) + p.add_argument("--label", help="experiment label in comparisons") + p.add_argument( + "--compare", + type=Path, + nargs="+", + default=[], + help="previous experiment directories", + ) + p.add_argument( + "--report-only", + action="store_true", + help="rebuild a comparison of --compare directories without Docker", + ) + p.add_argument( + "--interval", + type=positive, + default=0.5, + help="memory sampling interval in seconds", + ) + p.add_argument( + "--memory", default="16g", help="Docker memory and memory+swap limit" + ) + p.add_argument( + "--shm-size", default="8g", help="/dev/shm capacity; charged against --memory" + ) + p.add_argument( + "--platform", + choices=("linux/arm64", "linux/amd64"), + help="defaults to Docker native architecture", + ) + return p + + +def mount(source, target, readonly=False): + source = str(source.resolve()) + if "," in source: + raise ValueError("Docker bind paths cannot contain commas") + return [ + "--mount", + f"type=bind,src={source},dst={target}" + (",readonly" if readonly else ""), + ] + + +def container_phase(spec, output, data, image, phase_name): + """Retain Docker exit/OOM state even when the supervisor cannot finish.""" + phase = output / phase_name + phase.mkdir() + name = "abench-" + uuid.uuid4().hex[:12] + args = [ + "docker", + "run", + "--name", + name, + "--cgroupns=private", + "--memory", + spec["memory"], + "--memory-swap", + spec["memory"], + "--shm-size", + spec["shm_size"], + "--network=none", + ] + if spec["platform"]: + args += ["--platform", spec["platform"]] + args += mount(output / "model", "/model", True) + args += mount(output / "runner", "/benchmark", True) + args += mount(data, "/data", True) + args += mount(output, "/results") + args += [image, "supervise", f"/results/{phase_name}"] + try: + command(args, phase / "console.log") + finally: + try: + state = json.loads( + command(["docker", "inspect", name, "--format", "{{json .State}}"]) + ) + write_json(phase / "docker-state.json", state) + finally: + subprocess.run( + ["docker", "rm", "-f", name], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + check=False, + ) + + +def main(argv=None): + p = parser() + argv = list(sys.argv[1:] if argv is None else argv) + action = argv.pop(0) if argv and argv[0] in ("run", "report", "validate") else "run" + args = p.parse_args(argv) + if action == "report": + args.report_only = True + if action != "validate" and args.output_dir is None: + p.error("--output-dir is required") + comparisons = [path.expanduser().resolve() for path in args.compare] + for previous in comparisons: + load_run(previous) + output = args.output_dir.expanduser().resolve() if args.output_dir else None + if args.report_only: + if output is None: + p.error("--output-dir is required for reporting") + if not comparisons: + p.error("--report-only requires --compare") + report(comparisons, output) + print(output) + return 0 + root = args.model_dir.expanduser().resolve() + try: + profile = load_profile(args.profile, root) + sources = resolve_sources( + profile.get("sources", []), + args.source, + args.activitysim_commit, + args.sharrow_commit, + ) + data = ( + args.data_dir.expanduser().resolve() + if args.data_dir + else (root / profile.get("data_dir", "data")).resolve() + ) + validate_model(profile, root, data, args.households) + except ValueError as error: + p.error(str(error)) + args.activitysim_commit = next( + s["commit"] for s in sources if s["name"] == "activitysim" + ) + args.sharrow_commit = next( + (s["commit"] for s in sources if s["name"] == "sharrow"), None + ) + if args.sharrow and not args.sharrow_commit: + p.error("Sharrow enabled: provide a sharrow source override") + if args.households < 0: + p.error("--households must be nonnegative") + if args.multiprocess and (args.processes is None or args.processes < 1): + p.error("--multiprocess requires --processes >= 1") + if not args.multiprocess and args.processes not in (None, 1): + p.error("--processes > 1 requires --multiprocess") + for value in (args.memory, args.shm_size): + if not re.fullmatch(r"[1-9][0-9]*[bkmgBKMG]?", value): + p.error("memory sizes must be positive integer Docker sizes, such as 16g") + args.data_dir = data + overlays = [(root / path.expanduser()).resolve() for path in args.config_overlay] + for path in overlays: + if not path.is_dir() or (output is not None and output.is_relative_to(path)): + p.error("config overlays must be existing directories outside --output-dir") + seed = args.cache_from.expanduser().resolve() if args.cache_from else None + if seed: + prior = read_json(seed / "experiment.json", {}) + for key in ("activitysim_commit", "sharrow_commit"): + if prior.get(key) != getattr(args, key): + p.error(f"--cache-from must use the same {key}") + if not (seed / "cache/flows").is_dir(): + p.error("--cache-from has no flow cache") + for source in profile["snapshot"]: + if output is not None and output.is_relative_to(root / source): + p.error("--output-dir must be outside snapshot source directories") + if seed and prior.get("sources") != sources: + p.error("--cache-from must use the same complete source dependency manifest") + if "," in str(output) or "," in str(data): + p.error("Docker bind paths cannot contain commas") + if output is not None and output.exists(): + p.error( + "--output-dir must not already exist; each experiment owns a fresh cache" + ) + docker = json.loads(command(["docker", "info", "--format", "{{json .}}"])) + if docker.get("OSType") != "linux" or str(docker.get("CgroupVersion")) != "2": + p.error("Docker must run Linux containers using cgroup v2") + if action == "validate": + print( + json.dumps( + { + "profile": profile, + "sources": sources, + "data_dir": str(data), + "docker": docker, + }, + indent=2, + ) + ) + return 0 + output.mkdir(parents=True) + spec = vars(args).copy() + spec.update( + schema_version=2, + abench_version=__version__, + profile=profile, + profile_name=profile["name"], + sources=sources, + label=args.label or output.name, + processes=args.processes or 1, + created_at=datetime.now(timezone.utc).isoformat(), + model_commit=git_info(root, ["rev-parse", "HEAD"]), + model_git_status=git_info(root, ["status", "--porcelain"]), + docker={ + key: docker.get(key) + for key in ( + "ServerVersion", + "Architecture", + "NCPU", + "MemTotal", + "KernelVersion", + "CgroupVersion", + ) + }, + ) + spec = json.loads(json.dumps(spec, default=str)) + (output / "model").mkdir() + for config in profile["snapshot"]: + source, target = root / config, output / "model" / config + target.parent.mkdir(parents=True, exist_ok=True) + if source.is_dir(): + shutil.copytree( + source, target, ignore=shutil.ignore_patterns("__pycache__", "*.pyc") + ) + else: + shutil.copy2(source, target) + for i, path in enumerate(overlays): + shutil.copytree(path, output / "model" / f"overlay-{i}") + shutil.copytree( + PACKAGE / "runtime", + output / "runner", + ignore=shutil.ignore_patterns("__pycache__"), + ) + shutil.copytree( + PACKAGE, output / "harness", ignore=shutil.ignore_patterns("__pycache__") + ) + shutil.copy2(PACKAGE / "Dockerfile", output / "production-benchmark.Dockerfile") + spec["config_sha256"] = { + str(path.relative_to(output / "model")): hashlib.sha256( + path.read_bytes() + ).hexdigest() + for path in sorted((output / "model").rglob("*")) + if path.is_file() + } + spec["harness_sha256"] = { + str(path.relative_to(output / "harness")): hashlib.sha256( + path.read_bytes() + ).hexdigest() + for path in sorted((output / "harness").rglob("*")) + if path.is_file() + } + spec["input_files"] = { + str(path.relative_to(data)): { + "bytes": path.stat().st_size, + "mtime_ns": path.stat().st_mtime_ns, + } + for path in sorted(data.rglob("*")) + if path.is_file() + } + write_json(output / "experiment.json", spec) + image = "abench:" + uuid.uuid4().hex[:12] + stage = "build" + try: + print(f"Building pinned packages; log: {output / 'build.log'}", flush=True) + with tempfile.TemporaryDirectory() as context: + shutil.copy2( + output / "production-benchmark.Dockerfile", Path(context) / "Dockerfile" + ) + shutil.copy2( + PACKAGE / "runtime" / "build_sources.py", + Path(context) / "build_sources.py", + ) + write_json( + Path(context) / "dependencies.json", + { + "sources": sources, + "requirements": profile.get("requirements", []), + "constraints": profile.get("constraints", []), + }, + ) + build = [ + "docker", + "build", + "-t", + image, + "--build-arg", + f"PYTHON_IMAGE={profile.get('python_image', 'python:3.11-slim-bookworm')}", + ] + if args.platform: + build += ["--platform", args.platform] + command(build + [context], output / "build.log") + spec["image_id"] = command( + ["docker", "image", "inspect", image, "--format", "{{.Id}}"] + ) + write_json(output / "experiment.json", spec) + freeze = command( + [ + "docker", + "run", + "--rm", + "--network=none", + "--entrypoint", + "cat", + image, + "/opt/pip-freeze.txt", + ] + ) + (output / "pip-freeze.txt").write_text(freeze + "\n") + provenance = command( + [ + "docker", + "run", + "--rm", + "--network=none", + "--entrypoint", + "cat", + image, + "/opt/source-provenance.json", + ] + ) + (output / "source-provenance.json").write_text(provenance + "\n") + if seed: + if (seed / "pip-freeze.txt").read_text().strip() != freeze.strip(): + raise ValueError("Cannot seed cache: installed dependencies differ") + # Only flow artifacts are reused. Full warmup still checks coverage + # under the new configuration before compilation-blocked measurement. + shutil.copytree(seed / "cache/flows", output / "cache/flows") + if args.sharrow: + stage = "warmup" + print("Preparing Sharrow cache with a complete matching run…", flush=True) + container_phase(spec, output, data, image, "warmup") + stage = "measured" + print("Running measured model…", flush=True) + container_phase(spec, output, data, image, "measured") + except (Exception, KeyboardInterrupt) as error: + spec["failure"] = {"phase": stage, "error": str(error)} + write_json(output / "experiment.json", spec) + raise + finally: + report(comparisons + [output], output / "report.html") + print(f"Report: {output / 'report.html'}", flush=True) + return 0 if load_run(output)["valid"] else 1 + + +def entrypoint(): + """Expose CLI errors without an unnecessary Python traceback.""" + try: + sys.exit(main()) + except (ValueError, OSError, subprocess.CalledProcessError) as error: + print(f"Benchmark failed: {error}", file=sys.stderr) + sys.exit(1) diff --git a/src/abench/common.py b/src/abench/common.py new file mode 100644 index 0000000..ab78a72 --- /dev/null +++ b/src/abench/common.py @@ -0,0 +1,11 @@ +"""Small artifact helpers shared by the host tools.""" + +import json + + +def read_json(path, default=None): + return json.loads(path.read_text()) if path.exists() else default + + +def write_json(path, value): + path.write_text(json.dumps(value, indent=2) + "\n") diff --git a/src/abench/profiles.py b/src/abench/profiles.py new file mode 100644 index 0000000..3815953 --- /dev/null +++ b/src/abench/profiles.py @@ -0,0 +1,139 @@ +"""Declarative model layouts and inexpensive preflight checks.""" + +import csv +from pathlib import Path, PurePosixPath + +import yaml + +FIELDS = { + "schema_version", + "name", + "configs", + "mp_configs", + "snapshot", + "data_dir", + "required_inputs", + "household_table", + "extensions", + "adapter", + "settings", + "models_from", + "exclude_models", + "mp_settings", + "input_tables", + "output_tables", + "output_prefix", + "zone_label", + "sources", + "requirements", + "constraints", + "python_image", +} + + +def relative_path(value): + """Snapshot and config paths must resolve within the supplied model checkout.""" + if ( + not isinstance(value, str) + or not value + or PurePosixPath(value).is_absolute() + or ".." in PurePosixPath(value).parts + or "\\" in value + ): + raise ValueError(f"expected a relative model path: {value!r}") + return value + + +def load_profile(value, root): + """Load a model-owned YAML profile or one of the shipped example profiles.""" + path = Path(value).expanduser() + if value in ("mtc", "sandag"): + path = Path(__file__).parent / "profiles" / f"{value}.yaml" + elif not path.is_absolute(): + path = root / path + if not path.is_file(): + raise ValueError(f"profile not found: {path}") + profile = yaml.safe_load(path.read_text()) + if not isinstance(profile, dict) or profile.get("schema_version") != 1: + raise ValueError("profile must be a mapping with schema_version: 1") + unknown = set(profile) - FIELDS + if unknown: + raise ValueError(f"unknown profile fields: {sorted(unknown)}") + if not isinstance(profile.get("name"), str) or not profile["name"]: + raise ValueError("profile needs a name") + for key in ("configs", "snapshot"): + if not isinstance(profile.get(key), list) or not profile[key]: + raise ValueError(f"profile needs a nonempty {key} list") + for key in ("configs", "mp_configs", "snapshot"): + for item in profile.get(key, []): + relative_path(item) + for key in ("models_from", "mp_settings"): + if profile.get(key): + relative_path(profile[key]) + for key in ("extensions", "exclude_models", "requirements", "constraints"): + if not isinstance(profile.get(key, []), list) or any( + not isinstance(v, str) for v in profile.get(key, []) + ): + raise ValueError(f"{key} must be a list of strings") + if not isinstance(profile.get("settings", {}), dict): + raise ValueError("settings must be a mapping") + if profile.get("adapter") and ( + not isinstance(profile["adapter"], str) or profile["adapter"].count(":") != 1 + ): + raise ValueError("adapter must be module:function") + for key in ("input_tables", "output_tables"): + tables = profile.get(key, {}) + if not isinstance(tables, dict) or any( + not isinstance(opts, dict) for opts in tables.values() + ): + raise ValueError(f"{key} must map table names to summary options") + return profile + + +def validate_model(profile, root, data, households): + """Check the snapshot closure and required files before starting Docker builds.""" + if not data.is_dir(): + raise ValueError(f"data directory not found: {data}") + snapshots = [root / item for item in profile["snapshot"]] + for path in snapshots: + if not path.exists() or not path.resolve().is_relative_to(root): + raise ValueError(f"missing or external snapshot path: {path}") + for i, path in enumerate(snapshots): + if any( + path.is_relative_to(other) or other.is_relative_to(path) + for other in snapshots[:i] + ): + raise ValueError("snapshot paths must not overlap") + # Do not accidentally copy external data through a directory symlink. + if path.is_dir() and any(p.is_symlink() for p in path.rglob("*")): + raise ValueError(f"snapshot contains symlinks; use ordinary files: {path}") + for name in ( + profile["configs"] + + profile.get("mp_configs", []) + + [profile[k] for k in ("models_from", "mp_settings") if profile.get(k)] + ): + path = root / name + if not path.exists() or not any( + path.is_relative_to(item) for item in snapshots + ): + raise ValueError(f"config must exist and be included in snapshot: {name}") + for group in profile.get("required_inputs", []): + choices = [group] if isinstance(group, str) else group + if not isinstance(choices, list) or not choices: + raise ValueError( + "required_inputs entries must be paths or nonempty alternative lists" + ) + if not any((data / relative_path(name)).is_file() for name in choices): + raise ValueError(f"missing required input; expected one of: {choices}") + # Parquet counts are checked against realized outputs after the run. CSV + # counting uses streaming and no host pandas dependency. + household_file = data / profile.get("household_table", "households.csv") + if households > 0 and household_file.suffix == ".csv" and household_file.is_file(): + with household_file.open(newline="") as stream: + reader = csv.reader(stream) + next(reader, None) + count = sum(1 for row in reader if row) + if households > count: + raise ValueError( + f"requested {households:,} households but only {count:,} are available" + ) diff --git a/src/abench/profiles/mtc.yaml b/src/abench/profiles/mtc.yaml new file mode 100644 index 0000000..9ab542e --- /dev/null +++ b/src/abench/profiles/mtc.yaml @@ -0,0 +1,46 @@ +schema_version: 1 +settings: + chunk_size: 0 + use_shadow_pricing: false + trace_hh_id: null + trace_od: null + resume_after: null + instrument: false + memory_profile: false + expression_profile: false + checkpoints: true + rng_base_seed: 0 + benchmarking: false +exclude_models: +- track_skim_usage +constraints: +- multimethod<2 +- pandas<3 +input_tables: + households: {} + persons: {} + land_use: {} +output_tables: + households: {} + persons: {} + land_use: {} + tours: {} + trips: {} + joint_tour_participants: {} + vehicles: {} +name: MTC prototype +configs: +- configs +mp_configs: +- configs_mp +snapshot: +- configs +- configs_mp +data_dir: data_full +models_from: configs/settings.yaml +zone_label: zones +required_inputs: +- households.csv +- persons.csv +- land_use.csv +- skims.omx diff --git a/src/abench/profiles/sandag.yaml b/src/abench/profiles/sandag.yaml new file mode 100644 index 0000000..7db7438 --- /dev/null +++ b/src/abench/profiles/sandag.yaml @@ -0,0 +1,65 @@ +schema_version: 1 +settings: + chunk_size: 0 + use_shadow_pricing: false + trace_hh_id: null + trace_od: null + resume_after: null + instrument: false + memory_profile: false + expression_profile: false + checkpoints: true + rng_base_seed: 0 + benchmarking: false + recode_pipeline_columns: true +exclude_models: +- track_skim_usage +constraints: +- multimethod<2 +- pandas<3 +input_tables: + households: {} + persons: {} + land_use: {} +output_tables: + households: {} + persons: {} + land_use: {} + tours: {} + trips: {} + joint_tour_participants: {} + vehicles: {} +name: SANDAG ABM3 +configs: +- configs/common +- configs/resident +snapshot: +- configs +- extensions +data_dir: benchmarking-data +models_from: configs/resident/settings.yaml +mp_settings: configs/resident/settings_mp.yaml +extensions: +- extensions +zone_label: MAZs +required_inputs: +- - households.csv + - households.parquet +- - persons.csv + - persons.parquet +- - land_use.csv + - land_use.parquet +- - maz_maz_walk.csv + - maz_maz_walk.parquet +- - maz_maz_bike.csv + - maz_maz_bike.parquet +- traffic_skims_EA.omx +- transit_skims_EA.omx +- traffic_skims_AM.omx +- transit_skims_AM.omx +- traffic_skims_MD.omx +- transit_skims_MD.omx +- traffic_skims_PM.omx +- transit_skims_PM.omx +- traffic_skims_EV.omx +- transit_skims_EV.omx diff --git a/src/abench/report.py b/src/abench/report.py new file mode 100644 index 0000000..ced5be9 --- /dev/null +++ b/src/abench/report.py @@ -0,0 +1,436 @@ +"""Offline readers and interactive reports, including legacy experiment support.""" + +import csv +import html +import json +import math +import re +import statistics + +from .common import read_json, write_json + +COLORS = ("#0072b2", "#d55e00", "#009e73", "#cc79a7", "#e69f00", "#56b4e9") + + +def clock_seconds(value): + """Parse elapsed MM:SS or HH:MM:SS values from ActivitySim logs.""" + total = 0.0 + for part in value.split(":"): + total = total * 60 + float(part) + return total + + +def legacy_windows(phase): + """Recover approximate windows for artifacts predating monotonic timestamps. + + MP completion messages arrive at the parent after execution and can include + checkpoint time. Their relative logging clock also starts slightly after the + sampler. Keep this fallback explicitly approximate instead of inventing exact + timestamps by summing worker durations across unmeasured checkpoint gaps. + """ + path = phase / "console.log" + windows = {} + if not path.exists(): + return windows + pattern = re.compile( + r"^\[(?P[\d:.]+)\].*?\b(?Pmp_\w+) " + r"(?P\w+) : (?P[\d.]+) seconds\b" + ) + serial = re.compile( + r"^\[(?P[\d:.]+)\].*?time to execute run\." + r"(?P\w+) : (?P[\d:.]+)(?: seconds)?\s*$" + ) + for line in path.read_text(errors="replace").splitlines(): + match = pattern.search(line) or serial.search(line) + if match: + end = clock_seconds(match["end"]) + duration = clock_seconds(match["duration"]) + process = match.groupdict().get("process") or "MainProcess" + key = (process, match["component"]) + windows.setdefault(key, []).append( + { + "start_seconds": max(0.0, end - duration), + "end_seconds": end, + "source": "approximate completion log", + } + ) + return windows + + +def component_windows(observations, phase): + """Retain each worker interval separately, including overlaps and gaps.""" + fallback = ( + legacy_windows(phase) + if any( + "start_seconds" not in row or "end_seconds" not in row + for row in observations + ) + else {} + ) + windows = [] + for row in observations: + key = (row.get("process", "MainProcess"), row["component"]) + # Consume a matching fallback even for a timestamped observation so a + # mixed-format artifact cannot assign it to a later repeated execution. + matches = fallback.get(key, []) + approximate = matches.pop(0) if matches else None + if "start_seconds" in row and "end_seconds" in row: + interval = {k: row[k] for k in ("start_seconds", "end_seconds")} + interval["source"] = "recorded monotonic clock" + elif approximate: + interval = approximate + else: + continue + start, end = interval["start_seconds"], interval["end_seconds"] + if not (math.isfinite(start) and math.isfinite(end) and 0 <= start <= end): + continue + windows.append( + dict( + interval, + component=row["component"], + process=key[0], + pid=row.get("pid"), + succeeded=row["succeeded"], + ) + ) + return windows + + +def load_run(directory): + """Use raw worker observations, never duplicate ActivitySim's locutor CSV.""" + spec = read_json(directory / "experiment.json") + if not spec or spec.get("schema_version") not in (1, 2): + raise ValueError(f"Not a supported benchmark experiment: {directory}") + phase = directory / "measured" + grouped = {} + observations = [] + for path in sorted(phase.glob("components-*.jsonl")): + for line in path.read_text().splitlines(): + row = json.loads(line) + observations.append(row) + if row["succeeded"]: + grouped.setdefault(row["component"], []).append(row["seconds"]) + components = { + key: { + "n": len(values), + "mean": statistics.fmean(values), + "sd": statistics.pstdev(values), + "maximum": max(values), + } + for key, values in grouped.items() + } + memory = [] + if (phase / "memory.csv").exists(): + with (phase / "memory.csv").open() as stream: + memory = [ + {key: float(value) for key, value in row.items()} + for row in csv.DictReader(stream) + ] + status = read_json(phase / "status.json", {}) + docker = read_json(phase / "docker-state.json", {}) + outputs = read_json(phase / "output-summary.json", {}) + requested = spec.get("households", 0) + sample_matches = ( + not requested or outputs.get("households", {}).get("rows") == requested + ) + valid = ( + sample_matches + and status.get("returncode") == 0 + and docker.get("ExitCode") == 0 + and bool(memory) + and bool(components) + and not docker.get("OOMKilled") + and not list(phase.glob("cache-miss-*.txt")) + ) + return { + "spec": spec, + "components": components, + "component_windows": component_windows(observations, phase), + "memory": memory, + "status": status, + "valid": valid, + "inputs": read_json(phase / "input-summary.json", {}), + "outputs": outputs, + "peak": max((r["peak_bytes"] for r in memory), default=0), + "docker": docker, + "path": str(directory), + } + + +def escape(value): + return html.escape(str(value), quote=True) + + +def memory_chart(run, xmax, ymax): + """Standalone SVG uses common axes across experiments for fair comparison.""" + points = " ".join( + f"{55 + 620 * r['elapsed_seconds'] / xmax:.2f},{235 - 200 * r['current_bytes'] / ymax:.2f}" + for r in run["memory"] + ) + # Mapped skim pages can fill the cgroup with reclaimable file cache. Show + # anonymous plus shared allocations separately without calling them total RAM. + allocation_line = "" + allocation_legend = "" + if run["memory"] and all("anonymous_bytes" in row for row in run["memory"]): + allocations = " ".join( + f"{55 + 620 * row['elapsed_seconds'] / xmax:.2f},{235 - 200 * (row['anonymous_bytes'] + row.get('shared_bytes', 0)) / ymax:.2f}" + for row in run["memory"] + ) + allocation_line = f'' + allocation_legend = "

Blue: total container memory. Green dashed: anonymous + shared memory (excludes mapped skim pages and other file cache).

" + ticks = "".join( + f'{ymax * i / 4 / 2**30:.1f}' + f'' + f'{xmax * i / 4:.0f}' + for i in range(5) + ) + bands = [] + for window in run["component_windows"]: + start, end = window["start_seconds"], window["end_seconds"] + # Clip to the sampled chart domain; never bridge gaps between workers. + left, right = min(start, xmax), min(end, xmax) + title = ( + f"{window['process']}: {start:.3f}–{end:.3f} s " + f"({window['source']}; {'completed' if window['succeeded'] else 'failed'})" + ) + bands.append( + f'" + ) + return ( + '
' + f'{ticks}{"".join(bands)}' + f'{allocation_line}' + 'GiBElapsed seconds' + f"{allocation_legend}" + '

Choose a component to highlight its worker windows.

' + ) + + +def runtime_chart(runs, components): + """Grouped horizontal bars compare component means with population SD whiskers.""" + maximum = ( + max( + (c["mean"] + c["sd"] for run in runs for c in run["components"].values()), + default=1, + ) + or 1 + ) + rows = [] + y = 30 + for name in components: + rows.append(f'{escape(name)}') + for i, run in enumerate(runs): + value = run["components"].get(name) + if value: + scale = 530 / maximum + width = value["mean"] * scale + lo, hi = ( + max(0, value["mean"] - value["sd"]) * scale, + (value["mean"] + value["sd"]) * scale, + ) + rows.append( + f'{escape(run["spec"]["label"])}: {value["mean"]:.3f} ± {value["sd"]:.3f} s' + ) + y += 17 + y += 10 + return f'0 seconds{maximum:.1f} s{"".join(rows)}' + + +def experiment_card(run, xmax, ymax): + """Present the primary settings and counts before detailed provenance.""" + spec = run["spec"] + settings = "".join( + f"{escape(key.replace('_', ' '))}{escape(spec.get(key, 'unavailable'))}" + for key in ( + "sources", + "profile_name", + "abench_version", + "activitysim_commit", + "sharrow_commit", + "multiprocess", + "processes", + "sharrow", + "households", + "data_dir", + "output_dir", + "memory", + "shm_size", + "interval", + "platform", + "compare", + "config_overlay", + "cache_from", + ) + ) + counts = [] + for name in dict.fromkeys( + [ + "households", + "persons", + "land_use", + "tours", + "trips", + *run["inputs"], + *run["outputs"], + ] + ): + values = [run[key].get(name, {}).get("rows") for key in ("inputs", "outputs")] + cells = "".join( + f"{value:,}" if value is not None else "—" + for value in values + ) + counts.append( + f"{escape(spec.get('profile', {}).get('zone_label', 'zones') if name == 'land_use' else name)}{cells}" + ) + taz_cells = "".join( + f"{run[key].get('land_use', {}).get('taz_count', '—')}" + for key in ("inputs", "outputs") + ) + counts.append(f"TAZs{taz_cells}") + elapsed = run["status"].get("elapsed_seconds") + elapsed = f"{elapsed:.3f}" if elapsed is not None else "unavailable" + failure = f"

{escape(spec['failure'])}

" if spec.get("failure") else "" + details = "".join( + f"
{title}
{escape(json.dumps(value, indent=2))}
" + for title, value in ( + ("Complete settings and provenance", spec), + ("Input totals and categories", run["inputs"]), + ("Output totals and categories", run["outputs"]), + ("Container exit and OOM status", run["docker"]), + ) + ) + return ( + f"

{escape(spec['label'])}

{'SUCCEEDED' if run['valid'] else 'FAILED / INCOMPLETE'}" + f" · elapsed: {elapsed} s · peak: {run['peak'] / 2**30:.3f} GiB

{failure}" + f"{memory_chart(run, xmax, ymax)}

Experiment settings

{settings}
" + f"

Population and outputs

{''.join(counts)}
TableInput rowsOutput rows
" + f"

Output households and persons are the realized sample.

{details}
" + ) + + +def comparison_notes(runs): + """Surface workload/environment differences without claiming causal speedups.""" + notes = [] + fields = ( + "profile_name", + "sources", + "households", + "multiprocess", + "processes", + "sharrow", + "config_sha256", + "input_files", + "docker", + ) + for field in fields: + values = [json.dumps(run["spec"].get(field), sort_keys=True) for run in runs] + if len(set(values)) > 1: + notes.append(field.replace("_", " ")) + if not notes: + return "" + return ( + "

Comparison differences: " + + escape(", ".join(notes)) + + ". Inspect provenance before attributing differences to a source revision.

" + ) + + +def report(directories, destination): + """Generate a portable, offline HTML report and normalized comparison data.""" + runs = [load_run(path) for path in directories] + components = list(dict.fromkeys(name for run in runs for name in run["components"])) + headers = "".join(f"{escape(r['spec']['label'])}" for r in runs) + table = [] + for name in components: + eligible = [ + r["components"][name]["mean"] + for r in runs + if r["valid"] and name in r["components"] + ] + fastest = min(eligible, default=None) + cells = [] + for run in runs: + c = run["components"].get(name) + if c is None: + cells.append("—") + continue + winner = ( + run["valid"] + and fastest is not None + and math.isclose(c["mean"], fastest, rel_tol=1e-9) + ) + cells.append( + f'{c["mean"]:.3f} ± {c["sd"]:.3f} s
n={c["n"]}; max={c["maximum"]:.3f} s' + ) + table.append(f"{escape(name)}{''.join(cells)}") + xmax = ( + max((row["elapsed_seconds"] for r in runs for row in r["memory"]), default=1) + or 1 + ) + ymax = ( + max((row["current_bytes"] for r in runs for row in r["memory"]), default=1) or 1 + ) + cards = [experiment_card(run, xmax, ymax) for run in runs] + legend = " ".join( + f'■ {escape(r["spec"]["label"])}' + for i, r in enumerate(runs) + ) + selectable = list( + dict.fromkeys( + components + + [ + window["component"] + for run in runs + for window in run["component_windows"] + ] + ) + ) + options = "".join( + f'' for name in selectable + ) + # Component names live only in escaped HTML attributes/text, never in JS. + selector_script = """ + +""" + document = f"""abench report + +

abench report

Whole-container cgroup v2 memory counts shared pages once, including file cache, kernel memory and the supervisor. Swap is recorded separately in memory.csv. Cache preparation and post-run summaries are excluded. Peak is the kernel high-water mark sampled during the model lifetime, including container startup. Memory panels use identical axes.

+{comparison_notes(runs)} + + +

Selection applies to every memory chart. Each translucent band is one worker execution; darker overlaps indicate concurrent workers. Gaps remain unshaded. Hover over a band for its worker and time range.

+ +
{"".join(cards)}

Component runtimes

Mean ± population standard deviation across worker executions, with observation count and maximum. These describe worker imbalance, not uncertainty across repeated experiments. Component timings exclude pipeline checkpoint writes; elapsed time includes startup, I/O and coordination. Parallel component times must not be summed to estimate wall time. Green cells mark the fastest successful experiment's mean; failed runs are excluded from winners.

{headers}{"".join(table)}
Component

Runtime comparison

{legend}

{runtime_chart(runs, components)}
{selector_script}""" + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_text(document) + write_json(destination.with_suffix(".json"), runs) diff --git a/src/abench/runtime/__init__.py b/src/abench/runtime/__init__.py new file mode 100644 index 0000000..bdc50ef --- /dev/null +++ b/src/abench/runtime/__init__.py @@ -0,0 +1 @@ +"""Standalone files copied into each experiment and mounted in its containers.""" diff --git a/src/abench/runtime/build_sources.py b/src/abench/runtime/build_sources.py new file mode 100644 index 0000000..a531de6 --- /dev/null +++ b/src/abench/runtime/build_sources.py @@ -0,0 +1,109 @@ +"""Build exact source checkouts into wheels and resolve them together in Docker.""" + +import email +import importlib.metadata +import json +import re +import subprocess +import sys +import zipfile +from pathlib import Path + + +def command(args): + """Use argument arrays throughout; profile values are never shell fragments.""" + return subprocess.check_output(args, text=True).strip() + + +def normalized(name): + return re.sub(r"[-_.]+", "-", name).lower() + + +def install(manifest, directory=Path("/opt")): + """Verify Git objects and wheel identities before installing all overrides.""" + directory.mkdir(parents=True, exist_ok=True) + constraints = directory / "constraints.txt" + constraints.write_text("\n".join(manifest.get("constraints", [])) + "\n") + wheels = [] + provenance = [] + for item in manifest["sources"]: + checkout = directory / "sources" / item["name"] + checkout.mkdir(parents=True) + command(["git", "init", str(checkout)]) + url = f"https://github.com/{item['repository']}.git" + command(["git", "-C", str(checkout), "remote", "add", "origin", url]) + command( + ["git", "-C", str(checkout), "fetch", "--tags", "origin", item["commit"]] + ) + command(["git", "-C", str(checkout), "checkout", "--detach", "FETCH_HEAD"]) + resolved = command(["git", "-C", str(checkout), "rev-parse", "HEAD"]) + if resolved != item["commit"]: + raise ValueError(f"Git commit mismatch for {item['name']}") + package = (checkout / item.get("subdirectory", "")).resolve() + if not package.is_relative_to(checkout.resolve()): + raise ValueError("package subdirectory escapes checkout") + wheel_dir = directory / "wheels" / item["name"] + command( + [ + sys.executable, + "-m", + "pip", + "wheel", + "--no-deps", + "--wheel-dir", + str(wheel_dir), + str(package), + ] + ) + built = list(wheel_dir.glob("*.whl")) + if len(built) != 1: + raise ValueError(f"expected one wheel for {item['name']}") + with zipfile.ZipFile(built[0]) as archive: + metadata_path = next( + n for n in archive.namelist() if n.endswith(".dist-info/METADATA") + ) + metadata = email.message_from_bytes(archive.read(metadata_path)) + if normalized(metadata["Name"]) != item["name"]: + raise ValueError( + f"wrong distribution: expected {item['name']}, got {metadata['Name']}" + ) + extras = "[" + ",".join(item["extras"]) + "]" if item.get("extras") else "" + wheels.append(str(built[0]) + extras) + provenance.append( + dict(item, version=metadata["Version"], resolved_commit=resolved) + ) + # Supplying all local wheels in one transaction prevents dependency resolution + # from quietly replacing a requested source package with a registry release. + requirements = ["pyyaml", "pyarrow", *manifest.get("requirements", [])] + command( + [ + sys.executable, + "-m", + "pip", + "install", + "--no-cache-dir", + "-c", + str(constraints), + *wheels, + *requirements, + ] + ) + command([sys.executable, "-m", "pip", "check"]) + for item in provenance: + dist = importlib.metadata.distribution(item["name"]) + direct = json.loads(dist.read_text("direct_url.json") or "{}") + expected = (directory / "wheels" / item["name"]).resolve().as_uri() + "/" + if dist.version != item["version"] or not direct.get("url", "").startswith( + expected + ): + raise ValueError(f"installed source identity mismatch: {item['name']}") + (directory / "source-provenance.json").write_text( + json.dumps(provenance, indent=2) + "\n" + ) + (directory / "pip-freeze.txt").write_text( + command([sys.executable, "-m", "pip", "freeze"]) + "\n" + ) + + +if __name__ == "__main__": + install(json.loads(Path(sys.argv[1]).read_text())) diff --git a/src/abench/runtime/instrumentation.py b/src/abench/runtime/instrumentation.py new file mode 100644 index 0000000..fdd7199 --- /dev/null +++ b/src/abench/runtime/instrumentation.py @@ -0,0 +1,61 @@ +"""Process-local component timing and a strict Sharrow disk-cache guard.""" + +import json +import multiprocessing +import os +import time +from pathlib import Path + + +def install(flow_cache=Path("/results/cache/flows")): + """Install in the model parent and every spawned worker, before model imports.""" + from activitysim.core.workflow import runner + + original = runner.run_named_step + records = Path(os.environ["BENCH_PHASE_DIR"]) + + def timed(name, context, **kwargs): + started = time.perf_counter() + succeeded = False + try: + result = original(name, context, **kwargs) + succeeded = True + return result + finally: + finished = time.perf_counter() + row = { + "component": name, + "seconds": finished - started, + "process": multiprocessing.current_process().name, + "pid": os.getpid(), + "succeeded": succeeded, + } + # Linux's monotonic clock is shared across processes. Use the + # supervisor's origin so worker windows align with memory samples. + if "BENCH_STARTED_MONOTONIC" in os.environ: + origin = float(os.environ["BENCH_STARTED_MONOTONIC"]) + row.update( + start_seconds=started - origin, end_seconds=finished - origin + ) + # Separate files avoid interleaving writes from concurrent workers. + with (records / f"components-{os.getpid()}.jsonl").open("a") as stream: + stream.write(json.dumps(row) + "\n") + + runner.run_named_step = timed + if os.environ.get("BENCH_STRICT_CACHE") == "1": + from numba.core.dispatcher import _FunctionCompiler + + compile_original = _FunctionCompiler.compile + flow_cache = flow_cache.resolve() + + def compile_checked(self, *args, **kwargs): + # Numba reaches this method only after failing to load a compiled + # overload from disk. Ordinary ActivitySim/Numba JIT is still allowed. + filename = Path(self.py_func.__code__.co_filename).resolve() + if filename.is_relative_to(flow_cache): + with (records / f"cache-miss-{os.getpid()}.txt").open("a") as stream: + stream.write(str(filename) + "\n") + raise RuntimeError(f"Measured Sharrow flow cache miss: {filename}") + return compile_original(self, *args, **kwargs) + + _FunctionCompiler.compile = compile_checked diff --git a/src/abench/runtime/worker.py b/src/abench/runtime/worker.py new file mode 100644 index 0000000..3cedcc6 --- /dev/null +++ b/src/abench/runtime/worker.py @@ -0,0 +1,255 @@ +"""Linux container supervisor; the model runs in a separate process tree.""" + +import csv +import json +import os +import subprocess +import sys +import time +from pathlib import Path + +# multiprocessing's spawn imports this file again as __mp_main__. +if os.environ.get("BENCH_MODEL") == "1": + from instrumentation import install + + install() + + +def write_json(path, value): + path.write_text(json.dumps(value, indent=2) + "\n") + + +def make_state( + spec, + phase, + model_root=Path("/model"), + data_root=Path("/data"), + results_root=Path("/results"), +): + """Resolve model configs < profile defaults < overlays < explicit CLI controls. + + A generated inheriting config carries defaults through worker reconstruction; + direct settings overrides are reserved for the experiment's required controls. + """ + import activitysim.abm # noqa: F401 -- register standard components + import yaml + from activitysim.core.workflow import State + + profile = spec["profile"] + defaults = dict(profile.get("settings", {})) + if profile.get("models_from"): + base = yaml.safe_load((model_root / profile["models_from"]).read_text()) + defaults["models"] = [ + name + for name in base["models"] + if name not in profile.get("exclude_models", []) + ] + if spec["multiprocess"] and profile.get("mp_settings"): + mp = yaml.safe_load((model_root / profile["mp_settings"]).read_text()) + defaults["multiprocess_steps"] = mp["multiprocess_steps"] + generated = phase / "profile-config" + generated.mkdir(exist_ok=True) + for name in ("settings.yaml", "settings_mp.yaml", "settings_mp_sharrow.yaml"): + (generated / name).write_text( + yaml.safe_dump(dict(defaults, inherit_settings=True)) + ) + configs = [ + model_root / f"overlay-{i}" for i in range(len(spec.get("config_overlay", []))) + ] + configs += [generated] + if spec["multiprocess"]: + configs += [model_root / name for name in profile.get("mp_configs", [])] + configs += [model_root / name for name in profile["configs"]] + state = State.make_default( + working_dir=model_root, + configs_dir=configs, + data_dir=data_root, + output_dir=phase / "output", + cache_dir=results_root / "cache/model", + settings=dict( + households_sample_size=spec["households"], + multiprocess=spec["multiprocess"], + num_processes=spec["processes"], + sharrow="require" if spec["sharrow"] else False, + fail_fast=True, + ), + ) + # Every sliced phase honors the requested count, including phases that have + # their own worker count in production configs or explicit chunk overlays. + for step in state.settings.multiprocess_steps or []: + if isinstance(step, dict): + if "slice" in step: + step["num_processes"] = spec["processes"] + elif getattr(step, "slice", None) is not None: + step.num_processes = spec["processes"] + state.filesystem.sharrow_cache_dir = results_root / "cache/flows" + state.settings.sharrow_cache_dir = str(state.filesystem.sharrow_cache_dir) + sys.path.insert(0, str(model_root)) + state.set("imported_extensions", []) + for extension in profile.get("extensions", []): + state.import_extensions(extension) + if profile.get("adapter"): + import importlib + + module, function = profile["adapter"].split(":") + getattr(importlib.import_module(module), function)(state, spec, phase) + state.set("run_timestamp", "benchmark") + state.set("run_id", str(state.tracing.run_id)) + return state + + +def run_model(spec, phase): + """Use identical samples, process layout and cache paths in both phases.""" + state = make_state(spec, phase) + state.logging.config_logger() + write_json( + phase / "effective-settings.json", state.settings.model_dump(mode="json") + ) + state.run.all(resume_after=None) + if not spec["multiprocess"]: + state.checkpoint.close_store() + + +def table_summary(directory, prefix="", tables=None): + """Stream input/output tables so full-population summaries need bounded RAM.""" + import pandas as pd + import pyarrow.parquet as pq + + result = {} + tables = tables or { + name: {} + for name in ( + "households", + "persons", + "land_use", + "tours", + "trips", + "joint_tour_participants", + "vehicles", + ) + } + for name, options in tables.items(): + path = directory / f"{prefix}{options.get('file', name)}.csv" + if path.exists(): + chunks = pd.read_csv(path, chunksize=100_000) + else: + path = directory / f"{prefix}{options.get('file', name)}.parquet" + if not path.exists(): + continue + chunks = ( + batch.to_pandas() for batch in pq.ParquetFile(path).iter_batches() + ) + info = {"rows": 0, "totals": {}, "categories": {}} + zones = set() + for chunk in chunks: + if name == "land_use": + for col in options.get("zone_columns", ["taz", "TAZ"]): + if col in chunk: + zones.update(chunk[col].dropna().unique().tolist()) + info["rows"] += len(chunk) + for col in options.get( + "totals", ["TOTPOP", "TOTHH", "TOTEMP", "pop", "hh", "emp_total"] + ): + if col in chunk: + info["totals"][col] = info["totals"].get(col, 0) + float( + chunk[col].sum() + ) + for col in options.get( + "categories", + ( + "tour_category", + "tour_type", + "tour_mode", + "trip_mode", + "primary_purpose", + ), + ): + if col in chunk: + counts = info["categories"].setdefault(col, {}) + for value, count in chunk[col].value_counts(dropna=False).items(): + counts[str(value)] = counts.get(str(value), 0) + int(count) + if zones: + info["taz_count"] = len(zones) + result[name] = info + return result + + +def sample_memory(root, elapsed): + """cgroup v2 charges shared pages once; file includes shmem, not vice versa.""" + stat = dict( + line.split() for line in (root / "memory.stat").read_text().splitlines() + ) + return { + "elapsed_seconds": elapsed, + "current_bytes": int((root / "memory.current").read_text()), + "peak_bytes": int((root / "memory.peak").read_text()), + "swap_bytes": int((root / "memory.swap.current").read_text()), + "anonymous_bytes": int(stat.get("anon", 0)), + "file_bytes": int(stat.get("file", 0)), + "shared_bytes": int(stat.get("shmem", 0)), + } + + +def supervise(spec, phase): + """Measure only the model subprocess lifetime; summarize after sampling ends.""" + root = Path("/sys/fs/cgroup") + if not (root / "memory.peak").exists(): + raise RuntimeError("A private cgroup v2 with memory.peak is required") + (phase / "output").mkdir() + for name in ("flows", "model"): + Path(f"/results/cache/{name}").mkdir(parents=True, exist_ok=True) + env = dict(os.environ, BENCH_MODEL="1", BENCH_PHASE_DIR=str(phase)) + env["BENCH_STRICT_CACHE"] = ( + "1" if spec["sharrow"] and phase.name == "measured" else "0" + ) + started = time.perf_counter() + env["BENCH_STARTED_MONOTONIC"] = str(started) + with (phase / "memory.csv").open("w", buffering=1) as stream: + writer = csv.DictWriter(stream, fieldnames=sample_memory(root, 0).keys()) + writer.writeheader() + writer.writerow(sample_memory(root, 0)) + process = subprocess.Popen( + [sys.executable, __file__, "model", str(phase)], env=env + ) + while True: + writer.writerow(sample_memory(root, time.perf_counter() - started)) + if process.poll() is not None: + break + try: + process.wait(timeout=spec["interval"]) + except subprocess.TimeoutExpired: + pass + write_json( + phase / "status.json", + { + "returncode": process.returncode, + "elapsed_seconds": time.perf_counter() - started, + }, + ) + # This separate process keeps pandas/Arrow and summary allocations out of the + # measured lifetime; the host report uses only the recorded cgroup samples. + subprocess.run([sys.executable, __file__, "summary", str(phase)], check=True) + return process.returncode + + +if __name__ == "__main__": + mode, phase_arg = sys.argv[1:] + phase = Path(phase_arg) + spec = json.loads(Path("/results/experiment.json").read_text()) + if mode == "model": + run_model(spec, phase) + elif mode == "summary": + write_json( + phase / "input-summary.json", + table_summary(Path("/data"), tables=spec["profile"].get("input_tables")), + ) + write_json( + phase / "output-summary.json", + table_summary( + phase / "output", + spec["profile"].get("output_prefix", "final_"), + spec["profile"].get("output_tables"), + ), + ) + else: + sys.exit(supervise(spec, phase)) diff --git a/src/abench/sources.py b/src/abench/sources.py new file mode 100644 index 0000000..d10ef0e --- /dev/null +++ b/src/abench/sources.py @@ -0,0 +1,93 @@ +"""Normalize exact GitHub source overrides before executing any build commands.""" + +import re +from pathlib import PurePosixPath + + +def canonical_name(name): + """Compare distribution names using Python packaging's normalization rules.""" + return re.sub(r"[-_.]+", "-", name).lower() + + +def source(value): + """Accept a CLI shorthand or a profile mapping, including extras/subdirectory.""" + if isinstance(value, str): + match = re.fullmatch( + r"([A-Za-z0-9][A-Za-z0-9_.-]*)(?:\[([A-Za-z0-9_,.-]+)\])?=([A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+)@([0-9a-fA-F]{40})(?:#subdirectory=([^\s]+))?", + value, + ) + if not match: + raise ValueError( + "source must be distribution[extras]=organization/repository@full-SHA[#subdirectory=path]" + ) + name, extras, repository, commit, subdirectory = match.groups() + value = dict( + name=name, + repository=repository, + commit=commit, + extras=extras.split(",") if extras else [], + subdirectory=subdirectory or "", + ) + if not isinstance(value, dict) or set(value) - { + "name", + "repository", + "commit", + "extras", + "subdirectory", + }: + raise ValueError("invalid source fields") + result = dict(value) + for key, pattern in [ + ("name", r"[A-Za-z0-9][A-Za-z0-9_.-]*"), + ("repository", r"[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+"), + ("commit", r"[0-9a-fA-F]{40}"), + ]: + if not isinstance(result.get(key), str) or not re.fullmatch( + pattern, result[key] + ): + raise ValueError(f"invalid source {key}: {result.get(key)!r}") + result["name"] = canonical_name(result["name"]) + result["commit"] = result["commit"].lower() + result.setdefault("subdirectory", "") + result.setdefault("extras", []) + subdir = result["subdirectory"] + if ( + not isinstance(subdir, str) + or PurePosixPath(subdir).is_absolute() + or ".." in PurePosixPath(subdir).parts + or "\\" in subdir + ): + raise ValueError("source subdirectory must stay inside the checkout") + if not isinstance(result["extras"], list) or any( + not isinstance(e, str) or not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]*", e) + for e in result["extras"] + ): + raise ValueError("invalid package extras") + result["extras"] = sorted(set(result["extras"])) + return result + + +def resolve_sources(defaults, overrides, activitysim=None, sharrow=None): + """CLI entries replace profile pins; aliases cannot silently contradict them.""" + selected = {} + for value in defaults: + item = source(value) + if item["name"] in selected: + raise ValueError(f"duplicate profile source: {item['name']}") + selected[item["name"]] = item + explicit = {} + for value in overrides: + item = source(value) + if item["name"] in explicit: + raise ValueError(f"duplicate CLI source: {item['name']}") + explicit[item["name"]] = item + for name, sha in [("activitysim", activitysim), ("sharrow", sharrow)]: + if sha: + item = source(f"{name}=ActivitySim/{name}@{sha}") + if name in explicit and explicit[name] != item: + raise ValueError(f"conflicting --source and --{name}-commit") + explicit[name] = item + selected.update(explicit) + if "activitysim" not in selected: + raise ValueError("provide an exact ActivitySim source commit") + return [selected[name] for name in sorted(selected)] diff --git a/tests/fixtures/tiny/benchmark.yaml b/tests/fixtures/tiny/benchmark.yaml new file mode 100644 index 0000000..59a5211 --- /dev/null +++ b/tests/fixtures/tiny/benchmark.yaml @@ -0,0 +1,20 @@ +schema_version: 1 +name: CI extension model +configs: [configs] +snapshot: [configs, tiny_extension.py] +extensions: [tiny_extension] +data_dir: data +required_inputs: [households.csv] +settings: + models: [bench_initialize, bench_compute, bench_export] + checkpoints: true + rng_base_seed: 0 + chunk_size: 0 + trace_hh_id: null + trace_od: null +constraints: ['pandas<3', 'multimethod<2'] +input_tables: + households: {} +output_tables: + households: {} +adapter: tiny_extension:prepare diff --git a/tests/fixtures/tiny/configs/settings.yaml b/tests/fixtures/tiny/configs/settings.yaml new file mode 100644 index 0000000..052f22d --- /dev/null +++ b/tests/fixtures/tiny/configs/settings.yaml @@ -0,0 +1,15 @@ +models: [bench_initialize, bench_compute, bench_export] +multiprocess_steps: +- name: mp_initialize + begin: bench_initialize +- name: mp_households + begin: bench_compute + num_processes: 20 + slice: + tables: [households] +- name: mp_finalize + begin: bench_export +input_table_list: +- tablename: households + filename: households.csv + index_col: household_id diff --git a/tests/fixtures/tiny/data/households.csv b/tests/fixtures/tiny/data/households.csv new file mode 100644 index 0000000..21f2c9b --- /dev/null +++ b/tests/fixtures/tiny/data/households.csv @@ -0,0 +1,5 @@ +household_id,value +1,10 +2,20 +3,30 +4,40 diff --git a/tests/fixtures/tiny/tiny_extension.py b/tests/fixtures/tiny/tiny_extension.py new file mode 100644 index 0000000..9ea406b --- /dev/null +++ b/tests/fixtures/tiny/tiny_extension.py @@ -0,0 +1,54 @@ +"""Tiny real ActivitySim workflow: exercises registration, slicing, and caching.""" + +import importlib +import sys +from pathlib import Path + +import pandas as pd +from activitysim.core import workflow + + +@workflow.step +def bench_initialize(state: workflow.State): + """Seed four rows and a generated Numba module in the guarded flow directory.""" + households = pd.read_csv(Path("/data/households.csv"), index_col="household_id") + state.add_table("households", households) + cache = Path(state.settings.sharrow_cache_dir) + cache.mkdir(parents=True, exist_ok=True) + module = cache / "tiny_generated.py" + if not module.exists(): + module.write_text( + "from numba import njit\n@njit(cache=True)\ndef twice(x):\n return x * 2\n" + ) + + +@workflow.step +def bench_compute(state: workflow.State, households: pd.DataFrame): + """Run independently in each slice; measured cache misses must be rejected.""" + sys.path.insert(0, state.settings.sharrow_cache_dir) + twice = importlib.import_module("tiny_generated").twice + households = households.copy() + households["value"] = [twice(int(value)) for value in households["value"]] + state.add_table("households", households) + + +@workflow.step +def bench_export(state: workflow.State, households: pd.DataFrame): + """Export merged results so the harness validates the realized population.""" + households.to_csv(Path(state.filesystem.output_dir) / "final_households.csv") + + +def prepare(state, spec, phase): + """Supply tiny shared skims; this fixture has no destination/shadow models.""" + state.set("shadow_pricing_info", None) + state.set("shadow_pricing_choice_info", None) + state.set("network_los_preload", None) + if spec["sharrow"]: + import numpy as np + import sharrow as sh + + dataset = sh.Dataset({"distance": (("otaz", "dtaz"), np.ones((1, 1)))}) + state.set( + "skim_dataset", + dataset.shm.to_shared_memory("skim_dataset", mode="r", load=True), + ) diff --git a/tests/requirements.txt b/tests/requirements.txt new file mode 100644 index 0000000..99da3b9 --- /dev/null +++ b/tests/requirements.txt @@ -0,0 +1,6 @@ +# The measurement hook uses ActivitySim/Numba internals; pin real implementations +# in CI so unrelated upstream changes do not silently change the test contract. +activitysim @ git+https://github.com/ActivitySim/activitysim.git@5c6fae24a91a57a2d6dfc2e1dbe062a61d94545a +sharrow @ git+https://github.com/ActivitySim/sharrow.git@fc175b27d8e0c5d202721c67d96b050e6117b235 +multimethod<2 +pandas<3 diff --git a/tests/test_builder.py b/tests/test_builder.py new file mode 100644 index 0000000..932af91 --- /dev/null +++ b/tests/test_builder.py @@ -0,0 +1,82 @@ +"""Exercise source verification without network access or host package mutation.""" + +import json +import zipfile +from types import SimpleNamespace + +import pytest + +from abench.runtime import build_sources +from abench.sources import source + +SHA = "a" * 40 + + +def fake_commands(tmp_path, monkeypatch, *, resolved=SHA, metadata_name="addon"): + """Emulate Git/pip boundaries but pass real wheel metadata through the builder.""" + calls = [] + + def command(args): + calls.append(args) + if args[-2:] == ["rev-parse", "HEAD"]: + return resolved + if "wheel" in args: + wheel_dir = tmp_path / "wheels/addon" + wheel_dir.mkdir(parents=True) + with zipfile.ZipFile( + wheel_dir / "addon-1.0-py3-none-any.whl", "w" + ) as archive: + archive.writestr( + "addon-1.0.dist-info/METADATA", + f"Name: {metadata_name}\nVersion: 1.0\n", + ) + return "" + + monkeypatch.setattr(build_sources, "command", command) + monkeypatch.setattr( + build_sources.importlib.metadata, + "distribution", + lambda name: SimpleNamespace( + version="1.0", + read_text=lambda key: json.dumps( + {"url": (tmp_path / "wheels/addon/addon-1.0-py3-none-any.whl").as_uri()} + ), + ), + ) + return calls + + +def test_builder_verifies_and_installs_wheels_together(tmp_path, monkeypatch): + calls = fake_commands(tmp_path, monkeypatch) + item = source(f"addon[fast]=org/addon@{SHA}") + build_sources.install( + dict(sources=[item], requirements=["numpy"], constraints=["numpy<3"]), tmp_path + ) + install = next(c for c in calls if "install" in c) + assert any(value.endswith(".whl[fast]") for value in install) + assert "numpy" in install + assert (tmp_path / "constraints.txt").read_text() == "numpy<3\n" + assert ( + json.loads((tmp_path / "source-provenance.json").read_text())[0][ + "resolved_commit" + ] + == SHA + ) + + +@pytest.mark.parametrize( + "options,match", + [ + ({"resolved": "b" * 40}, "Git commit mismatch"), + ({"metadata_name": "different"}, "wrong distribution"), + ], +) +def test_builder_rejects_wrong_identity_before_install( + tmp_path, monkeypatch, options, match +): + calls = fake_commands(tmp_path, monkeypatch, **options) + with pytest.raises(ValueError, match=match): + build_sources.install( + dict(sources=[source(f"addon=org/addon@{SHA}")]), tmp_path + ) + assert not any("install" in c for c in calls) diff --git a/tests/test_docker.py b/tests/test_docker.py new file mode 100644 index 0000000..add8287 --- /dev/null +++ b/tests/test_docker.py @@ -0,0 +1,61 @@ +"""Opt-in integration through the real Docker supervisor and ActivitySim workflow.""" + +import json +import os +import shutil +from pathlib import Path + +import pytest + +from abench import cli +from abench.report import load_run + +pytestmark = [ + pytest.mark.docker, + pytest.mark.skipif( + os.environ.get("ABENCH_DOCKER_TESTS") != "1", + reason="set ABENCH_DOCKER_TESTS=1 to run Docker integration", + ), +] +ACTIVITYSIM = "5c6fae24a91a57a2d6dfc2e1dbe062a61d94545a" +SHARROW = "fc175b27d8e0c5d202721c67d96b050e6117b235" + + +@pytest.mark.parametrize("multiprocess,sharrow", [(False, False), (True, True)]) +def test_tiny_model(tmp_path, multiprocess, sharrow): + """Build pinned sources, run a full warmup where enabled, then measure.""" + root = tmp_path / "model" + shutil.copytree(Path(__file__).parent / "fixtures/tiny", root) + output = tmp_path / "experiment" + args = [ + "run", + "--model-dir", + str(root), + "--source", + f"activitysim=ActivitySim/activitysim@{ACTIVITYSIM}", + "--source", + f"sharrow=ActivitySim/sharrow@{SHARROW}", + "--households", + "4", + "--memory", + "3g", + "--shm-size", + "256m", + "--output-dir", + str(output), + ] + args += ( + ["--multiprocess", "--processes", "2"] if multiprocess else ["--single-process"] + ) + args += ["--sharrow"] if sharrow else ["--no-sharrow"] + assert cli.main(args) == 0 + run = load_run(output) + assert run["valid"] + assert run["components"]["bench_compute"]["n"] == (2 if multiprocess else 1) + assert run["outputs"]["households"]["rows"] == 4 + assert run["memory"] and all(row["current_bytes"] > 0 for row in run["memory"]) + assert (output / "warmup").exists() == sharrow + assert len(json.loads((output / "source-provenance.json").read_text())) == 2 + assert not list((output / "measured").glob("cache-miss-*")) + lines = (output / "measured/output/final_households.csv").read_text().splitlines() + assert set(lines[1:]) == {"1,20", "2,40", "3,60", "4,80"} diff --git a/tests/test_measurement.py b/tests/test_measurement.py new file mode 100644 index 0000000..9dca0d7 --- /dev/null +++ b/tests/test_measurement.py @@ -0,0 +1,364 @@ +"""Measurement/report regression tests; no running Docker daemon required.""" + +import csv +import json +import os +import re +import shutil +import subprocess +import sys +import tempfile +import time +import unittest +from html.parser import HTMLParser +from pathlib import Path + +from abench import report as benchmark +from abench.cli import commit, positive +from abench.runtime import worker + +benchmark.commit = commit +benchmark.positive = positive +HERE = Path(worker.__file__).parent + + +class BenchmarkTests(unittest.TestCase): + def make_run(self, root, label, times, success=True): + """Construct raw artifacts with the same schema emitted by workers.""" + root.mkdir() + phase = root / "measured" + phase.mkdir() + benchmark.write_json( + root / "experiment.json", {"schema_version": 1, "label": label} + ) + benchmark.write_json( + phase / "status.json", {"returncode": 0, "elapsed_seconds": 9} + ) + benchmark.write_json( + phase / "docker-state.json", {"ExitCode": 0, "OOMKilled": not success} + ) + for i, duration in enumerate(times): + (phase / f"components-{i}.jsonl").write_text( + json.dumps( + { + "component": "auto_ownership", + "seconds": duration, + "succeeded": True, + } + ) + + "\n" + ) + # A duplicate locutor CSV must never contribute another observation. + (phase / "timing_log.csv").write_text( + "model_name,seconds\nauto_ownership,999\n" + ) + with (phase / "memory.csv").open("w") as stream: + writer = csv.DictWriter( + stream, fieldnames=["elapsed_seconds", "current_bytes", "peak_bytes"] + ) + writer.writeheader() + writer.writerows( + [ + {"elapsed_seconds": 0, "current_bytes": 100, "peak_bytes": 100}, + {"elapsed_seconds": 9, "current_bytes": 200, "peak_bytes": 300}, + ] + ) + return root + + def test_requested_sample_must_be_realized(self): + """Do not call a fallback to a smaller input population a valid trial.""" + with tempfile.TemporaryDirectory() as tmp: + run = self.make_run(Path(tmp) / "sample", "sample", [1]) + benchmark.write_json( + run / "experiment.json", {"schema_version": 1, "households": 500000} + ) + benchmark.write_json( + run / "measured/output-summary.json", {"households": {"rows": 28365}} + ) + self.assertFalse(benchmark.load_run(run)["valid"]) + + def test_worker_statistics_and_failed_run_winners(self): + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + a = self.make_run(root / "a", "A ", document, re.DOTALL)[1] + controller_test = r""" +const assert = require('node:assert/strict'); +const vm = require('node:vm'); +const input = JSON.parse(require('node:fs').readFileSync(0, 'utf8')); +let change; +const selector = {value: '', addEventListener(type, fn) {assert.equal(type, 'change'); change = fn;}}; +const panels = input.panels.map(attrs => { + const bands = attrs.map(a => ({dataset: {component: a['data-component'], source: a['data-source']}, style: {display: 'none'}})); + const status = {textContent: ''}; + return {bands, status, querySelectorAll() {return bands;}, querySelector() {return status;}}; +}); +vm.runInNewContext(input.script, {document: { + getElementById(id) {assert.equal(id, 'memory-component'); return selector;}, + querySelectorAll() {return panels;} +}}); +selector.value = input.name; +change(); +assert.equal(panels[0].bands.filter(b => b.style.display === '').length, 3); +assert.match(panels[0].status.textContent, /3 worker execution windows/); +assert.match(panels[1].status.textContent, /No execution-window data/); +selector.value = ''; +change(); +assert.ok(panels.every(p => p.bands.every(b => b.style.display === 'none'))); +""" + subprocess.run( + [shutil.which("node"), "-e", controller_test], + input=json.dumps( + {"panels": parsed.panels, "name": name, "script": script} + ), + text=True, + check=True, + ) + + def test_commit_and_interval_validation(self): + import argparse + + for value in ("main", "1234567", "x" * 40): + with self.assertRaises(argparse.ArgumentTypeError): + benchmark.commit(value) + for value in ("nan", "inf", "0", "-1"): + with self.assertRaises(argparse.ArgumentTypeError): + benchmark.positive(value) + + def test_table_summaries(self): + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + (root / "households.csv").write_text("HHID,PERSONS\n1,2\n2,1\n") + (root / "land_use.csv").write_text("TAZ,TOTPOP,TOTEMP\n1,20,10\n2,30,15\n") + (root / "trips.csv").write_text( + "trip_id,trip_mode\n1,WALK\n2,WALK\n3,SOV\n" + ) + summary = worker.table_summary(root) + self.assertEqual(summary["households"]["rows"], 2) + self.assertEqual(summary["land_use"]["totals"]["TOTPOP"], 50) + self.assertEqual( + summary["trips"]["categories"]["trip_mode"], {"WALK": 2, "SOV": 1} + ) + + def test_real_numba_disk_hit_allowed_and_miss_blocked(self): + """A fresh interpreter must load warmed overloads but reject new ones.""" + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + (root / "generated_flow.py").write_text( + "from numba import njit\n@njit(cache=True)\ndef f(x):\n return x + 1\n" + ) + env = dict( + os.environ, + PYTHONPATH=os.pathsep.join([str(HERE), str(root)]), + BENCH_PHASE_DIR=str(root), + BENCH_STRICT_CACHE="1", + ) + warm = subprocess.run( + [ + sys.executable, + "-c", + "from generated_flow import f; assert f(1) == 2", + ], + env=env, + check=False, + capture_output=True, + text=True, + ) + self.assertEqual(warm.returncode, 0, warm.stderr) + script = f"from pathlib import Path; from instrumentation import install; install(Path({str(root)!r})); from generated_flow import f; " + hit = subprocess.run( + [sys.executable, "-c", script + "assert f(1) == 2"], + env=env, + check=False, + capture_output=True, + text=True, + ) + self.assertEqual(hit.returncode, 0, hit.stderr) + miss = subprocess.run( + [sys.executable, "-c", script + "f(1.5)"], + env=env, + check=False, + capture_output=True, + text=True, + ) + self.assertNotEqual(miss.returncode, 0) + self.assertIn("Measured Sharrow flow cache miss", miss.stderr) + self.assertTrue(list(root.glob("cache-miss-*.txt"))) + + def test_spawned_workers_each_record_components(self): + """Exercise the same top-level hook installation used by MP workers.""" + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + script = root / "spawn_check.py" + script.write_text( + "from activitysim.core.workflow import runner\n" + "runner.run_named_step = lambda name, context: 'ok'\n" + "import worker\n" + "import multiprocessing\n" + "def task():\n" + " assert runner.run_named_step('smoke', {}) == 'ok'\n" + "if __name__ == '__main__':\n" + " ctx = multiprocessing.get_context('spawn')\n" + " processes = [ctx.Process(target=task) for _ in range(2)]\n" + " for p in processes: p.start()\n" + " for p in processes:\n" + " p.join()\n" + " assert p.exitcode == 0\n" + ) + env = dict( + os.environ, + PYTHONPATH=str(HERE), + BENCH_MODEL="1", + BENCH_PHASE_DIR=str(root), + BENCH_STRICT_CACHE="0", + BENCH_STARTED_MONOTONIC=str(time.perf_counter()), + ) + result = subprocess.run( + [sys.executable, str(script)], + env=env, + capture_output=True, + text=True, + check=False, + ) + self.assertEqual(result.returncode, 0, result.stderr) + records = [ + json.loads(path.read_text()) for path in root.glob("components-*.jsonl") + ] + self.assertEqual(len({row["pid"] for row in records}), 2) + self.assertTrue( + all(row["succeeded"] and row["component"] == "smoke" for row in records) + ) + for row in records: + self.assertGreater(row["start_seconds"], 0) + self.assertAlmostEqual( + row["end_seconds"] - row["start_seconds"], row["seconds"] + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_profiles_sources.py b/tests/test_profiles_sources.py new file mode 100644 index 0000000..5ba4a9e --- /dev/null +++ b/tests/test_profiles_sources.py @@ -0,0 +1,165 @@ +"""Contracts that let a new model or source dependency reuse the same harness.""" + +import json + +import pytest +import yaml + +from abench import cli +from abench.profiles import load_profile, validate_model +from abench.runtime.worker import make_state +from abench.sources import resolve_sources, source + +SHA = "a" * 40 +OTHER = "b" * 40 + + +def test_source_overrides_and_alias_conflicts(): + defaults = [f"activitysim=ActivitySim/activitysim@{SHA}", f"My_Ext=org/old@{SHA}"] + resolved = resolve_sources( + defaults, [f"my-ext[fast]=org/new@{OTHER}#subdirectory=python/package"] + ) + assert resolved[1] == dict( + name="my-ext", + repository="org/new", + commit=OTHER, + subdirectory="python/package", + extras=["fast"], + ) + with pytest.raises(ValueError, match="conflicting"): + resolve_sources([], [f"activitysim=fork/activitysim@{SHA}"], SHA) + with pytest.raises(ValueError, match="duplicate"): + resolve_sources([], [f"activitysim=org/a@{SHA}"] * 2) + + +@pytest.mark.parametrize( + "value", + [ + "x=org/repo@main", + f"x=org/repo@{SHA}#subdirectory=../bad", + f"x=org/repo@{SHA}#subdirectory=/bad", + f"x=org/repo;echo@{SHA}", + dict(name="x", repository="org/repo", commit=SHA, extras=["x;echo"]), + ], +) +def test_invalid_source_rejected(value): + with pytest.raises(ValueError): + source(value) + + +def make_model(root): + """A self-contained profile is independent of either example repository.""" + (root / "configs").mkdir() + (root / "data").mkdir() + (root / "data/households.csv").write_text("household_id\n1\n2\n") + (root / "configs/settings.yaml").write_text( + "chunk_size: 100\nmodels: []\nmultiprocess_steps:\n- name: mp_households\n begin: test_step\n num_processes: 20\n slice:\n tables: [households]\n" + ) + profile = dict( + schema_version=1, + name="tiny", + configs=["configs"], + snapshot=["configs"], + settings={"chunk_size": 200}, + required_inputs=["households.csv"], + sources=[f"activitysim=ActivitySim/activitysim@{SHA}"], + ) + (root / "benchmark.yaml").write_text(yaml.safe_dump(profile)) + return profile + + +def test_profiles_preflight_and_oversampling(tmp_path): + make_model(tmp_path) + profile = load_profile("benchmark.yaml", tmp_path) + validate_model(profile, tmp_path, tmp_path / "data", 2) + with pytest.raises(ValueError, match="only 2"): + validate_model(profile, tmp_path, tmp_path / "data", 3) + (tmp_path / "data/households.csv").unlink() + with pytest.raises(ValueError, match="missing required"): + validate_model(profile, tmp_path, tmp_path / "data", 0) + for name in ("mtc", "sandag"): + assert load_profile(name, tmp_path)["configs"] + + +def test_config_overlay_and_process_precedence(tmp_path): + profile = make_model(tmp_path) + overlay = tmp_path / "overlay-0" + overlay.mkdir() + (overlay / "settings.yaml").write_text( + "inherit_settings: true\nchunk_size: 300\nchunk_training_mode: explicit\nnum_processes: 99\n" + ) + phase = tmp_path / "measured" + phase.mkdir() + (phase / "output").mkdir() + spec = dict( + profile=profile, + households=2, + multiprocess=True, + processes=4, + sharrow=False, + config_overlay=[str(overlay)], + ) + state = make_state(spec, phase, tmp_path, tmp_path / "data", tmp_path) + assert state.settings.chunk_size == 300 + assert state.settings.chunk_training_mode == "explicit" + assert state.settings.num_processes == 4 + assert state.settings.multiprocess_steps[0].num_processes == 4 + assert "chunk_size: 100" in (tmp_path / "configs/settings.yaml").read_text() + + +def test_validate_does_not_build_or_write(tmp_path, monkeypatch, capsys): + make_model(tmp_path) + calls = [] + + def command(args, log=None): + calls.append(args) + return json.dumps(dict(OSType="linux", CgroupVersion="2")) + + monkeypatch.setattr(cli, "command", command) + assert ( + cli.main( + [ + "validate", + "--model-dir", + str(tmp_path), + "--no-sharrow", + "--households", + "2", + ] + ) + == 0 + ) + assert calls == [["docker", "info", "--format", "{{json .}}"]] + assert json.loads(capsys.readouterr().out)["sources"][0]["commit"] == SHA + assert not (tmp_path / "measured").exists() + + +def test_cli_failure_keeps_report(tmp_path, monkeypatch): + make_model(tmp_path) + output = tmp_path / "experiment" + + def command(args, log=None): + if args[:2] == ["docker", "info"]: + return json.dumps(dict(OSType="linux", CgroupVersion="2")) + raise RuntimeError("deliberate build failure") + + monkeypatch.setattr(cli, "command", command) + with pytest.raises(RuntimeError, match="deliberate"): + cli.main( + [ + "run", + "--model-dir", + str(tmp_path), + "--no-sharrow", + "--households", + "2", + "--output-dir", + str(output), + ] + ) + assert (output / "report.html").is_file() + assert ( + json.loads((output / "experiment.json").read_text())["failure"]["phase"] + == "build" + ) + assert (output / "runner/build_sources.py").is_file() From 6a9c35f5024c5c479f538826ce64526f1b93fbb7 Mon Sep 17 00:00:00 2001 From: Jeff Newman Date: Sun, 13 Sep 2026 15:13:00 -0500 Subject: [PATCH 2/8] Add named YAML experiment suites with shared defaults Support variable expansion, per-run overrides, suite preflight, sequential execution, and comparison reports. Include a SANDAG example, tests, documentation, and expanded CI artifacts. --- .github/workflows/ci.yml | 15 +- MANIFEST.in | 1 + README.md | 73 ++++++++++ examples/sandag-chunked.yaml | 30 ++++ src/abench/cli.py | 19 ++- src/abench/experiments.py | 258 +++++++++++++++++++++++++++++++++++ tests/test_docker.py | 37 +++++ tests/test_experiments.py | 180 ++++++++++++++++++++++++ 8 files changed, 604 insertions(+), 9 deletions(-) create mode 100644 examples/sandag-chunked.yaml create mode 100644 src/abench/experiments.py create mode 100644 tests/test_experiments.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index df4d8d4..23f2a24 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -58,12 +58,11 @@ jobs: with: name: docker-experiments path: | - experiments/ci/**/experiment/*.html - experiments/ci/**/experiment/*.json - experiments/ci/**/experiment/*.txt - experiments/ci/**/experiment/*.log - experiments/ci/**/experiment/measured/*.json* - experiments/ci/**/experiment/measured/*.csv - experiments/ci/**/experiment/measured/*.log - experiments/ci/**/experiment/warmup/*.log + experiments/ci/**/*.html + experiments/ci/**/*.json + experiments/ci/**/*.jsonl + experiments/ci/**/*.txt + experiments/ci/**/*.log + experiments/ci/**/*.csv + experiments/ci/**/*.yaml if-no-files-found: ignore diff --git a/MANIFEST.in b/MANIFEST.in index eb5bcdb..d930cc1 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,3 +1,4 @@ include LICENSE README.md .pre-commit-config.yaml recursive-include tests *.py *.yaml *.csv *.txt recursive-include .github *.yml +recursive-include examples *.yaml diff --git a/README.md b/README.md index 17f1df6..141e758 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,79 @@ Bookworm/Python 3.11. Current instrumentation requires ActivitySim's `workflow.State` API (1.4-era or newer); arbitrary historical revisions are not promised to work. Build/runtime failures retain diagnostics and a failure report. +## Named experiment files + +Write common options once and override only what differs between runs: + +```yaml +schema_version: 1 +vars: + model: /path/to/sandag-abm3-example +output_root: ./results/sandag-${timestamp} +defaults: + model_dir: ${model} + profile: sandag + data_dir: ${model}/benchmarking-data + config_overlay: ["${model}/configs_explicit_chunk"] + multiprocess: true + processes: 4 + sharrow: true + households: 28365 + memory: 80g + shm_size: 8g + sources: + - sharrow=ActivitySim/sharrow@fc175b27d8e0c5d202721c67d96b050e6117b235 +runs: + main: + sources: + - activitysim=ActivitySim/activitysim@5c6fae24a91a57a2d6dfc2e1dbe062a61d94545a + pr1110: + sources: + - activitysim=ActivitySim/activitysim@51e298a84276813946e1d623c9a5785e078e022f +``` + +Save it as `sandag.yaml`, then run: + +```bash +abench sandag.yaml +# Or check all runs without building images or running models: +abench validate sandag.yaml +``` + +A ready-to-use [SANDAG chunked suite](examples/sandag-chunked.yaml) is included +in the repository. Its paths assume abench and the SANDAG repository are siblings. +The pinned `main` revision is the one used in the earlier trials, not a moving +branch reference. + +- `defaults` accepts CLI options using underscores (`shm_size`, `config_overlay`, + etc.). Use `multiprocess: false` for serial execution and `sharrow: false` to + disable Sharrow. `sources` accepts the same strings/mappings as model profiles. +- `runs` is an ordered mapping of names to overrides. Each run inherits defaults; + ordinary values and lists are replaced. **Sources merge by normalized package + name**, so changing ActivitySim does not discard the shared Sharrow pin. +- `${name}` substitutes a reusable scalar from `vars`; terms can reference other + terms. A whole-value reference preserves its type, including numbers/booleans. + Undefined references and cycles are errors. No shell or environment expansion + is performed. `${timestamp}` is a built-in UTC launch identifier shared by all + runs, with microseconds to avoid reusing output directories. +- All explicit paths in the suite are relative to the YAML file, independent of + the terminal's current directory. This includes overlays and custom profile + paths. Built-in `mtc`/`sandag` profile names retain their meaning. When omitted, + `model_dir` defaults to the YAML file's directory; the model profile still + supplies its usual default data/config paths. +- The suite owns output locations: `output_root//`. Set `output_root` + once instead of `output_dir` in each run. Existing roots are rejected. +- All runs are preflighted before the first starts, then run sequentially in file + order. Failure stops the suite and retains partial results. The combined report + is `output_root/comparison.html`; individual runs retain their own reports. + `experiments.yaml` and `suite.json` record the original file and expanded plan. +- File invocations do not accept additional CLI overrides. Edit `defaults` or the + relevant run to keep the file a complete description of the experiment. + +This experiment file describes **which tests to run**. A model profile such as +`benchmark.yaml` describes **how to configure a model**, and remains reusable +across suites. + ## Run controls - `--single-process` (default), or `--multiprocess --processes N`. The count applies diff --git a/examples/sandag-chunked.yaml b/examples/sandag-chunked.yaml new file mode 100644 index 0000000..a808e23 --- /dev/null +++ b/examples/sandag-chunked.yaml @@ -0,0 +1,30 @@ +# Run from any directory: abench /path/to/abench/examples/sandag-chunked.yaml +schema_version: 1 +vars: + model: ../../sandag-abm3-example + activitysim_main: 5c6fae24a91a57a2d6dfc2e1dbe062a61d94545a + activitysim_pr1110: 51e298a84276813946e1d623c9a5785e078e022f +output_root: ../experiments/sandag-chunked-${timestamp} +defaults: + model_dir: ${model} + profile: sandag + data_dir: ${model}/benchmarking-data + config_overlay: ["${model}/configs_explicit_chunk"] + multiprocess: true + processes: 4 + sharrow: true + households: 28365 + memory: 80g + shm_size: 8g + platform: linux/arm64 + sources: + - sharrow=ActivitySim/sharrow@fc175b27d8e0c5d202721c67d96b050e6117b235 +runs: + main: + label: SANDAG main — explicit chunking + sources: + - activitysim=ActivitySim/activitysim@${activitysim_main} + pr1110: + label: SANDAG PR1110 — explicit chunking + sources: + - activitysim=ActivitySim/activitysim@${activitysim_pr1110} diff --git a/src/abench/cli.py b/src/abench/cli.py index de67962..1ae6f50 100644 --- a/src/abench/cli.py +++ b/src/abench/cli.py @@ -53,7 +53,10 @@ def positive(value): def parser(): - p = argparse.ArgumentParser(description=__doc__) + p = argparse.ArgumentParser( + description=__doc__, + epilog="Named experiments: abench experiments.yaml; preflight: abench validate experiments.yaml", + ) p.add_argument("--model-dir", type=Path, default=Path.cwd()) p.add_argument("--profile", default="benchmark.yaml") p.add_argument( @@ -181,6 +184,20 @@ def container_phase(spec, output, data, image, phase_name): def main(argv=None): p = parser() argv = list(sys.argv[1:] if argv is None else argv) + # A file invocation stays separate from model profiles and ordinary flags. + candidate = argv[1:] if argv and argv[0] in ("run", "validate") else argv + if ( + candidate + and not candidate[0].startswith("-") + and candidate[0] not in ("run", "report", "validate") + ): + if len(candidate) != 1: + p.error( + "an experiment file cannot be mixed with command-line overrides; edit its defaults or runs" + ) + from .experiments import run_suite + + return run_suite(Path(candidate[0]), main, validate_only=argv[0] == "validate") action = argv.pop(0) if argv and argv[0] in ("run", "report", "validate") else "run" args = p.parse_args(argv) if action == "report": diff --git a/src/abench/experiments.py b/src/abench/experiments.py new file mode 100644 index 0000000..b5d380a --- /dev/null +++ b/src/abench/experiments.py @@ -0,0 +1,258 @@ +"""Named experiment suites with shared options and safe string substitution.""" + +import io +import re +from contextlib import redirect_stdout +from datetime import datetime, timezone +from pathlib import Path + +import yaml + +from .common import write_json +from .report import report +from .sources import source + +OPTIONS = { + "model_dir", + "profile", + "sources", + "activitysim_commit", + "sharrow_commit", + "multiprocess", + "processes", + "sharrow", + "households", + "data_dir", + "config_overlay", + "cache_from", + "label", + "compare", + "interval", + "memory", + "shm_size", + "platform", +} +PATHS = {"model_dir", "data_dir", "cache_from", "config_overlay", "compare"} +LISTS = {"config_overlay", "compare"} +TOKEN = re.compile(r"\$\{([^{}]+)\}") + + +class SuiteLoader(yaml.SafeLoader): + """Reject duplicate keys so an accidental repeated default cannot disappear.""" + + +def unique_mapping(loader, node): + result = {} + for key_node, value_node in node.value: + key = loader.construct_object(key_node) + if not isinstance(key, str): + raise ValueError("experiment configuration keys must be strings") + if key in result: + raise ValueError(f"duplicate configuration key: {key}") + result[key] = loader.construct_object(value_node) + return result + + +SuiteLoader.add_constructor( + yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, unique_mapping +) + + +def expand_variables(document, timestamp): + """Resolve named terms recursively; never execute shell code or read env vars.""" + terms = document.get("vars", {}) + if not isinstance(terms, dict): + raise ValueError("vars must be a mapping") + if "timestamp" in terms: + raise ValueError("timestamp is a reserved variable") + resolved = {"timestamp": timestamp} + + def term(name, stack): + if name in resolved: + return resolved[name] + if name in stack: + raise ValueError(f"cyclic variable: {' -> '.join((*stack, name))}") + if name not in terms: + raise ValueError(f"undefined variable: {name}") + value = terms[name] + if not isinstance(value, (str, int, float, bool)): + raise ValueError(f"variable {name} must be a scalar") + resolved[name] = expand(value, (*stack, name)) + return resolved[name] + + def expand(value, stack=()): + if isinstance(value, str): + match = TOKEN.fullmatch(value) + if match: + return term(match[1], stack) + return TOKEN.sub(lambda m: str(term(m[1], stack)), value) + if isinstance(value, list): + return [expand(v, stack) for v in value] + if isinstance(value, dict): + return {k: expand(v, stack) for k, v in value.items()} + return value + + for name in terms: + term(name, ()) + return expand(document) + + +def merge_options(defaults, overrides): + """Runs replace ordinary defaults; source pins merge by distribution name.""" + if not isinstance(overrides, dict): + raise ValueError("defaults and each run must be option mappings") + unknown = set(overrides) - OPTIONS + if unknown: + raise ValueError(f"unknown experiment options: {sorted(unknown)}") + merged = dict(defaults, **overrides) + if "sources" in overrides: + if not isinstance(overrides["sources"], list): + raise ValueError("sources must be a list") + pins = {item["name"]: item for item in defaults.get("sources", [])} + seen = set() + for value in overrides["sources"]: + item = source(value) + if item["name"] in seen: + raise ValueError(f"duplicate source: {item['name']}") + seen.add(item["name"]) + pins[item["name"]] = item + merged["sources"] = list(pins.values()) + return merged + + +def arguments(options, base): + """Translate typed YAML options into the existing CLI's validation interface.""" + argv = [] + options = dict(options) + options.setdefault("model_dir", str(base)) + for key, value in options.items(): + if value is None: + continue + if key in ("multiprocess", "sharrow"): + if not isinstance(value, bool): + raise ValueError(f"{key} must be a YAML boolean") + argv.append( + ("--multiprocess" if value else "--single-process") + if key == "multiprocess" + else ("--sharrow" if value else "--no-sharrow") + ) + continue + if key == "sources": + for item in value: + extras = "[" + ",".join(item["extras"]) + "]" if item["extras"] else "" + subdir = ( + "#subdirectory=" + item["subdirectory"] + if item["subdirectory"] + else "" + ) + argv += [ + "--source", + f"{item['name']}{extras}={item['repository']}@{item['commit']}{subdir}", + ] + continue + values = value if key in LISTS else [value] + if not isinstance(values, list) or any( + not isinstance(v, (str, int, float)) or isinstance(v, bool) for v in values + ): + raise ValueError(f"invalid value for {key}") + if not values: + continue + if key in PATHS or (key == "profile" and value not in ("mtc", "sandag")): + values = [str((base / Path(v).expanduser()).resolve()) for v in values] + argv += ["--" + key.replace("_", "-"), *map(str, values)] + return argv + + +def load_suite(path): + """Expand an entire suite before creating output or running any experiments.""" + path = path.expanduser().resolve() + try: + raw = path.read_text() + document = yaml.load(raw, Loader=SuiteLoader) + except yaml.YAMLError as error: + raise ValueError(f"invalid experiment YAML: {error}") from error + if not isinstance(document, dict) or document.get("schema_version") != 1: + raise ValueError("experiment file requires schema_version: 1") + unknown = set(document) - { + "schema_version", + "vars", + "defaults", + "runs", + "output_root", + } + if unknown: + raise ValueError(f"unknown experiment file fields: {sorted(unknown)}") + document = expand_variables( + document, datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S-%f") + ) + defaults = merge_options({}, document.get("defaults", {})) + runs = document.get("runs") + if not isinstance(runs, dict) or not runs: + raise ValueError("runs must be a nonempty mapping of run names to options") + output = document.get("output_root") + if not isinstance(output, str) or not output: + raise ValueError("output_root must be a path") + output = (path.parent / Path(output).expanduser()).resolve() + if output.exists(): + raise ValueError( + f"output_root already exists: {output}; use ${{timestamp}} for repeatable launches" + ) + plan = [] + for name, overrides in runs.items(): + if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_-]*", name): + raise ValueError(f"invalid run name: {name!r}") + options = merge_options(defaults, overrides) + options.setdefault("label", name) + argv = arguments(options, path.parent) + destination = output / name + plan.append( + { + "name": name, + "argv": argv + ["--output-dir", str(destination)], + "output_dir": str(destination), + } + ) + return { + "source_file": str(path), + "original_yaml": raw, + "configuration": document, + "output_root": str(output), + "runs": plan, + } + + +def run_suite(path, invoke, validate_only=False): + """Preflight every run, execute serially, and preserve partial failure reports.""" + plan = load_suite(path) + root = Path(plan["output_root"]) + # Validation uses the same CLI checks as individual runs and creates nothing. + # Run it for the whole suite first, so a typo in run two cannot waste run one. + for run in plan["runs"]: + with redirect_stdout(io.StringIO()): + code = invoke(["validate", *run["argv"]]) + if code: + return code + if validate_only: + print(f"Validated {len(plan['runs'])} experiments; output root: {root}") + return 0 + root.mkdir(parents=True) + (root / "experiments.yaml").write_text(plan["original_yaml"]) + write_json(root / "suite.json", plan) + completed = [] + try: + for run in plan["runs"]: + print(f"Running experiment {run['name']}…", flush=True) + try: + code = invoke(["run", *run["argv"]]) + finally: + # Failed builds/models still have an experiment record and belong + # in the comparison, with failure status rather than winner badges. + if (Path(run["output_dir"]) / "experiment.json").is_file(): + completed.append(Path(run["output_dir"])) + if code: + return code + finally: + if completed: + report(completed, root / "comparison.html") + print(f"Comparison: {root / 'comparison.html'}", flush=True) + return 0 diff --git a/tests/test_docker.py b/tests/test_docker.py index add8287..e430bb0 100644 --- a/tests/test_docker.py +++ b/tests/test_docker.py @@ -59,3 +59,40 @@ def test_tiny_model(tmp_path, multiprocess, sharrow): assert not list((output / "measured").glob("cache-miss-*")) lines = (output / "measured/output/final_households.csv").read_text().splitlines() assert set(lines[1:]) == {"1,20", "2,40", "3,60", "4,80"} + + +def test_named_suite(tmp_path): + """Exercise file dispatch, shared defaults, serial/MP overrides, and comparison.""" + import yaml + + root = tmp_path / "model" + shutil.copytree(Path(__file__).parent / "fixtures/tiny", root) + path = tmp_path / "experiments.yaml" + path.write_text( + yaml.safe_dump( + dict( + schema_version=1, + output_root="results", + defaults=dict( + model_dir="model", + sharrow=True, + households=4, + memory="3g", + shm_size="256m", + sources=[ + f"activitysim=ActivitySim/activitysim@{ACTIVITYSIM}", + f"sharrow=ActivitySim/sharrow@{SHARROW}", + ], + ), + runs={"serial": {}, "parallel": {"multiprocess": True, "processes": 2}}, + ), + sort_keys=False, + ) + ) + assert cli.main([str(path)]) == 0 + root = tmp_path / "results" + runs = json.loads((root / "comparison.json").read_text()) + assert len(runs) == 2 and all(run["valid"] for run in runs) + assert [run["components"]["bench_compute"]["n"] for run in runs] == [1, 2] + assert (root / "experiments.yaml").read_text() == path.read_text() + assert (root / "suite.json").is_file() diff --git a/tests/test_experiments.py b/tests/test_experiments.py new file mode 100644 index 0000000..e2fc2b5 --- /dev/null +++ b/tests/test_experiments.py @@ -0,0 +1,180 @@ +"""Named suites reuse CLI validation/execution without requiring Docker in tests.""" + +import json +from pathlib import Path + +import pytest +import yaml + +from abench import cli, experiments +from abench.experiments import load_suite, run_suite + +SHA = "a" * 40 +OTHER = "b" * 40 + + +def suite(tmp_path, **updates): + """Write a portable two-run suite with common paths and source dependencies.""" + document = dict( + schema_version=1, + vars={"model": "model", "sample": 4}, + output_root="results-${timestamp}", + defaults=dict( + model_dir="${model}", + profile="sandag", + data_dir="${model}/data", + config_overlay=["${model}/chunks"], + households="${sample}", + multiprocess=True, + processes=2, + sharrow=True, + sources=[ + f"sharrow=ActivitySim/sharrow@{SHA}", + f"activitysim=ActivitySim/activitysim@{SHA}", + ], + ), + runs={ + "main": {}, + "pr": {"sources": [f"activitysim=ActivitySim/activitysim@{OTHER}"]}, + }, + ) + document.update(updates) + path = tmp_path / "experiment.yaml" + path.write_text(yaml.safe_dump(document, sort_keys=False)) + return path + + +def test_defaults_variables_and_source_overrides(tmp_path, monkeypatch): + path = suite(tmp_path) + monkeypatch.chdir("/") + plan = load_suite(path) + first, second = plan["runs"] + assert first["name"] == "main" + assert second["name"] == "pr" + for run in plan["runs"]: + args = cli.parser().parse_args(run["argv"]) + assert args.households == 4 + assert args.model_dir == tmp_path / "model" + assert args.data_dir == tmp_path / "model/data" + assert args.config_overlay == [tmp_path / "model/chunks"] + assert f"sharrow=ActivitySim/sharrow@{SHA}" in args.source + assert f"activitysim=ActivitySim/activitysim@{OTHER}" in second["argv"] + assert f"activitysim=ActivitySim/activitysim@{SHA}" not in second["argv"] + assert not Path(plan["output_root"]).exists() + + +@pytest.mark.parametrize( + "updates,match", + [ + ({"vars": {"a": "${b}", "b": "${a}"}}, "cyclic"), + ({"vars": {}}, "undefined"), + ({"vars": {"timestamp": "x"}}, "reserved"), + ({"runs": {"../bad": {}}}, "invalid run name"), + ({"runs": {}}, "nonempty"), + ({"runs": {"bad": {"household": 1}}}, "unknown experiment options"), + ({"runs": {"bad": {"multiprocess": "false"}}}, "YAML boolean"), + ({"runs": {"bad": {"output_dir": "somewhere"}}}, "unknown experiment options"), + ], +) +def test_bad_suite_rejected_before_execution(tmp_path, updates, match): + with pytest.raises(ValueError, match=match): + load_suite(suite(tmp_path, **updates)) + + +def test_duplicate_yaml_and_existing_root(tmp_path): + path = suite(tmp_path, output_root="results") + (tmp_path / "results").mkdir() + with pytest.raises(ValueError, match="already exists"): + load_suite(path) + path.write_text("schema_version: 1\nruns: {}\nruns: {}\n") + with pytest.raises(ValueError, match="duplicate"): + load_suite(path) + + +def test_suite_preflights_all_then_runs_and_reports(tmp_path, monkeypatch): + path = suite(tmp_path) + calls, reports = [], [] + + def invoke(argv): + calls.append(argv) + if argv[0] == "run": + output = Path(argv[argv.index("--output-dir") + 1]) + output.mkdir() + (output / "experiment.json").write_text("{}") + return 0 + + monkeypatch.setattr( + experiments, + "report", + lambda paths, destination: reports.append((paths, destination)), + ) + assert run_suite(path, invoke) == 0 + assert [args[0] for args in calls] == ["validate", "validate", "run", "run"] + assert len(reports[0][0]) == 2 + root = reports[0][1].parent + assert (root / "experiments.yaml").read_text() == path.read_text() + assert len(json.loads((root / "suite.json").read_text())["runs"]) == 2 + + +def test_invalid_later_run_prevents_first_run(tmp_path): + calls = [] + + def invoke(argv): + calls.append(argv[0]) + if len(calls) == 2: + raise ValueError("bad input") + return 0 + + with pytest.raises(ValueError, match="bad input"): + run_suite(suite(tmp_path), invoke) + assert calls == ["validate", "validate"] + assert not list(tmp_path.glob("results-*")) + + +def test_model_failure_stops_suite_and_reports_partial_results(tmp_path, monkeypatch): + calls, reports = [], [] + + def invoke(argv): + calls.append(argv[0]) + if argv[0] == "run": + output = Path(argv[argv.index("--output-dir") + 1]) + output.mkdir() + (output / "experiment.json").write_text("{}") + raise RuntimeError("model failed") + return 0 + + monkeypatch.setattr( + experiments, "report", lambda paths, dest: reports.append(paths) + ) + with pytest.raises(RuntimeError, match="model failed"): + run_suite(suite(tmp_path), invoke) + assert calls == ["validate", "validate", "run"] + assert len(reports[0]) == 1 + + +def test_cli_file_dispatch_and_validation(tmp_path, monkeypatch): + calls = [] + monkeypatch.setattr( + experiments, + "run_suite", + lambda path, invoke, validate_only: calls.append((path, validate_only)) or 0, + ) + path = tmp_path / "named.yaml" + assert cli.main([str(path)]) == 0 + assert cli.main(["run", str(path)]) == 0 + assert cli.main(["validate", str(path)]) == 0 + assert calls == [(path, False), (path, False), (path, True)] + with pytest.raises(SystemExit): + cli.main([str(path), "--households", "5"]) + + +def test_shipped_sandag_suite(): + path = Path(__file__).parents[1] / "examples/sandag-chunked.yaml" + plan = load_suite(path) + for run in plan["runs"]: + args = cli.parser().parse_args(run["argv"]) + assert args.households == 28365 + assert args.processes == 4 + assert args.sharrow + assert args.data_dir.name == "benchmarking-data" + assert args.config_overlay[0].name == "configs_explicit_chunk" From ed37819cfa239968a4f1ffb1b7a39690937df3ed Mon Sep 17 00:00:00 2001 From: Jeff Newman Date: Sun, 13 Sep 2026 19:34:02 -0500 Subject: [PATCH 3/8] Use capped serial Sharrow warmups and improve failure diagnostics Add configurable warmup household limits while preserving measured settings and strict cache coverage checks. Record phase settings and missing flow signatures, and surface actionable failures in CLI output and reports. --- README.md | 29 +++++-- examples/sandag-chunked.yaml | 6 +- src/abench/cli.py | 25 +++++- src/abench/experiments.py | 1 + src/abench/failures.py | 82 ++++++++++++++++++++ src/abench/report.py | 23 +++++- src/abench/runtime/instrumentation.py | 14 +++- src/abench/runtime/worker.py | 63 ++++++++++++++- tests/fixtures/tiny/tiny_extension.py | 2 + tests/test_docker.py | 15 ++++ tests/test_experiments.py | 1 + tests/test_failure_diagnostics.py | 106 ++++++++++++++++++++++++++ tests/test_profiles_sources.py | 26 +++++++ tests/test_warmup.py | 63 +++++++++++++++ 14 files changed, 442 insertions(+), 14 deletions(-) create mode 100644 src/abench/failures.py create mode 100644 tests/test_failure_diagnostics.py create mode 100644 tests/test_warmup.py diff --git a/README.md b/README.md index 141e758..db37522 100644 --- a/README.md +++ b/README.md @@ -33,6 +33,8 @@ Write common options once and override only what differs between runs: ```yaml schema_version: 1 vars: + households: 28365 + warmup_households: 5000 model: /path/to/sandag-abm3-example output_root: ./results/sandag-${timestamp} defaults: @@ -43,7 +45,8 @@ defaults: multiprocess: true processes: 4 sharrow: true - households: 28365 + households: ${households} + warmup_households: ${warmup_households} memory: 80g shm_size: 8g sources: @@ -111,7 +114,7 @@ across suites. - `--memory 16g`, `--shm-size 8g`, `--interval 0.5`, and optional `--platform`. - `--output-dir` must be new. `--label` names an experiment, and `--compare` accepts earlier experiment directories. `--cache-from` seeds compatible generated flows; - a complete warmup still runs. + the small serial warmup still runs. `abench validate` accepts the same experiment arguments without `--output-dir`. It checks the profile, required inputs, CSV population size, source pin syntax, @@ -227,9 +230,25 @@ model outputs. ## Measurement and reports -Sharrow runs first complete a full matching warmup in a separate container. The -measured run uses the same sample, seed, layout, and cache path. Numba compilation -of generated flow overloads is rejected during measurement; ordinary non-flow +Sharrow runs first execute the model in a separate **single-process** warmup, +using **min(target households, 500)** households by default. For `--households 0`, +the target is the full available population, so warmup uses at most 500 of those +households. Set `--warmup-households N` (or `warmup_households: N` in experiment +YAML) to change this positive cap. Warmup always uses one process; measured runs +retain their requested sample and worker count. Model config directories, seed, +chunk overlays, and flow cache path are retained from the target experiment. + +A smaller serial warmup may not exercise every flow signature required by the +measured run. **Cache misses still invalidate the measured experiment**: abench +does not silently compile, retry with compilation enabled, or accept incomplete +cache coverage. Increase the warmup cap when needed. +The CLI and failure report identify the phase, failed component, cache flow, and +log path; `cache-miss-details-*.jsonl` records the exact required Numba signature. +`phase-settings.json` and +`effective-settings.json` preserve the actual settings for each phase, and the +report includes the cache-build settings separately from the measured target. + +Numba compilation of generated flow overloads is rejected during measurement; ordinary non-flow compilation and disk-cache loading remain included. This guard uses private Numba internals and is covered by real disk-cache hit/miss tests. diff --git a/examples/sandag-chunked.yaml b/examples/sandag-chunked.yaml index a808e23..a4774e6 100644 --- a/examples/sandag-chunked.yaml +++ b/examples/sandag-chunked.yaml @@ -1,6 +1,7 @@ # Run from any directory: abench /path/to/abench/examples/sandag-chunked.yaml schema_version: 1 vars: + households: 28365 model: ../../sandag-abm3-example activitysim_main: 5c6fae24a91a57a2d6dfc2e1dbe062a61d94545a activitysim_pr1110: 51e298a84276813946e1d623c9a5785e078e022f @@ -13,7 +14,10 @@ defaults: multiprocess: true processes: 4 sharrow: true - households: 28365 + households: ${households} + # The 500-household cache build misses rare at-work flows and dtype signatures. + # This model-specific override keeps the same sample, but warmup is still serial. + warmup_households: 5000 memory: 80g shm_size: 8g platform: linux/arm64 diff --git a/src/abench/cli.py b/src/abench/cli.py index 1ae6f50..d10f101 100644 --- a/src/abench/cli.py +++ b/src/abench/cli.py @@ -15,6 +15,7 @@ from . import __version__ from .common import read_json, write_json +from .failures import BenchmarkFailure, describe_failure from .profiles import load_profile, validate_model from .report import load_run, report from .sources import resolve_sources @@ -76,6 +77,12 @@ def parser(): p.add_argument( "--households", type=int, default=1000, help="0 means full population" ) + p.add_argument( + "--warmup-households", + type=int, + default=500, + help="maximum cache-build households (default: 500); warmup is always single-process", + ) p.add_argument("--data-dir", type=Path, default=None) p.add_argument( "--config-overlay", @@ -241,6 +248,8 @@ def main(argv=None): ) if args.sharrow and not args.sharrow_commit: p.error("Sharrow enabled: provide a sharrow source override") + if args.warmup_households < 1: + p.error("--warmup-households must be positive") if args.households < 0: p.error("--households must be nonnegative") if args.multiprocess and (args.processes is None or args.processes < 1): @@ -424,17 +433,27 @@ def main(argv=None): if seed: if (seed / "pip-freeze.txt").read_text().strip() != freeze.strip(): raise ValueError("Cannot seed cache: installed dependencies differ") - # Only flow artifacts are reused. Full warmup still checks coverage - # under the new configuration before compilation-blocked measurement. + # Only flow artifacts are reused. A small serial warmup prepares flows; + # measurement remains responsible for rejecting missing signatures. shutil.copytree(seed / "cache/flows", output / "cache/flows") if args.sharrow: stage = "warmup" - print("Preparing Sharrow cache with a complete matching run…", flush=True) + print( + f"Preparing Sharrow cache in single process (up to {args.warmup_households} households)…", + flush=True, + ) container_phase(spec, output, data, image, "warmup") stage = "measured" print("Running measured model…", flush=True) container_phase(spec, output, data, image, "measured") + if not load_run(output)["valid"]: + raise BenchmarkFailure(describe_failure(output, "measured")) except (Exception, KeyboardInterrupt) as error: + if isinstance(error, subprocess.CalledProcessError): + error = BenchmarkFailure(describe_failure(output, stage)) + spec["failure"] = {"phase": stage, "error": str(error)} + write_json(output / "experiment.json", spec) + raise error from None spec["failure"] = {"phase": stage, "error": str(error)} write_json(output / "experiment.json", spec) raise diff --git a/src/abench/experiments.py b/src/abench/experiments.py index b5d380a..fc60461 100644 --- a/src/abench/experiments.py +++ b/src/abench/experiments.py @@ -22,6 +22,7 @@ "processes", "sharrow", "households", + "warmup_households", "data_dir", "config_overlay", "cache_from", diff --git a/src/abench/failures.py b/src/abench/failures.py new file mode 100644 index 0000000..0307713 --- /dev/null +++ b/src/abench/failures.py @@ -0,0 +1,82 @@ +"""Translate container exit codes into useful model failure messages.""" + +import json +import re +from pathlib import Path + +from .common import read_json + + +class BenchmarkFailure(ValueError): + """An experiment failed with diagnostics retained beside its report.""" + + +def tail(path, limit=131072): + """Read a bounded log tail even when a full model log is many gigabytes.""" + if not path.is_file(): + return "" + with path.open("rb") as stream: + stream.seek(0, 2) + stream.seek(max(0, stream.tell() - limit)) + return stream.read().decode("utf-8", errors="replace") + + +def describe_failure(output, phase): + """Prefer explicit cache/OOM evidence over a generic worker exit exception.""" + output = Path(output) + directory = output / phase + log = output / "build.log" if phase == "build" else directory / "console.log" + docker = read_json(directory / "docker-state.json", {}) + components = set() + for path in directory.glob("components-*.jsonl"): + for line in tail(path).splitlines(): + try: + row = json.loads(line) + except (ValueError, TypeError): + continue + if not row.get("succeeded", True): + components.add(row["component"]) + where = f" in {', '.join(sorted(components))}" if components else "" + misses = sorted( + { + line.strip() + for path in directory.glob("cache-miss-*.txt") + for line in tail(path).splitlines() + if line.strip() + } + ) + if misses: + reason = f"Sharrow flow cache miss{where}: " + ", ".join( + Path(name).parent.name for name in misses + ) + remedy = "Results rejected; no flow compilation was allowed. The warmup did not cover the required flow/type signature. Increase warmup_households (up to the target sample) and start a new experiment." + elif docker.get("OOMKilled"): + reason = f"Docker killed the container for exceeding its memory limit{where}." + remedy = ( + "Increase the container/VM memory budget or reduce component chunk sizes." + ) + elif docker.get("ExitCode") == 0 and phase == "measured": + spec = read_json(output / "experiment.json", {}) + summaries = read_json(directory / "output-summary.json", {}) + actual = summaries.get("households", {}).get("rows") + requested = spec.get("households", 0) + if requested and actual != requested: + reason = f"Household sample mismatch: requested {requested}, output contains {actual if actual is not None else 'no household table'}." + else: + reason = "Required runtime or memory measurements are missing." + remedy = "Results rejected even though the container exited successfully." + else: + errors = re.findall( + r"^([\w.]+(?:Error|Exception): .+)$", tail(log), re.MULTILINE + ) + useful = [e for e in errors if "SubprocessError: Process " not in e] + reason = ( + useful + or errors + or [ + docker.get("Error") + or f"Container exited with code {docker.get('ExitCode', 'unavailable')}" + ] + )[0] + remedy = "See the retained log for the complete traceback." + return f"{output.name}: {phase} failed{where if not misses else ''}. {reason}\n{remedy}\nLog: {log}\nReport: {output / 'report.html'}" diff --git a/src/abench/report.py b/src/abench/report.py index ced5be9..7e1b2cc 100644 --- a/src/abench/report.py +++ b/src/abench/report.py @@ -8,6 +8,7 @@ import statistics from .common import read_json, write_json +from .failures import describe_failure COLORS = ("#0072b2", "#d55e00", "#009e73", "#cc79a7", "#e69f00", "#56b4e9") @@ -142,6 +143,14 @@ def load_run(directory): and not docker.get("OOMKilled") and not list(phase.glob("cache-miss-*.txt")) ) + warmup_settings = read_json(directory / "warmup/phase-settings.json", {}) + effective = read_json(directory / "warmup/effective-settings.json", {}) + if effective: + warmup_settings = { + "households": effective.get("households_sample_size"), + "multiprocess": effective.get("multiprocess"), + "processes": effective.get("num_processes"), + } return { "spec": spec, "components": components, @@ -151,6 +160,12 @@ def load_run(directory): "valid": valid, "inputs": read_json(phase / "input-summary.json", {}), "outputs": outputs, + "failure_reason": describe_failure( + directory, spec.get("failure", {}).get("phase", "measured") + ) + if not valid + else None, + "warmup_settings": warmup_settings, "peak": max((r["peak_bytes"] for r in memory), default=0), "docker": docker, "path": str(directory), @@ -255,6 +270,7 @@ def experiment_card(run, xmax, ymax): "processes", "sharrow", "households", + "warmup_households", "data_dir", "output_dir", "memory", @@ -293,11 +309,16 @@ def experiment_card(run, xmax, ymax): counts.append(f"TAZs{taz_cells}") elapsed = run["status"].get("elapsed_seconds") elapsed = f"{elapsed:.3f}" if elapsed is not None else "unavailable" - failure = f"

{escape(spec['failure'])}

" if spec.get("failure") else "" + failure = ( + f"
{escape(run['failure_reason'])}
" + if run.get("failure_reason") + else "" + ) details = "".join( f"
{title}
{escape(json.dumps(value, indent=2))}
" for title, value in ( ("Complete settings and provenance", spec), + ("Actual cache-build settings", run.get("warmup_settings", {})), ("Input totals and categories", run["inputs"]), ("Output totals and categories", run["outputs"]), ("Container exit and OOM status", run["docker"]), diff --git a/src/abench/runtime/instrumentation.py b/src/abench/runtime/instrumentation.py index fdd7199..e329e75 100644 --- a/src/abench/runtime/instrumentation.py +++ b/src/abench/runtime/instrumentation.py @@ -55,7 +55,19 @@ def compile_checked(self, *args, **kwargs): if filename.is_relative_to(flow_cache): with (records / f"cache-miss-{os.getpid()}.txt").open("a") as stream: stream.write(str(filename) + "\n") - raise RuntimeError(f"Measured Sharrow flow cache miss: {filename}") + detail = { + "flow": str(filename), + "function": self.py_func.__name__, + "signature": str(args[0] if args else kwargs), + "process": multiprocessing.current_process().name, + } + with (records / f"cache-miss-details-{os.getpid()}.jsonl").open( + "a" + ) as stream: + stream.write(json.dumps(detail) + "\n") + raise RuntimeError( + f"Measured Sharrow flow cache miss: {filename} ({self.py_func.__name__}); required signature: {detail['signature']}" + ) return compile_original(self, *args, **kwargs) _FunctionCompiler.compile = compile_checked diff --git a/src/abench/runtime/worker.py b/src/abench/runtime/worker.py index 3cedcc6..897d5fd 100644 --- a/src/abench/runtime/worker.py +++ b/src/abench/runtime/worker.py @@ -19,6 +19,57 @@ def write_json(path, value): path.write_text(json.dumps(value, indent=2) + "\n") +def phase_spec(spec, phase_name, data_root=Path("/data")): + """Shrink only warmup; retain target config selection and the measured spec.""" + result = dict(spec) + if phase_name != "warmup": + return result + cap = spec.get("warmup_households", 500) + target = spec["households"] + if target == 0: + # Zero denotes the full population, not an empty sample. Count only as + # far as the cap for CSV, or use Parquet metadata without loading rows. + profile = spec["profile"] + filename = profile.get("household_table") + if filename is None: + stem = ( + profile.get("input_tables", {}) + .get("households", {}) + .get("file", "households") + ) + filename = stem + ".csv" + path = data_root / filename + if not path.exists() and path.suffix == ".csv": + path = path.with_suffix(".parquet") + if path.suffix == ".csv" and path.is_file(): + with path.open(newline="") as stream: + reader = csv.reader(stream) + next(reader, None) + target = 0 + for row in reader: + if row: + target += 1 + if target >= cap: + break + elif path.suffix == ".parquet" and path.is_file(): + import pyarrow.parquet as pq + + target = pq.ParquetFile(path).metadata.num_rows + else: + raise ValueError( + "Cannot determine full-population warmup size; configure household_table or set a positive --households target" + ) + if target < 1: + raise ValueError("Cannot build a flow cache from an empty household table") + result.update( + households=min(target, cap), + multiprocess=False, + processes=1, + _target_multiprocess=spec["multiprocess"], + ) + return result + + def make_state( spec, phase, @@ -36,6 +87,7 @@ def make_state( from activitysim.core.workflow import State profile = spec["profile"] + config_multiprocess = spec.get("_target_multiprocess", spec["multiprocess"]) defaults = dict(profile.get("settings", {})) if profile.get("models_from"): base = yaml.safe_load((model_root / profile["models_from"]).read_text()) @@ -44,7 +96,7 @@ def make_state( for name in base["models"] if name not in profile.get("exclude_models", []) ] - if spec["multiprocess"] and profile.get("mp_settings"): + if config_multiprocess and profile.get("mp_settings"): mp = yaml.safe_load((model_root / profile["mp_settings"]).read_text()) defaults["multiprocess_steps"] = mp["multiprocess_steps"] generated = phase / "profile-config" @@ -57,7 +109,7 @@ def make_state( model_root / f"overlay-{i}" for i in range(len(spec.get("config_overlay", []))) ] configs += [generated] - if spec["multiprocess"]: + if config_multiprocess: configs += [model_root / name for name in profile.get("mp_configs", [])] configs += [model_root / name for name in profile["configs"]] state = State.make_default( @@ -99,7 +151,12 @@ def make_state( def run_model(spec, phase): - """Use identical samples, process layout and cache paths in both phases.""" + """Warm flows with a small serial run; measure the untouched target settings.""" + spec = phase_spec(spec, phase.name) + write_json( + phase / "phase-settings.json", + {key: spec[key] for key in ("households", "multiprocess", "processes")}, + ) state = make_state(spec, phase) state.logging.config_logger() write_json( diff --git a/tests/fixtures/tiny/tiny_extension.py b/tests/fixtures/tiny/tiny_extension.py index 9ea406b..26d46c9 100644 --- a/tests/fixtures/tiny/tiny_extension.py +++ b/tests/fixtures/tiny/tiny_extension.py @@ -12,6 +12,8 @@ def bench_initialize(state: workflow.State): """Seed four rows and a generated Numba module in the guarded flow directory.""" households = pd.read_csv(Path("/data/households.csv"), index_col="household_id") + if state.settings.households_sample_size: + households = households.head(state.settings.households_sample_size) state.add_table("households", households) cache = Path(state.settings.sharrow_cache_dir) cache.mkdir(parents=True, exist_ok=True) diff --git a/tests/test_docker.py b/tests/test_docker.py index e430bb0..2da1a16 100644 --- a/tests/test_docker.py +++ b/tests/test_docker.py @@ -55,6 +55,11 @@ def test_tiny_model(tmp_path, multiprocess, sharrow): assert run["outputs"]["households"]["rows"] == 4 assert run["memory"] and all(row["current_bytes"] > 0 for row in run["memory"]) assert (output / "warmup").exists() == sharrow + if sharrow: + warmup = json.loads((output / "warmup/effective-settings.json").read_text()) + assert warmup["households_sample_size"] == 4 + assert warmup["multiprocess"] is False + assert warmup["num_processes"] == 1 assert len(json.loads((output / "source-provenance.json").read_text())) == 2 assert not list((output / "measured").glob("cache-miss-*")) lines = (output / "measured/output/final_households.csv").read_text().splitlines() @@ -77,6 +82,7 @@ def test_named_suite(tmp_path): model_dir="model", sharrow=True, households=4, + warmup_households=2, memory="3g", shm_size="256m", sources=[ @@ -96,3 +102,12 @@ def test_named_suite(tmp_path): assert [run["components"]["bench_compute"]["n"] for run in runs] == [1, 2] assert (root / "experiments.yaml").read_text() == path.read_text() assert (root / "suite.json").is_file() + for name in ("serial", "parallel"): + warmup = json.loads( + (root / name / "warmup/effective-settings.json").read_text() + ) + assert warmup["households_sample_size"] == 2 + assert warmup["multiprocess"] is False + assert warmup["num_processes"] == 1 + summaries = json.loads((root / name / "warmup/output-summary.json").read_text()) + assert summaries["households"]["rows"] == 2 diff --git a/tests/test_experiments.py b/tests/test_experiments.py index e2fc2b5..0c4de6c 100644 --- a/tests/test_experiments.py +++ b/tests/test_experiments.py @@ -174,6 +174,7 @@ def test_shipped_sandag_suite(): for run in plan["runs"]: args = cli.parser().parse_args(run["argv"]) assert args.households == 28365 + assert args.warmup_households == args.households assert args.processes == 4 assert args.sharrow assert args.data_dir.name == "benchmarking-data" diff --git a/tests/test_failure_diagnostics.py b/tests/test_failure_diagnostics.py new file mode 100644 index 0000000..f684e9b --- /dev/null +++ b/tests/test_failure_diagnostics.py @@ -0,0 +1,106 @@ +"""Container failures should identify the model cause and useful artifact paths.""" + +import json + +import pytest + +from abench import cli +from abench.failures import BenchmarkFailure, describe_failure + + +def test_cache_miss_message_names_component_and_remedy(tmp_path): + phase = tmp_path / "measured" + phase.mkdir() + (phase / "cache-miss-1.txt").write_text( + "/results/cache/flows/flow_ABC/__init__.py\n" + ) + (phase / "components-1.jsonl").write_text( + json.dumps({"succeeded": False, "component": "school_escorting"}) + "\n" + ) + message = describe_failure(tmp_path, "measured") + assert "school_escorting" in message and "flow_ABC" in message + assert "Results rejected" in message + assert str(phase / "console.log") in message + assert "warmup_households" in message + + +def test_oom_and_python_exception(tmp_path): + phase = tmp_path / "warmup" + phase.mkdir() + (phase / "docker-state.json").write_text('{"OOMKilled": true, "ExitCode": 137}') + assert "memory limit" in describe_failure(tmp_path, "warmup") + (phase / "docker-state.json").write_text('{"ExitCode": 1}') + (phase / "console.log").write_text( + "KeyError: missing input column\nSubprocessError: Process worker failed\n" + ) + assert "KeyError: missing input column" in describe_failure(tmp_path, "warmup") + + +def test_entrypoint_prints_readable_failure_without_traceback(monkeypatch, capsys): + monkeypatch.setattr( + cli, + "main", + lambda: (_ for _ in ()).throw( + BenchmarkFailure("measured failed: cache miss\nLog: example.log") + ), + ) + with pytest.raises(SystemExit) as error: + cli.entrypoint() + assert error.value.code == 1 + assert "Log: example.log" in capsys.readouterr().err + + +def test_model_container_error_is_wrapped_with_root_cause(tmp_path, monkeypatch): + from test_profiles_sources import make_model + + make_model(tmp_path) + output = tmp_path / "experiment" + + def command(args, log=None): + if args[:2] == ["docker", "info"]: + return json.dumps({"OSType": "linux", "CgroupVersion": "2"}) + if "/opt/source-provenance.json" in args: + return "[]" + return "test-image" + + def container(spec, output, data, image, phase_name): + import subprocess + + phase = output / phase_name + phase.mkdir() + (phase / "cache-miss-1.txt").write_text( + "/results/cache/flows/flow_ABC/__init__.py\n" + ) + raise subprocess.CalledProcessError(1, ["docker", "run", "test-image"]) + + monkeypatch.setattr(cli, "command", command) + monkeypatch.setattr(cli, "container_phase", container) + with pytest.raises(BenchmarkFailure, match="flow_ABC") as error: + cli.main( + [ + "run", + "--model-dir", + str(tmp_path), + "--households", + "2", + "--no-sharrow", + "--output-dir", + str(output), + ] + ) + assert "returned non-zero exit status" not in str(error.value) + assert (output / "report.html").is_file() + failure = json.loads((output / "experiment.json").read_text())["failure"] + assert failure["phase"] == "measured" + assert "Results rejected" in failure["error"] + + +def test_successful_exit_with_wrong_sample_is_explained(tmp_path): + phase = tmp_path / "measured" + phase.mkdir() + (tmp_path / "experiment.json").write_text('{"households": 1000}') + (phase / "docker-state.json").write_text('{"ExitCode": 0}') + (phase / "output-summary.json").write_text('{"households": {"rows": 500}}') + message = describe_failure(tmp_path, "measured") + assert "requested 1000, output contains 500" in message + assert "Results rejected" in message diff --git a/tests/test_profiles_sources.py b/tests/test_profiles_sources.py index 5ba4a9e..a667369 100644 --- a/tests/test_profiles_sources.py +++ b/tests/test_profiles_sources.py @@ -163,3 +163,29 @@ def command(args, log=None): == "build" ) assert (output / "runner/build_sources.py").is_file() + + +def test_serial_warmup_retains_target_mp_configs(tmp_path): + """Only execution layout changes; target-only config layers must still apply.""" + from abench.runtime.worker import phase_spec + + profile = make_model(tmp_path) + mp_configs = tmp_path / "mp_configs" + mp_configs.mkdir() + (mp_configs / "settings.yaml").write_text( + "inherit_settings: true\nrng_base_seed: 123\n" + ) + profile["mp_configs"] = ["mp_configs"] + phase = tmp_path / "warmup" + phase.mkdir() + (phase / "output").mkdir() + spec = dict( + profile=profile, households=1000, multiprocess=True, processes=4, sharrow=True + ) + state = make_state( + phase_spec(spec, "warmup"), phase, tmp_path, tmp_path / "data", tmp_path + ) + assert state.settings.rng_base_seed == 123 + assert state.settings.households_sample_size == 500 + assert state.settings.multiprocess is False + assert state.settings.num_processes == 1 diff --git a/tests/test_warmup.py b/tests/test_warmup.py new file mode 100644 index 0000000..4346070 --- /dev/null +++ b/tests/test_warmup.py @@ -0,0 +1,63 @@ +"""Cache preparation is small and serial without relaxing measurement validity.""" + +import copy + +import pandas as pd +import pytest + +from abench.cli import parser +from abench.runtime.worker import phase_spec + + +@pytest.mark.parametrize( + "target,cap,expected", + [(28365, 500, 500), (200, 500, 200), (500, 500, 500), (28365, 1000, 1000)], +) +def test_warmup_cap_and_measured_isolation(target, cap, expected): + spec = dict( + households=target, + warmup_households=cap, + multiprocess=True, + processes=4, + profile={"configs": ["configs"]}, + ) + original = copy.deepcopy(spec) + warmup = phase_spec(spec, "warmup") + assert warmup["households"] == expected + assert not warmup["multiprocess"] + assert warmup["processes"] == 1 + assert warmup["_target_multiprocess"] + assert phase_spec(spec, "measured") == original + assert spec == original + + +@pytest.mark.parametrize("count", [3, 600]) +@pytest.mark.parametrize("format", ["csv", "parquet"]) +def test_full_population_caps_available_households(tmp_path, count, format): + data = pd.DataFrame({"household_id": range(count)}) + if format == "csv": + data.to_csv(tmp_path / "households.csv", index=False) + else: + data.to_parquet(tmp_path / "households.parquet", index=False) + spec = dict(households=0, multiprocess=True, processes=4, profile={}) + assert phase_spec(spec, "warmup", tmp_path)["households"] == min(count, 500) + assert phase_spec(spec, "measured", tmp_path)["households"] == 0 + + +def test_custom_household_table_and_empty_population(tmp_path): + spec = dict( + households=0, + multiprocess=False, + processes=1, + profile={"household_table": "population.csv"}, + ) + (tmp_path / "population.csv").write_text("id\n") + with pytest.raises(ValueError, match="empty household"): + phase_spec(spec, "warmup", tmp_path) + (tmp_path / "population.csv").write_text("id\n1\n2\n") + assert phase_spec(spec, "warmup", tmp_path)["households"] == 2 + + +def test_default_and_override_parser(): + assert parser().parse_args([]).warmup_households == 500 + assert parser().parse_args(["--warmup-households", "100"]).warmup_households == 100 From 5027f4c7288df9b3dee03d22971af7bedeb6fe33 Mon Sep 17 00:00:00 2001 From: Jeff Newman Date: Sun, 13 Sep 2026 19:46:44 -0500 Subject: [PATCH 4/8] Increase default warmup cap to 5,000 households --- README.md | 4 ++-- examples/sandag-chunked.yaml | 2 +- src/abench/cli.py | 4 ++-- src/abench/runtime/worker.py | 2 +- tests/test_experiments.py | 2 +- tests/test_profiles_sources.py | 2 +- tests/test_warmup.py | 17 ++++++++++++++--- 7 files changed, 22 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index db37522..425abef 100644 --- a/README.md +++ b/README.md @@ -231,8 +231,8 @@ model outputs. ## Measurement and reports Sharrow runs first execute the model in a separate **single-process** warmup, -using **min(target households, 500)** households by default. For `--households 0`, -the target is the full available population, so warmup uses at most 500 of those +using **min(target households, 5000)** households by default. For `--households 0`, +the target is the full available population, so warmup uses at most 5000 of those households. Set `--warmup-households N` (or `warmup_households: N` in experiment YAML) to change this positive cap. Warmup always uses one process; measured runs retain their requested sample and worker count. Model config directories, seed, diff --git a/examples/sandag-chunked.yaml b/examples/sandag-chunked.yaml index a4774e6..abb042b 100644 --- a/examples/sandag-chunked.yaml +++ b/examples/sandag-chunked.yaml @@ -16,7 +16,7 @@ defaults: sharrow: true households: ${households} # The 500-household cache build misses rare at-work flows and dtype signatures. - # This model-specific override keeps the same sample, but warmup is still serial. + # Use 5,000 households for the serial warmup; measured cache misses remain errors. warmup_households: 5000 memory: 80g shm_size: 8g diff --git a/src/abench/cli.py b/src/abench/cli.py index d10f101..9800182 100644 --- a/src/abench/cli.py +++ b/src/abench/cli.py @@ -80,8 +80,8 @@ def parser(): p.add_argument( "--warmup-households", type=int, - default=500, - help="maximum cache-build households (default: 500); warmup is always single-process", + default=5000, + help="maximum cache-build households (default: 5000); warmup is always single-process", ) p.add_argument("--data-dir", type=Path, default=None) p.add_argument( diff --git a/src/abench/runtime/worker.py b/src/abench/runtime/worker.py index 897d5fd..9da4844 100644 --- a/src/abench/runtime/worker.py +++ b/src/abench/runtime/worker.py @@ -24,7 +24,7 @@ def phase_spec(spec, phase_name, data_root=Path("/data")): result = dict(spec) if phase_name != "warmup": return result - cap = spec.get("warmup_households", 500) + cap = spec.get("warmup_households", 5000) target = spec["households"] if target == 0: # Zero denotes the full population, not an empty sample. Count only as diff --git a/tests/test_experiments.py b/tests/test_experiments.py index 0c4de6c..d205b87 100644 --- a/tests/test_experiments.py +++ b/tests/test_experiments.py @@ -174,7 +174,7 @@ def test_shipped_sandag_suite(): for run in plan["runs"]: args = cli.parser().parse_args(run["argv"]) assert args.households == 28365 - assert args.warmup_households == args.households + assert args.warmup_households == 5000 assert args.processes == 4 assert args.sharrow assert args.data_dir.name == "benchmarking-data" diff --git a/tests/test_profiles_sources.py b/tests/test_profiles_sources.py index a667369..c98f9df 100644 --- a/tests/test_profiles_sources.py +++ b/tests/test_profiles_sources.py @@ -186,6 +186,6 @@ def test_serial_warmup_retains_target_mp_configs(tmp_path): phase_spec(spec, "warmup"), phase, tmp_path, tmp_path / "data", tmp_path ) assert state.settings.rng_base_seed == 123 - assert state.settings.households_sample_size == 500 + assert state.settings.households_sample_size == 1000 assert state.settings.multiprocess is False assert state.settings.num_processes == 1 diff --git a/tests/test_warmup.py b/tests/test_warmup.py index 4346070..3990030 100644 --- a/tests/test_warmup.py +++ b/tests/test_warmup.py @@ -31,7 +31,7 @@ def test_warmup_cap_and_measured_isolation(target, cap, expected): assert spec == original -@pytest.mark.parametrize("count", [3, 600]) +@pytest.mark.parametrize("count", [3, 6000]) @pytest.mark.parametrize("format", ["csv", "parquet"]) def test_full_population_caps_available_households(tmp_path, count, format): data = pd.DataFrame({"household_id": range(count)}) @@ -40,7 +40,7 @@ def test_full_population_caps_available_households(tmp_path, count, format): else: data.to_parquet(tmp_path / "households.parquet", index=False) spec = dict(households=0, multiprocess=True, processes=4, profile={}) - assert phase_spec(spec, "warmup", tmp_path)["households"] == min(count, 500) + assert phase_spec(spec, "warmup", tmp_path)["households"] == min(count, 5000) assert phase_spec(spec, "measured", tmp_path)["households"] == 0 @@ -59,5 +59,16 @@ def test_custom_household_table_and_empty_population(tmp_path): def test_default_and_override_parser(): - assert parser().parse_args([]).warmup_households == 500 + assert parser().parse_args([]).warmup_households == 5000 assert parser().parse_args(["--warmup-households", "100"]).warmup_households == 100 + + +@pytest.mark.parametrize("target", [200, 5000, 28365]) +def test_default_warmup_cap(target): + """The default bounds cache preparation without increasing small samples.""" + spec = dict(households=target, multiprocess=True, processes=4, profile={}) + warmup = phase_spec(spec, "warmup") + assert warmup["households"] == min(target, 5000) + assert warmup["multiprocess"] is False + assert warmup["processes"] == 1 + assert phase_spec(spec, "measured") == spec From 9a7ab1ea9c81ef976f1f7bedc53986d23bb60f4d Mon Sep 17 00:00:00 2001 From: Jeff Newman Date: Sun, 13 Sep 2026 19:59:53 -0500 Subject: [PATCH 5/8] Automatically reuse compatible compiled flows across experiments Add a persistent flow cache keyed by compiler and platform identity, with isolated copies, serialized warmups, and atomic updates. Expose CLI and YAML controls and add documentation and reuse tests. --- .gitignore | 2 +- README.md | 27 +++++++- examples/sandag-chunked.yaml | 5 +- src/abench/cli.py | 58 +++++++++++++++- src/abench/experiments.py | 15 +++- src/abench/flow_cache.py | 59 ++++++++++++++++ src/abench/report.py | 3 + src/abench/runtime/cache_identity.py | 37 ++++++++++ tests/fixtures/tiny/tiny_extension.py | 14 ++++ tests/test_docker.py | 12 ++++ tests/test_flow_cache.py | 98 +++++++++++++++++++++++++++ 11 files changed, 319 insertions(+), 11 deletions(-) create mode 100644 src/abench/flow_cache.py create mode 100644 src/abench/runtime/cache_identity.py create mode 100644 tests/test_flow_cache.py diff --git a/.gitignore b/.gitignore index f7bfe9e..aaa90af 100644 --- a/.gitignore +++ b/.gitignore @@ -186,7 +186,7 @@ cython_debug/ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. -# .idea/ +.idea/ # Abstra # Abstra is an AI-powered process automation framework. diff --git a/README.md b/README.md index 425abef..a266644 100644 --- a/README.md +++ b/README.md @@ -113,8 +113,8 @@ across suites. priority order. Relative overlay paths are relative to the model directory. - `--memory 16g`, `--shm-size 8g`, `--interval 0.5`, and optional `--platform`. - `--output-dir` must be new. `--label` names an experiment, and `--compare` accepts - earlier experiment directories. `--cache-from` seeds compatible generated flows; - the small serial warmup still runs. + earlier experiment directories. Compatible compiled flows are reused automatically; + the serial warmup still runs. `abench validate` accepts the same experiment arguments without `--output-dir`. It checks the profile, required inputs, CSV population size, source pin syntax, @@ -238,6 +238,29 @@ YAML) to change this positive cap. Warmup always uses one process; measured runs retain their requested sample and worker count. Model config directories, seed, chunk overlays, and flow cache path are retained from the target experiment. +Compiled flows are automatically reused across runs and suites, including changes +in ActivitySim revisions, sample sizes, process counts, and model configs. The +persistent host cache defaults to `~/.cache/abench/flows`. Change it with +`--flow-cache-dir PATH` (`flow_cache_dir` in YAML), or disable automatic reads and +writes with `--no-reuse-flows` (`reuse_flows: false`). `--cache-from` remains an +explicit seed option with its existing stricter dependency checks. + +Compatibility uses the **installed** Sharrow, Numba, llvmlite, and NumPy versions, +plus their source repository/commit identities when applicable, Python version, +and container architecture/CPU features. ActivitySim and model settings are +excluded from this key: Sharrow identifies generated flows by their contents, +and Numba checks cached signatures. Changed flows can compile during warmup. +An existing cache does **not** guarantee that warmup will need no compilation. + +Each experiment receives a private copy with source timestamps preserved. Warmup +always runs, and successful warmups atomically update the persistent cache before +measurement starts. Compatible simultaneous warmups wait for each other to avoid +losing compiled signatures; measured runs remain independent. Only `cache/flows` +is shared, never model data, outputs, or shared-memory artifacts. Cache identity +and reuse counts are recorded in `flow-cache-identity.json` and `experiment.json`. +The persistent cache can be deleted between runs to reclaim disk space; older +experiments created before this feature are not automatically imported. + A smaller serial warmup may not exercise every flow signature required by the measured run. **Cache misses still invalidate the measured experiment**: abench does not silently compile, retry with compilation enabled, or accept incomplete diff --git a/examples/sandag-chunked.yaml b/examples/sandag-chunked.yaml index abb042b..e7751ea 100644 --- a/examples/sandag-chunked.yaml +++ b/examples/sandag-chunked.yaml @@ -15,9 +15,8 @@ defaults: processes: 4 sharrow: true households: ${households} - # The 500-household cache build misses rare at-work flows and dtype signatures. - # Use 5,000 households for the serial warmup; measured cache misses remain errors. - warmup_households: 5000 + # The 5000-household cache build misses some cache builds, so this test uses more + warmup_households: ${households} memory: 80g shm_size: 8g platform: linux/arm64 diff --git a/src/abench/cli.py b/src/abench/cli.py index 9800182..7a044f4 100644 --- a/src/abench/cli.py +++ b/src/abench/cli.py @@ -10,12 +10,14 @@ import sys import tempfile import uuid +from contextlib import nullcontext from datetime import datetime, timezone from pathlib import Path from . import __version__ from .common import read_json, write_json from .failures import BenchmarkFailure, describe_failure +from .flow_cache import reuse_flows from .profiles import load_profile, validate_model from .report import load_run, report from .sources import resolve_sources @@ -91,6 +93,18 @@ def parser(): default=[], help="extra config directories, highest priority first", ) + p.add_argument( + "--flow-cache-dir", + type=Path, + default=Path.home() / ".cache/abench/flows", + help="persistent compiled-flow cache (default: ~/.cache/abench/flows)", + ) + p.add_argument( + "--reuse-flows", + action=argparse.BooleanOptionalAction, + default=True, + help="automatically reuse and update compatible flows; warmup still runs", + ) p.add_argument( "--cache-from", type=Path, @@ -264,6 +278,12 @@ def main(argv=None): for path in overlays: if not path.is_dir() or (output is not None and output.is_relative_to(path)): p.error("config overlays must be existing directories outside --output-dir") + args.flow_cache_dir = args.flow_cache_dir.expanduser().resolve() + if args.sharrow and args.reuse_flows and output is not None: + if args.flow_cache_dir.is_relative_to(output) or output.is_relative_to( + args.flow_cache_dir + ): + p.error("--flow-cache-dir and --output-dir must be separate directories") seed = args.cache_from.expanduser().resolve() if args.cache_from else None if seed: prior = read_json(seed / "experiment.json", {}) @@ -281,7 +301,7 @@ def main(argv=None): p.error("Docker bind paths cannot contain commas") if output is not None and output.exists(): p.error( - "--output-dir must not already exist; each experiment owns a fresh cache" + "--output-dir must not already exist; each experiment needs a separate output directory" ) docker = json.loads(command(["docker", "info", "--format", "{{json .}}"])) if docker.get("OSType") != "linux" or str(docker.get("CgroupVersion")) != "2": @@ -442,7 +462,41 @@ def main(argv=None): f"Preparing Sharrow cache in single process (up to {args.warmup_households} households)…", flush=True, ) - container_phase(spec, output, data, image, "warmup") + cache = nullcontext(None) + if args.reuse_flows: + identity = json.loads( + command( + [ + "docker", + "run", + "--rm", + "--network=none", + "--entrypoint", + "python", + image, + "-c", + (PACKAGE / "runtime/cache_identity.py").read_text(), + ] + ) + ) + write_json(output / "flow-cache-identity.json", identity) + cache = reuse_flows( + args.flow_cache_dir, identity, output / "cache/flows" + ) + print( + "Checking compatible flow cache (waiting for any active warmup)…", + flush=True, + ) + with cache as cache_info: + if cache_info is not None: + spec["flow_cache"] = cache_info + write_json(output / "experiment.json", spec) + print( + f"Reused {cache_info['restored_files']} flow-cache files.", + flush=True, + ) + container_phase(spec, output, data, image, "warmup") + write_json(output / "experiment.json", spec) stage = "measured" print("Running measured model…", flush=True) container_phase(spec, output, data, image, "measured") diff --git a/src/abench/experiments.py b/src/abench/experiments.py index fc60461..aa4bfd4 100644 --- a/src/abench/experiments.py +++ b/src/abench/experiments.py @@ -26,6 +26,8 @@ "data_dir", "config_overlay", "cache_from", + "flow_cache_dir", + "reuse_flows", "label", "compare", "interval", @@ -33,7 +35,14 @@ "shm_size", "platform", } -PATHS = {"model_dir", "data_dir", "cache_from", "config_overlay", "compare"} +PATHS = { + "model_dir", + "data_dir", + "cache_from", + "flow_cache_dir", + "config_overlay", + "compare", +} LISTS = {"config_overlay", "compare"} TOKEN = re.compile(r"\$\{([^{}]+)\}") @@ -129,13 +138,13 @@ def arguments(options, base): for key, value in options.items(): if value is None: continue - if key in ("multiprocess", "sharrow"): + if key in ("multiprocess", "sharrow", "reuse_flows"): if not isinstance(value, bool): raise ValueError(f"{key} must be a YAML boolean") argv.append( ("--multiprocess" if value else "--single-process") if key == "multiprocess" - else ("--sharrow" if value else "--no-sharrow") + else ("--" if value else "--no-") + key.replace("_", "-") ) continue if key == "sources": diff --git a/src/abench/flow_cache.py b/src/abench/flow_cache.py new file mode 100644 index 0000000..798a6ce --- /dev/null +++ b/src/abench/flow_cache.py @@ -0,0 +1,59 @@ +"""Persistent flow artifacts, isolated from measured runs and model data.""" + +import fcntl +import hashlib +import json +import shutil +import tempfile +from contextlib import contextmanager +from pathlib import Path + +from .common import read_json, write_json + + +def cache_key(identity): + """Use installed compiler identities rather than the entire experiment spec.""" + return hashlib.sha256(json.dumps(identity, sort_keys=True).encode()).hexdigest() + + +@contextmanager +def reuse_flows(root, identity, destination): + """Serialize compatible warmups and atomically publish successful snapshots. + + Keeping the lock through warmup prevents concurrent writers from losing Numba + signature indexes. Measurement uses its own copy and never holds this lock. + copy2 preserves source mtimes, which Numba checks when loading compiled code. + """ + bucket = Path(root).expanduser().resolve() / cache_key(identity) + bucket.mkdir(parents=True, exist_ok=True) + with (bucket / "lock").open("a") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + current = read_json(bucket / "current.json", {}) + previous = bucket / current["snapshot"] if current else None + restored = 0 + if previous: + restored = sum(path.is_file() for path in previous.rglob("*")) + shutil.copytree(previous, destination, dirs_exist_ok=True) + metadata = { + "key": bucket.name, + "directory": str(bucket), + "restored_files": restored, + "published": False, + } + yield metadata + # A failed warmup leaves the last successful snapshot intact. Publishing + # before measurement also retains useful flows if measurement later fails. + snapshot = Path(tempfile.mkdtemp(prefix="snapshot-", dir=bucket)) + try: + if destination.exists(): + shutil.copytree(destination, snapshot, dirs_exist_ok=True) + write_json(bucket / "identity.json", identity) + pointer = bucket / "current.tmp" + write_json(pointer, {"snapshot": snapshot.name}) + pointer.replace(bucket / "current.json") + except BaseException: + shutil.rmtree(snapshot) + raise + metadata["published"] = True + if previous: + shutil.rmtree(previous) diff --git a/src/abench/report.py b/src/abench/report.py index 7e1b2cc..4d95fd0 100644 --- a/src/abench/report.py +++ b/src/abench/report.py @@ -280,6 +280,9 @@ def experiment_card(run, xmax, ymax): "compare", "config_overlay", "cache_from", + "reuse_flows", + "flow_cache_dir", + "flow_cache", ) ) counts = [] diff --git a/src/abench/runtime/cache_identity.py b/src/abench/runtime/cache_identity.py new file mode 100644 index 0000000..d32c445 --- /dev/null +++ b/src/abench/runtime/cache_identity.py @@ -0,0 +1,37 @@ +"""Describe the installed compiler environment inside the benchmark container.""" + +import importlib.metadata +import json +import platform +import sys +from pathlib import Path + + +def identity(provenance_path=Path("/opt/source-provenance.json")): + """Exclude model/ActivitySim settings: Sharrow hashes generated flow contents.""" + import llvmlite.binding as llvm + + names = ("sharrow", "numba", "llvmlite", "numpy") + provenance = json.loads(provenance_path.read_text()) + return { + "schema_version": 1, + "python": [sys.implementation.name, *sys.version_info[:3]], + "machine": platform.machine(), + "system": platform.system(), + "libc": platform.libc_ver(), + "cpu": str(llvm.get_host_cpu_name()), + "cpu_features": llvm.get_host_cpu_features().flatten(), + "versions": {name: importlib.metadata.version(name) for name in names}, + "sources": { + item["name"]: { + key: item.get(key) + for key in ("repository", "resolved_commit", "subdirectory") + } + for item in provenance + if item["name"] in names + }, + } + + +if __name__ == "__main__": + print(json.dumps(identity(), sort_keys=True)) diff --git a/tests/fixtures/tiny/tiny_extension.py b/tests/fixtures/tiny/tiny_extension.py index 26d46c9..4d295a7 100644 --- a/tests/fixtures/tiny/tiny_extension.py +++ b/tests/fixtures/tiny/tiny_extension.py @@ -42,6 +42,20 @@ def bench_export(state: workflow.State, households: pd.DataFrame): def prepare(state, spec, phase): """Supply tiny shared skims; this fixture has no destination/shadow models.""" + if phase.name == "warmup": + from numba.core.dispatcher import _FunctionCompiler + + original = _FunctionCompiler.compile + + def record_compile(compiler, *args, **kwargs): + # Record actual generated-code compilation, not dispatcher creation. + if "tiny_generated.py" in compiler.py_func.__code__.co_filename: + (Path(state.filesystem.output_dir) / "flow-compiled.txt").write_text( + "compiled" + ) + return original(compiler, *args, **kwargs) + + _FunctionCompiler.compile = record_compile state.set("shadow_pricing_info", None) state.set("shadow_pricing_choice_info", None) state.set("network_los_preload", None) diff --git a/tests/test_docker.py b/tests/test_docker.py index 2da1a16..6134d50 100644 --- a/tests/test_docker.py +++ b/tests/test_docker.py @@ -35,6 +35,8 @@ def test_tiny_model(tmp_path, multiprocess, sharrow): f"activitysim=ActivitySim/activitysim@{ACTIVITYSIM}", "--source", f"sharrow=ActivitySim/sharrow@{SHARROW}", + "--flow-cache-dir", + str(tmp_path / "shared-flows"), "--households", "4", "--memory", @@ -80,6 +82,7 @@ def test_named_suite(tmp_path): output_root="results", defaults=dict( model_dir="model", + flow_cache_dir="shared-flows", sharrow=True, households=4, warmup_households=2, @@ -111,3 +114,12 @@ def test_named_suite(tmp_path): assert warmup["num_processes"] == 1 summaries = json.loads((root / name / "warmup/output-summary.json").read_text()) assert summaries["households"]["rows"] == 2 + + # Both warmups execute, but the second must load compiled code from the first. + serial = json.loads((root / "serial/experiment.json").read_text())["flow_cache"] + parallel = json.loads((root / "parallel/experiment.json").read_text())["flow_cache"] + assert serial["restored_files"] == 0 + assert parallel["restored_files"] > 0 + assert serial["key"] == parallel["key"] + assert (root / "serial/warmup/output/flow-compiled.txt").exists() + assert not (root / "parallel/warmup/output/flow-compiled.txt").exists() diff --git a/tests/test_flow_cache.py b/tests/test_flow_cache.py new file mode 100644 index 0000000..808b1f8 --- /dev/null +++ b/tests/test_flow_cache.py @@ -0,0 +1,98 @@ +"""Persistent reuse must preserve compiled artifacts and survive failed warmups.""" + +import os +from pathlib import Path + +import pytest + +from abench import cli +from abench.experiments import arguments +from abench.flow_cache import cache_key, reuse_flows + + +def test_reuse_accumulates_and_preserves_mtime(tmp_path): + root = tmp_path / "shared" + first, second, third = [tmp_path / name for name in ("first", "second", "third")] + identity = {"sharrow": "abc", "numba": "1"} + with reuse_flows(root, identity, first) as info: + assert info["restored_files"] == 0 + first.mkdir() + (first / "generated.py").write_text("code") + os.utime(first / "generated.py", (1234567890, 1234567890)) + (first / "compiled.nbc").write_bytes(b"compiled") + assert info["published"] + with reuse_flows(root, identity, second) as info: + assert info["restored_files"] == 2 + assert (second / "generated.py").stat().st_mtime_ns == ( + first / "generated.py" + ).stat().st_mtime_ns + (second / "additional.nbc").write_bytes(b"new signature") + with reuse_flows(root, identity, third) as info: + assert info["restored_files"] == 3 + assert (third / "compiled.nbc").read_bytes() == b"compiled" + assert len(list((root / cache_key(identity)).glob("snapshot-*"))) == 1 + + +def test_failed_warmup_does_not_publish(tmp_path): + root = tmp_path / "shared" + with reuse_flows(root, {}, tmp_path / "first"): + (tmp_path / "first").mkdir() + (tmp_path / "first/good").write_text("original") + with pytest.raises(RuntimeError): + with reuse_flows(root, {}, tmp_path / "failed"): + (tmp_path / "failed/good").write_text("damaged") + raise RuntimeError("warmup failed") + with reuse_flows(root, {}, tmp_path / "last"): + assert (tmp_path / "last/good").read_text() == "original" + + +def test_incompatible_environments_start_empty(tmp_path): + with reuse_flows(tmp_path / "shared", {"numba": "1"}, tmp_path / "first"): + (tmp_path / "first").mkdir() + (tmp_path / "first/code").write_text("old") + with reuse_flows(tmp_path / "shared", {"numba": "2"}, tmp_path / "second") as info: + assert info["restored_files"] == 0 + assert not (tmp_path / "second").exists() + + +def test_cache_options_in_yaml(tmp_path): + args = cli.parser().parse_args( + arguments({"flow_cache_dir": "shared", "reuse_flows": False}, tmp_path) + ) + assert args.flow_cache_dir == tmp_path / "shared" + assert args.reuse_flows is False + assert cli.parser().parse_args([]).reuse_flows is True + assert ( + cli.parser().parse_args([]).flow_cache_dir + == Path.home() / ".cache/abench/flows" + ) + + +def test_identity_tracks_compiler_sources_but_not_activitysim(tmp_path, monkeypatch): + """Equal release versions must not hide a changed source commit.""" + import json + + from abench.runtime.cache_identity import identity + + monkeypatch.setattr("importlib.metadata.version", lambda name: "1.0") + provenance = tmp_path / "sources.json" + pins = [ + { + "name": "sharrow", + "repository": "ActivitySim/sharrow", + "resolved_commit": "a", + }, + { + "name": "activitysim", + "repository": "ActivitySim/activitysim", + "resolved_commit": "b", + }, + ] + provenance.write_text(json.dumps(pins)) + original = identity(provenance) + pins[1]["resolved_commit"] = "c" + provenance.write_text(json.dumps(pins)) + assert cache_key(identity(provenance)) == cache_key(original) + pins[0]["resolved_commit"] = "d" + provenance.write_text(json.dumps(pins)) + assert cache_key(identity(provenance)) != cache_key(original) From 847cabb7fd1d0ec4475d956c622c863f58c93e75 Mon Sep 17 00:00:00 2001 From: Jeff Newman Date: Fri, 18 Sep 2026 16:27:12 -0500 Subject: [PATCH 6/8] Retry benchmark attempts that compile Sharrow flows Add configurable cache retries, preserve attempt diagnostics, and publish newly compiled flows. Accept only compilation-free measurements and stop immediately on model or validation failures. --- README.md | 39 ++++++++----- examples/sandag-chunked.yaml | 2 +- src/abench/attempts.py | 66 +++++++++++++++++++++ src/abench/cli.py | 61 ++++++++++++++----- src/abench/common.py | 13 ++++- src/abench/experiments.py | 1 + src/abench/failures.py | 27 ++++++++- src/abench/flow_cache.py | 44 ++++++++------ src/abench/report.py | 35 +++++++++-- src/abench/runtime/instrumentation.py | 16 +++-- src/abench/runtime/worker.py | 14 +++-- tests/fixtures/tiny/tiny_extension.py | 10 +++- tests/test_attempts.py | 84 +++++++++++++++++++++++++++ tests/test_docker.py | 35 ++++++++++- tests/test_experiments.py | 2 +- tests/test_failure_diagnostics.py | 9 ++- 16 files changed, 381 insertions(+), 77 deletions(-) create mode 100644 src/abench/attempts.py create mode 100644 tests/test_attempts.py diff --git a/README.md b/README.md index a266644..3ebfd57 100644 --- a/README.md +++ b/README.md @@ -254,26 +254,35 @@ An existing cache does **not** guarantee that warmup will need no compilation. Each experiment receives a private copy with source timestamps preserved. Warmup always runs, and successful warmups atomically update the persistent cache before -measurement starts. Compatible simultaneous warmups wait for each other to avoid -losing compiled signatures; measured runs remain independent. Only `cache/flows` +measurement starts. Compatible simultaneous experiments wait for each other to avoid losing +compiled signatures when attempts update Numba cache indexes. Only `cache/flows` is shared, never model data, outputs, or shared-memory artifacts. Cache identity and reuse counts are recorded in `flow-cache-identity.json` and `experiment.json`. The persistent cache can be deleted between runs to reclaim disk space; older experiments created before this feature are not automatically imported. -A smaller serial warmup may not exercise every flow signature required by the -measured run. **Cache misses still invalidate the measured experiment**: abench -does not silently compile, retry with compilation enabled, or accept incomplete -cache coverage. Increase the warmup cap when needed. -The CLI and failure report identify the phase, failed component, cache flow, and -log path; `cache-miss-details-*.jsonl` records the exact required Numba signature. -`phase-settings.json` and -`effective-settings.json` preserve the actual settings for each phase, and the -report includes the cache-build settings separately from the measured target. - -Numba compilation of generated flow overloads is rejected during measurement; ordinary non-flow -compilation and disk-cache loading remain included. This guard uses private -Numba internals and is covered by real disk-cache hit/miss tests. +A smaller serial warmup may not exercise every flow/type signature needed by the +measured run. Each measured attempt therefore records flow compilation and allows +it to finish. If compilation occurred, the completed attempt becomes **cache +preparation**, and none of its runtime or memory results qualify as benchmark +results. Its diagnostics and outputs are retained under `attempts/attempt-001`, +`attempts/attempt-002`, etc. Newly compiled flows are published to the shared cache. + +The model then restarts in a fresh container with fresh outputs and model caches, +using the same settings and expanded flow cache. Only an attempt with **zero flow +compilations** is accepted. All attempts keep permanent directories under +`attempts/`; `measured/` links to the accepted attempt. By default abench allows +**two additional attempts** (three total). Set `--cache-retries N` or +`cache_retries: N` in YAML; zero allows no retries. If compilation persists, the +experiment fails with the final attempt's diagnostics retained. Ordinary model +errors, OOMs, and output validation failures stop immediately and are never +retried as cache preparation. + +The report and `experiment.json` include attempt history. Each attempt retains +`cache-miss-details-*.jsonl`, its settings, component timings, and memory samples. +Ordinary non-flow compilation and disk-cache loading remain included in accepted +measurements. Flow tracking uses private Numba internals and is covered by real +compilation, cache-hit, and Docker retry tests. Memory is the whole-container cgroup v2 charge, counting shared pages once. Blue is `memory.current` (including file cache, shared memory, and kernel costs). diff --git a/examples/sandag-chunked.yaml b/examples/sandag-chunked.yaml index e7751ea..0ed86e2 100644 --- a/examples/sandag-chunked.yaml +++ b/examples/sandag-chunked.yaml @@ -16,7 +16,7 @@ defaults: sharrow: true households: ${households} # The 5000-household cache build misses some cache builds, so this test uses more - warmup_households: ${households} + warmup_households: 5000 memory: 80g shm_size: 8g platform: linux/arm64 diff --git a/src/abench/attempts.py b/src/abench/attempts.py new file mode 100644 index 0000000..4808c01 --- /dev/null +++ b/src/abench/attempts.py @@ -0,0 +1,66 @@ +"""Retry completed runs that prepared new flow signatures, never model errors.""" + +import shutil + +from .common import write_json +from .failures import BenchmarkFailure, describe_failure +from .report import load_run + + +def measured_attempts(spec, output, run_phase, publish_cache): + """Select only a compilation-free attempt for the report. + + Each attempt starts in a new container with fresh outputs and model caches. + Only compiled flows survive; failed preparation timings never enter reports. + """ + spec["attempts"] = [] + for number in range(1, spec.get("cache_retries", 2) + 2): + model_cache = output / "cache/model" + if model_cache.exists(): + shutil.rmtree(model_cache) + staging = output / "attempts" / f"attempt-{number:03d}" + relative = str(staging.relative_to(output)) + record = {"number": number, "directory": relative, "status": "running"} + spec["attempts"].append(record) + spec["measured_directory"] = relative + write_json(output / "experiment.json", spec) + print(f"Running measured attempt {number}…", flush=True) + try: + run_phase(relative) + # Check all ordinary validity conditions before deciding to retry. + if not load_run(output, allow_cache_misses=True)["valid"]: + raise BenchmarkFailure( + describe_failure(output, "measured", ignore_cache=True) + ) + except BaseException: + record["status"] = "failed" + write_json(output / "experiment.json", spec) + raise + phase = staging + misses = sum( + len(path.read_text().splitlines()) + for path in phase.glob("cache-miss-*.txt") + ) + record.update( + compilations=misses, + status="cache preparation" if misses else "accepted", + ) + publish_cache() + write_json(output / "experiment.json", spec) + if not misses: + # Convenience alias only after containers finish; attempt directories + # never move or change identity while Docker may cache their paths. + (output / "measured").symlink_to(relative, target_is_directory=True) + return + if number > spec.get("cache_retries", 2): + raise BenchmarkFailure( + f"Flow compilation persisted after {number} completed attempts. " + "No valid benchmark was produced. Inspect attempts/*/cache-miss-details-*.jsonl " + "for changing signatures; increase --cache-retries if appropriate. " + f"Diagnostics: {output}" + ) + print( + f"Attempt {number} compiled {misses} flow signatures; retained as cache " + f"preparation at {staging}. Retrying with a fresh model state…", + flush=True, + ) diff --git a/src/abench/cli.py b/src/abench/cli.py index 7a044f4..1b5032a 100644 --- a/src/abench/cli.py +++ b/src/abench/cli.py @@ -15,9 +15,10 @@ from pathlib import Path from . import __version__ +from .attempts import measured_attempts from .common import read_json, write_json from .failures import BenchmarkFailure, describe_failure -from .flow_cache import reuse_flows +from .flow_cache import publish_flows, reuse_flows from .profiles import load_profile, validate_model from .report import load_run, report from .sources import resolve_sources @@ -85,6 +86,12 @@ def parser(): default=5000, help="maximum cache-build households (default: 5000); warmup is always single-process", ) + p.add_argument( + "--cache-retries", + type=int, + default=2, + help="additional attempts after completed runs compile flows (default: 2)", + ) p.add_argument("--data-dir", type=Path, default=None) p.add_argument( "--config-overlay", @@ -162,13 +169,22 @@ def mount(source, target, readonly=False): def container_phase(spec, output, data, image, phase_name): """Retain Docker exit/OOM state even when the supervisor cannot finish.""" phase = output / phase_name - phase.mkdir() + phase.mkdir(parents=True) + phase_kind = "warmup" if phase_name == "warmup" else "measured" name = "abench-" + uuid.uuid4().hex[:12] + # Each container reads an immutable, uniquely named settings snapshot. Reusing + # the mutable root manifest across Docker Desktop mounts can expose stale data. + snapshot_name = f"spec-{name}.json" + write_json(phase / snapshot_name, spec) args = [ "docker", "run", "--name", name, + "--env", + f"BENCH_SPEC_PATH=/results/{phase_name}/{snapshot_name}", + "--env", + f"BENCH_PHASE_NAME={phase_kind}", "--cgroupns=private", "--memory", spec["memory"], @@ -262,6 +278,8 @@ def main(argv=None): ) if args.sharrow and not args.sharrow_commit: p.error("Sharrow enabled: provide a sharrow source override") + if args.cache_retries < 0: + p.error("--cache-retries must be nonnegative") if args.warmup_households < 1: p.error("--warmup-households must be positive") if args.households < 0: @@ -456,6 +474,8 @@ def main(argv=None): # Only flow artifacts are reused. A small serial warmup prepares flows; # measurement remains responsible for rejecting missing signatures. shutil.copytree(seed / "cache/flows", output / "cache/flows") + cache = nullcontext(None) + identity = None if args.sharrow: stage = "warmup" print( @@ -487,24 +507,35 @@ def main(argv=None): "Checking compatible flow cache (waiting for any active warmup)…", flush=True, ) - with cache as cache_info: + with cache as cache_info: + + def publish_cache(): if cache_info is not None: - spec["flow_cache"] = cache_info - write_json(output / "experiment.json", spec) - print( - f"Reused {cache_info['restored_files']} flow-cache files.", - flush=True, - ) + publish_flows(identity, output / "cache/flows", cache_info) + + if cache_info is not None: + spec["flow_cache"] = cache_info + write_json(output / "experiment.json", spec) + print( + f"Reused {cache_info['restored_files']} flow-cache files.", + flush=True, + ) + if args.sharrow: container_phase(spec, output, data, image, "warmup") + publish_cache() write_json(output / "experiment.json", spec) - stage = "measured" - print("Running measured model…", flush=True) - container_phase(spec, output, data, image, "measured") - if not load_run(output)["valid"]: - raise BenchmarkFailure(describe_failure(output, "measured")) + stage = "measured" + measured_attempts( + spec, + output, + lambda phase_name: container_phase( + spec, output, data, image, phase_name + ), + publish_cache, + ) except (Exception, KeyboardInterrupt) as error: if isinstance(error, subprocess.CalledProcessError): - error = BenchmarkFailure(describe_failure(output, stage)) + error = BenchmarkFailure(describe_failure(output, stage, ignore_cache=True)) spec["failure"] = {"phase": stage, "error": str(error)} write_json(output / "experiment.json", spec) raise error from None diff --git a/src/abench/common.py b/src/abench/common.py index ab78a72..5330de7 100644 --- a/src/abench/common.py +++ b/src/abench/common.py @@ -1,6 +1,8 @@ """Small artifact helpers shared by the host tools.""" import json +import os +import tempfile def read_json(path, default=None): @@ -8,4 +10,13 @@ def read_json(path, default=None): def write_json(path, value): - path.write_text(json.dumps(value, indent=2) + "\n") + """Atomically replace metadata so Docker bind readers never see a rewrite.""" + content = json.dumps(value, indent=2) + "\n" + descriptor, temporary = tempfile.mkstemp(prefix=f".{path.name}-", dir=path.parent) + try: + with os.fdopen(descriptor, "w") as stream: + stream.write(content) + os.replace(temporary, path) + finally: + if os.path.exists(temporary): + os.unlink(temporary) diff --git a/src/abench/experiments.py b/src/abench/experiments.py index aa4bfd4..9257638 100644 --- a/src/abench/experiments.py +++ b/src/abench/experiments.py @@ -23,6 +23,7 @@ "sharrow", "households", "warmup_households", + "cache_retries", "data_dir", "config_overlay", "cache_from", diff --git a/src/abench/failures.py b/src/abench/failures.py index 0307713..d76e025 100644 --- a/src/abench/failures.py +++ b/src/abench/failures.py @@ -21,10 +21,13 @@ def tail(path, limit=131072): return stream.read().decode("utf-8", errors="replace") -def describe_failure(output, phase): +def describe_failure(output, phase, *, ignore_cache=False): """Prefer explicit cache/OOM evidence over a generic worker exit exception.""" output = Path(output) - directory = output / phase + spec = read_json(output / "experiment.json", {}) + directory = output / ( + spec.get("measured_directory", "measured") if phase == "measured" else phase + ) log = output / "build.log" if phase == "build" else directory / "console.log" docker = read_json(directory / "docker-state.json", {}) components = set() @@ -45,7 +48,25 @@ def describe_failure(output, phase): if line.strip() } ) - if misses: + spec = read_json(output / "experiment.json", {}) + retry_mode = "cache_retries" in spec + completed = ( + docker.get("ExitCode") == 0 + and read_json(directory / "status.json", {}).get("returncode") == 0 + ) + if ( + misses + and not ignore_cache + and retry_mode + and completed + and not docker.get("OOMKilled") + ): + reason = "Completed attempt compiled Sharrow flows; its runtime and memory are excluded." + remedy = ( + spec.get("failure", {}).get("error") + or "See attempt history and cache-miss-details-*.jsonl for required signatures." + ) + elif misses and not ignore_cache and not retry_mode: reason = f"Sharrow flow cache miss{where}: " + ", ".join( Path(name).parent.name for name in misses ) diff --git a/src/abench/flow_cache.py b/src/abench/flow_cache.py index 798a6ce..f582fd2 100644 --- a/src/abench/flow_cache.py +++ b/src/abench/flow_cache.py @@ -18,10 +18,10 @@ def cache_key(identity): @contextmanager def reuse_flows(root, identity, destination): - """Serialize compatible warmups and atomically publish successful snapshots. + """Serialize compatible experiments and atomically publish successful snapshots. - Keeping the lock through warmup prevents concurrent writers from losing Numba - signature indexes. Measurement uses its own copy and never holds this lock. + Keeping the lock through all attempts prevents concurrent writers from losing + Numba signature indexes when a measured attempt compiles additional flows. copy2 preserves source mtimes, which Numba checks when loading compiled code. """ bucket = Path(root).expanduser().resolve() / cache_key(identity) @@ -41,19 +41,25 @@ def reuse_flows(root, identity, destination): "published": False, } yield metadata - # A failed warmup leaves the last successful snapshot intact. Publishing - # before measurement also retains useful flows if measurement later fails. - snapshot = Path(tempfile.mkdtemp(prefix="snapshot-", dir=bucket)) - try: - if destination.exists(): - shutil.copytree(destination, snapshot, dirs_exist_ok=True) - write_json(bucket / "identity.json", identity) - pointer = bucket / "current.tmp" - write_json(pointer, {"snapshot": snapshot.name}) - pointer.replace(bucket / "current.json") - except BaseException: - shutil.rmtree(snapshot) - raise - metadata["published"] = True - if previous: - shutil.rmtree(previous) + publish_flows(identity, destination, metadata) + + +def publish_flows(identity, destination, metadata): + """Publish under the lock held by reuse_flows, including completed attempts.""" + bucket = Path(metadata["directory"]) + current = read_json(bucket / "current.json", {}) + previous = bucket / current["snapshot"] if current else None + snapshot = Path(tempfile.mkdtemp(prefix="snapshot-", dir=bucket)) + try: + if destination.exists(): + shutil.copytree(destination, snapshot, dirs_exist_ok=True) + write_json(bucket / "identity.json", identity) + pointer = bucket / "current.tmp" + write_json(pointer, {"snapshot": snapshot.name}) + pointer.replace(bucket / "current.json") + except BaseException: + shutil.rmtree(snapshot) + raise + metadata["published"] = True + if previous: + shutil.rmtree(previous) diff --git a/src/abench/report.py b/src/abench/report.py index 4d95fd0..a6fb681 100644 --- a/src/abench/report.py +++ b/src/abench/report.py @@ -97,12 +97,12 @@ def component_windows(observations, phase): return windows -def load_run(directory): +def load_run(directory, *, allow_cache_misses=False): """Use raw worker observations, never duplicate ActivitySim's locutor CSV.""" spec = read_json(directory / "experiment.json") if not spec or spec.get("schema_version") not in (1, 2): raise ValueError(f"Not a supported benchmark experiment: {directory}") - phase = directory / "measured" + phase = directory / spec.get("measured_directory", "measured") grouped = {} observations = [] for path in sorted(phase.glob("components-*.jsonl")): @@ -136,12 +136,13 @@ def load_run(directory): ) valid = ( sample_matches + and not spec.get("failure") and status.get("returncode") == 0 and docker.get("ExitCode") == 0 and bool(memory) and bool(components) and not docker.get("OOMKilled") - and not list(phase.glob("cache-miss-*.txt")) + and (allow_cache_misses or not list(phase.glob("cache-miss-*.txt"))) ) warmup_settings = read_json(directory / "warmup/phase-settings.json", {}) effective = read_json(directory / "warmup/effective-settings.json", {}) @@ -160,8 +161,11 @@ def load_run(directory): "valid": valid, "inputs": read_json(phase / "input-summary.json", {}), "outputs": outputs, - "failure_reason": describe_failure( - directory, spec.get("failure", {}).get("phase", "measured") + "failure_reason": ( + spec.get("failure", {}).get("error") + or describe_failure( + directory, spec.get("failure", {}).get("phase", "measured") + ) ) if not valid else None, @@ -271,6 +275,7 @@ def experiment_card(run, xmax, ymax): "sharrow", "households", "warmup_households", + "cache_retries", "data_dir", "output_dir", "memory", @@ -317,6 +322,24 @@ def experiment_card(run, xmax, ymax): if run.get("failure_reason") else "" ) + attempts = spec.get("attempts", []) + attempt_table = "" + if attempts: + rows = "".join( + "" + + "".join( + f"{escape(attempt.get(key, '—'))}" + for key in ("number", "status", "compilations", "directory") + ) + + "" + for attempt in attempts + ) + attempt_table = ( + "

Attempt history

Cache-preparation attempts are excluded from benchmark results.

" + "" + + rows + + "
AttemptResultFlow compilationsArtifacts
" + ) details = "".join( f"
{title}
{escape(json.dumps(value, indent=2))}
" for title, value in ( @@ -329,7 +352,7 @@ def experiment_card(run, xmax, ymax): ) return ( f"

{escape(spec['label'])}

{'SUCCEEDED' if run['valid'] else 'FAILED / INCOMPLETE'}" - f" · elapsed: {elapsed} s · peak: {run['peak'] / 2**30:.3f} GiB

{failure}" + f" · elapsed: {elapsed} s · peak: {run['peak'] / 2**30:.3f} GiB

{failure}{attempt_table}" f"{memory_chart(run, xmax, ymax)}

Experiment settings

{settings}
" f"

Population and outputs

{''.join(counts)}
TableInput rowsOutput rows
" f"

Output households and persons are the realized sample.

{details}
" diff --git a/src/abench/runtime/instrumentation.py b/src/abench/runtime/instrumentation.py index e329e75..413ae05 100644 --- a/src/abench/runtime/instrumentation.py +++ b/src/abench/runtime/instrumentation.py @@ -1,4 +1,4 @@ -"""Process-local component timing and a strict Sharrow disk-cache guard.""" +"""Process-local component timing and Sharrow compilation tracking.""" import json import multiprocessing @@ -42,7 +42,10 @@ def timed(name, context, **kwargs): stream.write(json.dumps(row) + "\n") runner.run_named_step = timed - if os.environ.get("BENCH_STRICT_CACHE") == "1": + if ( + os.environ.get("BENCH_TRACK_CACHE") == "1" + or os.environ.get("BENCH_STRICT_CACHE") == "1" + ): from numba.core.dispatcher import _FunctionCompiler compile_original = _FunctionCompiler.compile @@ -65,9 +68,12 @@ def compile_checked(self, *args, **kwargs): "a" ) as stream: stream.write(json.dumps(detail) + "\n") - raise RuntimeError( - f"Measured Sharrow flow cache miss: {filename} ({self.py_func.__name__}); required signature: {detail['signature']}" - ) + # Preserve strict mode for old runners and diagnostic tools. + # New runs compile, finish, and let the host reject this attempt. + if os.environ.get("BENCH_STRICT_CACHE") == "1": + raise RuntimeError( + f"Measured Sharrow flow cache miss: {filename} ({self.py_func.__name__}); required signature: {detail['signature']}" + ) return compile_original(self, *args, **kwargs) _FunctionCompiler.compile = compile_checked diff --git a/src/abench/runtime/worker.py b/src/abench/runtime/worker.py index 9da4844..16acb5e 100644 --- a/src/abench/runtime/worker.py +++ b/src/abench/runtime/worker.py @@ -152,7 +152,7 @@ def make_state( def run_model(spec, phase): """Warm flows with a small serial run; measure the untouched target settings.""" - spec = phase_spec(spec, phase.name) + spec = phase_spec(spec, os.environ.get("BENCH_PHASE_NAME", phase.name)) write_json( phase / "phase-settings.json", {key: spec[key] for key in ("households", "multiprocess", "processes")}, @@ -256,8 +256,12 @@ def supervise(spec, phase): for name in ("flows", "model"): Path(f"/results/cache/{name}").mkdir(parents=True, exist_ok=True) env = dict(os.environ, BENCH_MODEL="1", BENCH_PHASE_DIR=str(phase)) - env["BENCH_STRICT_CACHE"] = ( - "1" if spec["sharrow"] and phase.name == "measured" else "0" + env["BENCH_STRICT_CACHE"] = "0" + env["BENCH_TRACK_CACHE"] = ( + "1" + if spec["sharrow"] + and os.environ.get("BENCH_PHASE_NAME", phase.name) == "measured" + else "0" ) started = time.perf_counter() env["BENCH_STARTED_MONOTONIC"] = str(started) @@ -292,7 +296,9 @@ def supervise(spec, phase): if __name__ == "__main__": mode, phase_arg = sys.argv[1:] phase = Path(phase_arg) - spec = json.loads(Path("/results/experiment.json").read_text()) + spec = json.loads( + Path(os.environ.get("BENCH_SPEC_PATH", "/results/experiment.json")).read_text() + ) if mode == "model": run_model(spec, phase) elif mode == "summary": diff --git a/tests/fixtures/tiny/tiny_extension.py b/tests/fixtures/tiny/tiny_extension.py index 4d295a7..c381b6a 100644 --- a/tests/fixtures/tiny/tiny_extension.py +++ b/tests/fixtures/tiny/tiny_extension.py @@ -1,6 +1,7 @@ """Tiny real ActivitySim workflow: exercises registration, slicing, and caching.""" import importlib +import os import sys from pathlib import Path @@ -26,11 +27,16 @@ def bench_initialize(state: workflow.State): @workflow.step def bench_compute(state: workflow.State, households: pd.DataFrame): - """Run independently in each slice; measured cache misses must be rejected.""" + """Exercise compiled overload reuse, including a measured-only test signature.""" sys.path.insert(0, state.settings.sharrow_cache_dir) twice = importlib.import_module("tiny_generated").twice households = households.copy() - households["value"] = [twice(int(value)) for value in households["value"]] + measured_signature = ( + Path("/model/force-measured-signature").exists() + and os.environ.get("BENCH_PHASE_NAME") == "measured" + ) + cast = float if measured_signature else int + households["value"] = [int(twice(cast(value))) for value in households["value"]] state.add_table("households", households) diff --git a/tests/test_attempts.py b/tests/test_attempts.py new file mode 100644 index 0000000..6fcf145 --- /dev/null +++ b/tests/test_attempts.py @@ -0,0 +1,84 @@ +"""Completed compilation attempts retry; ordinary failures never do.""" + +import json +import subprocess + +import pytest + +from abench.attempts import measured_attempts +from abench.common import write_json +from abench.failures import BenchmarkFailure +from abench.report import load_run + + +def run_fixture(tmp_path, misses, retries=2, failure=None): + """Provide real report artifacts while faking only the container boundary.""" + spec = dict(schema_version=2, cache_retries=retries, households=1) + calls, publications = [], [] + + def run(phase_name): + assert not (tmp_path / "cache/model").exists() + (tmp_path / "cache/model").mkdir(parents=True) + phase = tmp_path / phase_name + phase.mkdir(parents=True) + calls.append(True) + if failure == "process": + raise subprocess.CalledProcessError(1, "docker") + write_json(phase / "status.json", dict(returncode=0)) + write_json(phase / "docker-state.json", dict(ExitCode=0, OOMKilled=False)) + write_json( + phase / "output-summary.json", + dict(households=dict(rows=2 if failure else 1)), + ) + (phase / "components-1.jsonl").write_text( + json.dumps(dict(component="test", seconds=len(calls), succeeded=True)) + + "\n" + ) + (phase / "memory.csv").write_text( + "elapsed_seconds,current_bytes,peak_bytes\n0,1,1\n" + ) + if misses[len(calls) - 1]: + (phase / "cache-miss-1.txt").write_text( + "/results/cache/flows/flow_a/__init__.py\n" + ) + + measured_attempts(spec, tmp_path, run, lambda: publications.append(True)) + return spec, calls, publications + + +def test_preparation_excluded_and_clean_attempt_accepted(tmp_path): + spec, calls, publications = run_fixture(tmp_path, [True, False]) + assert len(calls) == len(publications) == 2 + assert [a["status"] for a in spec["attempts"]] == ["cache preparation", "accepted"] + assert (tmp_path / "attempts/attempt-001/cache-miss-1.txt").exists() + assert load_run(tmp_path)["valid"] + assert load_run(tmp_path)["components"]["test"]["mean"] == 2 + + +@pytest.mark.parametrize("retries", [0, 2]) +def test_persistent_compilation_exhausts_budget(tmp_path, retries): + with pytest.raises(BenchmarkFailure, match="No valid benchmark"): + run_fixture(tmp_path, [True] * (retries + 1), retries) + spec = json.loads((tmp_path / "experiment.json").read_text()) + assert len(spec["attempts"]) == retries + 1 + assert not load_run(tmp_path)["valid"] + + +@pytest.mark.parametrize("failure", ["process", "sample"]) +def test_ordinary_failure_not_retried_even_with_misses(tmp_path, failure): + with pytest.raises((BenchmarkFailure, subprocess.CalledProcessError)): + run_fixture(tmp_path, [True], failure=failure) + spec = json.loads((tmp_path / "experiment.json").read_text()) + assert len(spec["attempts"]) == 1 + assert spec["attempts"][0]["status"] == "failed" + + +def test_retry_options_in_yaml(tmp_path): + from abench.cli import parser + from abench.experiments import arguments + + assert parser().parse_args([]).cache_retries == 2 + assert ( + parser().parse_args(arguments({"cache_retries": 4}, tmp_path)).cache_retries + == 4 + ) diff --git a/tests/test_docker.py b/tests/test_docker.py index 6134d50..1026761 100644 --- a/tests/test_docker.py +++ b/tests/test_docker.py @@ -21,11 +21,22 @@ SHARROW = "fc175b27d8e0c5d202721c67d96b050e6117b235" -@pytest.mark.parametrize("multiprocess,sharrow", [(False, False), (True, True)]) -def test_tiny_model(tmp_path, multiprocess, sharrow): +@pytest.mark.parametrize( + "multiprocess,sharrow,retry", + [(False, False, False), (True, True, False), (True, True, True)], +) +def test_tiny_model(tmp_path, multiprocess, sharrow, retry): """Build pinned sources, run a full warmup where enabled, then measure.""" root = tmp_path / "model" shutil.copytree(Path(__file__).parent / "fixtures/tiny", root) + if retry: + (root / "force-measured-signature").touch() + profile = root / "benchmark.yaml" + import yaml + + settings = yaml.safe_load(profile.read_text()) + settings["snapshot"].append("force-measured-signature") + profile.write_text(yaml.safe_dump(settings)) output = tmp_path / "experiment" args = [ "run", @@ -53,6 +64,16 @@ def test_tiny_model(tmp_path, multiprocess, sharrow): assert cli.main(args) == 0 run = load_run(output) assert run["valid"] + attempts = run["spec"]["attempts"] + assert len(attempts) == (2 if retry else 1) + assert attempts[-1]["status"] == "accepted" + assert attempts[-1]["compilations"] == 0 + if retry: + assert attempts[0]["status"] == "cache preparation" + archived = output / attempts[0]["directory"] + assert list(archived.glob("cache-miss-*.txt")) + assert json.loads((archived / "status.json").read_text())["returncode"] == 0 + assert (archived / "output/final_households.csv").exists() assert run["components"]["bench_compute"]["n"] == (2 if multiprocess else 1) assert run["outputs"]["households"]["rows"] == 4 assert run["memory"] and all(row["current_bytes"] > 0 for row in run["memory"]) @@ -66,6 +87,16 @@ def test_tiny_model(tmp_path, multiprocess, sharrow): assert not list((output / "measured").glob("cache-miss-*")) lines = (output / "measured/output/final_households.csv").read_text().splitlines() assert set(lines[1:]) == {"1,20", "2,40", "3,60", "4,80"} + if retry: + # The signature prepared by measurement must survive in the shared cache. + repeated = tmp_path / "repeat" + args[args.index("--output-dir") + 1] = str(repeated) + assert cli.main(args) == 0 + reused = load_run(repeated) + assert reused["valid"] + assert len(reused["spec"]["attempts"]) == 1 + assert reused["spec"]["flow_cache"]["restored_files"] > 0 + assert "cache preparation" in (output / "report.html").read_text() def test_named_suite(tmp_path): diff --git a/tests/test_experiments.py b/tests/test_experiments.py index d205b87..ae39d53 100644 --- a/tests/test_experiments.py +++ b/tests/test_experiments.py @@ -174,7 +174,7 @@ def test_shipped_sandag_suite(): for run in plan["runs"]: args = cli.parser().parse_args(run["argv"]) assert args.households == 28365 - assert args.warmup_households == 5000 + assert args.warmup_households == 28365 assert args.processes == 4 assert args.sharrow assert args.data_dir.name == "benchmarking-data" diff --git a/tests/test_failure_diagnostics.py b/tests/test_failure_diagnostics.py index f684e9b..15e2d3b 100644 --- a/tests/test_failure_diagnostics.py +++ b/tests/test_failure_diagnostics.py @@ -67,15 +67,18 @@ def container(spec, output, data, image, phase_name): import subprocess phase = output / phase_name - phase.mkdir() + phase.mkdir(parents=True) (phase / "cache-miss-1.txt").write_text( "/results/cache/flows/flow_ABC/__init__.py\n" ) + (phase / "console.log").write_text( + "ValueError: ordinary model error after compilation\n" + ) raise subprocess.CalledProcessError(1, ["docker", "run", "test-image"]) monkeypatch.setattr(cli, "command", command) monkeypatch.setattr(cli, "container_phase", container) - with pytest.raises(BenchmarkFailure, match="flow_ABC") as error: + with pytest.raises(BenchmarkFailure, match="ordinary model error") as error: cli.main( [ "run", @@ -92,7 +95,7 @@ def container(spec, output, data, image, phase_name): assert (output / "report.html").is_file() failure = json.loads((output / "experiment.json").read_text())["failure"] assert failure["phase"] == "measured" - assert "Results rejected" in failure["error"] + assert "ordinary model error" in failure["error"] def test_successful_exit_with_wrong_sample_is_explained(tmp_path): From c4a1b9c4d17964d4439eeba825600849150a080c Mon Sep 17 00:00:00 2001 From: Jeff Newman Date: Fri, 18 Sep 2026 17:09:05 -0500 Subject: [PATCH 7/8] Configure PyPI trusted publishing and uvx releases --- .github/workflows/ci.yml | 1 + .github/workflows/publish.yml | 61 +++++++++++++++++++++++++++++++++++ MANIFEST.in | 2 +- README.md | 20 +++++++++++- RELEASING.md | 52 +++++++++++++++++++++++++++++ pyproject.toml | 21 +++++++++++- src/abench/cli.py | 1 + tests/test_experiments.py | 2 +- 8 files changed, 156 insertions(+), 4 deletions(-) create mode 100644 .github/workflows/publish.yml create mode 100644 RELEASING.md diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 23f2a24..2867518 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -3,6 +3,7 @@ on: push: pull_request: workflow_dispatch: + workflow_call: permissions: contents: read concurrency: diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 0000000..0610176 --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,61 @@ +name: Publish to PyPI +on: + release: + types: [published] +permissions: + contents: read +concurrency: + group: pypi-${{ github.event.release.tag_name }} + cancel-in-progress: false +jobs: + tests: + uses: ./.github/workflows/ci.yml + build: + needs: tests + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.11' + - uses: astral-sh/setup-uv@v6 + - run: python -m pip install build twine + # Build the wheel from the sdist, so both published formats are exercised. + - run: python -m build + - run: python -m twine check --strict dist/* + - run: python -m pip install dist/*.whl + - name: Match release tag to package version + env: + RELEASE_TAG: ${{ github.event.release.tag_name }} + run: | + python - <<'PY' + import os + from importlib.metadata import version + expected = 'v' + version('abench') + if os.environ['RELEASE_TAG'] != expected: + raise SystemExit(f'Release tag must be {expected}') + PY + - name: Smoke test isolated uvx installation + run: | + uvx --from "$PWD"/dist/abench-*.whl abench --version + uvx --from "$PWD"/dist/abench-*.whl abench --help + - uses: actions/upload-artifact@v4 + with: + name: python-distributions + path: dist/* + if-no-files-found: error + publish: + needs: build + runs-on: ubuntu-latest + environment: + name: pypi + url: https://pypi.org/p/abench + permissions: + id-token: write + steps: + - uses: actions/download-artifact@v4 + with: + name: python-distributions + path: dist/ + - name: Publish distributions using Trusted Publishing + uses: pypa/gh-action-pypi-publish@release/v1 diff --git a/MANIFEST.in b/MANIFEST.in index d930cc1..e3854d2 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,4 +1,4 @@ -include LICENSE README.md .pre-commit-config.yaml +include LICENSE README.md RELEASING.md .pre-commit-config.yaml recursive-include tests *.py *.yaml *.csv *.txt recursive-include .github *.yml recursive-include examples *.yaml diff --git a/README.md b/README.md index 3ebfd57..189495b 100644 --- a/README.md +++ b/README.md @@ -4,8 +4,21 @@ Run reproducible ActivitySim runtime and memory experiments in Linux Docker, from macOS or Linux. One runner supports MTC, SANDAG ABM3, and other models through small YAML profiles. ActivitySim itself does not need to be installed on the host. +With [uv](https://docs.astral.sh/uv/), run without managing a Python environment: + +```bash +uvx abench --help +uvx abench experiments.yaml +``` + +For a specific release use `uvx abench@0.1.0 experiments.yaml`; use +`uvx abench@latest` to refresh to the latest release. Docker and model data must +still be available locally. macOS and Linux hosts are supported. + +Alternatively, install with pip: + ```bash -python -m pip install /path/to/abench +python -m pip install abench abench run --model-dir /path/to/sandag-abm3-example --profile sandag \ --source activitysim=ActivitySim/activitysim@ \ --source sharrow=ActivitySim/sharrow@ \ @@ -331,3 +344,8 @@ and reports. They do not require either example repository or large datasets. Adapted from the MTC and SANDAG benchmark harnesses developed in this workspace. The original measurement approach was informed by WSP's Lighthouse production benchmark. See LICENSE for the retained BSD license. + +## Releases + +See [RELEASING.md](https://github.com/ActivitySim/abench/blob/main/RELEASING.md) +for Trusted Publishing setup and release instructions. diff --git a/RELEASING.md b/RELEASING.md new file mode 100644 index 0000000..3fbd873 --- /dev/null +++ b/RELEASING.md @@ -0,0 +1,52 @@ +# Publishing abench + +The package supports macOS and Linux hosts with Python 3.10 or newer. Docker is +an external prerequisite; ActivitySim and Sharrow run inside its Linux containers. +The console entry point allows `uvx abench experiments.yaml` once published. + +## One-time PyPI setup + +On the owning PyPI account's [Publishing page](https://pypi.org/manage/account/publishing/), +add a pending GitHub publisher with exactly these values: + +| Field | Value | +| --- | --- | +| PyPI project name | `abench` | +| Owner | `ActivitySim` | +| Repository name | `abench` | +| Workflow name | `publish.yml` | +| Environment name | `pypi` | + +Create the `pypi` environment in the GitHub repository settings. No PyPI API token +or GitHub Actions secret is needed. The first successful upload creates the PyPI +project; registering a pending publisher does not reserve the name. Configure +project owners/organization membership on PyPI as appropriate after creation. + +See the official [pending publisher instructions](https://docs.pypi.org/trusted-publishers/creating-a-project-through-oidc/). + +## Release procedure + +1. Update `__version__` in `src/abench/__init__.py`. Package metadata reads this + value automatically; there is no second version field to synchronize. +2. Run pre-commit and the test suite. Review the code intended for the release. +3. Commit and push the release changes on a working branch, then merge through + the repository's normal review process. +4. Create a GitHub release with tag `v` targeting that reviewed commit + (for example, `v0.1.0`). Publishing the release triggers `publish.yml`. +5. The workflow runs unit and Docker tests, builds and checks the sdist and wheel, + verifies the tag, and runs an isolated `uvx` smoke test before uploading to PyPI. +6. Verify from outside the checkout: + + ```bash + uvx abench@0.1.0 --version + uvx abench@0.1.0 --help + uvx abench@latest /absolute/path/to/experiments.yaml + ``` + +For prereleases, use a PEP 440 version such as `0.1.0rc1` and tag `v0.1.0rc1`. +A published GitHub prerelease also triggers publishing; users should explicitly +request its version with `uvx abench@0.1.0rc1`. + +PyPI release files cannot be replaced. If published code needs changing, release +a new version. A failed workflow can be rerun after fixing publisher configuration +if no distributions have been uploaded yet. diff --git a/pyproject.toml b/pyproject.toml index 4f18484..7cbd69c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,13 +4,32 @@ build-backend = "setuptools.build_meta" [project] name = "abench" -version = "0.1.0" +dynamic = ["version"] description = "Reproducible ActivitySim runtime and memory benchmarks in Linux containers" readme = "README.md" requires-python = ">=3.10" license = "BSD-3-Clause" license-files = ["LICENSE"] dependencies = ["PyYAML>=6"] +keywords = ["activitysim", "benchmark", "transportation", "memory", "sharrow"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Environment :: Console", + "Operating System :: MacOS", + "Operating System :: POSIX :: Linux", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Topic :: System :: Benchmark", +] + +[project.urls] +Homepage = "https://github.com/ActivitySim/abench" +Documentation = "https://github.com/ActivitySim/abench#readme" +Repository = "https://github.com/ActivitySim/abench" +Issues = "https://github.com/ActivitySim/abench/issues" + +[tool.setuptools.dynamic] +version = {attr = "abench.__version__"} [project.optional-dependencies] test = ["pytest>=8", "pandas<3", "pyarrow", "numba", "ruff", "build"] diff --git a/src/abench/cli.py b/src/abench/cli.py index 1b5032a..b98f72e 100644 --- a/src/abench/cli.py +++ b/src/abench/cli.py @@ -61,6 +61,7 @@ def parser(): description=__doc__, epilog="Named experiments: abench experiments.yaml; preflight: abench validate experiments.yaml", ) + p.add_argument("--version", action="version", version=f"abench {__version__}") p.add_argument("--model-dir", type=Path, default=Path.cwd()) p.add_argument("--profile", default="benchmark.yaml") p.add_argument( diff --git a/tests/test_experiments.py b/tests/test_experiments.py index ae39d53..d205b87 100644 --- a/tests/test_experiments.py +++ b/tests/test_experiments.py @@ -174,7 +174,7 @@ def test_shipped_sandag_suite(): for run in plan["runs"]: args = cli.parser().parse_args(run["argv"]) assert args.households == 28365 - assert args.warmup_households == 28365 + assert args.warmup_households == 5000 assert args.processes == 4 assert args.sharrow assert args.data_dir.name == "benchmarking-data" From e4f3dcdbf02a95f11fe1eec3c086bf5db0fc1a5e Mon Sep 17 00:00:00 2001 From: Jeff Newman Date: Fri, 18 Sep 2026 17:10:22 -0500 Subject: [PATCH 8/8] Run model containers as the host user for writable Linux caches --- src/abench/cli.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/abench/cli.py b/src/abench/cli.py index b98f72e..6d976d0 100644 --- a/src/abench/cli.py +++ b/src/abench/cli.py @@ -4,6 +4,7 @@ import hashlib import json import math +import os import re import shutil import subprocess @@ -182,6 +183,10 @@ def container_phase(spec, output, data, image, phase_name): "run", "--name", name, + # Native Linux bind mounts preserve ownership. Run as the invoking user + # so outputs and caches remain writable between attempts and experiments. + "--user", + f"{os.getuid()}:{os.getgid()}", "--env", f"BENCH_SPEC_PATH=/results/{phase_name}/{snapshot_name}", "--env",