已合并
[inductor] Add akg backend ut. #38510
huawuyi创建于 6月15日
[inductor] Add akg backend ut. #38510
已合并
从已删除 :v2.10.0合入到Ascend/pytorchv2.10.0
共 1 个文件变更+453-0
| @@ -0,0 +1,453 @@ | |||
| 1 | +# Owner(s): ["module: inductor"] | ||
| 2 | +"""Install latest akg/mfusion wheels, then smoke-test TORCHINDUCTOR_USE_AKG=1. | ||
| 3 | + | ||
| 4 | +Repository: https://repo.mindspore.cn/mindspore/akg/newest/ | ||
| 5 | + - akg -> scheduler/{arch}/ | ||
| 6 | + - mfusion -> fusion/{arch}/ | ||
| 7 | + | ||
| 8 | +When akg/torch_mlir is installed, verify AkgScheduling and akg_auto_fallback. | ||
| 9 | +When akg is missing, verify the same entry falls back to NpuMlirScheduling and | ||
| 10 | +mlir_auto_fallback, matching npu_inductor_plugin.register_mlir_codegen_backend(). | ||
| 11 | + | ||
| 12 | +Also reuses test_torch_mlir.TestAdd pointwise cases for end-to-end smoke coverage. | ||
| 13 | + | ||
| 14 | +Usage: | ||
| 15 | + python test_akg_torch_mlir.py # try install akg, then run tests | ||
| 16 | + python test_akg_torch_mlir.py --insecure # bypass SSL verification | ||
| 17 | + python test_akg_torch_mlir.py --download-only # download wheels only (fail on error) | ||
| 18 | + python test_akg_torch_mlir.py --arch aarch64 --py-tag cp311 | ||
| 19 | + | ||
| 20 | +If wheel download/install fails (e.g. network), tests still run using MLIR fallback. | ||
| 21 | +""" | ||
| 22 | + | ||
| 23 | +from __future__ import annotations | ||
| 24 | + | ||
| 25 | +import sys | ||
| 26 | + | ||
| 27 | +# Install-only flags must be stripped before importing common_utils (it reads sys.argv at import). | ||
| 28 | +_INSTALL_FLAGS = frozenset({"--download-only", "--insecure"}) | ||
| 29 | +_INSTALL_VALUE_FLAGS = frozenset({"--arch", "--py-tag", "--download-dir"}) | ||
| 30 | +SAVED_INSTALL_ARGV: list[str] = [] | ||
| 31 | + | ||
| 32 | + | ||
| 33 | +def _strip_install_cli_args() -> None: | ||
| 34 | + """Remove install-only flags from sys.argv; save them for install_akg_wheels().""" | ||
| 35 | + global SAVED_INSTALL_ARGV | ||
| 36 | + install_tokens: list[str] = [] | ||
| 37 | + test_argv: list[str] = [] | ||
| 38 | + index = 1 | ||
| 39 | + while index < len(sys.argv): | ||
| 40 | + arg = sys.argv[index] | ||
| 41 | + if arg in _INSTALL_FLAGS: | ||
| 42 | + install_tokens.append(arg) | ||
| 43 | + index += 1 | ||
| 44 | + elif arg in _INSTALL_VALUE_FLAGS and index + 1 < len(sys.argv): | ||
| 45 | + install_tokens.extend([arg, sys.argv[index + 1]]) | ||
| 46 | + index += 2 | ||
| 47 | + else: | ||
| 48 | + test_argv.append(arg) | ||
| 49 | + index += 1 | ||
| 50 | + SAVED_INSTALL_ARGV = install_tokens | ||
| 51 | + sys.argv = [sys.argv[0], *test_argv] | ||
| 52 | + | ||
| 53 | + | ||
| 54 | +if __name__ == "__main__": | ||
| 55 | + _strip_install_cli_args() | ||
| 56 | + | ||
| 57 | +import argparse | ||
| 58 | +import hashlib | ||
| 59 | +import http.client | ||
| 60 | +import importlib | ||
| 61 | +import logging | ||
| 62 | +import os | ||
| 63 | +import platform | ||
| 64 | +import re | ||
| 65 | +import ssl | ||
| 66 | +import subprocess | ||
| 67 | +import time | ||
| 68 | +import urllib.error | ||
| 69 | +import urllib.request | ||
| 70 | +from collections.abc import Callable | ||
| 71 | +from contextlib import contextmanager | ||
| 72 | +from pathlib import Path | ||
| 73 | + | ||
| 74 | +import torch | ||
| 75 | +from torch._inductor.codegen.common import device_codegens | ||
| 76 | +from torch._inductor.utils import run_and_get_code | ||
| 77 | +from torch.testing._internal.common_utils import run_tests | ||
| 78 | + | ||
| 79 | +logger = logging.getLogger(__name__) | ||
| 80 | + | ||
| 81 | +# --------------------------------------------------------------------------- | ||
| 82 | +# AKG wheel download / install (from install_akg.py) | ||
| 83 | +# --------------------------------------------------------------------------- | ||
| 84 | + | ||
| 85 | +BASE_URL = "https://repo.mindspore.cn/mindspore/akg/newest/" | ||
| 86 | +PACKAGES = {"akg": "scheduler/{arch}/", "mfusion": "fusion/{arch}/"} | ||
| 87 | +ARCH_MAP = {"x86_64": "x86_64", "amd64": "x86_64", "aarch64": "aarch64", "arm64": "aarch64"} | ||
| 88 | +PIP_INDEX = "https://mirrors.aliyun.com/pypi/simple" | ||
| 89 | +TRUSTED_HOSTS = ("repo.mindspore.cn", "mirrors.aliyun.com", "pypi.org", "files.pythonhosted.org") | ||
| 90 | +MAX_RETRIES, RETRY_DELAY_SEC, CHUNK_SIZE = 3, 3, 1024 * 1024 | ||
| 91 | +_NET_ERRORS = ( | ||
| 92 | + urllib.error.URLError, | ||
| 93 | + urllib.error.ContentTooShortError, | ||
| 94 | + http.client.HTTPException, | ||
| 95 | + TimeoutError, | ||
| 96 | + ConnectionResetError, | ||
| 97 | + BrokenPipeError, | ||
| 98 | + OSError, | ||
| 99 | +) | ||
| 100 | + | ||
| 101 | +_ssl_ctx: ssl.SSLContext | None = None | ||
| 102 | +_insecure = False | ||
| 103 | + | ||
| 104 | + | ||
| 105 | +class IntegrityError(Exception): | ||
| 106 | + pass | ||
| 107 | + | ||
| 108 | + | ||
| 109 | +def _with_retry(label: str, fn: Callable[[], object], errors: tuple[type[BaseException], ...] = _NET_ERRORS): | ||
| 110 | + last_error: BaseException | None = None | ||
| 111 | + for attempt in range(1, MAX_RETRIES + 1): | ||
| 112 | + try: | ||
| 113 | + return fn() | ||
| 114 | + except RuntimeError: | ||
| 115 | + raise | ||
| 116 | + except errors as exc: | ||
| 117 | + last_error = exc | ||
| 118 | + if attempt >= MAX_RETRIES: | ||
| 119 | + break | ||
| 120 | + logger.info("%s failed (attempt %d/%d): %s", label, attempt, MAX_RETRIES, exc) | ||
| 121 | + time.sleep(RETRY_DELAY_SEC) | ||
| 122 | + raise RuntimeError(f"{label} failed after {MAX_RETRIES} attempts") from last_error | ||
| 123 | + | ||
| 124 | + | ||
| 125 | +def _open(url: str, timeout: int): | ||
| 126 | + kwargs = {"timeout": timeout, "context": _ssl_ctx} if _ssl_ctx else {"timeout": timeout} | ||
| 127 | + try: | ||
| 128 | + return urllib.request.urlopen( | ||
| 129 | + urllib.request.Request(url, headers={"User-Agent": "test_akg_torch_mlir.py/1.0"}), **kwargs | ||
| 130 | + ) | ||
| 131 | + except urllib.error.URLError as exc: | ||
| 132 | + if "CERTIFICATE_VERIFY_FAILED" in str(exc): | ||
| 133 | + raise RuntimeError( | ||
| 134 | + "SSL certificate verification failed. Retry with: python test_akg_torch_mlir.py --insecure" | ||
| 135 | + ) from exc | ||
| 136 | + raise | ||
| 137 | + | ||
| 138 | + | ||
| 139 | +def fetch(url: str, timeout: int = 120) -> bytes: | ||
| 140 | + def _read() -> bytes: | ||
| 141 | + with _open(url, timeout) as resp: | ||
| 142 | + return resp.read() | ||
| 143 | + | ||
| 144 | + return _with_retry(f"fetch {url}", _read) | ||
| 145 | + | ||
| 146 | + | ||
| 147 | +def _sha256(path: Path) -> str: | ||
| 148 | + return hashlib.sha256(path.read_bytes()).hexdigest() | ||
| 149 | + | ||
| 150 | + | ||
| 151 | +def _fetch_expected_hash(sha256_url: str) -> str | None: | ||
| 152 | + try: | ||
| 153 | + return fetch(sha256_url).decode().strip().split()[0] | ||
| 154 | + except urllib.error.HTTPError: | ||
| 155 | + return None | ||
| 156 | + | ||
| 157 | + | ||
| 158 | +def download(url: str, dest: Path, sha256_url: str) -> None: | ||
| 159 | + dest.parent.mkdir(parents=True, exist_ok=True) | ||
| 160 | + expected = _fetch_expected_hash(sha256_url) | ||
| 161 | + if not expected: | ||
| 162 | + logger.warning("No checksum file found; skipping verification for %s", dest.name) | ||
| 163 | + elif dest.exists(): | ||
| 164 | + if _sha256(dest) == expected: | ||
| 165 | + logger.info("%s already present and verified, skipping download", dest.name) | ||
| 166 | + return | ||
| 167 | + dest.unlink() | ||
| 168 | + | ||
| 169 | + tmp = dest.with_suffix(dest.suffix + ".part") | ||
| 170 | + | ||
| 171 | + def _once() -> None: | ||
| 172 | + with _open(url, 600) as resp, tmp.open("wb") as out: | ||
| 173 | + nbytes = 0 | ||
| 174 | + while chunk := resp.read(CHUNK_SIZE): | ||
| 175 | + out.write(chunk) | ||
| 176 | + nbytes += len(chunk) | ||
| 177 | + if resp.length is not None and nbytes != resp.length: | ||
| 178 | + raise IntegrityError(f"incomplete download for {dest.name}") | ||
| 179 | + if expected and _sha256(tmp) != expected: | ||
| 180 | + raise IntegrityError(f"SHA256 mismatch for {dest.name}") | ||
| 181 | + tmp.replace(dest) | ||
| 182 | + logger.info("Downloaded %s (%.2f MiB)", dest.name, dest.stat().st_size / 1024 / 1024) | ||
| 183 | + | ||
| 184 | + _with_retry(f"download {dest.name}", _once, (*_NET_ERRORS, IntegrityError)) | ||
| 185 | + if expected: | ||
| 186 | + logger.info("SHA256 verified for %s", dest.name) | ||
| 187 | + | ||
| 188 | + | ||
| 189 | +def resolve_target(arch_arg: str | None, py_tag_arg: str | None) -> tuple[str, str]: | ||
| 190 | + arch = arch_arg or ARCH_MAP.get(platform.machine().lower()) | ||
| 191 | + if not arch: | ||
| 192 | + raise RuntimeError(f"Unsupported machine type {platform.machine()!r}; use --arch") | ||
| 193 | + version = sys.version_info[:2] | ||
| 194 | + if py_tag_arg: | ||
| 195 | + return arch, py_tag_arg | ||
| 196 | + if version not in {(3, 10), (3, 11), (3, 12)}: | ||
| 197 | + raise RuntimeError(f"No wheels for Python {version[0]}.{version[1]}; use 3.10-3.12 or --py-tag") | ||
| 198 | + return arch, f"cp{version[0]}{version[1]}" | ||
| 199 | + | ||
| 200 | + | ||
| 201 | +def pick_wheel(subdir: str, package: str, py_tag: str, arch: str) -> str: | ||
| 202 | + token = f"linux_{arch}" | ||
| 203 | + wheels = sorted( | ||
| 204 | + name | ||
| 205 | + for name in ( | ||
| 206 | + href.split("/")[-1] | ||
| 207 | + for href in re.findall(r'href="([^"]+\.whl)"', fetch(BASE_URL + subdir).decode("utf-8", "replace")) | ||
| 208 | + if not href.endswith(".whl.sha256") | ||
| 209 | + ) | ||
| 210 | + if name.startswith(f"{package}-") and f"-{py_tag}-{py_tag}-" in name and token in name | ||
| 211 | + ) | ||
| 212 | + if not wheels: | ||
| 213 | + raise RuntimeError(f"No {package} wheel found for py={py_tag}, arch={arch}") | ||
| 214 | + return wheels[-1] | ||
| 215 | + | ||
| 216 | + | ||
| 217 | +def pip_install(wheels: list[Path]) -> None: | ||
| 218 | + cmd = [sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps", "--extra-index-url", PIP_INDEX] | ||
| 219 | + if _insecure: | ||
| 220 | + for host in TRUSTED_HOSTS: | ||
| 221 | + cmd.extend(["--trusted-host", host]) | ||
| 222 | + cmd.extend(map(str, wheels)) | ||
| 223 | + logger.info("Running: %s", " ".join(cmd)) | ||
| 224 | + _with_retry("pip install", lambda: subprocess.run(cmd, check=True), (subprocess.CalledProcessError,)) | ||
| 225 | + | ||
| 226 | + | ||
| 227 | +def cleanup_downloads(download_dir: Path) -> None: | ||
| 228 | + if not download_dir.exists(): | ||
| 229 | + return | ||
| 230 | + for path in download_dir.glob("*.whl*"): | ||
| 231 | + path.unlink(missing_ok=True) | ||
| 232 | + if not any(download_dir.iterdir()): | ||
| 233 | + download_dir.rmdir() | ||
| 234 | + logger.info("Removed downloaded wheels from %s", download_dir.resolve()) | ||
| 235 | + | ||
| 236 | + | ||
| 237 | +def log_release_info() -> None: | ||
| 238 | + try: | ||
| 239 | + text = fetch(BASE_URL + "release_info.yaml").decode() | ||
| 240 | + info = { | ||
| 241 | + k: (m.group(1)[:12] if k == "commit_id" else m.group(1)) | ||
| 242 | + for k in ("branch", "commit_id", "date") | ||
| 243 | + if (m := re.search(rf"{k}:\s*(\S+)", text)) | ||
| 244 | + } | ||
| 245 | + if info: | ||
| 246 | + logger.info( | ||
| 247 | + "Latest build: branch=%s, commit=%s, date=%s", | ||
| 248 | + info.get("branch"), | ||
| 249 | + info.get("commit_id"), | ||
| 250 | + info.get("date"), | ||
| 251 | + ) | ||
| 252 | + except (urllib.error.URLError, RuntimeError) as exc: | ||
| 253 | + logger.warning("Failed to read release_info.yaml: %s", exc) | ||
| 254 | + | ||
| 255 | + | ||
| 256 | +def install_akg_wheels(args: argparse.Namespace) -> bool: | ||
| 257 | + """Download and optionally install akg/mfusion wheels. | ||
| 258 | + | ||
| 259 | + Returns True when wheels are installed successfully. | ||
| 260 | + Returns False when install fails (network, checksum, pip, etc.). | ||
| 261 | + Raises on failure when --download-only is set. | ||
| 262 | + """ | ||
| 263 | + global _ssl_ctx, _insecure | ||
| 264 | + | ||
| 265 | + if args.insecure: | ||
| 266 | + ctx = ssl.create_default_context() | ||
| 267 | + ctx.check_hostname, ctx.verify_mode = False, ssl.CERT_NONE | ||
| 268 | + _ssl_ctx, _insecure = ctx, True | ||
| 269 | + logger.warning("HTTPS verification disabled; wheel SHA256 verification remains enabled") | ||
| 270 | + | ||
| 271 | + arch, py_tag = resolve_target(args.arch, args.py_tag) | ||
| 272 | + logger.info("Repository: %s", BASE_URL) | ||
| 273 | + logger.info("Target arch=%s, py_tag=%s", arch, py_tag) | ||
| 274 | + log_release_info() | ||
| 275 | + | ||
| 276 | + wheels: list[Path] = [] | ||
| 277 | + try: | ||
| 278 | + for package, subdir_tpl in PACKAGES.items(): | ||
| 279 | + subdir = subdir_tpl.format(arch=arch) | ||
| 280 | + name = pick_wheel(subdir, package, py_tag, arch) | ||
| 281 | + dest = args.download_dir / name | ||
| 282 | + logger.info("Processing package %s: %s", package, name) | ||
| 283 | + download(BASE_URL + subdir + name, dest, BASE_URL + subdir + name + ".sha256") | ||
| 284 | + wheels.append(dest) | ||
| 285 | + if args.download_only: | ||
| 286 | + logger.info("Wheels saved to %s", args.download_dir.resolve()) | ||
| 287 | + return True | ||
| 288 | + pip_install(wheels) | ||
| 289 | + logger.info("Installation completed") | ||
| 290 | + return True | ||
| 291 | + except Exception as exc: | ||
| 292 | + if args.download_only: | ||
| 293 | + raise | ||
| 294 | + logger.warning("AKG wheel installation failed (%s); continuing with MLIR fallback tests", exc) | ||
| 295 | + return False | ||
| 296 | + finally: | ||
| 297 | + if not args.download_only: | ||
| 298 | + cleanup_downloads(args.download_dir) | ||
| 299 | + | ||
| 300 | + | ||
| 301 | +def _build_install_parser() -> argparse.ArgumentParser: | ||
| 302 | + parser = argparse.ArgumentParser(add_help=False) | ||
| 303 | + parser.add_argument("--arch", choices=["x86_64", "aarch64"]) | ||
| 304 | + parser.add_argument("--py-tag") | ||
| 305 | + parser.add_argument("--download-dir", type=Path, default=Path.cwd() / "akg_wheels") | ||
| 306 | + parser.add_argument("--download-only", action="store_true") | ||
| 307 | + parser.add_argument("--insecure", action="store_true", help="Disable HTTPS certificate verification") | ||
| 308 | + return parser | ||
| 309 | + | ||
| 310 | + | ||
| 311 | +# --------------------------------------------------------------------------- | ||
| 312 | +# AKG / MLIR smoke tests | ||
| 313 | +# --------------------------------------------------------------------------- | ||
| 314 | + | ||
| 315 | +AKG_ENV = { | ||
| 316 | + "TORCHINDUCTOR_NPU_BACKEND": "mlir", | ||
| 317 | + "TORCHINDUCTOR_USE_AKG": "1", | ||
| 318 | +} | ||
| 319 | + | ||
| 320 | +HAS_AKG_STACK = False | ||
| 321 | +EXPECTED_SCHEDULING = "" | ||
| 322 | +EXPECTED_COMPILE_API = "" | ||
| 323 | +EXPECTED_SCHEDULING_CLS = None | ||
| 324 | +TestAkgTorchMlir = None | ||
| 325 | + | ||
| 326 | + | ||
| 327 | +def _set_akg_env() -> dict[str, str | None]: | ||
| 328 | + original = {name: os.environ.get(name) for name in AKG_ENV} | ||
| 329 | + os.environ.update(AKG_ENV) | ||
| 330 | + return original | ||
| 331 | + | ||
| 332 | + | ||
| 333 | +def _restore_env(original: dict[str, str | None]) -> None: | ||
| 334 | + for name, value in original.items(): | ||
| 335 | + if value is None: | ||
| 336 | + os.environ.pop(name, None) | ||
| 337 | + else: | ||
| 338 | + os.environ[name] = value | ||
| 339 | + | ||
| 340 | + | ||
| 341 | + | ||
| 342 | +def _temporary_akg_env(): | ||
| 343 | + original = _set_akg_env() | ||
| 344 | + try: | ||
| 345 | + yield | ||
| 346 | + finally: | ||
| 347 | + _restore_env(original) | ||
| 348 | + | ||
| 349 | + | ||
| 350 | +def _bootstrap_tests() -> None: | ||
| 351 | + """Probe akg stack and import torch_npu._inductor after optional wheel install.""" | ||
| 352 | + global HAS_AKG_STACK, EXPECTED_SCHEDULING, EXPECTED_COMPILE_API | ||
| 353 | + global EXPECTED_SCHEDULING_CLS, TestAkgTorchMlir | ||
| 354 | + | ||
| 355 | + try: | ||
| 356 | + importlib.import_module("akg") | ||
| 357 | + importlib.import_module("torch_mlir") | ||
| 358 | + HAS_AKG_STACK = True | ||
| 359 | + except ImportError: | ||
| 360 | + HAS_AKG_STACK = False | ||
| 361 | + | ||
| 362 | + EXPECTED_SCHEDULING = "AkgScheduling" if HAS_AKG_STACK else "NpuMlirScheduling" | ||
| 363 | + EXPECTED_COMPILE_API = "akg_auto_fallback" if HAS_AKG_STACK else "mlir_auto_fallback" | ||
| 364 | + | ||
| 365 | + # Backend selection is resolved when torch_npu._inductor is first imported. | ||
| 366 | + with _temporary_akg_env(): | ||
| 367 | + importlib.import_module("torch_npu._inductor") | ||
| 368 | + from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.codegen.akg import AkgScheduling | ||
| 369 | + from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.codegen.mlir import NpuMlirScheduling | ||
| 370 | + | ||
| 371 | + # isort: off | ||
| 372 | + import test_torch_mlir as torch_mlir_tests | ||
| 373 | + # isort: on | ||
| 374 | + | ||
| 375 | + EXPECTED_SCHEDULING_CLS = AkgScheduling if HAS_AKG_STACK else NpuMlirScheduling | ||
| 376 | + | ||
| 377 | + class _TestAkgTorchMlir(torch_mlir_tests.TestAdd): | ||
| 378 | + """Verify AKG enablement, or MLIR fallback when akg packages are unavailable.""" | ||
| 379 | + | ||
| 380 | + | ||
| 381 | + def _fused_op_calc(first_element, second_element): | ||
| 382 | + return (first_element + second_element) * second_element | ||
| 383 | + | ||
| 384 | + def setUp(self): | ||
| 385 | + self._original_akg_env = _set_akg_env() | ||
| 386 | + try: | ||
| 387 | + super().setUp() | ||
| 388 | + except Exception: | ||
| 389 | + _restore_env(self._original_akg_env) | ||
| 390 | + raise | ||
| 391 | + | ||
| 392 | + def tearDown(self): | ||
| 393 | + try: | ||
| 394 | + super().tearDown() | ||
| 395 | + finally: | ||
| 396 | + _restore_env(self._original_akg_env) | ||
| 397 | + | ||
| 398 | + def test_backend_scheduling_registered(self): | ||
| 399 | + scheduling = getattr(device_codegens.get("npu"), "scheduling", None) | ||
| 400 | + self.assertIs( | ||
| 401 | + scheduling, | ||
| 402 | + EXPECTED_SCHEDULING_CLS, | ||
| 403 | + f"expected {EXPECTED_SCHEDULING} when HAS_AKG_STACK={HAS_AKG_STACK}", | ||
| 404 | + ) | ||
| 405 | + | ||
| 406 | + def test_fused_kernel_compile_path(self): | ||
| 407 | + """Fused multi-op subgraphs should use the selected backend compile API.""" | ||
| 408 | + shape = torch_mlir_tests.TestUtils._pointwise_demo_shapes[0] | ||
| 409 | + dtype = "float32" | ||
| 410 | + x = self._generate_tensor(shape, dtype) | ||
| 411 | + y = self._generate_tensor(shape, dtype) | ||
| 412 | + expected = self._fused_op_calc(x, y) | ||
| 413 | + | ||
| 414 | + compiled = torch.compile(self._fused_op_calc) | ||
| 415 | + result, codes = run_and_get_code(compiled, x, y) | ||
| 416 | + | ||
| 417 | + self.assertEqual(expected, result) | ||
| 418 | + self.assertIn( | ||
| 419 | + EXPECTED_COMPILE_API, | ||
| 420 | + codes[0], | ||
| 421 | + f"expected {EXPECTED_COMPILE_API} when HAS_AKG_STACK={HAS_AKG_STACK}", | ||
| 422 | + ) | ||
| 423 | + | ||
| 424 | + TestAkgTorchMlir = _TestAkgTorchMlir | ||
| 425 | + globals()["TestAkgTorchMlir"] = TestAkgTorchMlir | ||
| 426 | + | ||
| 427 | + | ||
| 428 | +def main() -> int: | ||
| 429 | + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") | ||
| 430 | + | ||
| 431 | + install_args = _build_install_parser().parse_args(SAVED_INSTALL_ARGV) | ||
| 432 | + | ||
| 433 | + try: | ||
| 434 | + installed = install_akg_wheels(install_args) | ||
| 435 | + except KeyboardInterrupt: | ||
| 436 | + logger.error("Interrupted by user") | ||
| 437 | + return 130 | ||
| 438 | + except Exception: | ||
| 439 | + logger.exception("Download failed") | ||
| 440 | + return 1 | ||
| 441 | + | ||
| 442 | + if install_args.download_only: | ||
| 443 | + return 0 if installed else 1 | ||
| 444 | + | ||
| 445 | + _bootstrap_tests() | ||
| 446 | + run_tests() | ||
| 447 | + return 0 | ||
| 448 | + | ||
| 449 | + | ||
| 450 | +if __name__ == "__main__": | ||
| 451 | + raise SystemExit(main()) | ||
| 452 | +else: | ||
| 453 | + _bootstrap_tests() | ||