From 7980d658990423053b5c815313af96aa9c02a075 Mon Sep 17 00:00:00 2001 From: Jeff Newman Date: Fri, 18 Sep 2026 18:19:27 -0500 Subject: [PATCH 1/4] Add typed suite inputs, source selectors, and data preparation Support --set overrides with validation and file-specific help. Resolve branch and PR sources to exact commits, prepare checksum-verified model assets in shared caches, and record suite provenance. --- README.md | 175 +++++++++++++++++++++++- examples/sandag-chunked.yaml | 21 ++- pyproject.toml | 2 +- src/abench/assets.py | 244 ++++++++++++++++++++++++++++++++++ src/abench/cli.py | 38 ++++-- src/abench/experiments.py | 54 ++++++-- src/abench/inputs.py | 156 ++++++++++++++++++++++ src/abench/sources.py | 76 +++++++++++ tests/test_assets.py | 249 +++++++++++++++++++++++++++++++++++ tests/test_docker.py | 45 ++++++- tests/test_experiments.py | 153 ++++++++++++++++++++- tests/test_inputs.py | 161 ++++++++++++++++++++++ tests/test_sources.py | 96 ++++++++++++++ 13 files changed, 1429 insertions(+), 41 deletions(-) create mode 100644 src/abench/assets.py create mode 100644 src/abench/inputs.py create mode 100644 tests/test_assets.py create mode 100644 tests/test_inputs.py create mode 100644 tests/test_sources.py diff --git a/README.md b/README.md index 189495b..a2a0fe7 100644 --- a/README.md +++ b/README.md @@ -45,9 +45,16 @@ Write common options once and override only what differs between runs: ```yaml schema_version: 1 +inputs: + households: + type: integer + default: 28365 + minimum: 0 + warmup_households: + type: integer + default: 5000 + minimum: 1 vars: - households: 28365 - warmup_households: 5000 model: /path/to/sandag-abm3-example output_root: ./results/sandag-${timestamp} defaults: @@ -86,14 +93,124 @@ in the repository. Its paths assume abench and the SANDAG repository are sibling The pinned `main` revision is the one used in the earlier trials, not a moving branch reference. +To compare **current main against a PR**, named suites also accept source mappings +with `branch` or `pr` in place of `commit` (local development feature, not in +PyPI 0.1.0 yet): + +```yaml +schema_version: 1 +inputs: + activitysim_pr: + type: integer + required: true + minimum: 1 + description: ActivitySim PR number to compare against current main +output_root: benchmark-runs/mtc-${timestamp} +defaults: + profile: mtc + households: 500000 + multiprocess: true + processes: 4 + sharrow: true + sources: + - sharrow=ActivitySim/sharrow@fc175b27d8e0c5d202721c67d96b050e6117b235 +runs: + main: + sources: + - name: activitysim + repository: ActivitySim/activitysim + branch: main + pr: + label: ActivitySim PR ${activitysim_pr} + sources: + - name: activitysim + repository: ActivitySim/activitysim + pr: ${activitysim_pr} +``` + +Each mapping must provide exactly one of `commit`, `branch`, or `pr`. Branch/PR +selectors work for any source package in a suite, including extras/subdirectories. +Abench uses host Git and network access to resolve each selector once per suite, +before running models, and passes only exact commits to the builds. `pr` selects +`refs/pull//head` in the named repository, including PRs from forks; +it does **not** select the synthetic merge commit. Use a positive integer PR number. +`branch` names the branch without `refs/heads/`. + +The resolved commits are printed and recorded in `suite.json` alongside the +original YAML, and in each experiment's source provenance. A new invocation +resolves the selectors again, including `validate` and `prepare`. To reproduce an +earlier suite, replace selectors with `commit: `. CLI overrides +and model profiles still require exact commits. + +Override the example's PR without editing YAML: + +```bash +uvx abench ./activitysim-prototype-mtc/abench.yaml --set activitysim_pr=1110 +# Preflight the same selection without running benchmarks: +uvx abench validate ./activitysim-prototype-mtc/abench.yaml --set activitysim_pr=1110 +``` + +Declare user-facing settings in `inputs`, and keep internal reusable values and +expressions in `vars`. Only inputs can be overridden. For example: + +```yaml +inputs: + households: + type: integer + default: 500000 + minimum: 0 + description: Households to sample; 0 uses the full population + mode: + type: string + default: chunked + choices: [chunked, unchunked] + sharrow: + type: boolean + default: true +vars: + label: "${mode}, ${households} households" +``` + +Use repeatable `--set NAME=VALUE` with run, validate, or prepare commands. Quote +arguments containing spaces, for example `--set 'label=My experiment'` if `label` +is declared as a string input. `abench experiment.yaml --help` lists the file's +inputs, descriptions, defaults, and constraints without requiring input values, +contacting GitHub, downloading data, or creating outputs. + +- Every input requires a `type`: `string`, `integer`, `number`, or `boolean`. + Give it either a typed YAML `default` or `required: true`, but not both. +- Optional `choices` restricts allowed values; `minimum`/`maximum` are inclusive + bounds for numeric inputs. Numbers must be finite. Boolean overrides accept + `true` or `false` (case insensitive), not `yes`, `no`, `1`, or `0`. +- Values are converted according to their declared type, never parsed as YAML. + Strings retain literal text, including `=`, spaces, and `${...}`. Defaults and + input declarations are literal too; use `vars` for derived expressions. +- Defaults are applied first, followed by command-line overrides, then `${...}` + expansion. Inputs and vars share one namespace; duplicate names and the + reserved name `timestamp` are errors. Input names use letters, digits, and + underscores and cannot start with a digit. +- Unknown inputs, duplicate overrides, invalid values, and missing required + inputs fail before source resolution or downloads. There is no overriding vars. +- The YAML file is unchanged. `suite.json` records typed `cli_overrides`, effective + `input_values`, the expanded configuration, and exact source resolutions. + `experiments.yaml` preserves the original file, so replay its recorded overrides + too when reproducing a run. + +`branch: main` is freshly resolved on **every launch**; there is no saved branch +SHA to update manually. For the MTC example, use `--set activitysim_pr=1110` and +optionally `--set households=500000 --set processes=4`. + +These commands require a release containing this feature. Until then, use +`uvx --from /path/to/abench abench ...` with this local checkout. + - `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. +- `${name}` substitutes a scalar from `inputs` or `vars`; vars can reference + other vars and inputs. 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. @@ -108,8 +225,8 @@ branch reference. 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. +- File invocations accept only `--set` input overrides; other model options belong + in `defaults` or the relevant run. This experiment file describes **which tests to run**. A model profile such as `benchmark.yaml` describes **how to configure a model**, and remains reusable @@ -349,3 +466,49 @@ benchmark. See LICENSE for the retained BSD license. See [RELEASING.md](https://github.com/ActivitySim/abench/blob/main/RELEASING.md) for Trusted Publishing setup and release instructions. + +## Downloading model data + +Experiment suites may declare `data_assets` using ActivitySim's external-example +asset names, URLs, SHA-256 checksums, and `unpack` destinations: + +```yaml +data_assets: + name: prototype_mtc_extended + assets: + data_full.tar.zst: + url: https://github.com/ActivitySim/activitysim-prototype-mtc/releases/download/v1.3.4/data_full.tar.zst + sha256: b402506a61055e2d38621416dd9a5c7e3cf7517c0a9ae5869f6d760c03284ef3 + unpack: data_full +``` + +`abench experiment.yaml` prepares these assets before validating model inputs or +starting Docker. `abench prepare experiment.yaml` only prepares data. `abench +validate experiment.yaml` remains read-only: it checks declarations and requires +inputs to be present (use `prepare` first for a fresh clone). Asset destinations +are relative to the experiment YAML, regardless of `model_dir`; point each run's +`data_dir` to the appropriate destination. Variable substitutions also work here. + +The default download cache is exactly +`platformdirs.user_cache_dir("ActivitySim")/External-Examples//`—the same +layout used by `activitysim.examples.external.download_external_example`. +`name` is optional; `cache_dir` can override the root before appending `name`. +Keep names and checksums identical to ActivitySim's declarations to share files. +For assets cached by direct `download_asset(link=True)` calls, set `cache_dir` +to that call's `platformdirs.user_data_dir("ActivitySim")` and omit `name`. +No host ActivitySim installation is needed. + +Checksums are required and verified before reuse. For a `.gz` URL whose asset name +omits `.gz`, the checksum covers the decompressed file, as in ActivitySim. +Archives (`.tar.zst`, `.tar.gz`, `.zip`) use the archive checksum and preserve their +internal paths when unpacking. Verified extracted contents are cached under +`.abench-extracted//` beside the archive. Whole unpacked directories are +linked to the suite, and abench resolves `data_dir` before Docker mounts it. +Individual files are copied so external file symlinks cannot break inside Docker. +Existing destinations must match; modified inputs are never silently replaced. +Archive links, special files, and paths escaping the destination are rejected. +The original instructions and resolved cache locations are saved in `suite.json`. + +This feature is not present in abench 0.1.0. Until the next release, install the +updated checkout (`uv tool install /path/to/abench`) or run it with +`uvx --from /path/to/abench abench /path/to/model/abench.yaml`. diff --git a/examples/sandag-chunked.yaml b/examples/sandag-chunked.yaml index 0ed86e2..6376651 100644 --- a/examples/sandag-chunked.yaml +++ b/examples/sandag-chunked.yaml @@ -1,7 +1,20 @@ # Run from any directory: abench /path/to/abench/examples/sandag-chunked.yaml schema_version: 1 +inputs: + households: + type: integer + default: 28365 + minimum: 0 + description: Households to sample from benchmarking-data; 0 uses all + processes: + type: integer + default: 4 + minimum: 1 + warmup_households: + type: integer + default: 5000 + minimum: 1 vars: - households: 28365 model: ../../sandag-abm3-example activitysim_main: 5c6fae24a91a57a2d6dfc2e1dbe062a61d94545a activitysim_pr1110: 51e298a84276813946e1d623c9a5785e078e022f @@ -12,11 +25,11 @@ defaults: data_dir: ${model}/benchmarking-data config_overlay: ["${model}/configs_explicit_chunk"] multiprocess: true - processes: 4 + processes: ${processes} sharrow: true households: ${households} - # The 5000-household cache build misses some cache builds, so this test uses more - warmup_households: 5000 + # Cache misses in the measured run trigger a retry using the compiled flows. + warmup_households: ${warmup_households} memory: 80g shm_size: 8g platform: linux/arm64 diff --git a/pyproject.toml b/pyproject.toml index 7cbd69c..60205b9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ readme = "README.md" requires-python = ">=3.10" license = "BSD-3-Clause" license-files = ["LICENSE"] -dependencies = ["PyYAML>=6"] +dependencies = ["PyYAML>=6", "platformdirs>=3", "zstandard>=0.21"] keywords = ["activitysim", "benchmark", "transportation", "memory", "sharrow"] classifiers = [ "Development Status :: 3 - Alpha", diff --git a/src/abench/assets.py b/src/abench/assets.py new file mode 100644 index 0000000..9bc0dda --- /dev/null +++ b/src/abench/assets.py @@ -0,0 +1,244 @@ +"""Checksum-verified assets using ActivitySim's external-example cache layout.""" + +import fcntl +import gzip +import hashlib +import os +import shutil +import tarfile +import tempfile +import urllib.request +import zipfile +from pathlib import Path, PurePosixPath +from urllib.parse import urlparse + +import platformdirs + +from .common import read_json, write_json + + +def relative(value): + """Asset paths must stay within their chosen cache and destination roots.""" + if not isinstance(value, str) or not value or "\\" in value: + raise ValueError(f"invalid asset path: {value!r}") + path = PurePosixPath(value) + if path.is_absolute() or ".." in path.parts or str(path) == ".": + raise ValueError(f"asset path must be relative without '..': {value!r}") + return str(path) + + +def plan_assets(config, base): + """Resolve declarations without network access or filesystem mutations.""" + if config is None: + return [] + if not isinstance(config, dict) or set(config) - {"name", "cache_dir", "assets"}: + raise ValueError("data_assets accepts name, cache_dir, and assets") + cache = Path(platformdirs.user_cache_dir("ActivitySim")) / "External-Examples" + if "cache_dir" in config and ( + not isinstance(config["cache_dir"], str) or not config["cache_dir"] + ): + raise ValueError("data_assets.cache_dir must be a nonempty path string") + if config.get("cache_dir"): + cache = (base / Path(config["cache_dir"]).expanduser()).resolve() + if config.get("name"): + cache /= relative(config["name"]) + assets = config.get("assets") + if not isinstance(assets, dict) or not assets: + raise ValueError("data_assets.assets must be a nonempty mapping") + result = [] + for name, info in assets.items(): + name = relative(name) + if not isinstance(info, dict) or set(info) - {"url", "sha256", "unpack"}: + raise ValueError(f"asset {name} accepts url, sha256, and unpack") + url, digest = info.get("url"), info.get("sha256") + if not isinstance(url, str) or urlparse(url).scheme not in {"https", "http"}: + raise ValueError(f"asset {name} needs an HTTP(S) URL") + if ( + not isinstance(digest, str) + or len(digest) != 64 + or any(c not in "0123456789abcdefABCDEF" for c in digest) + ): + raise ValueError(f"asset {name} needs a full SHA-256 checksum") + unpack = relative(info["unpack"]) if info.get("unpack") else None + if unpack and not name.endswith((".zip", ".tar.gz", ".tar.zst")): + raise ValueError(f"unsupported asset archive: {name}") + destination = base / (unpack or name) + if not destination.parent.resolve().is_relative_to(base.resolve()): + raise ValueError( + f"asset destination parent escapes suite directory: {destination}" + ) + if any( + destination == Path(x["destination"]) + or destination in Path(x["destination"]).parents + or Path(x["destination"]) in destination.parents + for x in result + ): + raise ValueError("asset destinations must not overlap") + result.append( + dict( + name=name, + url=url, + sha256=digest.lower(), + unpack=unpack, + cache_file=str(cache / name), + destination=str(destination), + ) + ) + return result + + +def checksum(path): + digest = hashlib.sha256() + with path.open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def download(asset, cached): + """Never expose a partial or unverified download at ActivitySim's cache path.""" + if cached.is_file() and checksum(cached) == asset["sha256"]: + print(f"Using cached asset: {cached}", flush=True) + return + print(f"Downloading asset: {asset['url']}", flush=True) + with tempfile.TemporaryDirectory(dir=cached.parent) as temporary: + raw = Path(temporary) / "download" + with ( + urllib.request.urlopen(asset["url"], timeout=120) as response, + raw.open("wb") as stream, + ): + shutil.copyfileobj(response, stream) + ready = raw + # ActivitySim hashes the decompressed file for a .gz URL whose target + # filename lacks .gz. Archive checksums instead cover the archive bytes. + if asset["url"].endswith(".gz") and not cached.name.endswith(".gz"): + ready = Path(temporary) / "decoded" + with gzip.open(raw, "rb") as source, ready.open("wb") as stream: + shutil.copyfileobj(source, stream) + if checksum(ready) != asset["sha256"]: + raise ValueError( + f"asset checksum mismatch: {asset['name']}; expected {asset['sha256']}" + ) + os.replace(ready, cached) + + +def extract(archive, destination): + """Stream regular files only; reject links and archive path traversal.""" + + def target(name): + if name.rstrip("/") in {"", "."}: + return destination + return destination / relative(name.rstrip("/")) + + if archive.name.endswith(".zip"): + with zipfile.ZipFile(archive) as z: + for info in z.infolist(): + path = target(info.filename) + mode = (info.external_attr >> 16) & 0o170000 + if mode not in (0, 0o100000, 0o040000): + raise ValueError( + "asset archives cannot contain links or special files" + ) + if info.is_dir(): + path.mkdir(parents=True, exist_ok=True) + else: + path.parent.mkdir(parents=True, exist_ok=True) + with z.open(info) as source, path.open("wb") as stream: + shutil.copyfileobj(source, stream) + else: + import contextlib + + import zstandard + + with contextlib.ExitStack() as stack: + source = stack.enter_context(archive.open("rb")) + if archive.name.endswith(".tar.zst"): + source = stack.enter_context( + zstandard.ZstdDecompressor().stream_reader(source) + ) + tar = stack.enter_context(tarfile.open(fileobj=source, mode="r|*")) + for member in tar: + path = target(member.name) + if member.isdir(): + path.mkdir(parents=True, exist_ok=True) + elif member.isfile(): + path.parent.mkdir(parents=True, exist_ok=True) + with tar.extractfile(member) as incoming, path.open("wb") as stream: + shutil.copyfileobj(incoming, stream) + else: + raise ValueError( + "asset archives cannot contain links or special files" + ) + + +def matches(root, inventory): + """Validate installed content, including edits or incomplete extraction.""" + return all( + (root / name).is_file() and checksum(root / name) == digest + for name, digest in inventory.items() + ) + + +def prepare_assets(assets): + """Reuse ActivitySim archives and materialize inputs before suite preflight.""" + for asset in assets: + cached, target = Path(asset["cache_file"]), Path(asset["destination"]) + cached.parent.mkdir(parents=True, exist_ok=True) + with (cached.parent / (cached.name + ".abench.lock")).open("a") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + download(asset, cached) + if asset["unpack"]: + storage = cached.parent / ".abench-extracted" + storage.mkdir(exist_ok=True) + extracted = storage / asset["sha256"] + marker = storage / (asset["sha256"] + ".json") + inventory = read_json(marker) + if not inventory or not matches(extracted, inventory): + if extracted.exists(): + raise ValueError( + f"extracted cache was modified: {extracted}; remove it to rebuild" + ) + with tempfile.TemporaryDirectory(dir=storage) as temporary: + staging = Path(temporary) / "files" + staging.mkdir() + extract(cached, staging) + inventory = { + str(p.relative_to(staging)): checksum(p) + for p in staging.rglob("*") + if p.is_file() + } + if not inventory: + raise ValueError( + f"asset archive contains no files: {cached}" + ) + staging.rename(extracted) + write_json(marker, inventory) + if target.exists() or target.is_symlink(): + if not target.is_dir() or not matches(target, inventory): + raise ValueError( + f"existing asset destination differs; refusing to overwrite: {target}" + ) + else: + target.parent.mkdir(parents=True, exist_ok=True) + target.symlink_to(extracted, target_is_directory=True) + else: + if target.exists() or target.is_symlink(): + if not target.is_file() or checksum(target) != asset["sha256"]: + raise ValueError( + f"existing asset destination differs; refusing to overwrite: {target}" + ) + else: + target.parent.mkdir(parents=True, exist_ok=True) + # File symlinks pointing outside /data break in Docker. Plain + # files are copied; whole unpacked directories can be linked + # because abench resolves data_dir before mounting it. + with tempfile.NamedTemporaryFile( + dir=target.parent, delete=False + ) as stream: + temporary = Path(stream.name) + try: + shutil.copyfile(cached, temporary) + temporary.replace(target) + finally: + temporary.unlink(missing_ok=True) + print(f"Asset ready: {target}", flush=True) diff --git a/src/abench/cli.py b/src/abench/cli.py index 6d976d0..f9f126c 100644 --- a/src/abench/cli.py +++ b/src/abench/cli.py @@ -60,7 +60,7 @@ def positive(value): def parser(): p = argparse.ArgumentParser( description=__doc__, - epilog="Named experiments: abench experiments.yaml; preflight: abench validate experiments.yaml", + epilog="Named experiments: abench experiments.yaml [--set NAME=VALUE]; preflight: abench validate experiments.yaml [--set NAME=VALUE]; data only: abench prepare experiments.yaml", ) p.add_argument("--version", action="version", version=f"abench {__version__}") p.add_argument("--model-dir", type=Path, default=Path.cwd()) @@ -228,19 +228,39 @@ 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 + candidate = argv[1:] if argv and argv[0] in ("run", "validate", "prepare") else argv if ( candidate and not candidate[0].startswith("-") - and candidate[0] not in ("run", "report", "validate") + and candidate[0] not in ("run", "report", "validate", "prepare") ): - 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 + from .experiments import read_suite, run_suite + from .inputs import input_help - return run_suite(Path(candidate[0]), main, validate_only=argv[0] == "validate") + suite_parser = argparse.ArgumentParser( + prog="abench", + description="Run a named experiment suite.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=input_help(read_suite(Path(candidate[0]))[1]) + if any(flag in candidate[1:] for flag in ("--help", "-h")) + else None, + ) + suite_parser.add_argument("experiment_file", type=Path) + suite_parser.add_argument( + "--set", + action="append", + default=[], + metavar="NAME=VALUE", + help="override a declared experiment input (repeatable)", + ) + suite_args = suite_parser.parse_args(candidate) + return run_suite( + suite_args.experiment_file, + main, + prepare_only=argv[0] == "prepare", + validate_only=argv[0] == "validate", + assignments=suite_args.set, + ) 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 index 9257638..a18d010 100644 --- a/src/abench/experiments.py +++ b/src/abench/experiments.py @@ -8,9 +8,11 @@ import yaml +from .assets import plan_assets, prepare_assets from .common import write_json +from .inputs import resolve_inputs from .report import report -from .sources import source +from .sources import suite_source OPTIONS = { "model_dir", @@ -69,14 +71,14 @@ def unique_mapping(loader, node): ) -def expand_variables(document, timestamp): +def expand_variables(document, timestamp, inputs=None): """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} + resolved = {**(inputs or {}), "timestamp": timestamp} def term(name, stack): if name in resolved: @@ -105,10 +107,10 @@ def expand(value, stack=()): for name in terms: term(name, ()) - return expand(document) + return {k: v if k == "inputs" else expand(v) for k, v in document.items()} -def merge_options(defaults, overrides): +def merge_options(defaults, overrides, resolutions=None): """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") @@ -116,13 +118,15 @@ def merge_options(defaults, overrides): if unknown: raise ValueError(f"unknown experiment options: {sorted(unknown)}") merged = dict(defaults, **overrides) + if resolutions is None: + resolutions = {} 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) + item = suite_source(value, resolutions) if item["name"] in seen: raise ValueError(f"duplicate source: {item['name']}") seen.add(item["name"]) @@ -174,8 +178,8 @@ def arguments(options, base): return argv -def load_suite(path): - """Expand an entire suite before creating output or running any experiments.""" +def read_suite(path): + """Read and validate the file envelope without resolving inputs or sources.""" path = path.expanduser().resolve() try: raw = path.read_text() @@ -187,16 +191,28 @@ def load_suite(path): unknown = set(document) - { "schema_version", "vars", + "inputs", "defaults", "runs", "output_root", + "data_assets", } if unknown: raise ValueError(f"unknown experiment file fields: {sorted(unknown)}") + return raw, document + + +def load_suite(path, assignments=()): + """Resolve typed inputs and expand the suite before any external side effects.""" + path = path.expanduser().resolve() + raw, document = read_suite(path) + input_values, cli_overrides = resolve_inputs(document, assignments) document = expand_variables( - document, datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S-%f") + document, datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S-%f"), input_values ) - defaults = merge_options({}, document.get("defaults", {})) + assets = plan_assets(document.get("data_assets"), path.parent) + resolutions = {} + defaults = merge_options({}, document.get("defaults", {}), resolutions) 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") @@ -212,7 +228,7 @@ def load_suite(path): 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 = merge_options(defaults, overrides, resolutions) options.setdefault("label", name) argv = arguments(options, path.parent) destination = output / name @@ -226,17 +242,27 @@ def load_suite(path): return { "source_file": str(path), "original_yaml": raw, + "cli_overrides": cli_overrides, + "input_values": input_values, "configuration": document, "output_root": str(output), "runs": plan, + "data_assets": assets, + "source_resolutions": list(resolutions.values()), } -def run_suite(path, invoke, validate_only=False): +def run_suite(path, invoke, validate_only=False, prepare_only=False, assignments=()): """Preflight every run, execute serially, and preserve partial failure reports.""" - plan = load_suite(path) + plan = load_suite(path, assignments=assignments) root = Path(plan["output_root"]) - # Validation uses the same CLI checks as individual runs and creates nothing. + if not validate_only: + prepare_assets(plan["data_assets"]) + if prepare_only: + print(f"Prepared {len(plan['data_assets'])} assets; no experiments started") + return 0 + # Validation uses the same CLI checks as individual runs. Input preparation + # above runs only for execution/prepare; validate itself 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()): diff --git a/src/abench/inputs.py b/src/abench/inputs.py new file mode 100644 index 0000000..343a9f4 --- /dev/null +++ b/src/abench/inputs.py @@ -0,0 +1,156 @@ +"""Typed public inputs for experiment suites; internal vars remain read-only.""" + +import math +import re + +TYPES = { + "string": (str,), + "integer": (int,), + "number": (int, float), + "boolean": (bool,), +} + + +def check_value(name, spec, value): + """Validate YAML defaults and converted overrides using identical rules.""" + kind = spec["type"] + if type(value) not in TYPES[kind]: + raise ValueError(f"input {name} must be {kind}") + if kind == "number" and isinstance(value, float) and not math.isfinite(value): + raise ValueError(f"input {name} must be finite") + if "choices" in spec and value not in spec["choices"]: + raise ValueError(f"input {name} must be one of {spec['choices']!r}") + for key, invalid in ( + ("minimum", lambda a, b: a < b), + ("maximum", lambda a, b: a > b), + ): + if key in spec and invalid(value, spec[key]): + raise ValueError(f"input {name} must satisfy {key}: {spec[key]}") + return value + + +def input_schema(document): + """Validate declarations without requiring values, network, or file writes.""" + inputs = document.get("inputs", {}) + terms = document.get("vars", {}) + if not isinstance(inputs, dict) or not isinstance(terms, dict): + raise ValueError("inputs and vars must be mappings") + if "timestamp" in inputs or "timestamp" in terms: + raise ValueError("timestamp is a reserved variable") + if inputs.keys() & terms.keys(): + raise ValueError("inputs and vars must not declare the same name") + for name, spec in inputs.items(): + if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", name): + raise ValueError(f"invalid input name: {name!r}") + if not isinstance(spec, dict) or set(spec) - { + "type", + "default", + "required", + "description", + "choices", + "minimum", + "maximum", + }: + raise ValueError(f"invalid declaration for input {name}") + if not isinstance(spec.get("type"), str) or spec["type"] not in TYPES: + raise ValueError( + f"input {name} requires type: string, integer, number, or boolean" + ) + required = spec.get("required", False) + if type(required) is not bool: + raise ValueError(f"input {name} required must be boolean") + if required == ("default" in spec): + raise ValueError( + f"input {name} must have either a default or required: true" + ) + if "description" in spec and not isinstance(spec["description"], str): + raise ValueError(f"input {name} description must be a string") + for key in ("minimum", "maximum"): + if key in spec and ( + spec["type"] not in ("integer", "number") + or type(spec[key]) not in (int, float) + or (isinstance(spec[key], float) and not math.isfinite(spec[key])) + ): + raise ValueError( + f"input {name} {key} must be a finite numeric bound for a numeric input" + ) + if ( + "minimum" in spec + and "maximum" in spec + and spec["minimum"] > spec["maximum"] + ): + raise ValueError(f"input {name} minimum exceeds maximum") + if "choices" in spec: + if not isinstance(spec["choices"], list) or not spec["choices"]: + raise ValueError(f"input {name} choices must be a nonempty list") + for choice in spec["choices"]: + check_value(name, spec, choice) + if "default" in spec: + check_value(name, spec, spec["default"]) + return inputs + + +def resolve_inputs(document, assignments): + """Parse NAME=VALUE as declared scalar types, never as arbitrary YAML.""" + schema = input_schema(document) + explicit = {} + for assignment in assignments: + name, separator, text = assignment.partition("=") + if not separator or not name: + raise ValueError("--set requires NAME=VALUE") + if name not in schema: + raise ValueError( + f"unknown input {name!r}; --set can only override declared inputs" + ) + if name in explicit: + raise ValueError(f"duplicate --set input: {name}") + spec = schema[name] + kind = spec["type"] + try: + if kind == "integer": + if not re.fullmatch(r"[+-]?[0-9]+", text): + raise ValueError() + value = int(text) + elif kind == "number": + value = float(text) + elif kind == "boolean": + if text.lower() not in ("true", "false"): + raise ValueError() + value = text.lower() == "true" + else: + value = text + except ValueError as error: + raise ValueError(f"input {name} must be {kind}; got {text!r}") from error + explicit[name] = check_value(name, spec, value) + values = {} + for name, spec in schema.items(): + if name in explicit: + values[name] = explicit[name] + elif "default" in spec: + values[name] = spec["default"] + else: + raise ValueError( + f"missing required input {name}; supply --set {name}=VALUE" + ) + return values, explicit + + +def input_help(document): + """Describe the file's interface even when required inputs are missing.""" + lines = ["Experiment inputs (override with --set NAME=VALUE):"] + for name, spec in input_schema(document).items(): + status = "required" if spec.get("required") else f"default: {spec['default']!r}" + constraints = [ + f"{key}: {spec[key]}" + for key in ("choices", "minimum", "maximum") + if key in spec + ] + lines.append( + f" {name} ({spec['type']}; {status})" + + ("; " + "; ".join(constraints) if constraints else "") + ) + if spec.get("description"): + lines.append(f" {spec['description']}") + if len(lines) == 1: + lines.append(" No inputs declared.") + return "\n".join(lines) diff --git a/src/abench/sources.py b/src/abench/sources.py index d10ef0e..107a326 100644 --- a/src/abench/sources.py +++ b/src/abench/sources.py @@ -1,6 +1,7 @@ """Normalize exact GitHub source overrides before executing any build commands.""" import re +import subprocess from pathlib import PurePosixPath @@ -9,6 +10,81 @@ def canonical_name(name): return re.sub(r"[-_.]+", "-", name).lower() +def suite_source(value, resolutions): + """Pin a suite's GitHub branch or PR head once, before any builds begin. + + Only named suites accept moving selectors. The CLI, model profiles, and + Docker builder continue to receive immutable commits. A shared memo prevents + a branch update between runs from changing an inherited dependency. + """ + if not isinstance(value, dict) or not ({"branch", "pr"} & value.keys()): + return source(value) + if sum(key in value for key in ("commit", "branch", "pr")) != 1: + raise ValueError("source requires exactly one of commit, branch, or pr") + selector = "branch" if "branch" in value else "pr" + requested = value[selector] + if selector == "pr": + if type(requested) is not int or requested <= 0: + raise ValueError("source pr must be a positive integer") + ref = f"refs/pull/{requested}/head" + else: + if ( + not isinstance(requested, str) + or not requested + or any(c.isspace() or c in "~^:?*[\\" for c in requested) + or any(ord(c) < 32 or ord(c) == 127 for c in requested) + or ".." in requested + or "@{" in requested + or requested.endswith(".") + or any( + not p or p.startswith(".") or p.endswith(".lock") + for p in requested.split("/") + ) + ): + raise ValueError("invalid source branch") + ref = f"refs/heads/{requested}" + # Reuse exact-source validation before allowing any repository into Git. + pin = source( + {k: v for k, v in value.items() if k != selector} | {"commit": "0" * 40} + ) + repository = pin["repository"] + key = (repository.lower(), ref) + if key not in resolutions: + try: + result = subprocess.run( + [ + "git", + "ls-remote", + "--exit-code", + f"https://github.com/{repository}.git", + ref, + ], + capture_output=True, + text=True, + timeout=60, + check=True, + ) + except (OSError, subprocess.SubprocessError) as error: + raise ValueError( + f"Cannot resolve {repository} {ref}: {error}. " + "Check Git/network access and the branch or PR number, or use commit: ." + ) from error + matches = [line.split() for line in result.stdout.splitlines()] + commits = [parts[0] for parts in matches if len(parts) == 2 and parts[1] == ref] + if len(commits) != 1 or not re.fullmatch(r"[0-9a-fA-F]{40}", commits[0]): + raise ValueError( + f"GitHub did not return an exact commit for {repository} {ref}" + ) + resolutions[key] = { + "repository": repository, + "ref": ref, + "commit": commits[0].lower(), + } + print(f"Resolved {repository} {ref} → {commits[0].lower()}", flush=True) + pin["commit"] = resolutions[key]["commit"] + return pin + + def source(value): """Accept a CLI shorthand or a profile mapping, including extras/subdirectory.""" if isinstance(value, str): diff --git a/tests/test_assets.py b/tests/test_assets.py new file mode 100644 index 0000000..7441621 --- /dev/null +++ b/tests/test_assets.py @@ -0,0 +1,249 @@ +"""Asset integrity, archive handling, and ActivitySim cache interoperability.""" + +import gzip +import hashlib +import io +import tarfile +import zipfile +from pathlib import Path + +import pytest +import zstandard + +from abench import assets + + +def plan(tmp_path, content, name="data.csv", unpack=None): + info = dict( + url="https://example.test/" + name, sha256=hashlib.sha256(content).hexdigest() + ) + if unpack: + info["unpack"] = unpack + base = tmp_path / "model" + base.mkdir(exist_ok=True) + return assets.plan_assets( + dict(cache_dir=str(tmp_path / "cache"), name="example", assets={name: info}), + base, + ) + + +def mock_download(monkeypatch, data): + calls = [] + + def fetch(url, **kwargs): + calls.append(url) + return io.BytesIO(data) + + monkeypatch.setattr(assets.urllib.request, "urlopen", fetch) + return calls + + +def test_plain_download_and_reuse(tmp_path, monkeypatch): + instructions = plan(tmp_path, b"id\n1\n") + calls = mock_download(monkeypatch, b"id\n1\n") + assets.prepare_assets(instructions) + assets.prepare_assets(instructions) + assert len(calls) == 1 + assert (tmp_path / "model/data.csv").read_bytes() == b"id\n1\n" + assert not (tmp_path / "model/data.csv").is_symlink() + + +def test_gzip_checksum_is_decompressed(tmp_path, monkeypatch): + instructions = plan(tmp_path, b"data") + instructions[0]["url"] += ".gz" + mock_download(monkeypatch, gzip.compress(b"data")) + assets.prepare_assets(instructions) + assert Path(instructions[0]["cache_file"]).read_bytes() == b"data" + + +def test_bad_download_preserves_cache(tmp_path, monkeypatch): + instructions = plan(tmp_path, b"good") + cached = Path(instructions[0]["cache_file"]) + cached.parent.mkdir(parents=True) + cached.write_bytes(b"old") + mock_download(monkeypatch, b"wrong") + with pytest.raises(ValueError, match="checksum mismatch"): + assets.prepare_assets(instructions) + assert cached.read_bytes() == b"old" + assert not Path(instructions[0]["destination"]).exists() + + +@pytest.mark.parametrize("suffix", [".zip", ".tar.gz", ".tar.zst"]) +def test_unpack_reuses_verified_extraction(tmp_path, monkeypatch, suffix): + buffer = io.BytesIO() + if suffix == ".zip": + with zipfile.ZipFile(buffer, "w") as z: + z.writestr("households.csv", "id\n1\n") + else: + with tarfile.open(fileobj=buffer, mode="w") as t: + member = tarfile.TarInfo("households.csv") + member.size = 5 + t.addfile(member, io.BytesIO(b"id\n1\n")) + data = buffer.getvalue() + buffer = io.BytesIO( + gzip.compress(data) + if suffix == ".tar.gz" + else zstandard.ZstdCompressor().compress(data) + ) + instructions = plan(tmp_path, buffer.getvalue(), "data" + suffix, "data_full") + calls = mock_download(monkeypatch, buffer.getvalue()) + assets.prepare_assets(instructions) + monkeypatch.setattr( + assets, "extract", lambda *args: pytest.fail("unnecessary re-extraction") + ) + assets.prepare_assets(instructions) + assert len(calls) == 1 + target = Path(instructions[0]["destination"]) + assert target.is_symlink() + assert (target.resolve() / "households.csv").read_text() == "id\n1\n" + + +def test_existing_data_is_not_overwritten(tmp_path, monkeypatch): + instructions = plan(tmp_path, b"new") + target = Path(instructions[0]["destination"]) + target.write_text("user data") + mock_download(monkeypatch, b"new") + with pytest.raises(ValueError, match="refusing to overwrite"): + assets.prepare_assets(instructions) + assert target.read_text() == "user data" + + +def test_reject_archive_escape(tmp_path, monkeypatch): + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w") as z: + z.writestr("../escape", "bad") + instructions = plan(tmp_path, buffer.getvalue(), "data.zip", "data") + mock_download(monkeypatch, buffer.getvalue()) + with pytest.raises(ValueError, match="relative"): + assets.prepare_assets(instructions) + assert not Path(instructions[0]["destination"]).exists() + + +def test_shared_with_activitysim_in_both_directions(tmp_path, monkeypatch): + from activitysim.cli import create + from activitysim.examples.external import default_cache_dir + + assert assets.plan_assets( + dict(assets={"a": dict(url="https://example.test/a", sha256="a" * 64)}), + tmp_path, + )[0]["cache_file"] == str(default_cache_dir() / "a") + instructions = plan(tmp_path, b"shared") + item = instructions[0] + cached = Path(item["cache_file"]) + cached.parent.mkdir(parents=True) + # Simulate the file previously downloaded by ActivitySim. Its actual download + # function must accept abench's file without any requests call, and vice versa. + cached.write_bytes(b"shared") + mock_download(monkeypatch, b"wrong") + assets.prepare_assets(instructions) + cached.unlink() + Path(item["destination"]).unlink() + mock_download(monkeypatch, b"shared") + assets.prepare_assets(instructions) + monkeypatch.setattr( + create.requests, + "get", + lambda *a, **k: pytest.fail("ActivitySim redownloaded abench's cache"), + ) + output = tmp_path / "activitysim-copy/data.csv" + create.download_asset( + item["url"], + output, + sha256=item["sha256"], + link=cached.parent, + base_path=output.parent, + ) + assert output.read_bytes() == b"shared" + + +def test_suite_prepares_before_preflight_and_prepare_only_skips_runs( + tmp_path, monkeypatch +): + import yaml + + from abench.experiments import run_suite + + content = b"id\n1\n" + path = tmp_path / "abench.yaml" + path.write_text( + yaml.safe_dump( + dict( + schema_version=1, + output_root="results-${timestamp}", + data_assets=dict( + cache_dir="cache", + assets={ + "data/households.csv": dict( + url="https://example.test/households.csv", + sha256=hashlib.sha256(content).hexdigest(), + ) + }, + ), + runs={"example": {}}, + ) + ) + ) + calls = mock_download(monkeypatch, content) + assert ( + run_suite( + path, + lambda argv: pytest.fail("prepare invoked model validation"), + prepare_only=True, + ) + == 0 + ) + assert len(calls) == 1 + assert not list(tmp_path.glob("results-*")) + invoked = [] + + def invoke(argv): + assert (tmp_path / "data/households.csv").read_bytes() == content + invoked.append(argv[0]) + return 0 + + assert run_suite(path, invoke) == 0 + assert invoked == ["validate", "run"] + assert len(calls) == 1 + + +def test_validate_never_downloads(tmp_path, monkeypatch): + import yaml + + from abench.experiments import run_suite + + path = tmp_path / "abench.yaml" + path.write_text( + yaml.safe_dump( + dict( + schema_version=1, + output_root="results-${timestamp}", + runs={"x": {}}, + data_assets=dict( + cache_dir="cache", + assets={ + "data.csv": dict( + url="https://example.test/data", sha256="a" * 64 + ) + }, + ), + ) + ) + ) + monkeypatch.setattr( + assets.urllib.request, + "urlopen", + lambda *a, **k: pytest.fail("validate downloaded data"), + ) + assert run_suite(path, lambda argv: 0, validate_only=True) == 0 + assert not (tmp_path / "cache").exists() + + +@pytest.mark.parametrize("filename", ["../outside", "/absolute", "foo/../../bad"]) +def test_invalid_paths_rejected_before_download(tmp_path, filename): + with pytest.raises(ValueError, match="relative"): + assets.plan_assets( + dict( + assets={filename: dict(url="https://example.test/a", sha256="a" * 64)} + ), + tmp_path, + ) diff --git a/tests/test_docker.py b/tests/test_docker.py index 1026761..01cf5a5 100644 --- a/tests/test_docker.py +++ b/tests/test_docker.py @@ -99,23 +99,53 @@ def test_tiny_model(tmp_path, multiprocess, sharrow, retry): assert "cache preparation" in (output / "report.html").read_text() -def test_named_suite(tmp_path): +def test_named_suite(tmp_path, monkeypatch): """Exercise file dispatch, shared defaults, serial/MP overrides, and comparison.""" import yaml root = tmp_path / "model" shutil.copytree(Path(__file__).parent / "fixtures/tiny", root) + # Exercise archive preparation and a linked data directory through a real + # container mount, starting with no model inputs in the checkout. + import hashlib + import io + import zipfile + + from abench import assets + + archive = io.BytesIO() + with zipfile.ZipFile(archive, "w") as stream: + stream.write(root / "data/households.csv", "households.csv") + payload = archive.getvalue() + shutil.rmtree(root / "data") + monkeypatch.setattr( + assets.urllib.request, "urlopen", lambda *a, **k: io.BytesIO(payload) + ) path = tmp_path / "experiments.yaml" path.write_text( yaml.safe_dump( dict( schema_version=1, + inputs={ + "households": {"type": "integer", "default": 3, "minimum": 1}, + "sharrow": {"type": "boolean", "default": False}, + }, output_root="results", + data_assets=dict( + cache_dir="downloads", + assets={ + "data.zip": dict( + url="https://example.test/data.zip", + sha256=hashlib.sha256(payload).hexdigest(), + unpack="model/data", + ) + }, + ), defaults=dict( model_dir="model", flow_cache_dir="shared-flows", - sharrow=True, - households=4, + sharrow="${sharrow}", + households="${households}", warmup_households=2, memory="3g", shm_size="256m", @@ -129,13 +159,18 @@ def test_named_suite(tmp_path): sort_keys=False, ) ) - assert cli.main([str(path)]) == 0 + assert cli.main([str(path), "--set", "households=4", "--set", "sharrow=true"]) == 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() + suite = json.loads((root / "suite.json").read_text()) + assert ( + suite["input_values"] + == suite["cli_overrides"] + == {"households": 4, "sharrow": True} + ) for name in ("serial", "parallel"): warmup = json.loads( (root / name / "warmup/effective-settings.json").read_text() diff --git a/tests/test_experiments.py b/tests/test_experiments.py index d205b87..1e05ab5 100644 --- a/tests/test_experiments.py +++ b/tests/test_experiments.py @@ -1,7 +1,9 @@ """Named suites reuse CLI validation/execution without requiring Docker in tests.""" import json +import subprocess from pathlib import Path +from types import SimpleNamespace import pytest import yaml @@ -63,6 +65,59 @@ def test_defaults_variables_and_source_overrides(tmp_path, monkeypatch): assert not Path(plan["output_root"]).exists() +def test_suite_pins_moving_refs_before_execution(tmp_path, monkeypatch): + calls = [] + + def git(argv, **kwargs): + calls.append(argv[-1]) + sha = SHA if argv[-1] == "refs/heads/main" else OTHER + return SimpleNamespace(stdout=f"{sha}\t{argv[-1]}\n") + + monkeypatch.setattr(subprocess, "run", git) + path = suite(tmp_path) + document = yaml.safe_load(path.read_text()) + document["vars"]["pr"] = 1110 + document["defaults"]["sources"][1] = dict( + name="activitysim", repository="ActivitySim/activitysim", branch="main" + ) + document["runs"] = { + "main": {}, + "main_again": {}, + "pr": { + "sources": [ + dict( + name="activitysim", repository="ActivitySim/activitysim", pr="${pr}" + ) + ] + }, + } + path.write_text(yaml.safe_dump(document)) + executed = [] + monkeypatch.setattr(experiments, "report", lambda *a: None) + + def invoke(argv): + assert calls == ["refs/heads/main", "refs/pull/1110/head"] + executed.append(argv) + return 0 + + assert run_suite(path, invoke) == 0 + root = next(tmp_path.glob("results-*")) + recorded = json.loads((root / "suite.json").read_text()) + assert recorded["source_resolutions"] == [ + dict(repository="ActivitySim/activitysim", ref="refs/heads/main", commit=SHA), + dict( + repository="ActivitySim/activitysim", + ref="refs/pull/1110/head", + commit=OTHER, + ), + ] + assert "branch: main" in recorded["original_yaml"] + for run, sha in zip(recorded["runs"], [SHA, SHA, OTHER]): + assert f"activitysim=ActivitySim/activitysim@{sha}" in run["argv"] + assert f"sharrow=ActivitySim/sharrow@{SHA}" in run["argv"] + assert [args[0] for args in executed] == ["validate"] * 3 + ["run"] * 3 + + @pytest.mark.parametrize( "updates,match", [ @@ -157,17 +212,111 @@ def test_cli_file_dispatch_and_validation(tmp_path, monkeypatch): monkeypatch.setattr( experiments, "run_suite", - lambda path, invoke, validate_only: calls.append((path, validate_only)) or 0, + lambda path, invoke, **kwargs: calls.append((path, kwargs)) 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)] + assert cli.main([str(path), "--set", "activitysim_pr=1110"]) == 0 + assert cli.main(["validate", str(path), "--set=activitysim_pr=1110"]) == 0 + assert cli.main(["prepare", str(path), "--set", "activitysim_pr=1110"]) == 0 + assert calls == [ + (path, dict(validate_only=False, prepare_only=False, assignments=[])), + (path, dict(validate_only=False, prepare_only=False, assignments=[])), + (path, dict(validate_only=True, prepare_only=False, assignments=[])), + ( + path, + dict( + validate_only=False, + prepare_only=False, + assignments=["activitysim_pr=1110"], + ), + ), + ( + path, + dict( + validate_only=True, + prepare_only=False, + assignments=["activitysim_pr=1110"], + ), + ), + ( + path, + dict( + validate_only=False, + prepare_only=True, + assignments=["activitysim_pr=1110"], + ), + ), + ] with pytest.raises(SystemExit): cli.main([str(path), "--households", "5"]) +def test_pr_override_updates_labels_and_pins_and_main_is_fresh(tmp_path, monkeypatch): + path = suite(tmp_path) + document = yaml.safe_load(path.read_text()) + document["inputs"] = { + "activitysim_pr": {"type": "integer", "default": 100, "minimum": 1} + } + document["runs"] = { + "main": { + "sources": [ + dict( + name="activitysim", + repository="ActivitySim/activitysim", + branch="main", + ) + ] + }, + "pr": { + "label": "PR ${activitysim_pr}", + "sources": [ + dict( + name="activitysim", + repository="ActivitySim/activitysim", + pr="${activitysim_pr}", + ) + ], + }, + } + path.write_text(yaml.safe_dump(document)) + calls = [] + main_sha = SHA + + def git(argv, **kwargs): + calls.append(argv[-1]) + sha = main_sha if argv[-1] == "refs/heads/main" else OTHER + return SimpleNamespace(stdout=f"{sha}\t{argv[-1]}\n") + + monkeypatch.setattr(subprocess, "run", git) + first = load_suite(path, assignments=["activitysim_pr=1110"]) + assert first["cli_overrides"] == {"activitysim_pr": 1110} + assert first["input_values"]["activitysim_pr"] == 1110 + assert ( + yaml.safe_load(first["original_yaml"])["inputs"]["activitysim_pr"]["default"] + == 100 + ) + assert "PR 1110" in first["runs"][1]["argv"] + assert f"activitysim=ActivitySim/activitysim@{SHA}" in first["runs"][0]["argv"] + assert f"activitysim=ActivitySim/activitysim@{OTHER}" in first["runs"][1]["argv"] + main_sha = "c" * 40 + second = load_suite(path, assignments=["activitysim_pr=1110"]) + assert ( + f"activitysim=ActivitySim/activitysim@{main_sha}" in second["runs"][0]["argv"] + ) + assert calls == ["refs/heads/main", "refs/pull/1110/head"] * 2 + assert ( + yaml.safe_load(path.read_text())["inputs"]["activitysim_pr"]["default"] == 100 + ) + + +def test_pr_override_requires_declared_variable(tmp_path): + with pytest.raises(ValueError, match="unknown input"): + load_suite(suite(tmp_path), assignments=["activitysim_pr=1110"]) + + def test_shipped_sandag_suite(): path = Path(__file__).parents[1] / "examples/sandag-chunked.yaml" plan = load_suite(path) diff --git a/tests/test_inputs.py b/tests/test_inputs.py new file mode 100644 index 0000000..b9bcea6 --- /dev/null +++ b/tests/test_inputs.py @@ -0,0 +1,161 @@ +"""Public input validation happens before source resolution or data preparation.""" + +import subprocess + +import pytest +import yaml + +from abench import cli, experiments +from abench.inputs import input_schema, resolve_inputs + + +def test_typed_values_defaults_and_literal_strings(): + document = { + "inputs": { + "count": {"type": "integer", "default": 4, "minimum": 0, "maximum": 10}, + "scale": {"type": "number", "default": 1.5}, + "enabled": {"type": "boolean", "default": True}, + "text": {"type": "string", "default": "false"}, + } + } + values, overrides = resolve_inputs( + document, ["count=0", "enabled=false", "text=${count}=yes", "scale=1e-2"] + ) + assert ( + values + == overrides + == dict(count=0, enabled=False, text="${count}=yes", scale=0.01) + ) + expanded = experiments.expand_variables( + {**document, "vars": {"label": "value ${text}"}, "result": "${count}"}, + "now", + values, + ) + assert expanded["vars"]["label"] == "value ${count}=yes" + assert expanded["result"] == 0 + assert expanded["inputs"] == document["inputs"] + defaults, explicit = resolve_inputs(document, []) + assert defaults == dict(count=4, scale=1.5, enabled=True, text="false") + assert explicit == {} + + +@pytest.mark.parametrize( + "spec,value", + [ + ({"type": "integer", "default": 1}, "1.5"), + ({"type": "integer", "default": 1}, "true"), + ({"type": "number", "default": 1}, "NaN"), + ({"type": "number", "default": 1}, "inf"), + ({"type": "boolean", "default": True}, "yes"), + ({"type": "integer", "default": 1, "minimum": 1}, "0"), + ({"type": "integer", "default": 1, "maximum": 2}, "3"), + ({"type": "string", "default": "a", "choices": ["a", "b"]}, "c"), + ], +) +def test_invalid_overrides(spec, value): + with pytest.raises(ValueError, match="input x"): + resolve_inputs({"inputs": {"x": spec}}, [f"x={value}"]) + + +@pytest.mark.parametrize( + "spec", + [ + {"type": "integer", "default": True}, + {"type": "number", "default": float("nan")}, + {"type": "string", "default": 1}, + {"type": "integer"}, + {"type": "integer", "required": True, "default": 1}, + {"type": "integer", "required": "true"}, + {"type": "string", "default": "x", "minimum": 1}, + {"type": "integer", "default": 1, "minimum": 2, "maximum": 1}, + {"type": "integer", "default": 1, "choices": [True]}, + {"type": "integer", "default": 1, "choices": []}, + {"type": "integer", "default": 1, "minimum": float("inf")}, + {"type": "integer", "default": 1, "maximum": False}, + {"type": "integer", "default": 1, "typo": 2}, + {"type": "list", "default": []}, + ], +) +def test_bad_schema(spec): + with pytest.raises(ValueError, match="input x"): + input_schema({"inputs": {"x": spec}}) + + +@pytest.mark.parametrize( + "assignments,match", + [ + (["x=1", "x=2"], "duplicate"), + (["x"], "NAME=VALUE"), + (["internal=1"], "unknown input"), + ([], "missing required input x"), + ], +) +def test_missing_unknown_duplicate(assignments, match): + with pytest.raises(ValueError, match=match): + resolve_inputs( + { + "inputs": {"x": {"type": "integer", "required": True}}, + "vars": {"internal": 1}, + }, + assignments, + ) + + +@pytest.mark.parametrize( + "document", + [ + {"inputs": {"timestamp": {"type": "string", "default": "x"}}}, + {"inputs": {"x": {"type": "integer", "default": 1}}, "vars": {"x": 1}}, + {"inputs": {"bad name": {"type": "integer", "default": 1}}}, + ], +) +def test_invalid_names(document): + with pytest.raises(ValueError): + input_schema(document) + + +def test_help_and_invalid_input_never_resolve_sources(tmp_path, monkeypatch, capsys): + def unexpected(*args, **kwargs): + pytest.fail("unexpected external work") + + monkeypatch.setattr(subprocess, "run", unexpected) + monkeypatch.setattr(experiments, "prepare_assets", unexpected) + path = tmp_path / "suite.yaml" + path.write_text( + yaml.safe_dump( + { + "schema_version": 1, + "inputs": { + "pr": { + "type": "integer", + "required": True, + "minimum": 1, + "description": "PR to benchmark", + } + }, + "output_root": "results", + "runs": { + "main": { + "sources": [ + { + "name": "activitysim", + "repository": "ActivitySim/activitysim", + "branch": "main", + } + ] + } + }, + } + ) + ) + with pytest.raises(SystemExit) as error: + cli.main([str(path), "--help"]) + assert error.value.code == 0 + help_text = capsys.readouterr().out + assert "pr (integer; required)" in help_text + assert "PR to benchmark" in help_text + assert "minimum: 1" in help_text + for extra in ([], ["--set", "pr=0"], ["--set", "typo=1"]): + with pytest.raises(ValueError): + cli.main([str(path), *extra]) + assert not (tmp_path / "results").exists() diff --git a/tests/test_sources.py b/tests/test_sources.py new file mode 100644 index 0000000..e53f13a --- /dev/null +++ b/tests/test_sources.py @@ -0,0 +1,96 @@ +"""Moving suite selectors resolve to immutable, reproducible build inputs.""" + +import subprocess +from types import SimpleNamespace + +import pytest + +from abench.sources import source, suite_source + +SHA = "a" * 40 + + +def test_branch_and_pr_heads_resolve_once(monkeypatch): + calls = [] + + def git(argv, **kwargs): + calls.append(argv) + return SimpleNamespace(stdout=f"{SHA}\t{argv[-1]}\n") + + monkeypatch.setattr(subprocess, "run", git) + memo = {} + base = dict( + name="My_Addon", repository="Org/Repo", extras=["fast"], subdirectory="python" + ) + branch = suite_source(base | {"branch": "feature/test"}, memo) + assert branch == source(base | {"commit": SHA}) + assert suite_source(base | {"branch": "feature/test"}, memo) == branch + suite_source(base | {"pr": 1110}, memo) + assert [args[-1] for args in calls] == [ + "refs/heads/feature/test", + "refs/pull/1110/head", + ] + assert calls[1][-2] == "https://github.com/Org/Repo.git" + assert len(memo) == 2 + + +@pytest.mark.parametrize( + "fields", + [ + {"branch": "main", "commit": SHA}, + {"branch": "main", "pr": 1}, + {"pr": True}, + {"pr": "1110"}, + {"pr": 0}, + {"pr": -1}, + {"branch": "*"}, + {"branch": "../main"}, + {"branch": "main\n"}, + {"branch": "main", "repository": "https://evil.example/repo"}, + {"branch": "main", "typo": 1}, + ], +) +def test_invalid_selector_never_contacts_github(monkeypatch, fields): + def unexpected(*args, **kwargs): + pytest.fail("invalid source reached Git") + + monkeypatch.setattr(subprocess, "run", unexpected) + with pytest.raises(ValueError): + suite_source( + dict(name="activitysim", repository="ActivitySim/activitysim") | fields, {} + ) + + +@pytest.mark.parametrize( + "failure", + [ + FileNotFoundError("git"), + subprocess.TimeoutExpired("git", 60), + subprocess.CalledProcessError(2, "git"), + ], +) +def test_resolution_failure_explains_source(monkeypatch, failure): + def fail(*args, **kwargs): + raise failure + + monkeypatch.setattr(subprocess, "run", fail) + with pytest.raises(ValueError, match="Cannot resolve Org/Repo refs/pull/1110/head"): + suite_source(dict(name="addon", repository="Org/Repo", pr=1110), {}) + + +def test_unexpected_ref_response_rejected(monkeypatch): + monkeypatch.setattr( + subprocess, + "run", + lambda *a, **kw: SimpleNamespace(stdout=f"{SHA}\trefs/pull/1/merge\n"), + ) + with pytest.raises(ValueError, match="exact commit"): + suite_source(dict(name="addon", repository="Org/Repo", pr=1), {}) + + +def test_exact_pin_does_not_need_git(monkeypatch): + def unexpected(*args, **kwargs): + pytest.fail("exact pin reached Git") + + monkeypatch.setattr(subprocess, "run", unexpected) + assert suite_source(f"addon=Org/Repo@{SHA}", {})["commit"] == SHA From 8490fae71f759b5e9b6ac24026fcaefce6bfe17d Mon Sep 17 00:00:00 2001 From: Jeff Newman Date: Fri, 18 Sep 2026 18:40:45 -0500 Subject: [PATCH 2/4] Add interactive suite input prompts and benchmark progress updates --- README.md | 34 +++++++++---- src/abench/cli.py | 13 ++++- src/abench/experiments.py | 31 +++++++++--- src/abench/inputs.py | 80 ++++++++++++++++++++++-------- src/abench/progress.py | 62 +++++++++++++++++++++++ tests/test_experiments.py | 5 ++ tests/test_inputs.py | 101 ++++++++++++++++++++++++++++++++++++++ tests/test_progress.py | 34 +++++++++++++ 8 files changed, 322 insertions(+), 38 deletions(-) create mode 100644 src/abench/progress.py create mode 100644 tests/test_progress.py diff --git a/README.md b/README.md index a2a0fe7..b6a4ae4 100644 --- a/README.md +++ b/README.md @@ -142,12 +142,12 @@ resolves the selectors again, including `validate` and `prepare`. To reproduce a earlier suite, replace selectors with `commit: `. CLI overrides and model profiles still require exact commits. -Override the example's PR without editing YAML: +Start the suite and answer its input prompts: ```bash -uvx abench ./activitysim-prototype-mtc/abench.yaml --set activitysim_pr=1110 +uvx abench ./activitysim-prototype-mtc/abench.yaml # Preflight the same selection without running benchmarks: -uvx abench validate ./activitysim-prototype-mtc/abench.yaml --set activitysim_pr=1110 +uvx abench validate ./activitysim-prototype-mtc/abench.yaml ``` Declare user-facing settings in `inputs`, and keep internal reusable values and @@ -171,7 +171,16 @@ vars: label: "${mode}, ${households} households" ``` -Use repeatable `--set NAME=VALUE` with run, validate, or prepare commands. Quote +In a terminal, run, validate, and prepare prompt for each input in YAML order. +Press Enter to accept a displayed default. Required inputs have no default and +must be entered; empty or invalid answers prompt again with an explanation. +Ctrl-C cancels before source resolution or downloads. Selected values are printed +and saved in `suite.json` as `input_values`. + +For scripting, `--non-interactive` uses defaults and supplied values without +prompting. Non-terminal stdin behaves the same way; missing required inputs fail +instead of hanging. Repeatable `--set NAME=VALUE` can supply values explicitly; +these inputs are not prompted. Quote arguments containing spaces, for example `--set 'label=My experiment'` if `label` is declared as a string input. `abench experiment.yaml --help` lists the file's inputs, descriptions, defaults, and constraints without requiring input values, @@ -189,16 +198,17 @@ contacting GitHub, downloading data, or creating outputs. expansion. Inputs and vars share one namespace; duplicate names and the reserved name `timestamp` are errors. Input names use letters, digits, and underscores and cannot start with a digit. -- Unknown inputs, duplicate overrides, invalid values, and missing required - inputs fail before source resolution or downloads. There is no overriding vars. +- Unknown inputs, duplicate overrides, and invalid explicit values fail before + source resolution or downloads. Missing required inputs prompt in a terminal + and fail in non-interactive mode. There is no overriding vars. - The YAML file is unchanged. `suite.json` records typed `cli_overrides`, effective `input_values`, the expanded configuration, and exact source resolutions. `experiments.yaml` preserves the original file, so replay its recorded overrides too when reproducing a run. `branch: main` is freshly resolved on **every launch**; there is no saved branch -SHA to update manually. For the MTC example, use `--set activitysim_pr=1110` and -optionally `--set households=500000 --set processes=4`. +SHA to update manually. For the MTC example, enter the PR number when prompted, +then press Enter twice to accept 500,000 households and 4 processes. These commands require a release containing this feature. Until then, use `uvx --from /path/to/abench abench ...` with this local checkout. @@ -225,13 +235,19 @@ These commands require a release containing this feature. Until then, use 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 accept only `--set` input overrides; other model options belong +- File invocations accept `--set` and `--non-interactive`; other model options belong in `defaults` or the relevant run. 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. +Terminal progress identifies the experiment number, image build, warmup, and +measured attempts/retries. Builds and model phases print elapsed time every +15 seconds; model phases also show current and peak cgroup memory when samples +are available. Full console output stays in the printed log paths. These are +status updates, not an estimated completion percentage. + ## Run controls - `--single-process` (default), or `--multiprocess --processes N`. The count applies diff --git a/src/abench/cli.py b/src/abench/cli.py index f9f126c..f53a19c 100644 --- a/src/abench/cli.py +++ b/src/abench/cli.py @@ -21,6 +21,7 @@ from .failures import BenchmarkFailure, describe_failure from .flow_cache import publish_flows, reuse_flows from .profiles import load_profile, validate_model +from .progress import run_logged from .report import load_run, report from .sources import resolve_sources @@ -30,8 +31,7 @@ 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) + run_logged(args, log) else: return subprocess.check_output(args, text=True).strip() @@ -253,6 +253,11 @@ def main(argv=None): metavar="NAME=VALUE", help="override a declared experiment input (repeatable)", ) + suite_parser.add_argument( + "--non-interactive", + action="store_true", + help="use defaults and --set values without prompting (required inputs must be supplied)", + ) suite_args = suite_parser.parse_args(candidate) return run_suite( suite_args.experiment_file, @@ -260,6 +265,7 @@ def main(argv=None): prepare_only=argv[0] == "prepare", validate_only=argv[0] == "validate", assignments=suite_args.set, + interactive=not suite_args.non_interactive and sys.stdin.isatty(), ) action = argv.pop(0) if argv and argv[0] in ("run", "report", "validate") else "run" args = p.parse_args(argv) @@ -578,6 +584,9 @@ def entrypoint(): """Expose CLI errors without an unnecessary Python traceback.""" try: sys.exit(main()) + except KeyboardInterrupt: + print("\nBenchmark cancelled.", file=sys.stderr) + sys.exit(130) except (ValueError, OSError, subprocess.CalledProcessError) as error: print(f"Benchmark failed: {error}", file=sys.stderr) sys.exit(1) diff --git a/src/abench/experiments.py b/src/abench/experiments.py index a18d010..a02b206 100644 --- a/src/abench/experiments.py +++ b/src/abench/experiments.py @@ -202,11 +202,13 @@ def read_suite(path): return raw, document -def load_suite(path, assignments=()): +def load_suite(path, assignments=(), interactive=False): """Resolve typed inputs and expand the suite before any external side effects.""" path = path.expanduser().resolve() raw, document = read_suite(path) - input_values, cli_overrides = resolve_inputs(document, assignments) + input_values, cli_overrides = resolve_inputs( + document, assignments, interactive=interactive + ) document = expand_variables( document, datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S-%f"), input_values ) @@ -252,11 +254,26 @@ def load_suite(path, assignments=()): } -def run_suite(path, invoke, validate_only=False, prepare_only=False, assignments=()): +def run_suite( + path, + invoke, + validate_only=False, + prepare_only=False, + assignments=(), + interactive=False, +): """Preflight every run, execute serially, and preserve partial failure reports.""" - plan = load_suite(path, assignments=assignments) + plan = load_suite(path, assignments=assignments, interactive=interactive) + if plan["input_values"]: + print( + "Selected inputs: " + + ", ".join(f"{k}={v}" for k, v in plan["input_values"].items()), + flush=True, + ) root = Path(plan["output_root"]) if not validate_only: + if plan["data_assets"]: + print("Preparing input data (checking shared cache)…", flush=True) prepare_assets(plan["data_assets"]) if prepare_only: print(f"Prepared {len(plan['data_assets'])} assets; no experiments started") @@ -277,8 +294,10 @@ def run_suite(path, invoke, validate_only=False, prepare_only=False, assignments write_json(root / "suite.json", plan) completed = [] try: - for run in plan["runs"]: - print(f"Running experiment {run['name']}…", flush=True) + for number, run in enumerate(plan["runs"], 1): + print( + f"Experiment {number}/{len(plan['runs'])}: {run['name']}…", flush=True + ) try: code = invoke(["run", *run["argv"]]) finally: diff --git a/src/abench/inputs.py b/src/abench/inputs.py index 343a9f4..4fd6e08 100644 --- a/src/abench/inputs.py +++ b/src/abench/inputs.py @@ -90,7 +90,28 @@ def input_schema(document): return inputs -def resolve_inputs(document, assignments): +def parse_value(name, spec, text): + """Convert terminal or CLI text using the declared scalar type.""" + kind = spec["type"] + try: + if kind == "integer": + if not re.fullmatch(r"[+-]?[0-9]+", text): + raise ValueError() + value = int(text) + elif kind == "number": + value = float(text) + elif kind == "boolean": + if text.lower() not in ("true", "false"): + raise ValueError() + value = text.lower() == "true" + else: + value = text + except ValueError as error: + raise ValueError(f"input {name} must be {kind}; got {text!r}") from error + return check_value(name, spec, value) + + +def resolve_inputs(document, assignments, interactive=False): """Parse NAME=VALUE as declared scalar types, never as arbitrary YAML.""" schema = input_schema(document) explicit = {} @@ -104,40 +125,28 @@ def resolve_inputs(document, assignments): ) if name in explicit: raise ValueError(f"duplicate --set input: {name}") - spec = schema[name] - kind = spec["type"] - try: - if kind == "integer": - if not re.fullmatch(r"[+-]?[0-9]+", text): - raise ValueError() - value = int(text) - elif kind == "number": - value = float(text) - elif kind == "boolean": - if text.lower() not in ("true", "false"): - raise ValueError() - value = text.lower() == "true" - else: - value = text - except ValueError as error: - raise ValueError(f"input {name} must be {kind}; got {text!r}") from error - explicit[name] = check_value(name, spec, value) + explicit[name] = parse_value(name, schema[name], text) values = {} for name, spec in schema.items(): if name in explicit: values[name] = explicit[name] + elif interactive: + values[name] = prompt_value(name, spec) elif "default" in spec: values[name] = spec["default"] else: raise ValueError( - f"missing required input {name}; supply --set {name}=VALUE" + f"missing required input {name}; run in a terminal to answer prompts " + f"or supply --set {name}=VALUE" ) return values, explicit def input_help(document): """Describe the file's interface even when required inputs are missing.""" - lines = ["Experiment inputs (override with --set NAME=VALUE):"] + lines = [ + "Experiment inputs (prompted at startup; optionally override with --set NAME=VALUE):" + ] for name, spec in input_schema(document).items(): status = "required" if spec.get("required") else f"default: {spec['default']!r}" constraints = [ @@ -154,3 +163,32 @@ def input_help(document): if len(lines) == 1: lines.append(" No inputs declared.") return "\n".join(lines) + + +def prompt_value(name, spec): + """Keep asking until the user supplies a valid value or accepts a default.""" + if spec.get("description"): + print(f"{name}: {spec['description']}", flush=True) + constraints = ", ".join( + f"{key}: {spec[key]}" + for key in ("choices", "minimum", "maximum") + if key in spec + ) + suffix = f" [{spec['default']}]" if "default" in spec else " (required)" + prompt = f" {name} ({spec['type']}{'; ' + constraints if constraints else ''}){suffix}: " + while True: + try: + text = input(prompt) + except EOFError as error: + raise ValueError( + f"Input ended while asking for {name}; experiment not started" + ) from error + if not text.strip(): + if "default" in spec: + return spec["default"] + print(f" {name} is required; enter a value.", flush=True) + continue + try: + return parse_value(name, spec, text) + except ValueError as error: + print(f" {error}. Try again.", flush=True) diff --git a/src/abench/progress.py b/src/abench/progress.py new file mode 100644 index 0000000..5a5ae3a --- /dev/null +++ b/src/abench/progress.py @@ -0,0 +1,62 @@ +"""Host-side progress for long commands without changing measured model work.""" + +import csv +import subprocess +import time + + +def memory_status(path): + """Read only the last complete sample; tolerate partial writes and startup.""" + try: + with path.open("rb") as stream: + header = stream.readline().decode().strip().split(",") + start = stream.tell() + stream.seek(0, 2) + stream.seek(max(start, stream.tell() - 4096)) + rows = stream.read().decode().split("\n") + # A final unterminated row may still be in flight from the container. + row = next(csv.reader([rows[-2]])) + sample = dict(zip(header, row)) + current = int(sample["current_bytes"]) / 2**30 + peak = int(sample["peak_bytes"]) / 2**30 + return f"; memory {current:.2f} GiB (peak {peak:.2f} GiB)" + except (OSError, UnicodeError, ValueError, KeyError, IndexError, csv.Error): + return "" + + +def run_logged(args, log, interval=15): + """Keep full logs on disk and print periodic stage status, including failures.""" + label = "Image build" if log.name == "build.log" else log.parent.name + started = time.monotonic() + print(f"{label}: started; log: {log}", flush=True) + with log.open("w") as stream: + process = subprocess.Popen(args, stdout=stream, stderr=subprocess.STDOUT) + try: + while True: + try: + code = process.wait(timeout=interval) + break + except subprocess.TimeoutExpired: + detail = memory_status(log.parent / "memory.csv") + print( + f"{label}: {time.monotonic() - started:.0f}s elapsed{detail}", + flush=True, + ) + except BaseException: + # Let container_phase's finally block remove a cancelled container. + # Reap the host Docker client so cancellation never leaves it running. + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + raise + status = "completed" if code == 0 else f"failed (exit {code})" + print( + f"{label}: {status} after {time.monotonic() - started:.0f}s" + + memory_status(log.parent / "memory.csv"), + flush=True, + ) + if code: + raise subprocess.CalledProcessError(code, args) diff --git a/tests/test_experiments.py b/tests/test_experiments.py index 1e05ab5..75ba1e0 100644 --- a/tests/test_experiments.py +++ b/tests/test_experiments.py @@ -208,6 +208,7 @@ def invoke(argv): def test_cli_file_dispatch_and_validation(tmp_path, monkeypatch): + monkeypatch.setattr("sys.stdin.isatty", lambda: False) calls = [] monkeypatch.setattr( experiments, @@ -221,6 +222,10 @@ def test_cli_file_dispatch_and_validation(tmp_path, monkeypatch): assert cli.main([str(path), "--set", "activitysim_pr=1110"]) == 0 assert cli.main(["validate", str(path), "--set=activitysim_pr=1110"]) == 0 assert cli.main(["prepare", str(path), "--set", "activitysim_pr=1110"]) == 0 + calls = [ + (path, {k: v for k, v in options.items() if k != "interactive"}) + for path, options in calls + ] assert calls == [ (path, dict(validate_only=False, prepare_only=False, assignments=[])), (path, dict(validate_only=False, prepare_only=False, assignments=[])), diff --git a/tests/test_inputs.py b/tests/test_inputs.py index b9bcea6..eaf25b0 100644 --- a/tests/test_inputs.py +++ b/tests/test_inputs.py @@ -159,3 +159,104 @@ def unexpected(*args, **kwargs): with pytest.raises(ValueError): cli.main([str(path), *extra]) assert not (tmp_path / "results").exists() + + +def test_interactive_defaults_required_and_validation(monkeypatch, capsys): + answers = iter(["", "", "bad", "0", "1110", "", ""]) + prompts = [] + + def respond(prompt): + prompts.append(prompt) + return next(answers) + + monkeypatch.setattr("builtins.input", respond) + values, explicit = resolve_inputs( + { + "inputs": { + "households": {"type": "integer", "default": 500000}, + "pr": {"type": "integer", "required": True, "minimum": 1}, + "sharrow": {"type": "boolean", "default": True}, + "label": {"type": "string", "default": ""}, + } + }, + [], + interactive=True, + ) + assert values == dict(households=500000, pr=1110, sharrow=True, label="") + assert explicit == {} + assert "[500000]" in prompts[0] + assert "(required)" in prompts[1] + output = capsys.readouterr().out + assert "pr is required" in output and "Try again" in output + + +def test_prompt_eof_cancels(monkeypatch): + def ended(prompt): + raise EOFError() + + monkeypatch.setattr("builtins.input", ended) + with pytest.raises(ValueError, match="experiment not started"): + resolve_inputs( + {"inputs": {"x": {"type": "integer", "required": True}}}, + [], + interactive=True, + ) + + +def test_explicit_values_skip_prompts(monkeypatch): + def unexpected(prompt): + pytest.fail("explicit input should not prompt") + + monkeypatch.setattr("builtins.input", unexpected) + assert resolve_inputs( + {"inputs": {"x": {"type": "integer", "required": True}}}, + ["x=1"], + interactive=True, + )[0] == {"x": 1} + + +def test_cli_terminal_detection(tmp_path, monkeypatch): + calls = [] + monkeypatch.setattr( + experiments, "run_suite", lambda *a, **kw: calls.append(kw) or 0 + ) + monkeypatch.setattr("sys.stdin.isatty", lambda: True) + path = str(tmp_path / "suite.yaml") + cli.main([path]) + cli.main([path, "--non-interactive"]) + monkeypatch.setattr("sys.stdin.isatty", lambda: False) + cli.main([path]) + assert [call["interactive"] for call in calls] == [True, False, False] + + +def test_prompted_inputs_are_saved_before_execution(tmp_path, monkeypatch): + import json + + path = tmp_path / "suite.yaml" + path.write_text( + yaml.safe_dump( + { + "schema_version": 1, + "inputs": {"households": {"type": "integer", "default": 4}}, + "output_root": "results", + "defaults": {"households": "${households}"}, + "runs": {"test": {}}, + } + ) + ) + answers = [] + monkeypatch.setattr("builtins.input", lambda prompt: answers.append(prompt) or "5") + monkeypatch.setattr(experiments, "report", lambda *a: None) + calls = [] + + def invoke(argv): + assert len(answers) == 1 + assert argv[argv.index("--households") + 1] == "5" + calls.append(argv[0]) + return 0 + + experiments.run_suite(path, invoke, interactive=True) + plan = json.loads((tmp_path / "results/suite.json").read_text()) + assert plan["input_values"] == {"households": 5} + assert plan["cli_overrides"] == {} + assert calls == ["validate", "run"] diff --git a/tests/test_progress.py b/tests/test_progress.py new file mode 100644 index 0000000..d4deae5 --- /dev/null +++ b/tests/test_progress.py @@ -0,0 +1,34 @@ +"""Progress is visible without streaming model logs or affecting exit codes.""" + +import subprocess +import sys + +import pytest + +from abench.progress import memory_status, run_logged + + +def test_memory_partial_sample(tmp_path): + path = tmp_path / "memory.csv" + assert memory_status(path) == "" + path.write_text( + "elapsed_seconds,current_bytes,peak_bytes\n1,1073741824,2147483648\n2,123" + ) + assert memory_status(path) == "; memory 1.00 GiB (peak 2.00 GiB)" + + +def test_logged_command_progress_and_exit(tmp_path, capsys): + log = tmp_path / "build.log" + run_logged( + [sys.executable, "-c", "import time; print('build output'); time.sleep(.15)"], + log, + interval=0.03, + ) + output = capsys.readouterr().out + assert "elapsed" in output and "completed" in output + assert "build output" not in output + assert "build output" in log.read_text() + with pytest.raises(subprocess.CalledProcessError) as error: + run_logged([sys.executable, "-c", "raise SystemExit(7)"], log) + assert error.value.returncode == 7 + assert "failed (exit 7)" in capsys.readouterr().out From 711cb7082ca40f0a14d3d98367e267edc6e8aafe Mon Sep 17 00:00:00 2001 From: Jeff Newman Date: Fri, 18 Sep 2026 18:47:47 -0500 Subject: [PATCH 3/4] Add experiment discovery and selection from model directories --- README.md | 25 +++++++++ src/abench/cli.py | 20 ++++++-- src/abench/discovery.py | 65 ++++++++++++++++++++++++ tests/test_discovery.py | 109 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 214 insertions(+), 5 deletions(-) create mode 100644 src/abench/discovery.py create mode 100644 tests/test_discovery.py diff --git a/README.md b/README.md index b6a4ae4..8448155 100644 --- a/README.md +++ b/README.md @@ -41,6 +41,31 @@ promised to work. Build/runtime failures retain diagnostics and a failure report ## Named experiment files +You can pass a model directory instead of a YAML file: + +```bash +abench /path/to/model +``` + +Abench looks for `.yaml` and `.yml` files directly inside `/path/to/model/.abench/` +(no recursive search). In a terminal it lists them alphabetically and asks which +experiment to run, then prompts for that experiment's inputs. Enter chooses the +first file; a single file still gets a selection prompt. Only the selected file +is loaded. `run`, `validate`, and `prepare` all support directory selection. + +`abench /path/to/model --help` lists available files without prompting or running +anything. Missing `.abench` directories and empty file lists produce clear errors. +Without a terminal, or with `--non-interactive`, one file is selected automatically; +multiple files require passing the desired YAML path directly. + +Relative paths remain relative to the selected YAML file. For a file inside +`.abench`, use `model_dir: ..` and, for example, `data_dir: ../data_full` and +`output_root: ../benchmark-runs/run-${timestamp}` to reference the parent model. +For declared downloads, keep unpack destinations inside `.abench` (for example, +`unpack: data_full` with `data_dir: data_full`); asset destinations cannot use `..`. +Moving an existing suite into `.abench` requires reviewing its paths. + + Write common options once and override only what differs between runs: ```yaml diff --git a/src/abench/cli.py b/src/abench/cli.py index f53a19c..e68ed20 100644 --- a/src/abench/cli.py +++ b/src/abench/cli.py @@ -234,18 +234,26 @@ def main(argv=None): and not candidate[0].startswith("-") and candidate[0] not in ("run", "report", "validate", "prepare") ): + from .discovery import directory_help, select_experiment from .experiments import read_suite, run_suite from .inputs import input_help + target = Path(candidate[0]).expanduser() suite_parser = argparse.ArgumentParser( prog="abench", - description="Run a named experiment suite.", + description="Run an experiment YAML file or choose from a directory’s .abench folder.", formatter_class=argparse.RawDescriptionHelpFormatter, - epilog=input_help(read_suite(Path(candidate[0]))[1]) + epilog=( + directory_help(target) + if target.is_dir() + else input_help(read_suite(target)[1]) + ) if any(flag in candidate[1:] for flag in ("--help", "-h")) else None, ) - suite_parser.add_argument("experiment_file", type=Path) + suite_parser.add_argument( + "experiment_file", type=Path, help="experiment YAML file or model directory" + ) suite_parser.add_argument( "--set", action="append", @@ -259,13 +267,15 @@ def main(argv=None): help="use defaults and --set values without prompting (required inputs must be supplied)", ) suite_args = suite_parser.parse_args(candidate) + interactive = not suite_args.non_interactive and sys.stdin.isatty() + selected = select_experiment(suite_args.experiment_file, interactive) return run_suite( - suite_args.experiment_file, + selected, main, prepare_only=argv[0] == "prepare", validate_only=argv[0] == "validate", assignments=suite_args.set, - interactive=not suite_args.non_interactive and sys.stdin.isatty(), + interactive=interactive, ) action = argv.pop(0) if argv and argv[0] in ("run", "report", "validate") else "run" args = p.parse_args(argv) diff --git a/src/abench/discovery.py b/src/abench/discovery.py new file mode 100644 index 0000000..7b6d381 --- /dev/null +++ b/src/abench/discovery.py @@ -0,0 +1,65 @@ +"""Discover and choose experiment instructions stored inside a model directory.""" + + +def instruction_files(directory): + """List immediate YAML files deterministically without loading any suite.""" + folder = directory.expanduser() / ".abench" + if not folder.is_dir(): + raise ValueError(f"No experiment directory found: {folder}") + files = sorted( + ( + p + for p in folder.iterdir() + if p.is_file() and p.suffix.lower() in (".yaml", ".yml") + ), + key=lambda p: (p.name.casefold(), p.name), + ) + if not files: + raise ValueError(f"No YAML experiment files found in {folder}") + return files + + +def directory_help(directory): + """Help lists available files without requiring an interactive selection.""" + files = instruction_files(directory) + return ( + "Available experiments:\n" + + "\n".join(f" {p.name}" for p in files) + + ( + "\n\nRun this directory to choose an experiment, or pass a file directly." + "\nRelative paths in each experiment are relative to its YAML file." + ) + ) + + +def select_experiment(path, interactive): + """Choose a suite before resolving its inputs, sources, or data assets.""" + path = path.expanduser() + if not path.is_dir(): + return path + files = instruction_files(path) + if not interactive: + if len(files) == 1: + return files[0] + raise ValueError( + f"Multiple experiments found in {path / '.abench'}: " + + ", ".join(p.name for p in files) + + ". Run in a terminal to choose, or pass the experiment YAML path directly." + ) + print(f"Experiments in {path / '.abench'}:", flush=True) + for number, file in enumerate(files, 1): + print(f" {number}. {file.name}", flush=True) + while True: + try: + answer = input(f"Choose experiment [1] (1–{len(files)}): ").strip() + except EOFError as error: + raise ValueError( + "Input ended while choosing an experiment; nothing started" + ) from error + if not answer: + answer = "1" + if answer.isascii() and answer.isdigit() and 1 <= int(answer) <= len(files): + chosen = files[int(answer) - 1] + print(f"Selected experiment: {chosen.name}", flush=True) + return chosen + print(f"Enter a number from 1 to {len(files)}.", flush=True) diff --git a/tests/test_discovery.py b/tests/test_discovery.py new file mode 100644 index 0000000..56891b3 --- /dev/null +++ b/tests/test_discovery.py @@ -0,0 +1,109 @@ +"""Directory selection precedes suite inputs and never runs unselected files.""" + +import pytest + +from abench import cli, experiments +from abench.discovery import instruction_files, select_experiment + + +def instructions(tmp_path): + folder = tmp_path / ".abench" + folder.mkdir() + for name in ("b.yml", "a.yaml", "ignored.txt"): + (folder / name).write_text("not loaded during selection") + (folder / "nested.yaml").mkdir() + return folder + + +def test_selection_order_and_invalid_answers(tmp_path, monkeypatch, capsys): + folder = instructions(tmp_path) + assert [p.name for p in instruction_files(tmp_path)] == ["a.yaml", "b.yml"] + answers = iter(["bad", "0", "3", "2"]) + monkeypatch.setattr("builtins.input", lambda prompt: next(answers)) + assert select_experiment(tmp_path, True) == folder / "b.yml" + assert "Enter a number" in capsys.readouterr().out + + +def test_single_file_still_prompts_and_accepts_enter(tmp_path, monkeypatch): + folder = instructions(tmp_path) + (folder / "b.yml").unlink() + prompts = [] + monkeypatch.setattr("builtins.input", lambda prompt: prompts.append(prompt) or "") + assert select_experiment(tmp_path, True) == folder / "a.yaml" + assert len(prompts) == 1 + assert select_experiment(tmp_path, False) == folder / "a.yaml" + + +def test_noninteractive_multiple_and_missing(tmp_path): + with pytest.raises(ValueError, match="No experiment directory"): + instruction_files(tmp_path) + folder = instructions(tmp_path) + with pytest.raises(ValueError, match="Multiple experiments"): + select_experiment(tmp_path, False) + (folder / "a.yaml").unlink() + (folder / "b.yml").unlink() + with pytest.raises(ValueError, match="No YAML"): + instruction_files(tmp_path) + + +def test_selection_eof(tmp_path, monkeypatch): + instructions(tmp_path) + + def ended(prompt): + raise EOFError() + + monkeypatch.setattr("builtins.input", ended) + with pytest.raises(ValueError, match="nothing started"): + select_experiment(tmp_path, True) + + +def test_directory_help_does_not_prompt_or_parse_suites(tmp_path, monkeypatch, capsys): + instructions(tmp_path) + + def unexpected(*args, **kwargs): + pytest.fail("help must not prompt or execute") + + monkeypatch.setattr("builtins.input", unexpected) + monkeypatch.setattr(experiments, "run_suite", unexpected) + with pytest.raises(SystemExit) as error: + cli.main([str(tmp_path), "--help"]) + assert error.value.code == 0 + assert "a.yaml" in capsys.readouterr().out + + +def test_choice_then_inputs_and_file_relative_paths(tmp_path, monkeypatch): + folder = instructions(tmp_path) + (folder / "b.yml").write_text("""schema_version: 1 +inputs: + households: + type: integer + required: true +output_root: ../results +runs: + test: + model_dir: .. + households: ${households} +""") + prompts = [] + answers = iter(["2", "42"]) + + def respond(prompt): + prompts.append(prompt) + return next(answers) + + monkeypatch.setattr("builtins.input", respond) + monkeypatch.setattr("sys.stdin.isatty", lambda: True) + plans = [] + + def run_suite(path, invoke, **options): + plans.append(experiments.load_suite(path, interactive=options["interactive"])) + return 0 + + monkeypatch.setattr(experiments, "run_suite", run_suite) + assert cli.main([str(tmp_path)]) == 0 + assert prompts[0].startswith("Choose experiment") + assert "households" in prompts[1] + assert plans[0]["input_values"] == {"households": 42} + args = plans[0]["runs"][0]["argv"] + assert args[args.index("--model-dir") + 1] == str(tmp_path) + assert plans[0]["output_root"] == str(tmp_path / "results") From 6bc8f74549d4cbc8ca929f39f9357b9cca0c29c8 Mon Sep 17 00:00:00 2001 From: Jeff Newman Date: Fri, 18 Sep 2026 19:00:55 -0500 Subject: [PATCH 4/4] Prepare abench 0.1.1 release --- README.md | 12 ++++-------- src/abench/__init__.py | 2 +- 2 files changed, 5 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index 8448155..be8bdeb 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ uvx abench --help uvx abench experiments.yaml ``` -For a specific release use `uvx abench@0.1.0 experiments.yaml`; use +For a specific release use `uvx abench@0.1.1 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. @@ -119,8 +119,7 @@ The pinned `main` revision is the one used in the earlier trials, not a moving branch reference. To compare **current main against a PR**, named suites also accept source mappings -with `branch` or `pr` in place of `commit` (local development feature, not in -PyPI 0.1.0 yet): +with `branch` or `pr` in place of `commit`: ```yaml schema_version: 1 @@ -235,8 +234,7 @@ contacting GitHub, downloading data, or creating outputs. SHA to update manually. For the MTC example, enter the PR number when prompted, then press Enter twice to accept 500,000 households and 4 processes. -These commands require a release containing this feature. Until then, use -`uvx --from /path/to/abench abench ...` with this local checkout. +These features are available in abench 0.1.1 and later. - `defaults` accepts CLI options using underscores (`shm_size`, `config_overlay`, etc.). Use `multiprocess: false` for serial execution and `sharrow: false` to @@ -550,6 +548,4 @@ Existing destinations must match; modified inputs are never silently replaced. Archive links, special files, and paths escaping the destination are rejected. The original instructions and resolved cache locations are saved in `suite.json`. -This feature is not present in abench 0.1.0. Until the next release, install the -updated checkout (`uv tool install /path/to/abench`) or run it with -`uvx --from /path/to/abench abench /path/to/model/abench.yaml`. +Data preparation is available in abench 0.1.1 and later. diff --git a/src/abench/__init__.py b/src/abench/__init__.py index 334e474..430437a 100644 --- a/src/abench/__init__.py +++ b/src/abench/__init__.py @@ -1,3 +1,3 @@ """Reproducible ActivitySim experiments in Linux containers.""" -__version__ = "0.1.0" +__version__ = "0.1.1"