| #!/usr/bin/env python3 | |
| """Utilities for CI. | |
| This dynamically prepares a list of routines that had a source file change based on | |
| git history. | |
| """ | |
| import json | |
| import os | |
| import pprint | |
| import re | |
| import subprocess as sp | |
| import sys | |
| from dataclasses import dataclass, field | |
| from functools import cache | |
| from glob import glob | |
| from inspect import cleandoc | |
| from os import getenv | |
| from pathlib import Path | |
| from typing import TypedDict, Self | |
| USAGE = cleandoc(""" | |
| usage: | |
| ./ci/ci-util.py <COMMAND> [flags] | |
| COMMAND: | |
| generate-matrix | |
| Calculate a matrix of which functions had source change, print that as | |
| a JSON object. | |
| locate-baseline [--download] [--extract] [--tag TAG] | |
| Locate the most recent benchmark baseline available in CI and, if flags | |
| specify, download and extract it. Never exits with nonzero status if | |
| downloading fails. | |
| `--tag` can be specified to look for artifacts with a specific tag, such as | |
| for a specific architecture. | |
| Note that `--extract` will overwrite files in `gungraun-home`. | |
| handle-bench-regressions PR_NUMBER | |
| Exit with success if the pull request contains a line starting with | |
| `ci: allow-regressions`, indicating that regressions in benchmarks should | |
| be accepted. Otherwise, exit 1. | |
| """) | |
| REPO_ROOT = Path(__file__).parent.parent | |
| GIT = ["git", "-C", REPO_ROOT] | |
| DEFAULT_BRANCH = "main" | |
| WORKFLOW_NAME = "CI" # Workflow that generates the benchmark artifacts | |
| ARTIFACT_PREFIX = "baseline-icount*" | |
| # Don't run exhaustive tests if these files change, even if they contaiin a function | |
| # definition. | |
| IGNORE_FILES = [ | |
| "libm/src/math/support/", | |
| "libm/src/libm_helper.rs", | |
| "libm/src/math/arch/intrinsics.rs", | |
| ] | |
| # libm PR CI takes a long time and doesn't need to run unless relevant files have been | |
| # changed. Anything matching this regex pattern will trigger a run. | |
| TRIGGER_LIBM_CI_FILE_PAT = ".*(libm|musl).*" | |
| TYPES = ["f16", "f32", "f64", "f128"] | |
| def eprint(*args, **kwargs): | |
| """Print to stderr.""" | |
| print(*args, file=sys.stderr, **kwargs) | |
| @dataclass(kw_only=True) | |
| class PrCfg: | |
| """Directives that we allow in the commit body to control test behavior. | |
| These are of the form `ci: foo`, at the start of a line. | |
| """ | |
| # The PR body | |
| body: str | |
| # Skip regression checks (must be at the start of a line). | |
| allow_regressions: bool = False | |
| # Don't run extensive tests | |
| skip_extensive: bool = False | |
| # Add these extensive tests to the list | |
| extra_extensive: list[str] = field(default_factory=list, init=False) | |
| # Allow running a large number of extensive tests. If not set, this script | |
| # will error out if a threshold is exceeded in order to avoid accidentally | |
| # spending huge amounts of CI time. | |
| allow_many_extensive: bool = False | |
| # Max number of extensive tests to run by default | |
| MANY_EXTENSIVE_THRESHOLD: int = 20 | |
| # Run tests for `libm` that may otherwise be skipped due to no changed files. | |
| always_test_libm: bool = False | |
| # String values of directive names | |
| DIR_ALLOW_REGRESSIONS: str = "allow-regressions" | |
| DIR_SKIP_EXTENSIVE: str = "skip-extensive" | |
| DIR_ALLOW_MANY_EXTENSIVE: str = "allow-many-extensive" | |
| DIR_TEST_LIBM: str = "test-libm" | |
| DIR_EXTRA_EXTENSIVE: str = "extra-extensive" | |
| def __post_init__(self): | |
| directives = re.finditer( | |
| r"^\s*ci:\s*(?P<dir_name>[^\s=]*)(?:\s*=\s*(?P<args>.*))?", | |
| self.body, | |
| re.MULTILINE, | |
| ) | |
| for dir in directives: | |
| name = dir.group("dir_name") | |
| args = dir.group("args") | |
| if name == self.DIR_ALLOW_REGRESSIONS: | |
| self.allow_regressions = True | |
| elif name == self.DIR_SKIP_EXTENSIVE: | |
| self.skip_extensive = True | |
| elif name == self.DIR_ALLOW_MANY_EXTENSIVE: | |
| self.allow_many_extensive = True | |
| elif name == self.DIR_TEST_LIBM: | |
| self.always_test_libm = True | |
| elif name == self.DIR_EXTRA_EXTENSIVE: | |
| self.extra_extensive = [x.strip() for x in args.split(",")] | |
| args = None | |
| else: | |
| eprint(f"Found unexpected directive `{name}`") | |
| exit(1) | |
| if args is not None: | |
| eprint("Found arguments where not expected") | |
| exit(1) | |
| eprint(pprint.pformat(self)) | |
| @dataclass | |
| class PrInfo: | |
| """GitHub response for PR query""" | |
| body: str | |
| commits: list[str] | |
| created_at: str | |
| number: int | |
| cfg: PrCfg | |
| @classmethod | |
| def from_env(cls) -> Self | None: | |
| """Create a PR object from the PR_NUMBER environment if set, `None` otherwise.""" | |
| pr_env = os.environ.get("PR_NUMBER") | |
| if pr_env is not None and len(pr_env) > 0: | |
| return cls.from_pr(pr_env) | |
| return None | |
| @classmethod | |
| @cache # Cache so we don't print info messages multiple times | |
| def from_pr(cls, pr_number: int | str) -> Self: | |
| """For a given PR number, query the body and commit list.""" | |
| pr_info = sp.check_output( | |
| [ | |
| "gh", | |
| "pr", | |
| "view", | |
| str(pr_number), | |
| "--json=number,commits,body,createdAt", | |
| # Flatten the commit list to only hashes, change a key to snake naming | |
| "--jq=.commits |= map(.oid) | .created_at = .createdAt | del(.createdAt)", | |
| ], | |
| text=True, | |
| ) | |
| pr_json = json.loads(pr_info) | |
| eprint("PR info:", json.dumps(pr_json, indent=4)) | |
| return cls(**pr_json, cfg=PrCfg(body=pr_json["body"])) | |
| class FunctionDef(TypedDict): | |
| """Type for an entry in `function-definitions.json`""" | |
| sources: list[str] | |
| type: str | |
| class Context: | |
| gh_ref: str | None | |
| changed: list[Path] | |
| defs: dict[str, FunctionDef] | |
| def __init__(self) -> None: | |
| self.gh_ref = getenv("GITHUB_REF") | |
| self.changed = [] | |
| self._init_change_list() | |
| with open(REPO_ROOT.joinpath("etc/function-definitions.json")) as f: | |
| defs = json.load(f) | |
| defs.pop("__comment", None) | |
| self.defs = defs | |
| def _init_change_list(self): | |
| """Create a list of files that have been changed. This uses GITHUB_REF if | |
| available, otherwise a diff between `HEAD` and `main`. | |
| """ | |
| # For pull requests, GitHub creates a ref `refs/pull/1234/merge` (1234 being | |
| # the PR number), and sets this as `GITHUB_REF`. | |
| ref = self.gh_ref | |
| eprint(f"using ref `{ref}`") | |
| if not self.is_pr(): | |
| # If the ref is not for `merge` then we are not in PR CI | |
| eprint("No diff available for ref") | |
| return | |
| # The ref is for a dummy merge commit. We can extract the merge base by | |
| # inspecting all parents (`^@`). | |
| merge_sha = sp.check_output( | |
| GIT + ["show-ref", "--hash", ref], text=True | |
| ).strip() | |
| merge_log = sp.check_output(GIT + ["log", "-1", merge_sha], text=True) | |
| eprint(f"Merge:\n{merge_log}\n") | |
| parents = ( | |
| sp.check_output(GIT + ["rev-parse", f"{merge_sha}^@"], text=True) | |
| .strip() | |
| .splitlines() | |
| ) | |
| assert len(parents) == 2, f"expected two-parent merge but got:\n{parents}" | |
| base = parents[0].strip() | |
| incoming = parents[1].strip() | |
| eprint(f"base: {base}, incoming: {incoming}") | |
| textlist = sp.check_output( | |
| GIT + ["diff", base, incoming, "--name-only"], text=True | |
| ) | |
| self.changed = [Path(p) for p in textlist.splitlines()] | |
| def is_pr(self) -> bool: | |
| """Check if we are looking at a PR rather than a push.""" | |
| return self.gh_ref is not None and "merge" in self.gh_ref | |
| @staticmethod | |
| def _ignore_file(fname: str) -> bool: | |
| return any(fname.startswith(pfx) for pfx in IGNORE_FILES) | |
| def changed_routines(self) -> dict[str, list[str]]: | |
| """Create a list of routines for which one or more files have been updated, | |
| separated by type. | |
| """ | |
| routines = set() | |
| for name, meta in self.defs.items(): | |
| # Don't update if changes to the file should be ignored | |
| sources = (f for f in meta["sources"] if not self._ignore_file(f)) | |
| # Select changed files | |
| changed = [f for f in sources if Path(f) in self.changed] | |
| if len(changed) > 0: | |
| eprint(f"changed files for {name}: {changed}") | |
| routines.add(name) | |
| ret: dict[str, list[str]] = {} | |
| for r in sorted(routines): | |
| ret.setdefault(self.defs[r]["type"], []).append(r) | |
| return ret | |
| def may_skip_libm_ci(self) -> bool: | |
| """If this is a PR and no libm files were changed, allow skipping libm | |
| jobs.""" | |
| # Always run on merge CI | |
| if not self.is_pr(): | |
| return False | |
| pr = PrInfo.from_env() | |
| assert pr is not None, "Is a PR but couldn't load PrInfo" | |
| # Allow opting in to libm tests | |
| if pr.cfg.always_test_libm: | |
| return False | |
| # By default, run if there are any changed files matching the pattern | |
| return all(not re.match(TRIGGER_LIBM_CI_FILE_PAT, str(f)) for f in self.changed) | |
| def emit_workflow_output(self): | |
| """Create a JSON object a list items for each type's changed files, if any | |
| did change, and the routines that were affected by the change. | |
| """ | |
| skip_tests = False | |
| error_on_many_tests = False | |
| extra_tests = {} | |
| pr = PrInfo.from_env() | |
| if pr is not None: | |
| skip_tests = pr.cfg.skip_extensive | |
| error_on_many_tests = not pr.cfg.allow_many_extensive | |
| for fn_name in pr.cfg.extra_extensive: | |
| extra_tests.setdefault(base_name(fn_name)[1], []).append(fn_name) | |
| if skip_tests: | |
| eprint("Skipping all extensive tests") | |
| changed = self.changed_routines() | |
| eprint(f"Changed: {changed}") | |
| matrix = [] | |
| total_to_test = 0 | |
| # Figure out which extensive tests need to run | |
| for ty in TYPES: | |
| ty_changed = changed.get(ty, []) | |
| ty_to_test = [] if skip_tests else ty_changed.copy() | |
| ty_to_test.extend(extra_tests.get(ty, [])) | |
| total_to_test += len(ty_to_test) | |
| item = { | |
| "ty": ty, | |
| "changed": ",".join(sorted(set(ty_changed))), | |
| "to_test": ",".join(sorted(set(ty_to_test))), | |
| } | |
| matrix.append(item) | |
| ext_matrix = json.dumps({"extensive_matrix": matrix}, separators=(",", ":")) | |
| may_skip = str(self.may_skip_libm_ci()).lower() | |
| print(f"extensive_matrix={ext_matrix}") | |
| print(f"may_skip_libm_ci={may_skip}") | |
| eprint(f"total extensive tests: {total_to_test}") | |
| if error_on_many_tests and total_to_test > PrCfg.MANY_EXTENSIVE_THRESHOLD: | |
| eprint( | |
| f"More than {PrCfg.MANY_EXTENSIVE_THRESHOLD} tests would be run; add" | |
| f" `{PrCfg.DIR_ALLOW_MANY_EXTENSIVE}` to the PR body if this is" | |
| " intentional. If this is refactoring that happens to touch a lot of" | |
| f" files, `{PrCfg.DIR_SKIP_EXTENSIVE}` can be used instead." | |
| ) | |
| exit(1) | |
| def base_name(name: str) -> tuple[str, str]: | |
| """Return the basename and type from a full function name. Keep in sync with Rust's | |
| `fn base_name`. | |
| """ | |
| known_mappings = [ | |
| ("erff", ("erf", "f32")), | |
| ("erf", ("erf", "f64")), | |
| ("modff", ("modf", "f32")), | |
| ("modf", ("modf", "f64")), | |
| ("lgammaf_r", ("lgamma_r", "f32")), | |
| ("lgamma_r", ("lgamma_r", "f64")), | |
| ] | |
| found = next((base for (full, base) in known_mappings if full == name), None) | |
| if found is not None: | |
| return found | |
| if name.endswith("f"): | |
| return (name.rstrip("f"), "f32") | |
| elif name.endswith("f16"): | |
| return (name.rstrip("f16"), "f16") | |
| elif name.endswith("f32"): | |
| return (name.rstrip("f32"), "f32") | |
| elif name.endswith("f64"): | |
| return (name.rstrip("f64"), "f64") | |
| elif name.endswith("f128"): | |
| return (name.rstrip("f128"), "f128") | |
| return (name, "f64") | |
| def locate_baseline(flags: list[str]) -> None: | |
| """Find the most recent baseline from CI, download it if specified. | |
| This returns rather than erroring, even if the `gh` commands fail. This is to avoid | |
| erroring in CI if the baseline is unavailable (artifact time limit exceeded, first | |
| run on the branch, etc). | |
| """ | |
| download = False | |
| extract = False | |
| tag = "" | |
| while len(flags) > 0: | |
| match flags[0]: | |
| case "--download": | |
| download = True | |
| case "--extract": | |
| extract = True | |
| case "--tag": | |
| tag = flags[1] | |
| flags = flags[1:] | |
| case _: | |
| eprint(USAGE) | |
| exit(1) | |
| flags = flags[1:] | |
| if extract and not download: | |
| eprint("cannot extract without downloading") | |
| exit(1) | |
| try: | |
| # Locate the most recent job to complete with success on our branch | |
| latest_job = sp.check_output( | |
| [ | |
| "gh", | |
| "run", | |
| "list", | |
| "--status=success", | |
| f"--branch={DEFAULT_BRANCH}", | |
| "--json=databaseId,url,headSha,conclusion,createdAt," | |
| "status,workflowDatabaseId,workflowName", | |
| # Return the first array element matching our workflow name. NB: cannot | |
| # just use `--limit=1`, jq filtering happens after limiting. We also | |
| # cannot just use `--workflow` because GH gets confused from | |
| # different file names in history. | |
| f'--jq=[.[] | select(.workflowName == "{WORKFLOW_NAME}")][0]', | |
| ], | |
| text=True, | |
| ) | |
| except sp.CalledProcessError as e: | |
| eprint(f"failed to run github command: {e}") | |
| return | |
| try: | |
| latest = json.loads(latest_job) | |
| eprint("latest job: ", json.dumps(latest, indent=4)) | |
| except json.JSONDecodeError as e: | |
| eprint(f"failed to decode json '{latest_job}', {e}") | |
| return | |
| if not download: | |
| eprint("--download not specified, returning") | |
| return | |
| job_id = latest.get("databaseId") | |
| if job_id is None: | |
| eprint("skipping download step") | |
| return | |
| artifact_glob = f"{ARTIFACT_PREFIX}{f"-{tag}" if tag else ""}*" | |
| # Skip checking because this will fail if the file already exists, which is fine. | |
| sp.run( | |
| ["gh", "run", "download", str(job_id), f"--pattern={artifact_glob}"], | |
| check=False, | |
| ) | |
| if not extract: | |
| eprint("skipping extraction step") | |
| return | |
| # Find the baseline with the most recent timestamp. GH downloads the files to e.g. | |
| # `some-dirname/some-dirname.tar.xz`, so just glob the whole thing together. | |
| candidate_baselines = glob(f"{artifact_glob}/{artifact_glob}") | |
| if len(candidate_baselines) == 0: | |
| eprint("no possible baseline directories found") | |
| return | |
| candidate_baselines.sort(reverse=True) | |
| baseline_archive = candidate_baselines[0] | |
| eprint(f"extracting {baseline_archive}") | |
| all_paths = sp.check_output(["tar", "tJf", baseline_archive], encoding="utf8") | |
| sp.run(["tar", "xJf", baseline_archive], check=True) | |
| # Print a short summary of paths, we don't use `tar v` since the list is huge | |
| short_paths = re.findall(r"^(?:[^/\n]+/?){1,3}", all_paths, re.MULTILINE) | |
| print("Extracted:") | |
| for path in sorted(set(short_paths)): | |
| print(f"* {path}") | |
| eprint("baseline extracted successfully") | |
| def handle_bench_regressions(args: list[str]): | |
| """Exit with error unless the PR message contains an ignore directive.""" | |
| match args: | |
| case [pr_number]: | |
| pr_number = pr_number | |
| case _: | |
| eprint(USAGE) | |
| exit(1) | |
| pr = PrInfo.from_pr(pr_number) | |
| if pr.cfg.allow_regressions: | |
| eprint("PR allows regressions") | |
| return | |
| eprint("Regressions were found; benchmark failed") | |
| exit(1) | |
| def main(): | |
| match sys.argv[1:]: | |
| case ["generate-matrix"]: | |
| ctx = Context() | |
| ctx.emit_workflow_output() | |
| case ["locate-baseline", *flags]: | |
| locate_baseline(flags) | |
| case ["handle-bench-regressions", *args]: | |
| handle_bench_regressions(args) | |
| case ["--help" | "-h"]: | |
| print(USAGE) | |
| exit() | |
| case _: | |
| eprint(USAGE) | |
| exit(1) | |
| if __name__ == "__main__": | |
| main() |