[ci] Build: auto install vulkan on Linux (taichi-dev#6969)

Issue: taichi-dev#6445 

### Brief Summary

.github/workflows/scripts/build.py

@@ -54,6 +54,26 @@ def setup_llvm(env_out: dict) -> None:
     env_out['LLVM_DIR'] = str(out)
 
 
+@banner('Setup Vulkan 1.3.236.0')
+def setup_vulkan(env: dict):
+    u = platform.uname()
+    if u.system == "Linux":
+        url = 'https://sdk.lunarg.com/sdk/download/1.3.236.0/linux/vulkansdk-linux-x86_64-1.3.236.0.tar.gz'
+        prefix = get_cache_home() / 'vulkan-1.3.236.0'
+        download_dep(url, prefix, strip=1)
+        sdk = prefix / 'x86_64'
+        env['VULKAN_SDK'] = str(sdk)
+        env['PATH'] = str(sdk / "bin") + ':' + env["PATH"]
+        env['LD_LIBRARY_PATH'] = str(sdk / "lib") + ':' + env.get(
+            "LD_LIBRARY_PATH", "")
+        env['VK_LAYER_PATH'] = str(sdk / 'etc' / 'vulkan' / 'explicit_layer.d')
+    # elif (u.system, u.machine) == ("Darwin", "arm64"):
+    # elif (u.system, u.machine) == ("Darwin", "x86_64"):
+    # elif u.system == "Windows":
+    else:
+        return
+
+
 @banner('Build Taichi Wheel')
 def build_wheel(python: Command, pip: Command, env: dict) -> None:
     '''
@@ -68,7 +88,7 @@ def build_wheel(python: Command, pip: Command, env: dict) -> None:
     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'
+        env['TAICHI_CMAKE_ARGS'] += ' -DTI_WITH_C_API=ON'
 
     if platform.system() == 'Linux':
         if is_manylinux2014():
@@ -85,15 +105,18 @@ def build_wheel(python: Command, pip: Command, env: dict) -> None:
 
 def main() -> None:
     env = {
+        'PATH': os.environ['PATH'],
+        'LD_LIBRARY_PATH': os.environ.get('LD_LIBRARY_PATH', ''),
         'TAICHI_CMAKE_ARGS': os.environ.get('TAICHI_CMAKE_ARGS', ''),
         'PROJECT_NAME': os.environ.get('PROJECT_NAME', 'taichi'),
     }
     setup_llvm(env)
+    setup_vulkan(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'])
+    python, pip = setup_python(env, os.environ['PY'])
     build_wheel(python, pip, env)
 
     sccache('-s')

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

@@ -69,7 +69,7 @@ def set_common_env():
 
 class _EnvironWrapper(_Environ):
     def __setitem__(self, name: str, value: str) -> None:
-        orig = self.get(name)
+        orig = self.get(name, '')
         _Environ.__setitem__(self, name, value)
         new = self[name]
 

.github/workflows/scripts/ci_common/python.py

@@ -27,7 +27,8 @@ def setup_miniforge3(prefix):
 
 
 @banner('Setup Python {version}')
-def setup_python(version: Optional[str] = None) -> Tuple[Command, Command]:
+def setup_python(env_out: dict,
+                 version: Optional[str] = None) -> Tuple[Command, Command]:
     '''
     Find the required Python environment and return the `python` and `pip` commands.
     '''
@@ -50,7 +51,7 @@ def setup_python(version: Optional[str] = None) -> Tuple[Command, Command]:
     if not exe.exists():
         conda.create('-y', '-n', version, f'python={version}')
 
-    os.environ['PATH'] = f'{env / "bin"}:{os.environ["PATH"]}'
+    env_out['PATH'] = f'{env / "bin"}:{env_out["PATH"]}'
     python = sh.bake(str(exe))
     pip = python.bake('-m', 'pip')
 

.github/workflows/scripts/ci_common/tinysh.py

@@ -1,15 +1,12 @@
 # -*- coding: utf-8 -*-
 
-# -- stdlib --
 import os
 import platform
+import sys
 from contextlib import contextmanager
 from typing import Any, Mapping, Sequence
 
-# -- third party --
-# -- own --
-
-# -- code --
+from .escapes import escape_codes
 
 # A minimal and naive imitiation of the sh library, which can work on Windows.
 # NOT written as a general purpose library, wild assumptions are made.
@@ -36,6 +33,9 @@ def __str__(self):
 ENVIRON_STACK = []
 PREFIX_STACK = []
 
+P = escape_codes['bold_purple']
+N = escape_codes['reset']
+
 
 class Command:
     def __init__(self, *args: str):
@@ -60,11 +60,23 @@ def __call__(self, *moreargs: Sequence[str]) -> None:
         for v in PREFIX_STACK:
             prefixes.extend(v)
 
-        env = os.environ.copy()
+        overlay = {}
         for v in ENVIRON_STACK:
-            env.update(v)
+            overlay.update(v)
 
         args = prefixes + args
+        cmd = ' '.join([quote(v) for v in args])
+
+        print(f'{P}:: RUN {cmd}{N}', file=sys.stderr, flush=True)
+        if overlay:
+            print(f'{P}>> WITH ADDITIONAL ENVS:{N}',
+                  file=sys.stderr,
+                  flush=True)
+            for k, v in overlay.items():
+                print(f'{P}       {k}={v}{N}', file=sys.stderr, flush=True)
+
+        env = os.environ.copy()
+        env.update(overlay)
 
         code = os.spawnvpe(os.P_WAIT, args[0], args, env)
         if code: