blob: 7f18900fbbf8342e21bce6acc931a708847f8652 [file] [log] [blame]
Stella Laurenzob76b6df2023-08-24 19:38:26 -07001#!/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 Todd8230f412024-12-02 10:34:38 -080010There are several modes in which to use this script:
Stella Laurenzob76b6df2023-08-24 19:38:26 -070011
Scott Todd38197022024-12-05 12:03:58 -080012* 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 Todd8230f412024-12-02 10:34:38 -080034* 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 Laurenzob76b6df2023-08-24 19:38:26 -070037
Scott Todd8230f412024-12-02 10:34:38 -080038 ```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 Laurenzob76b6df2023-08-24 19:38:26 -070070You must have the `gh` command line tool installed and authenticated if you
71will be fetching artifacts.
72"""
73
Scott Todd38197022024-12-05 12:03:58 -080074from glob import glob
75from pathlib import Path
Stella Laurenzob76b6df2023-08-24 19:38:26 -070076from typing import Optional, Dict, Tuple
77
78import argparse
79import functools
Stella Laurenzob76b6df2023-08-24 19:38:26 -070080import json
Stella Laurenzob76b6df2023-08-24 19:38:26 -070081import platform
82import subprocess
83import sys
84import tempfile
85import zipfile
86
Scott Todd8230f412024-12-02 10:34:38 -080087THIS_DIR = Path(__file__).parent.resolve()
88REPO_ROOT = THIS_DIR.parent.parent
Scott Todd38197022024-12-05 12:03:58 -080089BASE_API_PATH = "/repos/iree-org/iree"
Scott Todd8230f412024-12-02 10:34:38 -080090
91
92def 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 Todd38197022024-12-05 12:03:58 -0800101 "--fetch-gh-workflow",
102 help="Fetch artifacts from a specific GitHub workflow using its run ID, like `12125722686`",
Scott Todd8230f412024-12-02 10:34:38 -0800103 )
Scott Todd38197022024-12-05 12:03:58 -0800104 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 Todd8230f412024-12-02 10:34:38 -0800113
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 Todd38197022024-12-05 12:03:58 -0800128def 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 Toddb8ab7a22024-12-10 08:19:13 -0800151 print(f"\nFound workflow run: {latest_run['html_url']}")
Scott Todd38197022024-12-05 12:03:58 -0800152 return latest_run["id"]
153
154
Scott Todd8230f412024-12-02 10:34:38 -0800155def get_latest_workflow_run_id_for_ref(ref: str) -> int:
Scott Toddb8ab7a22024-12-10 08:19:13 -0800156 print(f"Finding workflow run for ref: {ref}")
Scott Todd8230f412024-12-02 10:34:38 -0800157 normalized_ref = (
158 subprocess.check_output(["git", "rev-parse", ref], cwd=REPO_ROOT)
159 .decode()
160 .strip()
161 )
162
Scott Toddb8ab7a22024-12-10 08:19:13 -0800163 print(f" Using normalized ref: {normalized_ref}")
Scott Todd8230f412024-12-02 10:34:38 -0800164 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 Todd38197022024-12-05 12:03:58 -0800171 f"{BASE_API_PATH}/actions/workflows/pkgci.yml/runs?head_sha={normalized_ref}",
Scott Todd8230f412024-12-02 10:34:38 -0800172 ]
Scott Toddb8ab7a22024-12-10 08:19:13 -0800173 print(f"\nRunning command to list workflow runs:\n {' '.join(workflow_run_args)}")
Scott Todd8230f412024-12-02 10:34:38 -0800174 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 Toddb8ab7a22024-12-10 08:19:13 -0800180 print(f"\nFound workflow run: {latest_run['html_url']}")
Scott Todd8230f412024-12-02 10:34:38 -0800181 return latest_run["id"]
182
Stella Laurenzob76b6df2023-08-24 19:38:26 -0700183
184@functools.lru_cache
185def list_gh_artifacts(run_id: str) -> Dict[str, str]:
Scott Todd8230f412024-12-02 10:34:38 -0800186 print(f"Fetching artifacts for workflow run: {run_id}")
Scott Todd38197022024-12-05 12:03:58 -0800187
Stella Laurenzob76b6df2023-08-24 19:38:26 -0700188 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 Todd38197022024-12-05 12:03:58 -0800196 f"{BASE_API_PATH}/actions/runs/{run_id}/artifacts",
Stella Laurenzob76b6df2023-08-24 19:38:26 -0700197 ]
198 )
199 data = json.loads(output)
200 # Uncomment to debug:
201 # print(json.dumps(data, indent=2))
202 artifacts = {
Scott Todd38197022024-12-05 12:03:58 -0800203 rec["name"]: f"{BASE_API_PATH}/actions/artifacts/{rec['id']}/zip"
Stella Laurenzob76b6df2023-08-24 19:38:26 -0700204 for rec in data["artifacts"]
205 }
Scott Toddb8ab7a22024-12-10 08:19:13 -0800206 print("\nFound artifacts:")
Stella Laurenzob76b6df2023-08-24 19:38:26 -0700207 for k, v in artifacts.items():
208 print(f" {k}: {v}")
209 return artifacts
210
211
212def 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
228def 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 Todd38197022024-12-05 12:03:58 -0800236def 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 Laurenzob76b6df2023-08-24 19:38:26 -0700282def main(args):
Scott Todd38197022024-12-05 12:03:58 -0800283 # 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 Todd8230f412024-12-02 10:34:38 -0800290 # 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 Laurenzob76b6df2023-08-24 19:38:26 -0700297 # 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 Laurenzob76b6df2023-08-24 19:38:26 -0700303 artifact_prefix = f"{platform.system().lower()}_{platform.machine()}"
304 wheels = []
305 for package_stem, variant in [
Marius Brehlerc651ba92024-11-07 21:44:30 +0100306 ("iree-base-compiler", args.compiler_variant),
307 ("iree-base-runtime", args.runtime_variant),
Stella Laurenzob76b6df2023-08-24 19:38:26 -0700308 ]:
309 wheels.append(
310 find_wheel_for_variants(args, artifact_prefix, package_stem, variant)
311 )
Scott Toddb8ab7a22024-12-10 08:19:13 -0800312 print("\nInstalling wheels:", wheels)
Stella Laurenzob76b6df2023-08-24 19:38:26 -0700313
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 Toddb8ab7a22024-12-10 08:19:13 -0800344 print(f"\nvenv setup complete at '{venv_path}'. Activate it with")
345 print(f" source {venv_path}/bin/activate")
346
Stella Laurenzob76b6df2023-08-24 19:38:26 -0700347 return 0
348
349
Stella Laurenzob76b6df2023-08-24 19:38:26 -0700350if __name__ == "__main__":
351 sys.exit(main(parse_arguments()))