[ci] Workflow Rewrite: Building on Linux (taichi-dev#6848)

Issue: taichi-dev#6445 

### Brief Summary

.github/workflows/scripts/build.py

@@ -0,0 +1,108 @@
+#!/usr/bin/env python3
+
+# -- prioritized --
+import ci_common  # isort: skip, early initialization happens here
+
+# -- stdlib --
+import glob
+import os
+import platform
+
+# -- third party --
+# -- own --
+from ci_common.dep import download_dep
+from ci_common.misc import banner, get_cache_home, is_manylinux2014
+from ci_common.python import setup_python
+from ci_common.sccache import setup_sccache
+from ci_common.tinysh import Command, environ, git, sh
+
+
+# -- code --
+@banner('Setup LLVM')
+def setup_llvm(env_out: dict) -> None:
+    '''
+    Download and install LLVM.
+    '''
+    u = platform.uname()
+    if u.system == 'Linux':
+        if 'AMDGPU_TEST' in os.environ:
+            # FIXME: AMDGPU bots are currently maintained separately,
+            #        we should unify them with the rest of the bots.
+            lnsf = sh.sudo.ln.bake('-sf')
+            lnsf('/usr/bin/clang++-10', '/usr/bin/clang++')
+            lnsf('/usr/bin/clang-10', '/usr/bin/clang')
+            lnsf('/usr/bin/ld.lld-10', '/usr/bin/ld.lld')
+            env_out['LLVM_DIR'] = '/taichi-llvm-15'
+            return
+        elif is_manylinux2014():
+            # FIXME: prebuilt llvm15 on ubuntu didn't work on manylinux2014 image of centos. Once that's fixed, remove this hack.
+            out = get_cache_home() / 'llvm15-manylinux2014'
+            url = 'https://github.com/ailzhang/torchhub_example/releases/download/0.3/taichi-llvm-15-linux.zip'
+        else:
+            out = get_cache_home() / 'llvm15'
+            url = 'https://github.com/taichi-dev/taichi_assets/releases/download/llvm15/taichi-llvm-15-linux.zip'
+    elif (u.system, u.machine) == ('Darwin', 'arm64'):
+        out = get_cache_home() / 'llvm15-m1'
+        url = 'https://github.com/taichi-dev/taichi_assets/releases/download/llvm15/taichi-llvm-15-m1.zip'
+    elif (u.system, u.machine) == ('Darwin', 'x86_64'):
+        out = get_cache_home() / 'llvm15-mac'
+        url = 'https://github.com/taichi-dev/taichi_assets/releases/download/llvm15/llvm-15-mac10.15.zip'
+    else:
+        raise RuntimeError(f'Unsupported platform: {u.system} {u.machine}')
+
+    download_dep(url, out, strip=1)
+    env_out['LLVM_DIR'] = str(out)
+
+
+@banner('Build Taichi Wheel')
+def build_wheel(python: Command, pip: Command, env: dict) -> None:
+    '''
+    Build the Taichi wheel
+    '''
+    pip.install('-r', 'requirements_dev.txt')
+    git.fetch('origin', 'master', '--tags')
+    proj = env['PROJECT_NAME']
+    proj_tags = []
+    extra = []
+
+    if proj == 'taichi-nightly':
+        proj_tags.extend(['egg_info', '--tag-date'])
+        # Include C-API in nightly builds
+        os.environ['TAICHI_CMAKE_ARGS'] += ' -DTI_WITH_C_API=ON'
+
+    if platform.system() == 'Linux':
+        if is_manylinux2014():
+            extra.extend(['-p', 'manylinux2014_x86_64'])
+        else:
+            extra.extend(['-p', 'manylinux_2_27_x86_64'])
+
+    python('misc/make_changelog.py', '--ver', 'origin/master', '--repo_dir',
+           './', '--save')
+
+    with environ(env):
+        python('setup.py', *proj_tags, 'bdist_wheel', *extra)
+
+
+def main() -> None:
+    env = {
+        'TAICHI_CMAKE_ARGS': os.environ.get('TAICHI_CMAKE_ARGS', ''),
+        'PROJECT_NAME': os.environ.get('PROJECT_NAME', 'taichi'),
+    }
+    setup_llvm(env)
+    sccache = setup_sccache(env)
+
+    # NOTE: We use conda/venv to build wheels, which may not be the same python
+    #       running this script.
+    python, pip = setup_python(os.environ['PY'])
+    build_wheel(python, pip, env)
+
+    sccache('-s')
+
+    distfiles = glob.glob('dist/*.whl')
+    if len(distfiles) != 1:
+        raise RuntimeError(
+            f'Failed to produce exactly one wheel file: {distfiles}')
+
+
+if __name__ == '__main__':
+    main()

.github/workflows/scripts/ci_common/__init__.py

@@ -0,0 +1,2 @@
+from .bootstrap import early_init
+early_init()

.github/workflows/scripts/ci_common/bootstrap.py

@@ -0,0 +1,103 @@
+# -*- coding: utf-8 -*-
+
+# -- stdlib --
+import importlib
+import os
+import sys
+from pathlib import Path
+
+# -- third party --
+# -- own --
+
+
+# -- code --
+def is_in_venv() -> bool:
+    '''
+    Are we in a virtual environment?
+    '''
+    return hasattr(sys, 'real_prefix') or (hasattr(sys, 'base_prefix')
+                                           and sys.base_prefix != sys.prefix)
+
+
+def ensure_dependencies():
+    '''
+    Automatically install dependencies if they are not installed.
+    '''
+    p = Path(__file__).parent.parent / 'requirements.txt'
+    if not p.exists():
+        raise RuntimeError(f'Cannot find {p}')
+
+    user = '' if is_in_venv() else '--user'
+
+    with open(p) as f:
+        deps = [i.strip().split('=')[0] for i in f.read().splitlines()]
+
+    try:
+        for dep in deps:
+            importlib.import_module(dep)
+    except ModuleNotFoundError:
+        print('Installing dependencies...')
+        if os.system(f'{sys.executable} -m pip install {user} -U pip'):
+            raise Exception('Unable to upgrade pip!')
+        if os.system(f'{sys.executable} -m pip install {user} -U -r {p}'):
+            raise Exception('Unable to install dependencies!')
+        os.execl(sys.executable, sys.executable, *sys.argv)
+
+
+def chdir_to_root():
+    '''
+    Change working directory to the root of the repository
+    '''
+    root = Path('/')
+    p = Path(__file__).resolve()
+    while p != root:
+        if (p / '.git').exists():
+            os.chdir(p)
+            break
+        p = p.parent
+
+
+def set_common_env():
+    '''
+    Set common environment variables.
+    '''
+    os.environ['TI_CI'] = '1'
+
+
+_Environ = os.environ.__class__
+
+
+class _EnvironWrapper(_Environ):
+    def __setitem__(self, name: str, value: str) -> None:
+        orig = self.get(name)
+        _Environ.__setitem__(self, name, value)
+        new = self[name]
+
+        if orig == new:
+            return
+
+        from .escapes import escape_codes
+
+        G = escape_codes['bold_green']
+        R = escape_codes['bold_red']
+        N = escape_codes['reset']
+        print(f'{R}:: ENV -{name}={orig}{N}', file=sys.stderr, flush=True)
+        print(f'{G}:: ENV +{name}={new}{N}', file=sys.stderr, flush=True)
+
+
+def monkey_patch_environ():
+    '''
+    Monkey patch os.environ to print changes.
+    '''
+    os.environ.__class__ = _EnvironWrapper
+
+
+def early_init():
+    '''
+    Do early initialization.
+    This must be called before any other non-stdlib imports.
+    '''
+    ensure_dependencies()
+    chdir_to_root()
+    monkey_patch_environ()
+    set_common_env()

.github/workflows/scripts/ci_common/dep.py

@@ -0,0 +1,100 @@
+# -*- coding: utf-8 -*-
+
+# -- stdlib --
+from pathlib import Path
+from urllib.parse import urlparse
+import zipfile
+import shutil
+
+# -- third party --
+import requests
+
+# -- own --
+from .misc import get_cache_home
+from .tinysh import tar
+
+
+# -- code --
+def unzip(filename, extract_dir, strip=0):
+    '''
+    Unpack zip `filename` to `extract_dir`, optionally stripping `strip` components.
+    '''
+    if not zipfile.is_zipfile(filename):
+        raise Exception(f"{filename} is not a zip file")
+
+    extract_dir = Path(extract_dir)
+
+    ar = zipfile.ZipFile(filename)
+    try:
+        for info in ar.infolist():
+            name = info.filename
+
+            # don't extract absolute paths or ones with .. in them
+            if name.startswith('/') or '..' in name:
+                continue
+
+            target = extract_dir.joinpath(*name.split('/')[strip:]).resolve()
+            if not target:
+                continue
+
+            target.parent.mkdir(parents=True, exist_ok=True)
+            if not name.endswith('/'):
+                # file
+                data = ar.read(info.filename)
+                f = open(target, 'wb')
+                try:
+                    f.write(data)
+                finally:
+                    f.close()
+                    del data
+    finally:
+        ar.close()
+
+
+def download_dep(url, outdir, *, strip=0, force=False):
+    '''
+    Download a dependency archive from `url` and expand it to `outdir`,
+    optionally stripping `strip` components.
+    '''
+    outdir = Path(outdir)
+    if outdir.exists() and len(list(outdir.glob('*'))) > 0 and not force:
+        return
+
+    shutil.rmtree(outdir, ignore_errors=True)
+
+    parsed = urlparse(url)
+    name = Path(parsed.path).name
+    escaped = url.replace('/', '_').replace(':', '_')
+    depcache = get_cache_home() / 'deps'
+    depcache.mkdir(parents=True, exist_ok=True)
+    local_cached = depcache / escaped
+
+    if not local_cached.exists():
+        cached_url = f'http://botmaster.tgr:9000/misc/depcache/{escaped}/{name}'
+        try:
+            resp = requests.head(cached_url, timeout=1)
+            if resp.ok:
+                print('Using near cache: ', cached_url)
+                url = cached_url
+        except Exception:
+            pass
+
+        import tqdm
+
+        with requests.get(url, stream=True) as r:
+            r.raise_for_status()
+            total_size = int(r.headers.get('content-length', 0))
+            prog = tqdm.tqdm(unit="B", unit_scale=True, unit_divisor=1024, total=total_size, desc=name)
+            with prog, open(local_cached, 'wb') as f:
+                for chunk in r.iter_content(chunk_size=8192):
+                    sz = f.write(chunk)
+                    prog.update(sz)
+
+    outdir.mkdir(parents=True, exist_ok=True)
+
+    if name.endswith('.zip'):
+        unzip(local_cached, outdir, strip=strip)
+    elif name.endswith('.tar.gz') or name.endswith('.tgz'):
+        tar('-xzf', local_cached, '-C', outdir, f'--strip-components={strip}')
+    else:
+        raise RuntimeError(f'Unknown file type: {name}')

.github/workflows/scripts/ci_common/escapes.py

@@ -0,0 +1,60 @@
+"""
+Generates a dictionary of ANSI escape codes.
+
+http://en.wikipedia.org/wiki/ANSI_escape_code
+
+Uses colorama as an optional dependency to support color on Windows
+
+---
+Copied from colorlog
+---
+
+This file is currently only used to display colored banner.
+"""
+
+__all__ = ('escape_codes', 'parse_colors')
+
+
+# Returns escape codes from format codes
+def esc(*x):
+    return '\033[' + ';'.join(x) + 'm'
+
+
+# The initial list of escape codes
+escape_codes = {
+    'reset': esc('0'),
+    'bold': esc('01'),
+    'thin': esc('02')
+}
+
+# The color names
+COLORS = [
+    'black',
+    'red',
+    'green',
+    'yellow',
+    'blue',
+    'purple',
+    'cyan',
+    'white'
+]
+
+PREFIXES = [
+    # Foreground without prefix
+    ('3', ''), ('01;3', 'bold_'), ('02;3', 'thin_'),
+
+    # Foreground with fg_ prefix
+    ('3', 'fg_'), ('01;3', 'fg_bold_'), ('02;3', 'fg_thin_'),
+
+    # Background with bg_ prefix - bold/light works differently
+    ('4', 'bg_'), ('10', 'bg_bold_'),
+]
+
+for prefix, prefix_name in PREFIXES:
+    for code, name in enumerate(COLORS):
+        escape_codes[prefix_name + name] = esc(prefix + str(code))
+
+
+def parse_colors(sequence):
+    """Return escape codes from a color sequence."""
+    return ''.join(escape_codes[n] for n in sequence.split(',') if n)