wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/scripts/ci.py
Executable
+672
View File
@@ -0,0 +1,672 @@
#!/usr/bin/env python3
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
# ruff: noqa: E501, RUF005
import argparse
import getpass
import grp
import inspect
import json
import multiprocessing
import os
import platform
import random
import shutil
import string
import subprocess
import sys
import textwrap
import typing
from collections.abc import Callable
from pathlib import Path
from typing import Any
REPO_ROOT = Path(__file__).resolve().parent.parent.parent
SCRIPT_DIR = REPO_ROOT / ".ci-py-scripts"
NPROC = multiprocessing.cpu_count()
class col:
BLUE = "\033[94m"
CYAN = "\033[96m"
GREEN = "\033[92m"
YELLOW = "\033[93m"
RED = "\033[91m"
RESET = "\033[0m"
BOLD = "\033[1m"
UNDERLINE = "\033[4m"
def print_color(color: str, msg: str, bold: bool, **kwargs: Any) -> None:
if hasattr(sys.stdout, "isatty") and sys.stdout.isatty():
bold_code = col.BOLD if bold else ""
print(bold_code + color + msg + col.RESET, **kwargs)
else:
print(msg, **kwargs)
warnings: list[str] = []
def clean_exit(msg: str) -> None:
print_color(col.RED, msg, bold=True, file=sys.stderr)
for warning in warnings:
print_color(col.YELLOW, warning, bold=False, file=sys.stderr)
exit(1)
def cmd(commands: list[Any], **kwargs: Any):
commands = [str(s) for s in commands]
command_str = " ".join(commands)
print_color(col.BLUE, command_str, bold=True)
proc = subprocess.run(commands, **kwargs)
if proc.returncode != 0:
raise RuntimeError(f"Command failed: '{command_str}'")
return proc
def get_build_dir(name: str) -> str:
build_dir = REPO_ROOT / f"build-{name}"
return str(build_dir.relative_to(REPO_ROOT))
def check_docker():
executable = shutil.which("docker")
if executable is None:
clean_exit("'docker' executable not found, install it first (e.g. 'apt install docker.io')")
if sys.platform == "linux":
# Check that the user is in the docker group before running
try:
group = grp.getgrnam("docker")
if getpass.getuser() not in group.gr_mem:
warnings.append(
f"Note: User '{getpass.getuser()}' is not in the 'docker' group, either:\n"
" * run with 'sudo'\n"
" * add user to 'docker': sudo usermod -aG docker $(whoami), then log out and back in",
)
except KeyError:
warnings.append("Note: 'docker' group does not exist")
def check_gpu():
if not (sys.platform == "linux" and shutil.which("lshw")):
# Can't check GPU on non-Linux platforms
return
# See if we can check if a GPU is present in case of later failures,
# but don't block on execution since this isn't critical
try:
proc = cmd(
["lshw", "-json", "-C", "display"],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
encoding="utf-8",
)
stdout = proc.stdout.strip().strip(",")
stdout = json.loads(stdout)
except (subprocess.CalledProcessError, json.decoder.JSONDecodeError):
# Do nothing if any step failed
return
if isinstance(stdout, dict):
# Sometimes lshw outputs a single item as a dict instead of a list of
# dicts, so wrap it up if necessary
stdout = [stdout]
if not isinstance(stdout, list):
return
vendors = [s.get("vendor", "").lower() for s in stdout]
if not any("nvidia" in vendor for vendor in vendors):
warnings.append(
"nvidia GPU not found in 'lshw', maybe use --cpu flag when running 'docs' command?"
)
def gen_name(s: str) -> str:
# random 4 letters
suffix = "".join([random.choice(string.ascii_lowercase) for i in range(5)])
return f"{s}-{suffix}"
def docker(
name: str,
image: str,
scripts: list[str],
env: dict[str, str],
interactive: bool,
additional_flags: dict[str, str] | None = None,
):
"""
Invoke a set of bash scripts through docker/bash.sh
name: container name
image: docker image name
scripts: list of bash commands to run
env: environment to set
"""
check_docker()
# As sccache is added to these images these can be uncommented
sccache_images = {
"ci_gpu",
"ci_cpu",
# "ci_wasm",
"ci_cortexm",
"ci_arm",
"ci_riscv",
"ci_adreno",
}
if image in sccache_images and os.getenv("USE_SCCACHE", "1") == "1":
scripts = [
"sccache --start-server",
] + scripts
# Set the C/C++ compiler so CMake picks them up in the build
env["CC"] = "/opt/sccache/cc"
env["CXX"] = "/opt/sccache/c++"
env["SCCACHE_CACHE_SIZE"] = os.getenv("SCCACHE_CACHE_SIZE", "50G")
env["SCCACHE_SERVER_PORT"] = os.getenv("SCCACHE_SERVER_PORT", "4226")
env["PLATFORM"] = name
docker_bash = REPO_ROOT / "docker" / "bash.sh"
command = [docker_bash]
if sys.stdout.isatty():
command.append("-t")
command.append("--name")
command.append(name)
if interactive:
command.append("-i")
scripts = ["interact() {", " bash", "}", "trap interact 0", ""] + scripts
for key, value in env.items():
command.append("--env")
command.append(f"{key}={value}")
if additional_flags is not None:
for key, value in additional_flags.items():
command.append(key)
command.append(value)
SCRIPT_DIR.mkdir(exist_ok=True)
script_file = SCRIPT_DIR / f"{name}.sh"
with open(script_file, "w") as f:
f.write("set -eux\n\n")
f.write("\n".join(scripts))
f.write("\n")
command += [image, "bash", str(script_file.relative_to(REPO_ROOT))]
try:
cmd(command)
except RuntimeError as e:
clean_exit(f"Error invoking Docker: {e}")
except KeyboardInterrupt:
cmd(["docker", "stop", "--time", "1", name])
finally:
if os.getenv("DEBUG", "0") != "1":
script_file.unlink()
def docs(
tutorial_pattern: str | None = None,
cpu: bool = False,
full: bool = False,
interactive: bool = False,
skip_build: bool = False,
docker_image: str | None = None,
) -> None:
"""
Build the documentation from gallery/ and docs/. By default this builds only
the Python docs without any tutorials.
arguments:
full -- Build all language docs, not just Python (cannot be used with --cpu)
cpu -- Use the 'ci_cpu' Docker image (useful for building docs on a machine without a GPU)
tutorial-pattern -- Regex for which tutorials to execute when building docs (cannot be used with --cpu)
skip_build -- skip build and setup scripts
interactive -- start a shell after running build / test scripts
docker-image -- manually specify the docker image to use
"""
build_dir = get_build_dir("gpu")
extra_setup = []
image = "ci_gpu" if docker_image is None else docker_image
if cpu:
# TODO: Change this to tlcpack/docs once that is uploaded
image = "ci_cpu" if docker_image is None else docker_image
build_dir = get_build_dir("cpu")
config_script = " && ".join(
[
f"mkdir -p {build_dir}",
f"pushd {build_dir}",
"cp ../cmake/config.cmake .",
"popd",
]
)
# These are taken from the ci-gpu image via pip freeze, consult that
# if there are any changes: https://github.com/apache/tvm/tree/main/docs#native
requirements = [
"Sphinx==8.1.3",
"sphinx_autodoc_annotation~=1.0",
"sphinx-gallery==0.20.0",
"sphinx-book-theme==1.1.4",
"pydata-sphinx-theme==0.15.4",
"autodocsumm==0.2.14",
"image==1.5.33",
"matplotlib==3.10.8",
"commonmark==0.9.1",
"docutils==0.21.2",
"Pillow==12.1.1",
]
extra_setup = [
"uv pip install " + " ".join(requirements),
]
else:
check_gpu()
config_script = f"./tests/scripts/task_config_build_gpu.sh {build_dir}"
scripts = extra_setup + [
config_script,
f"./tests/scripts/task_build.py --build-dir {build_dir}",
]
if skip_build:
scripts = []
scripts.append("./tests/scripts/task_python_docs.sh")
if tutorial_pattern is None:
tutorial_pattern = os.getenv("TVM_TUTORIAL_EXEC_PATTERN", ".py" if full else "none")
env = {
"TVM_TUTORIAL_EXEC_PATTERN": tutorial_pattern,
"PYTHON_DOCS_ONLY": "0" if full else "1",
"IS_LOCAL": "1",
"TVM_LIBRARY_PATH": str(REPO_ROOT / build_dir),
}
docker(name=gen_name("docs"), image=image, scripts=scripts, env=env, interactive=interactive)
print_color(
col.GREEN,
"Done building the docs. You can view them by running "
"'python3 tests/scripts/ci.py serve-docs' and visiting:"
" http://localhost:8000 in your browser.",
bold=True,
)
def serve_docs(directory: str = "_docs") -> None:
"""
Serve the docs using Python's http server
arguments:
directory -- Directory to serve from
"""
directory_path = Path(directory)
if not directory_path.exists():
clean_exit("Docs have not been built, run 'ci.py docs' first")
cmd([sys.executable, "-m", "http.server"], cwd=directory_path)
def lint(interactive: bool = False, docker_image: str | None = None) -> None:
"""
Run lint checks locally.
arguments:
interactive -- start a shell after running build / test scripts when using --docker-image
docker-image -- manually specify a docker image to use
"""
scripts = ["pre-commit run --all-files"]
if docker_image is None:
if interactive:
clean_exit("--interactive requires --docker-image")
cmd(scripts[0].split())
return
docker(
name=gen_name("lint"),
image=docker_image,
scripts=scripts,
env={},
interactive=interactive,
)
Option = tuple[str, list[str]]
def generate_command(
name: str,
options: dict[str, Option],
help: str,
precheck: Callable[[], None] | None = None,
post_build: list[str] | None = None,
additional_flags: dict[str, str] | None = None,
env: dict[str, str] | None = None,
):
"""
Helper to generate CLIs that:
1. Build a with a config matching a specific CI Docker image (e.g. 'cpu')
2. Run tests (either a pre-defined set from scripts or manually via invoking
pytest)
3. (optional) Drop down into a terminal into the Docker container
"""
def fn(
tests: list[str] | None = None,
skip_build: bool = False,
interactive: bool = False,
docker_image: str | None = None,
verbose: bool = False,
**kwargs,
) -> None:
"""
arguments:
tests -- pytest test IDs (e.g. tests/python or tests/python/a_file.py::a_test[param=1])
skip_build -- skip build and setup scripts
interactive -- start a shell after running build / test scripts
docker-image -- manually specify the docker image to use
verbose -- run verbose build
"""
if precheck is not None:
precheck()
build_dir = get_build_dir(name)
if skip_build:
scripts = []
else:
scripts = [
f"./tests/scripts/task_config_build_{name}.sh {build_dir}",
f"./tests/scripts/task_build.py --build-dir {build_dir}",
]
if post_build is not None:
scripts += post_build
# Check that a test suite was not used alongside specific test names
if any(v for v in kwargs.values()) and tests is not None:
option_flags = ", ".join([f"--{k}" for k in options.keys()])
clean_exit(f"{option_flags} cannot be used with --tests")
if tests is not None:
scripts.append(f"python3 -m pytest {' '.join(tests)}")
# Add named test suites
for option_name, (_, extra_scripts) in options.items():
if kwargs.get(option_name, False):
scripts.extend(script.format(build_dir=build_dir) for script in extra_scripts)
docker_env = {
# Need to specify the library path manually or else TVM can't
# determine which build directory to use (i.e. if there are
# multiple copies of libtvm_compiler.so / libtvm_runtime.so laying around)
"TVM_LIBRARY_PATH": str(REPO_ROOT / get_build_dir(name)),
"VERBOSE": "true" if verbose else "false",
}
if env is not None:
docker_env.update(env)
docker(
name=gen_name(f"ci-{name}"),
image=f"ci_{name}" if docker_image is None else docker_image,
scripts=scripts,
env=docker_env,
interactive=interactive,
additional_flags=additional_flags,
)
fn.__name__ = name
return fn, options, help
def check_arm_qemu() -> None:
"""
Check if a machine is ready to run an ARM Docker image
"""
machine = platform.machine().lower()
if "arm" in machine or "aarch64" in machine:
# No need to check anything if the machine runs ARM
return
binfmt = Path("/proc/sys/fs/binfmt_misc")
if not binfmt.exists() or len(list(binfmt.glob("qemu-*"))) == 0:
clean_exit(
textwrap.dedent(
"""
You must run a one-time setup to use ARM containers on x86 via QEMU:
sudo apt install -y qemu binfmt-support qemu-user-static
docker run --rm --privileged multiarch/qemu-user-static --reset -p yes
See https://www.stereolabs.com/docs/docker/building-arm-container-on-x86/ for details""".strip(
"\n"
)
)
)
def cli_name(s: str) -> str:
return s.replace("_", "-")
def typing_get_origin(annotation):
return typing.get_origin(annotation)
def typing_get_args(annotation):
return typing.get_args(annotation)
def is_optional_type(annotation):
return (
hasattr(annotation, "__origin__")
and (typing_get_origin(annotation) == typing.Union)
and (type(None) in typing_get_args(annotation))
)
def add_subparser(
func: Callable,
subparsers: Any,
options: dict[str, Option] | None = None,
help: str | None = None,
) -> Any:
"""
Utility function to make it so subparser commands can be defined locally
as a function rather than directly via argparse and manually dispatched
out.
"""
# Each function is intended follow the example for arguments in PEP257, so
# split apart the function documentation from the arguments
split = [s.strip() for s in func.__doc__.split("arguments:\n")]
if len(split) == 1:
args_help = None
command_help = split[0]
else:
command_help, args_help = split
if help is not None:
command_help = help
# Parse out the help text for each argument if present
arg_help_texts = {}
if args_help is not None:
for line in args_help.split("\n"):
line = line.strip()
name, help_text = [t.strip() for t in line.split(" -- ")]
arg_help_texts[cli_name(name)] = help_text
subparser = subparsers.add_parser(cli_name(func.__name__), help=command_help)
seen_prefixes = set()
# Add each parameter to the subparser
signature = inspect.signature(func)
for name, value in signature.parameters.items():
if name == "kwargs":
continue
arg_cli_name = cli_name(name)
kwargs: dict[str, str | bool] = {"help": arg_help_texts[arg_cli_name]}
is_optional = is_optional_type(value.annotation)
if is_optional:
arg_type = typing_get_args(value.annotation)[0]
else:
arg_type = value.annotation
# Grab the default value if present
has_default = False
if value.default is not value.empty:
kwargs["default"] = value.default
has_default = True
# Check if it should be a flag
if arg_type is bool:
kwargs["action"] = "store_true"
else:
kwargs["required"] = not is_optional and not has_default
if str(arg_type).startswith("typing.List"):
kwargs["action"] = "append"
if arg_cli_name[0] not in seen_prefixes:
subparser.add_argument(f"-{arg_cli_name[0]}", f"--{arg_cli_name}", **kwargs)
seen_prefixes.add(arg_cli_name[0])
else:
subparser.add_argument(f"--{arg_cli_name}", **kwargs)
if options is not None:
for option_name, (help, _) in options.items():
option_cli_name = cli_name(option_name)
if option_cli_name[0] not in seen_prefixes:
subparser.add_argument(
f"-{option_cli_name[0]}", f"--{option_cli_name}", action="store_true", help=help
)
seen_prefixes.add(option_cli_name[0])
else:
subparser.add_argument(f"--{option_cli_name}", action="store_true", help=help)
return subparser
CPP_UNITTEST = ("run c++ unitests", ["./tests/scripts/task_cpp_unittest.sh {build_dir}"])
generated = [
generate_command(
name="gpu",
help="Run GPU build and test(s)",
options={
"cpp": CPP_UNITTEST,
"unittest": (
"run unit tests",
[
"./tests/scripts/task_java_unittest.sh",
"./tests/scripts/task_python_unittest_gpuonly.sh",
],
),
},
),
generate_command(
name="cpu",
help="Run CPU build and test(s)",
options={
"cpp": CPP_UNITTEST,
"unittest": (
"run unit tests",
[
"./tests/scripts/task_python_unittest.sh",
],
),
},
),
generate_command(
name="wasm",
help="Run WASM build and test(s)",
options={
"cpp": CPP_UNITTEST,
"test": ("run WASM tests", ["./tests/scripts/task_web_wasm.sh"]),
},
),
generate_command(
name="arm",
help="Run ARM build and test(s) (native or via QEMU on x86)",
precheck=check_arm_qemu,
options={
"cpp": CPP_UNITTEST,
"python": (
"run full Python tests",
[
"./tests/scripts/task_python_unittest.sh",
],
),
},
),
]
def main():
description = """
Run CI jobs locally via Docker. This facilitates reproducing CI failures for
fast iteration. Note that many of the Docker images required are large (the
CPU and GPU images are both over 25GB) and may take some time to download on first use.
"""
parser = argparse.ArgumentParser(description=description)
subparsers = parser.add_subparsers(dest="command")
commands = {}
# Add manually defined commands
for func in [docs, serve_docs, lint]:
add_subparser(func, subparsers)
commands[cli_name(func.__name__)] = func
# Add generated commands
for func, options, help in generated:
add_subparser(func, subparsers, options, help)
commands[cli_name(func.__name__)] = func
args = parser.parse_args()
if args.command is None:
# Command not found in list, error out
parser.print_help()
exit(1)
func = commands[args.command]
# Extract out the parsed args and invoke the relevant function
kwargs = {k: getattr(args, k) for k in dir(args) if not k.startswith("_") and k != "command"}
func(**kwargs)
if __name__ == "__main__":
main()