Stella Laurenzo | b76b6df | 2023-08-24 19:38:26 -0700 | [diff] [blame] | 1 | #!/usr/bin/env python3 |
| 2 | # Copyright 2023 The IREE Authors |
| 3 | # |
| 4 | # Licensed under the Apache License v2.0 with LLVM Exceptions. |
| 5 | # See https://llvm.org/LICENSE.txt for license information. |
| 6 | # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| 7 | |
| 8 | """Sets up a Python venv with compiler/runtime from a workflow run. |
| 9 | |
Scott Todd | 8230f41 | 2024-12-02 10:34:38 -0800 | [diff] [blame] | 10 | There are several modes in which to use this script: |
Stella Laurenzo | b76b6df | 2023-08-24 19:38:26 -0700 | [diff] [blame] | 11 | |
Scott Todd | 3819702 | 2024-12-05 12:03:58 -0800 | [diff] [blame] | 12 | * Locally, `--fetch-gh-workflow=WORKFLOW_ID` can be used to download and |
| 13 | setup the venv from a specific workflow run in one step: |
| 14 | |
| 15 | |
| 16 | ```bash |
| 17 | python3.11 ./build_tools/pkgci/setup_venv.py /tmp/.venv --fetch-gh-workflow=11977414405 |
| 18 | ``` |
| 19 | |
| 20 | * Locally, `--fetch-git-ref=GIT_REF` can be used to download and setup the |
| 21 | venv from the latest workflow run for a given ref (commit) in one step: |
| 22 | |
| 23 | ```bash |
| 24 | python3.11 ./build_tools/pkgci/setup_venv.py /tmp/.venv --fetch-git-ref=main |
| 25 | ``` |
| 26 | |
| 27 | * Locally, `--fetch-latest-main` can be used to download and setup the |
| 28 | venv from the latest completed run from the `main` branch in one step: |
| 29 | |
| 30 | ```bash |
| 31 | python3.11 ./build_tools/pkgci/setup_venv.py /tmp/.venv --fetch-latest-main |
| 32 | ``` |
| 33 | |
Scott Todd | 8230f41 | 2024-12-02 10:34:38 -0800 | [diff] [blame] | 34 | * Within a workflow triggered by `workflow_call`, an artifact action will |
| 35 | typically be used to fetch relevant package artifacts. Specify the fetched |
| 36 | location with `--artifact-path=`: |
Stella Laurenzo | b76b6df | 2023-08-24 19:38:26 -0700 | [diff] [blame] | 37 | |
Scott Todd | 8230f41 | 2024-12-02 10:34:38 -0800 | [diff] [blame] | 38 | ```yml |
| 39 | - uses: actions/download-artifact@fa0a91b85d4f404e444e00e005971372dc801d16 # v4.1.8 |
| 40 | with: |
| 41 | name: linux_x86_64_release_packages |
| 42 | path: ${{ env.PACKAGE_DOWNLOAD_DIR }} |
| 43 | - name: Setup venv |
| 44 | run: | |
| 45 | ./build_tools/pkgci/setup_venv.py ${VENV_DIR} \ |
| 46 | --artifact-path=${PACKAGE_DOWNLOAD_DIR} |
| 47 | ``` |
| 48 | |
| 49 | * Within a workflow triggered by `workflow_dispatch`, pass `artifact_run_id` as |
| 50 | an input that developers must specify when running the workflow: |
| 51 | |
| 52 | ```yml |
| 53 | on: |
| 54 | workflow_dispatch: |
| 55 | inputs: |
| 56 | artifact_run_id: |
| 57 | type: string |
| 58 | default: "" |
| 59 | |
| 60 | ... |
| 61 | steps: |
| 62 | - name: Setup venv |
| 63 | run: | |
| 64 | ./build_tools/pkgci/setup_venv.py ${VENV_DIR} \ |
| 65 | --fetch-gh-workflow=${{ inputs.artifact_run_id }} |
| 66 | ``` |
| 67 | |
| 68 | (Note that these two modes are often combined to allow for workflow testing) |
| 69 | |
Stella Laurenzo | b76b6df | 2023-08-24 19:38:26 -0700 | [diff] [blame] | 70 | You must have the `gh` command line tool installed and authenticated if you |
| 71 | will be fetching artifacts. |
| 72 | """ |
| 73 | |
Scott Todd | 3819702 | 2024-12-05 12:03:58 -0800 | [diff] [blame] | 74 | from glob import glob |
| 75 | from pathlib import Path |
Stella Laurenzo | b76b6df | 2023-08-24 19:38:26 -0700 | [diff] [blame] | 76 | from typing import Optional, Dict, Tuple |
| 77 | |
| 78 | import argparse |
| 79 | import functools |
Stella Laurenzo | b76b6df | 2023-08-24 19:38:26 -0700 | [diff] [blame] | 80 | import json |
Stella Laurenzo | b76b6df | 2023-08-24 19:38:26 -0700 | [diff] [blame] | 81 | import platform |
| 82 | import subprocess |
| 83 | import sys |
| 84 | import tempfile |
| 85 | import zipfile |
| 86 | |
Scott Todd | 8230f41 | 2024-12-02 10:34:38 -0800 | [diff] [blame] | 87 | THIS_DIR = Path(__file__).parent.resolve() |
| 88 | REPO_ROOT = THIS_DIR.parent.parent |
Scott Todd | 3819702 | 2024-12-05 12:03:58 -0800 | [diff] [blame] | 89 | BASE_API_PATH = "/repos/iree-org/iree" |
Scott Todd | 8230f41 | 2024-12-02 10:34:38 -0800 | [diff] [blame] | 90 | |
| 91 | |
| 92 | def parse_arguments(argv=None): |
| 93 | parser = argparse.ArgumentParser(description="Setup venv") |
| 94 | parser.add_argument( |
| 95 | "venv_dir", type=Path, help="Directory in which to create the venv" |
| 96 | ) |
| 97 | parser.add_argument("--artifact-path", help="Path in which to find/fetch artifacts") |
| 98 | |
| 99 | fetch_group = parser.add_mutually_exclusive_group() |
| 100 | fetch_group.add_argument( |
Scott Todd | 3819702 | 2024-12-05 12:03:58 -0800 | [diff] [blame] | 101 | "--fetch-gh-workflow", |
| 102 | help="Fetch artifacts from a specific GitHub workflow using its run ID, like `12125722686`", |
Scott Todd | 8230f41 | 2024-12-02 10:34:38 -0800 | [diff] [blame] | 103 | ) |
Scott Todd | 3819702 | 2024-12-05 12:03:58 -0800 | [diff] [blame] | 104 | fetch_group.add_argument( |
| 105 | "--fetch-latest-main", |
| 106 | help="Fetch artifacts from the latest workflow run on the `main` branch", |
| 107 | action="store_true", |
| 108 | ) |
| 109 | fetch_group.add_argument( |
| 110 | "--fetch-git-ref", |
| 111 | help="Fetch artifacts for a specific git ref. Refs can be branch names (e.g. `main`), commit hashes (short like `abc123` or long), or tags (e.g. `iree-3.0.0`)", |
| 112 | ) |
Scott Todd | 8230f41 | 2024-12-02 10:34:38 -0800 | [diff] [blame] | 113 | |
| 114 | parser.add_argument( |
| 115 | "--compiler-variant", |
| 116 | default="", |
| 117 | help="Package variant to install for the compiler ('', 'asserts')", |
| 118 | ) |
| 119 | parser.add_argument( |
| 120 | "--runtime-variant", |
| 121 | default="", |
| 122 | help="Package variant to install for the runtime ('', 'asserts')", |
| 123 | ) |
| 124 | args = parser.parse_args(argv) |
| 125 | return args |
| 126 | |
| 127 | |
Scott Todd | 3819702 | 2024-12-05 12:03:58 -0800 | [diff] [blame] | 128 | def get_latest_workflow_run_id_for_main() -> int: |
| 129 | print(f"Looking up latest workflow run for main branch") |
| 130 | # Note: at a high level, we probably want to select one of these: |
| 131 | # A) The latest run that produced package artifacts |
| 132 | # B) The latest run that passed all checks |
| 133 | # Instead, we just check for the latest completed workflow. This can miss |
| 134 | # runs that are still pending (especially if jobs are queued waiting for |
| 135 | # runners) and can include jobs that failed tests (for better or worse). |
| 136 | workflow_run_args = [ |
| 137 | "gh", |
| 138 | "api", |
| 139 | "-H", |
| 140 | "Accept: application/vnd.github+json", |
| 141 | "-H", |
| 142 | "X-GitHub-Api-Version: 2022-11-28", |
| 143 | f"{BASE_API_PATH}/actions/workflows/pkgci.yml/runs?branch=main&event=push&status=completed&per_page=1", |
| 144 | ] |
| 145 | print( |
| 146 | f"Running command to find latest completed workflow run:\n {' '.join(workflow_run_args)}" |
| 147 | ) |
| 148 | workflow_run_output = subprocess.check_output(workflow_run_args) |
| 149 | workflow_run_json_output = json.loads(workflow_run_output) |
| 150 | latest_run = workflow_run_json_output["workflow_runs"][0] |
Scott Todd | b8ab7a2 | 2024-12-10 08:19:13 -0800 | [diff] [blame] | 151 | print(f"\nFound workflow run: {latest_run['html_url']}") |
Scott Todd | 3819702 | 2024-12-05 12:03:58 -0800 | [diff] [blame] | 152 | return latest_run["id"] |
| 153 | |
| 154 | |
Scott Todd | 8230f41 | 2024-12-02 10:34:38 -0800 | [diff] [blame] | 155 | def get_latest_workflow_run_id_for_ref(ref: str) -> int: |
Scott Todd | b8ab7a2 | 2024-12-10 08:19:13 -0800 | [diff] [blame] | 156 | print(f"Finding workflow run for ref: {ref}") |
Scott Todd | 8230f41 | 2024-12-02 10:34:38 -0800 | [diff] [blame] | 157 | normalized_ref = ( |
| 158 | subprocess.check_output(["git", "rev-parse", ref], cwd=REPO_ROOT) |
| 159 | .decode() |
| 160 | .strip() |
| 161 | ) |
| 162 | |
Scott Todd | b8ab7a2 | 2024-12-10 08:19:13 -0800 | [diff] [blame] | 163 | print(f" Using normalized ref: {normalized_ref}") |
Scott Todd | 8230f41 | 2024-12-02 10:34:38 -0800 | [diff] [blame] | 164 | workflow_run_args = [ |
| 165 | "gh", |
| 166 | "api", |
| 167 | "-H", |
| 168 | "Accept: application/vnd.github+json", |
| 169 | "-H", |
| 170 | "X-GitHub-Api-Version: 2022-11-28", |
Scott Todd | 3819702 | 2024-12-05 12:03:58 -0800 | [diff] [blame] | 171 | f"{BASE_API_PATH}/actions/workflows/pkgci.yml/runs?head_sha={normalized_ref}", |
Scott Todd | 8230f41 | 2024-12-02 10:34:38 -0800 | [diff] [blame] | 172 | ] |
Scott Todd | b8ab7a2 | 2024-12-10 08:19:13 -0800 | [diff] [blame] | 173 | print(f"\nRunning command to list workflow runs:\n {' '.join(workflow_run_args)}") |
Scott Todd | 8230f41 | 2024-12-02 10:34:38 -0800 | [diff] [blame] | 174 | workflow_run_output = subprocess.check_output(workflow_run_args) |
| 175 | workflow_run_json_output = json.loads(workflow_run_output) |
| 176 | if workflow_run_json_output["total_count"] == 0: |
| 177 | raise RuntimeError("Workflow did not run at this commit") |
| 178 | |
| 179 | latest_run = workflow_run_json_output["workflow_runs"][-1] |
Scott Todd | b8ab7a2 | 2024-12-10 08:19:13 -0800 | [diff] [blame] | 180 | print(f"\nFound workflow run: {latest_run['html_url']}") |
Scott Todd | 8230f41 | 2024-12-02 10:34:38 -0800 | [diff] [blame] | 181 | return latest_run["id"] |
| 182 | |
Stella Laurenzo | b76b6df | 2023-08-24 19:38:26 -0700 | [diff] [blame] | 183 | |
| 184 | @functools.lru_cache |
| 185 | def list_gh_artifacts(run_id: str) -> Dict[str, str]: |
Scott Todd | 8230f41 | 2024-12-02 10:34:38 -0800 | [diff] [blame] | 186 | print(f"Fetching artifacts for workflow run: {run_id}") |
Scott Todd | 3819702 | 2024-12-05 12:03:58 -0800 | [diff] [blame] | 187 | |
Stella Laurenzo | b76b6df | 2023-08-24 19:38:26 -0700 | [diff] [blame] | 188 | output = subprocess.check_output( |
| 189 | [ |
| 190 | "gh", |
| 191 | "api", |
| 192 | "-H", |
| 193 | "Accept: application/vnd.github+json", |
| 194 | "-H", |
| 195 | "X-GitHub-Api-Version: 2022-11-28", |
Scott Todd | 3819702 | 2024-12-05 12:03:58 -0800 | [diff] [blame] | 196 | f"{BASE_API_PATH}/actions/runs/{run_id}/artifacts", |
Stella Laurenzo | b76b6df | 2023-08-24 19:38:26 -0700 | [diff] [blame] | 197 | ] |
| 198 | ) |
| 199 | data = json.loads(output) |
| 200 | # Uncomment to debug: |
| 201 | # print(json.dumps(data, indent=2)) |
| 202 | artifacts = { |
Scott Todd | 3819702 | 2024-12-05 12:03:58 -0800 | [diff] [blame] | 203 | rec["name"]: f"{BASE_API_PATH}/actions/artifacts/{rec['id']}/zip" |
Stella Laurenzo | b76b6df | 2023-08-24 19:38:26 -0700 | [diff] [blame] | 204 | for rec in data["artifacts"] |
| 205 | } |
Scott Todd | b8ab7a2 | 2024-12-10 08:19:13 -0800 | [diff] [blame] | 206 | print("\nFound artifacts:") |
Stella Laurenzo | b76b6df | 2023-08-24 19:38:26 -0700 | [diff] [blame] | 207 | for k, v in artifacts.items(): |
| 208 | print(f" {k}: {v}") |
| 209 | return artifacts |
| 210 | |
| 211 | |
| 212 | def fetch_gh_artifact(api_path: str, file: Path): |
| 213 | print(f"Downloading artifact {api_path}") |
| 214 | contents = subprocess.check_output( |
| 215 | [ |
| 216 | "gh", |
| 217 | "api", |
| 218 | "-H", |
| 219 | "Accept: application/vnd.github+json", |
| 220 | "-H", |
| 221 | "X-GitHub-Api-Version: 2022-11-28", |
| 222 | api_path, |
| 223 | ] |
| 224 | ) |
| 225 | file.write_bytes(contents) |
| 226 | |
| 227 | |
| 228 | def find_venv_python(venv_path: Path) -> Optional[Path]: |
| 229 | paths = [venv_path / "bin" / "python", venv_path / "Scripts" / "python.exe"] |
| 230 | for p in paths: |
| 231 | if p.exists(): |
| 232 | return p |
| 233 | return None |
| 234 | |
| 235 | |
Scott Todd | 3819702 | 2024-12-05 12:03:58 -0800 | [diff] [blame] | 236 | def find_wheel_for_variants( |
| 237 | args, artifact_prefix: str, package_stem: str, variant: str |
| 238 | ) -> Tuple[Path, str]: |
| 239 | artifact_path = Path(args.artifact_path) |
| 240 | package_suffix = "" if variant == "" else f"-{variant}" |
| 241 | package_name = f"{package_stem}{package_suffix}" |
| 242 | |
| 243 | def has_package(): |
| 244 | norm_package_name = package_name.replace("-", "_") |
| 245 | pattern = str(artifact_path / f"{norm_package_name}-*.whl") |
| 246 | files = glob(pattern) |
| 247 | return bool(files) |
| 248 | |
| 249 | if has_package(): |
| 250 | return (artifact_path, package_name) |
| 251 | |
| 252 | if not args.fetch_gh_workflow: |
| 253 | raise RuntimeError( |
| 254 | f"Could not find package {package_name} to install from {artifact_path}" |
| 255 | ) |
| 256 | |
| 257 | # Fetch. |
| 258 | artifact_path.mkdir(parents=True, exist_ok=True) |
| 259 | artifact_suffix = "" if variant == "" else f"_{variant}" |
| 260 | artifact_name = f"{artifact_prefix}_release{artifact_suffix}_packages" |
| 261 | artifact_file = artifact_path / f"{artifact_name}.zip" |
| 262 | if not artifact_file.exists(): |
| 263 | print( |
| 264 | f"Package {package_name} not found in cache. Fetching from {artifact_name}..." |
| 265 | ) |
| 266 | artifacts = list_gh_artifacts(args.fetch_gh_workflow) |
| 267 | if artifact_name not in artifacts: |
| 268 | raise RuntimeError( |
| 269 | f"Could not find required artifact {artifact_name} in run {args.fetch_gh_workflow}" |
| 270 | ) |
| 271 | fetch_gh_artifact(artifacts[artifact_name], artifact_file) |
| 272 | print(f"Extracting {artifact_file}") |
| 273 | with zipfile.ZipFile(artifact_file) as zip_ref: |
| 274 | zip_ref.extractall(artifact_path) |
| 275 | |
| 276 | # Try again. |
| 277 | if not has_package(): |
| 278 | raise RuntimeError(f"Could not find {package_name} in {artifact_path}") |
| 279 | return (artifact_path, package_name) |
| 280 | |
| 281 | |
Stella Laurenzo | b76b6df | 2023-08-24 19:38:26 -0700 | [diff] [blame] | 282 | def main(args): |
Scott Todd | 3819702 | 2024-12-05 12:03:58 -0800 | [diff] [blame] | 283 | # Look up the workflow run for latest main. |
| 284 | if args.fetch_latest_main: |
| 285 | latest_gh_workflow = get_latest_workflow_run_id_for_main() |
| 286 | args.fetch_gh_workflow = str(latest_gh_workflow) |
| 287 | args.fetch_latest_main = "" |
| 288 | return main(args) |
| 289 | |
Scott Todd | 8230f41 | 2024-12-02 10:34:38 -0800 | [diff] [blame] | 290 | # Look up the workflow run for a ref. |
| 291 | if args.fetch_git_ref: |
| 292 | latest_gh_workflow = get_latest_workflow_run_id_for_ref(args.fetch_git_ref) |
| 293 | args.fetch_git_ref = "" |
| 294 | args.fetch_gh_workflow = str(latest_gh_workflow) |
| 295 | return main(args) |
| 296 | |
Stella Laurenzo | b76b6df | 2023-08-24 19:38:26 -0700 | [diff] [blame] | 297 | # Make sure we have an artifact path if fetching. |
| 298 | if not args.artifact_path and args.fetch_gh_workflow: |
| 299 | with tempfile.TemporaryDirectory() as td: |
| 300 | args.artifact_path = td |
| 301 | return main(args) |
| 302 | |
Stella Laurenzo | b76b6df | 2023-08-24 19:38:26 -0700 | [diff] [blame] | 303 | artifact_prefix = f"{platform.system().lower()}_{platform.machine()}" |
| 304 | wheels = [] |
| 305 | for package_stem, variant in [ |
Marius Brehler | c651ba9 | 2024-11-07 21:44:30 +0100 | [diff] [blame] | 306 | ("iree-base-compiler", args.compiler_variant), |
| 307 | ("iree-base-runtime", args.runtime_variant), |
Stella Laurenzo | b76b6df | 2023-08-24 19:38:26 -0700 | [diff] [blame] | 308 | ]: |
| 309 | wheels.append( |
| 310 | find_wheel_for_variants(args, artifact_prefix, package_stem, variant) |
| 311 | ) |
Scott Todd | b8ab7a2 | 2024-12-10 08:19:13 -0800 | [diff] [blame] | 312 | print("\nInstalling wheels:", wheels) |
Stella Laurenzo | b76b6df | 2023-08-24 19:38:26 -0700 | [diff] [blame] | 313 | |
| 314 | # Set up venv. |
| 315 | venv_path = args.venv_dir |
| 316 | python_exe = find_venv_python(venv_path) |
| 317 | |
| 318 | if not python_exe: |
| 319 | print(f"Creating venv at {str(venv_path)}") |
| 320 | subprocess.check_call([sys.executable, "-m", "venv", str(venv_path)]) |
| 321 | python_exe = find_venv_python(venv_path) |
| 322 | if not python_exe: |
| 323 | raise RuntimeError("Error creating venv") |
| 324 | |
| 325 | # Install each of the built wheels without deps or consulting an index. |
| 326 | # This is because we absolutely don't want this falling back to anything |
| 327 | # but what we said. |
| 328 | for artifact_path, package_name in wheels: |
| 329 | cmd = [ |
| 330 | str(python_exe), |
| 331 | "-m", |
| 332 | "pip", |
| 333 | "install", |
| 334 | "--no-deps", |
| 335 | "--no-index", |
| 336 | "-f", |
| 337 | str(artifact_path), |
| 338 | "--force-reinstall", |
| 339 | package_name, |
| 340 | ] |
| 341 | print(f"Running command: {' '.join([str(c) for c in cmd])}") |
| 342 | subprocess.check_call(cmd) |
| 343 | |
Scott Todd | b8ab7a2 | 2024-12-10 08:19:13 -0800 | [diff] [blame] | 344 | print(f"\nvenv setup complete at '{venv_path}'. Activate it with") |
| 345 | print(f" source {venv_path}/bin/activate") |
| 346 | |
Stella Laurenzo | b76b6df | 2023-08-24 19:38:26 -0700 | [diff] [blame] | 347 | return 0 |
| 348 | |
| 349 | |
Stella Laurenzo | b76b6df | 2023-08-24 19:38:26 -0700 | [diff] [blame] | 350 | if __name__ == "__main__": |
| 351 | sys.exit(main(parse_arguments())) |